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

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