sigw_loss

A single PyTorch loss function — sigw_loss — for relational alignment between text and vision tokens in a multimodal model. It takes attention matrices produced during the forward pass and returns a scalar loss you add to your training objective. No new parameters, no inference-time cost (use it during training only).

Source: sigw_loss.py.

Requirements

• Python 3
• PyTorch (torch)

The module has no other dependencies — just import torch.

Quick start

import torch
from sigw_loss import sigw_loss

n_t, n_v = 16, 64                                          # #text tokens, #vision tokens
A_text   = torch.softmax(torch.randn(n_t, n_t), dim=-1)    # text self-attention   [n_t, n_t]
A_vision = torch.softmax(torch.randn(n_v, n_v), dim=-1)    # vision self-attention [n_v, n_v]
C        = torch.softmax(torch.randn(n_t, n_v), dim=-1)    # text->vision attention [n_t, n_v]

loss, D_t, D_v, Dv_hat = sigw_loss(A_text, A_vision, C)
print(loss)        # scalar tensor

API

sigw_loss(A_text, A_vision, C, eps=1e-8) -> (loss, D_t, D_v, Dv_hat)

Inputs

Name	Shape	Description
A_text	[n_t, n_t]	Text self-attention (raw; need not be symmetric).
A_vision	[n_v, n_v]	Vision self-attention (raw; need not be symmetric).
C	[n_t, n_v]	Text→vision cross-attention. C[i, j] = text token i attending to vis j.
eps	float	Numerical-stability constant. Default 1e-8.

A_text / A_vision must be square, and C must be [n_t, n_v] matching them — otherwise an AssertionError is raised.

Returns

Name	Shape	Description
loss	scalar	The value to backprop (add to your LM loss).
D_t	[n_t, n_t]	Text distance matrix (for logging/visualization).
D_v	[n_v, n_v]	Current vision distance matrix (for logging).
Dv_hat	[n_v, n_v]	Target vision distance matrix (for logging/diagnostics).

Only loss is needed for training; the three matrices are returned for inspection (e.g. comparing D_v against the target Dv_hat).

Use in a training loop

lam = 0.1                                  # weight of the regularizer

loss_sigw, *_ = sigw_loss(A_text, A_vision, C)
total_loss = lm_loss + lam * loss_sigw     # lm_loss = your usual generation loss
total_loss.backward()

Gradients flow into the vision branch through A_vision, so make sure that tensor carries grad. A_text and C act as the fixed teacher / coupling; whether they receive gradient depends on how you extract them upstream.

What it computes

Briefly, given the three attention matrices it:

1. Turns each self-attention matrix into a distance matrix: D = -log(sym(A) + eps) (A is clamped ≥ 0, A_text row-normalized, then symmetrized).
2. Row-normalizes C and forms the visual marginals b = colsum(C).
3. Computes the target visual geometry in closed form: Dv_hat = (Cᵀ · D_t · C) / (b · bᵀ).
4. Returns the discrepancy loss = mean((D_v - Dv_hat) ** 2).

Cost is O(n_t²·n_v + n_v²·n_t) (two matrix multiplications).

Notes

• eps guards the three divisions/logs (row-normalization, -log, and the b·bᵀ denominator).
• The loss is a mean of squared errors, so its scale is roughly independent of token count; tune lam to balance it against lm_loss.
• Inputs are taken straight from attention maps — no need to pre-normalize them yourself; the function clamps, normalizes, and symmetrizes internally.

Reference

This loss (SI-GW) is the regularizer from the MIRROR method (ECCV 2026). A related variant relaxed_gwot_loss and the full objective gwot_loss_from_attentions live in gwot_loss.py.