Skip to content

losses

losses

Training objectives: conditional flow matching and delta (guidance) alignment.

ConditionalFlowMatchingLoss 

ConditionalFlowMatchingLoss(interpolant=None, coupling=None, loss_type='l2', time_scale=1.0)

Bases: BaseLoss

Regress a velocity field onto the conditional target velocity of a path.

Conditional flow matching trains \(v_\theta\) to match the per-pair target velocity \(u_t\) of a probability path by minimising

\[ \mathcal{L} = \mathbb{E}_{t \sim \mathcal{U}[0,1],\, x_0,\, x_1} \bigl\| v_\theta(x_t, t) - u_t \bigr\|^2, \]

where \((x_t, u_t)\) are produced by the chosen interpolant (for the linear path, \(x_t = (1-t)x_0 + t x_1\) and \(u_t = x_1 - x_0\)). Although the target is only defined conditionally on \((x_0, x_1)\), its regression minimiser is the marginal velocity field that transports noise onto data, which is exactly the field the sampler integrates.

References

Lipman et al., "Flow Matching for Generative Modeling" (2023), https://arxiv.org/abs/2210.02747.

Parameters:

Name Type Description Default
interpolant Optional[BaseInterpolant]

the probability path to regress against. Defaults to LinearInterpolant.

 None
coupling Optional[BaseCoupling]

optional train-time coupling that produces \((x_0, x_1)\) pairs from a batch of \(x_1\). See deltaflow.trainer.coupling. Kept separate from interpolant on purpose, since the two decisions (which path to use, and how to pair noise with data) are independent and swappable.

 None
loss_type str

one of "l2", "l1", "huber".

 'l2'
time_scale float

scales the continuous \(t \in [0, 1]\) before it reaches the model, e.g. to match a diffusion-style time embedding.

 1.0
Source code in deltaflow/losses/conditional_flow_matching.py 
51
52
53
54
55
56
57
58
59
60
61
def __init__(
    self,
    interpolant: Optional[BaseInterpolant] = None,
    coupling: Optional["BaseCoupling"] = None,
    loss_type: str = "l2",
    time_scale: float = 1.0,
):
    self.interpolant = interpolant or LinearInterpolant()
    self.coupling = coupling
    self.loss_type = loss_type
    self.time_scale = time_scale

DeltaAlignmentLoss 

DeltaAlignmentLoss(projector, lambda_flow=1.0, lambda_align=5.0)

Bases: Module

Combined flow-matching and delta-alignment loss.

The total objective linearly combines the velocity-regression term with the multi-scale guidance-alignment term,

\[ \mathcal{L} = \lambda_\text{flow}\,\mathcal{L}_\text{flow} + \lambda_\text{align}\,\mathcal{L}_\text{align}, \]

where, at each hierarchy level \(l\), the alignment term compares the guidance-difference embeddings \(z^{(i)}_l = g_l\bigl(\text{GAP}(\Delta h^{(i)}_l)\bigr)\) of two augmented views \(i \in \{1, 2\}\) via a cosine dissimilarity,

\[ \mathcal{L}_\text{align} = \frac{1}{L}\sum_{l=1}^{L} \Bigl(1 - \cos\bigl(z^{(1)}_l, z^{(2)}_l\bigr)\Bigr), \qquad \Delta h_l = h^\text{cond}_l - h^\text{uncond}_l. \]

During the alignment phase, \(\mathcal{L}_\text{flow}\) is computed over all four velocity predictions (two views by two conditioning modes) to preserve both the guided and unguided generative pathways while the alignment term shapes the guidance representation.

References

Di Via et al., "CDPM-Align: Multi-Scale Guidance-Aligned Diffusion Pretraining for Robust Few-Shot Anatomical Landmark Detection" (2026), https://arxiv.org/abs/2606.04898.

Source code in deltaflow/losses/delta_alignment.py 
123
124
125
126
127
128
129
130
131
132
def __init__(
    self,
    projector: MultiScaleProjector,
    lambda_flow: float = 1.0,
    lambda_align: float = 5.0,
):
    super().__init__()
    self.projector = projector
    self.lambda_flow = lambda_flow
    self.lambda_align = lambda_align

EquilibriumMatchingLoss 

EquilibriumMatchingLoss(interpolant=None, coupling=None, loss_type='l2')

Bases: BaseLoss

Regress a time-invariant field onto the Equilibrium Matching target.

Equilibrium Matching (EqM) trains a field \(f_\theta\) to match an energy-compatible target along the straight noise-to-data path. For an interpolation coefficient \(\gamma \in [0, 1]\) (a noise level), noise \(x_0\) and data \(x_1\), the corrupted sample and target are produced by an EquilibriumInterpolant,

\[ x_\gamma = (1 - \gamma)\,x_0 + \gamma\,x_1, \qquad u_\gamma = c(\gamma)\,(x_1 - x_0), \]

and the objective is (the paper's Eq. 3, up to the target-sign convention below)

\[ \mathcal{L}_{\text{EqM}} = \mathbb{E}_{\gamma,\, x_0,\, x_1} \bigl\| f_\theta(x_\gamma) - u_\gamma \bigr\|^2 . \]

\(\gamma\) is implicit, and there is no time at all. Unlike flow matching, where the model is conditioned on time \(t\), an EqM field is time-invariant and noise-unconditional, \(f_\theta(x)\). The coefficient \(\gamma\) is not seen by the model, and neither is any surrogate time. This loss therefore queries the model as model(x, **cond), passing no time (and no \(\gamma\)) at all, matching BaseEquilibriumField.

Target sign. The interpolant returns \(u_\gamma = c(\gamma)(x_1 - x_0)\) (data minus noise), so the field points from noise toward data and is sampled by gradient ascent \(x \leftarrow x + \eta f(x)\) in EquilibriumSolver. This matches the official EqM code. The paper writes the mirror-image \((\epsilon - x)c(\gamma)\) with descent, the two conventions are equivalent under a global sign flip.

References

Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.

Parameters:

Name Type Description Default
interpolant Optional[BaseEquilibriumInterpolant]

the energy-compatible path. Defaults to EquilibriumInterpolant.

 None
coupling Optional[BaseCoupling]

optional train-time coupling that produces \((x_0, x_1)\) pairs from a batch of \(x_1\). See deltaflow.trainer.coupling.

 None
loss_type str

one of "l2", "l1", "huber".

 'l2'
Source code in deltaflow/losses/equilibrium_matching.py 
66
67
68
69
70
71
72
73
74
def __init__(
    self,
    interpolant: Optional[BaseEquilibriumInterpolant] = None,
    coupling: Optional["BaseCoupling"] = None,
    loss_type: str = "l2",
):
    self.interpolant = interpolant or EquilibriumInterpolant()
    self.coupling = coupling
    self.loss_type = loss_type

delta_alignment_loss 

delta_alignment_loss(feats_u1, feats_c1, feats_u2, feats_c2, projector)

Multi-scale alignment loss on guidance-difference embeddings.

For each hierarchy level l and each of two augmented views i in {1, 2}::

delta_h_l_i = h_cond_i[l] - h_uncond_i[l]
z_l_i       = projector_l(GAP(delta_h_l_i)), L2-normalized
loss_l      = 1 - cos(z_l_1, z_l_2)

Returns the average of loss_l over all levels present in every feature dict and in projector.

Source code in deltaflow/losses/delta_alignment.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def delta_alignment_loss(
    feats_u1: Dict[str, torch.Tensor],
    feats_c1: Dict[str, torch.Tensor],
    feats_u2: Dict[str, torch.Tensor],
    feats_c2: Dict[str, torch.Tensor],
    projector: MultiScaleProjector,
) -> torch.Tensor:
    """Multi-scale alignment loss on guidance-difference embeddings.

    For each hierarchy level ``l`` and each of two augmented views ``i in {1, 2}``::

        delta_h_l_i = h_cond_i[l] - h_uncond_i[l]
        z_l_i       = projector_l(GAP(delta_h_l_i)), L2-normalized
        loss_l      = 1 - cos(z_l_1, z_l_2)

    Returns the average of ``loss_l`` over all levels present in every
    feature dict and in ``projector``.
    """
    layer_set = projector.layer_set
    total_loss = torch.tensor(0.0, device=next(iter(feats_u1.values())).device)
    num_layers = 0

    for layer in layer_set:
        if not all(layer in feats for feats in (feats_u1, feats_c1, feats_u2, feats_c2)):
            continue
        if layer not in projector.projectors:
            continue

        delta_h1 = feats_c1[layer] - feats_u1[layer]
        delta_h2 = feats_c2[layer] - feats_u2[layer]

        z1 = projector.pool_feature(delta_h1)
        z2 = projector.pool_feature(delta_h2)

        z1 = projector.projectors[layer](z1)
        z2 = projector.projectors[layer](z2)

        z1 = safe_normalize(z1, dim=-1)
        z2 = safe_normalize(z2, dim=-1)

        layer_loss = clamp_loss(_cosine_dissimilarity(z1, z2), max_val=10.0)
        total_loss = total_loss + layer_loss
        num_layers += 1

    if num_layers == 0:
        return total_loss
    return total_loss / num_layers