Smoother Flow Matching via Contrastive Trajectory Repulsion
Ziqi Jiang, Zhenqi He, Long Chen
cs.CV, cs.LG
2026-10-01
CoFlow adds contrastive repulsion to Flow Matching to lower local Lipschitz constants. SiT-XL 20-step ImageNet FID drops from 38.50 to 29.65 with no extra training cost.
Flow Matching draws nearly straight paths from noise to data, so it should need far fewer sampling steps than diffusion. High-fidelity images still often take 50 to 100 steps. The usual account is trajectory crossing: two paths sit close in the middle while pointing at different endpoints, the marginal velocity averages conflicting directions, and the field turns sharply in space. Earlier work treated that as velocity averaging in the objective. This paper turns it into a local smoothness number you can actually compute.
Prior fixes stay at the boundary. Reflow and distillation (ReFlow, DMD) skip intermediate integration after training. Coupling methods such as OT-CFM re-pair noise and data to cut crossings at the endpoints. The path in between is rarely constrained. This HKUST paper treats the same crossings as a learning problem. Near a collision, nearby points have very different velocity targets, so the local Lipschitz constant of the target field blows up. That constant measures how much velocity can differ between nearby points. A large L injects high-frequency content that networks fit poorly because of spectral bias. It also scales up the Euler truncation-error bound, so coarse few-step ODE solves accumulate large mistakes. Exact intersections are rare in high dimension; crossing here means trajectories that sit close together with divergent target velocities.
CoFlow brings contrastive learning into Flow Matching. Another trajectory in the mini-batch is a negative. Training pushes the positive path away from it, which is meant to bring the local Lipschitz constant down.
The physics comes from Doob's h-transform. A conditional bridge SDE adds a drift that pulls a particle toward a chosen endpoint. CoFlow reverses that force: it subtracts a repulsive drift b(xt, t, xrt) from the forward SDE so the particle leaves the negative path. Integrating the drift during training is expensive, so the paper applies a closed-form tweak to the stochastic interpolant, writing xt as a function of time with no SDE simulation. The usual line xt = (1-t)x0 + t x1 becomes, by default, x̂t = xt - λ t(1-t)(xrt - xt). The target velocity is the analytic time derivative of x̂t, and the network still trains with MSE. The schedule β(t)=t(1-t) zeros the correction at both endpoints, so the noise and data marginals stay unbiased. Training remains simulation-free. Inference matches standard FM, with no extra cost.
The state residual pushes along the relative position of the two paths at the same t, which matches the Lipschitz definition directly. Endpoint residuals have lower variance and velocity residuals echo DeltaFM; both lose to the state residual in the ablation. The other knobs are strength λ and the time schedule β(t). Negatives can be drawn at random, as nearest neighbors on endpoints, or as the pair that would produce the largest local Lipschitz constant. Defaults are λ=0.2, state residual, and random negatives.
DeltaFM, a concurrent method, also uses negatives, but it penalizes velocity similarity and leaves intermediate paths unchanged. The paper proves that DeltaFM's optimal field is an affine transform of the optimal FM field. That transform can be applied at inference as velocity centering (VC): center the velocity with the training-set mean in VAE latent space, with no retraining. VC stacks on a trained CoFlow model.
The main benchmark is class-conditional ImageNet 256×256 with SiT backbones and a frozen SD-VAE, 400K steps, global batch 256. Default eval uses a 20-step Euler-Maruyama sampler, 50K images, and no CFG.
| Setting | FID↓ | sFID↓ | IS↑ |
| SiT-B/2 FM | 61.43 | 35.18 | 27.46 |
| SiT-B/2 CoFlow | 52.79 | 21.71 | 35.00 |
| SiT-B/2 FM+REPA | 47.64 | 34.33 | 40.75 |
| SiT-B/2 CoFlow+REPA | 37.76 | 20.23 | 52.92 |
| SiT-XL/2 FM | 38.50 | 29.27 | 50.14 |
| SiT-XL/2 CoFlow | 29.65 | 17.41 | 64.12 |
| SiT-XL/2 FM+REPA | 35.35 | 27.98 | 63.13 |
| SiT-XL/2 CoFlow+REPA | 27.78 | 17.88 | 80.52 |
sFID moves even more: 29.27 to 17.41 on SiT-XL/2. IS rises from 50.14 to 64.12. The gain survives REPA, so repulsion is not spending the same budget as representation alignment.
A 2D Gaussian-to-two-moons run backs the Lipschitz story. CoFlow paths look straighter. Mean local Lipschitz of the learned field drops 32.1% versus standard FM. Relative Euler truncation error falls 30.4%, 25.8%, 24.0%, 23.5%, and 23.1% at 5, 10, 20, 30, and 50 steps.
On the NFE curve CoFlow leads at low step counts; by 250 NFE the two FIDs are close. A smoother field mainly cuts discretization error. Once the grid is fine, that edge fades.
The method stacks with neighbors. Adding CoFlow to OT-CFM moves FID from 69.90 to 63.92 and sFID from 38.87 to 25.96. On SiT-B/2 + REPA, CoFlow+VC takes FID from 37.76 to 27.74 and sFID to 8.68, beating a from-scratch DeltaFM model (FID 32.32) on all three metrics. At NFE=50, VC on a plain FM model already matches DeltaFM: SiT-B/2 scores 27.64, 20.40, and 19.97 for FM, DeltaFM, and FM+VC; SiT-XL/2 scores 11.40, 7.29, and 7.38.
Ablations use SiT-S/2. Random noise as the negative signal barely moves FID (86.99 to 86.49), so jitter alone is not the story. State residual wins at 79.55; endpoint residual hits 92.24 and velocity residual 122.19, both worse than the FM baseline. All three sampling policies help. Lipschitz sampling is best at 78.12 FID; random sits at 79.55 and is kept as the default. Sweeping λ through 0.05, 0.2, 0.5, and 1.0 gives FID 85.66, 79.55, 75.06, and 87.43. Moderate repulsion helps; λ=1.0 is slightly worse than standard FM.
Few-step generation is already crowded with distillation and coupling. CoFlow edits the middle of the training path. The implementation is a closed-form interpolant change with no extra compute, so anyone already training SiT or vanilla FM can drop it in as a trajectory regularizer. It composes with REPA, OT-CFM, and DeltaFM/VC.
The gain is incremental. The setup is still multi-step ODE sampling, not MeanFlow or Shortcut one-step generation. Twenty-step ImageNet FID without CFG remains in the 27 to 52 band, far from production high-res synthesis. What carries over is a diagnosis: some of the few-step collapse comes from a locally steep velocity field, and contrastive repulsion is a cheap way to flatten it.
The authors restrict the image study to class-conditional ImageNet 256 with SiT. Other datasets, higher resolutions, and other generation tasks are untested. The repulsion schedule is fixed, and negatives are chosen from endpoints, independent of t, to keep a closed-form velocity target.
The main tables are almost all 20-step, no CFG. There is no head-to-head against ReFlow, DMD, or MeanFlow under one protocol. Standalone OT-CFM scores FID 69.90, worse than standard FM at 61.43, so the stacking baseline is weak. The main runs freeze λ at 0.2, while SiT-S/2 prefers 0.5, so the headline tables may be under-tuned. λ=1.0 already hurts, and a new backbone will likely need a retune. On SiT-XL/2 + REPA, Recall slips from 0.55 to 0.53 with no discussion. The 2D Lipschitz and truncation-error drops are never remeasured in ImageNet latent space. How often local Lipschitz actually explodes in high dimension is still a motivation, not a measurement. The unbiased design pins the correction to zero at both ends; biased boundaries in the spirit of CFG are left for later.