PaperlayerView exact revision on arXiv

Source-linked AI summary

Smoother Flow Matching via Contrastive Trajectory Repulsion

Ziqi Jiang, Zhenqi He, Long Chen

arXiv:2610.01408v1cs.CVcs.LG

TL;DR

Trajectory crossings in Flow Matching can create abrupt velocity fields, making them difficult to learn and degrading few-step integration. CoFlow repels intermediate trajectories using negative samples and an equivalent stochastic-interpolant formulation. On ImageNet 256 × 256, the paper reports improved few-step generation; its experiments focus on class-conditional generation with SiT backbones.

  • Problem

    Crossing trajectories produce conflicting velocities and abrupt changes, while existing approaches address them through post-hoc distillation or endpoint coupling rather than direct regulation of intermediate trajectories.

    Source: §1 · p. 2

  • Method

    CoFlow uses negative samples to repel training trajectories through an SDE repulsive drift and its equivalent stochastic-interpolant formulation.

    Source: §1 · p. 2

  • Results

    CoFlow improves few-step generation on ImageNet 256 × 256 and is compatible with REPA and OT-CFM.

    Source: §5 · p. 9

  • Takeaways & Limitations

    When its repulsive drift vanishes at both endpoints, CoFlow preserves the transformation between prior noise and target data distributions.

    Source: §3.3 · p. 5

    Experiments focus on class-conditional ImageNet 256 × 256 with SiT backbones; effectiveness on other datasets, resolutions, and generation tasks remains unevaluated.

    Source: A LIMITATIONS · p. 13

Abstract

from arXiv · show

Trajectory crossing remains a critical bottleneck in Flow Matching (FM), and previous works typically view these crossings from a theoretical optimization perspective causing velocity averaging. They attempt to address it indirectly by post-hoc distillation or endpoint coupling, without explicitly regulating the intermediate trajectories. In this paper, we introduce a new network learning perspective: crossing points inherently induce large local Lipschitz constants in the target velocity field, leading to two drawbacks. First, high Lipschitz constants correspond to high-frequency signals in the velocity field that neural networks struggle to fit due to spectral bias. Second, they also imply drastic velocity variations, leading to severe numerical integration errors in few-step inference. To alleviate this, we propose CoFlow, a framework that introduces the contrastive learning paradigm into FM to explicitly repel trajectories during training, thereby lowering the local Lipschitz constants of the velocity field. Specifically, we formulate CoFlow from a Stochastic Differential Equation (SDE) perspective by injecting a repulsive drift term. This drift actively guides the forward process of positive samples away from negative trajectories, effectively reducing the local Lipschitz constant. Furthermore, we derive an equivalent stochastic interpolant formulation from this SDE, providing a simple and tractable design space to control the influence of negative samples. Extensive experiments on ImageNet 256x256 demonstrate that CoFlow significantly reduces FID compared to standard FM in few-step inference (e.g., 20 steps), with no added training overhead. The code can be accessed at: https://github.com/HKUST-LongGroup/CoFlow

1 INTRODUCTION

Trajectory crossings can make Flow Matching velocity targets abrupt, impairing learning and few-step integration. CoFlow explicitly repels intermediate trajectories, with experiments reporting improved few-step generation and a complementary inference-time connection to DeltaFM.

1 INTRODUCTION

  • Trajectory crossings create conflicting velocities and abrupt field changes, amplifying numerical solver errors and degrading few-step generation.Prior approaches address crossings through post-hoc distillation or endpoint coupling rather than directly adjusting intermediate trajectories during training.Source: §1 · p. 2
  • High local Lipschitz constants correspond to high-frequency velocity signals that neural networks struggle to fit and amplify ODE truncation error.The paper uses the local Lipschitz constant to quantify velocity-field smoothness.Source: §1 · p. 2
  • CoFlow uses negative samples to repel adjacent trajectories during training, aiming to lower local Lipschitz constants through a repulsive SDE drift.The authors derive an equivalent stochastic-interpolant formulation that controls negative-sample influence and preserves endpoint distributions when the drift vanishes at both endpoints.Source: §1 · p. 2
  • DeltaFM adds a contrastive loss to align velocities, whereas CoFlow adjusts intermediate trajectories; the paper characterizes DeltaFM’s optimal velocity as an inference-time affine transformation.That transformation can be combined with CoFlow at inference without extra training overhead.Source: §1 · p. 2
  • In a 2D Gaussian-to-Two-Moons example, CoFlow produces lower-curvature paths and reduces the learned field’s local Lipschitz constant and relative truncation error versus standard FM.Figure 3 presents the trajectory, Lipschitz-constant, and truncation-error comparisons.Source: Figure 3 · p. 3
  • On ImageNet-256, the paper reports that CoFlow improves few-step image-generation fidelity and is compatible with REPA, OT-CFM, and DeltaFM.The cited introduction describes extensive experiments but does not provide benchmark values.Source: §1 · p. 3

2 RELATED WORKS

The related work places CoFlow alongside approaches to few-step generation and contrastive learning in generative models. It presents direct adjustment of intermediate training trajectories as complementary to existing approaches.

2 RELATED WORKS

  • Few-step generation methods include coupling optimization, post-hoc rectification or distillation, and mean-velocity approaches using extra conditioning.CoFlow is presented as complementary: it adjusts intermediate training trajectories with negative samples to encourage smoother velocity fields and support few-step sampling.Source: §2 · p. 3
  • Contrastive learning relies on negative samples and is described as underexplored in continuous-time generative-model training.The text notes that classifier-free guidance acts as a contrastive mechanism at inference time, pushing conditional trajectories away from unconditional ones.Source: §2 · pp. 3–4
  • Drifting Models and DeltaFM are cited as pioneering work applying contrastive learning ideas to generative models.The passage names these methods without specifying their mechanisms.Source: §2 · p. 4

3 METHOD

The method connects stochastic differential equations with stochastic interpolants to derive CoFlow’s trajectory repulsion. Its design separates repulsion strength, timing, and direction, then trains a velocity field against a closed-form target.

3 METHOD

  • The method derives CoFlow first as an SDE with repulsive drift and then as an equivalent stochastic-interpolant formulation for training.The authors use the SDE view for physical intuition and the interpolant view as a flexible design space.Source: §3 · p. 4

3.1 PRELIMINARIES

  • The paper defines velocity-field smoothness by a Lipschitz bound and relates lower local Lipschitz constants to lower Euler-step truncation error.For fixed step size, suppressing the local Lipschitz constant effectively decreases the stated error bound and improves few-step sampling performance.Source: §3.1 · p. 4

3.2 FIXED ENDPOINTS VIA DOOB’S h-TRANSFORM

  • Doob’s h-transform conditions a forward SDE on a fixed terminal endpoint, adding an attractive drift that guides the trajectory toward that endpoint.In the stated linear configuration, the attractive force grows without bound as time approaches 1, guaranteeing termination at the target.Source: §3.2 · p. 5

3.3 COFLOW

  • CoFlow adds a repulsive force to the drift so a target trajectory moves away from a nearby negative trajectory with divergent target velocity.The modified SDE is equivalent to a modified stochastic interpolant; a reparameterization directly defines a tractable trajectory to avoid computationally intractable integral sampling during training.Source: §3.3 · p. 5
  • When the adjustment vanishes at both time boundaries, CoFlow guarantees an unbiased transformation between prior noise and target data distributions.The paper focuses on this regime for comparison with standard FM; it also identifies intentional boundary shifts as a broader design possibility.Source: §3.3 · p. 5

3.4 DESIGN SPACE

  • CoFlow’s design space separates repulsion into strength, time schedule, and negative signal; the schedule is set to zero at both endpoints.The paper describes state, endpoint, and velocity residuals as alternative negative-signal formulations.Source: §3.4 · p. 6
  • The training algorithm samples paired trajectories, adjusts the target interpolant using a negative trajectory, and regresses a velocity network to its target velocity.The paper summarizes this training pipeline in Algorithm 1.Source: §3.4 · p. 6
  • The method considers random, nearest, and largest-Lipschitz policies for selecting a reference trajectory from the mini-batch.Its defaults are linear interpolation, random reference sampling, and the state residual as the negative signal.Source: §3.4 · p. 6

4 EXPERIMENTS

Experiments evaluate CoFlow on class-conditional ImageNet 256×256, comparing it with standard flow matching and testing compatibility with other methods and design choices. The results report improved few-step generation and benefits from combining CoFlow with DeltaFM or OT-CFM.

4.1 COMPARISON WITH STANDARD FLOW MATCHING

  • At NFE=20 without REPA, FID falls from 61.43 to 52.79 on SiT-B/2 and from 38.50 to 29.65 on SiT-XL/2.The other four metrics also improve, and the improvements are retained with REPA.Source: §4.1 · p. 7, Table 1 · p. 8
  • CoFlow outperforms standard FM at low NFEs, but their FID is similar at 250 NFEs as the gap narrows with increased sampling budget.The passage attributes the trend to finer discretization reducing numerical truncation errors for both methods.Source: §4.1 · p. 7, Figure 5 · p. 8

4.2 COMPATIBILITY WITH OTHER METHODS

  • Velocity centering yields performance comparable to DeltaFM; on CoFlow, it reduces FID from 37.76 to 27.74 and outperforms DeltaFM on all three metrics.The velocity-centering operation is training-free and applies a transformation to a trained standard FM model.Source: §4.2 · p. 8, Table 2 · p. 8
  • With OT-CFM, CoFlow reduces FID from 69.90 to 63.92 and sFID from 38.87 to 25.96.The experiment used SiT-B/2 trained for 400K steps without REPA.Source: §4.2 · p. 8, Table 3 · p. 8

4.3 COMPONENT-WISE ABLATION

  • The state residual performs best across all three metrics, achieving FID 79.55; random-noise perturbation shifts FID only from 86.99 to 86.49.The authors hypothesize that state residuals better reduce local Lipschitz constants by using relative trajectory positions.Source: §4.3 · p. 9, Table 4 · p. 9
  • All three negative-sampling policies improve on standard FM; Lipschitz achieves the lowest FID of 78.12, while random sampling is competitive and chosen for simplicity.Nearest and Lipschitz negatives are selected using endpoint states to avoid time-varying selections that may cause interpolant discontinuities.Source: §4.3 · p. 9, Table 4 · p. 9
  • Increasing λ from 0.05 to 0.5 improves FID from 85.66 to 75.06, but λ=1.0 degrades FID to 87.43, slightly worse than standard FM.The results indicate that moderate repulsion improves generation quality, whereas an excessively strong adjustment can be detrimental.Source: §4.3 · p. 9, Table 4 · p. 9

5 CONCLUSION

CoFlow repels crossing training trajectories using negative samples and adjusts intermediate paths to encourage a smoother velocity field. The paper reports improved few-step ImageNet generation and compatibility with complementary methods.

5 CONCLUSION

  • CoFlow uses negative samples to repel crossing training trajectories and encourage a smoother velocity field.The framework is inspired by the Doob-h transform and derives a closed-form stochastic interpolant that preserves endpoint distributions and enables simulation-free training.Source: §5 · p. 9
  • ImageNet 256 × 256 experiments demonstrate improved few-step generation and compatibility with REPA and OT-CFM.The paper also reports that inference-time velocity adjustment equivalent to DeltaFM can further improve CoFlow without additional training.Source: §5 · p. 9
  • The authors conclude that contrastive learning can adjust training trajectories to improve few-step generation in flow matching.Source: §5 · p. 9

AI USE STATEMENT

The authors used ChatGPT only for language polishing and grammar improvements. They report reviewing and verifying all AI-assisted revisions and take responsibility for the final content.

AI USE STATEMENT

  • ChatGPT was used solely for language polishing and grammatical improvements.Source: AI USE STATEMENT · p. 10
  • The authors report no AI use for research ideation, data collection, mathematical proofs, experimental implementation, or scientific claims.Source: AI USE STATEMENT · p. 10
  • The authors say they reviewed and verified all AI-assisted revisions and take full responsibility for the final content, claims, and artifacts.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The paper points to its training algorithm and design settings, reports where evaluation settings and component-wise ablations appear, and identifies appendices for supporting derivations and assumptions.

REPRODUCIBILITY STATEMENT

A LIMITATIONS

The experiments are limited to class-conditional ImageNet at 256 × 256 with SiT backbones, while the appendix passages describe marginal representation and CoFlow equivalence derivations.

A LIMITATIONS

  • Experiments focus on class-conditional ImageNet 256 × 256 with SiT backbones; other datasets, resolutions, and generation tasks remain unevaluated.The authors attribute this scope to limited computational resources.Source: A LIMITATIONS · p. 13

B MARGINAL EQUIVALENCE BETWEEN SDE AND SI

  • The SDE admits a representation in distribution at each fixed time, and the derivation introduces an integrating factor.The passage notes that matching one-time marginals does not ensure the same joint distribution across time when a fixed x1 defines the interpolant.Source: B MARGINAL EQUIVALENCE BETWEEN SDE AND… · pp. 13–14

C EQUIVALENCE OF COFLOW

  • The CoFlow equivalence derivation adds a drift term while retaining the shared SDE terms, which can then be simplified.The passages also identify a simulation-free reparameterization of the modified drift.Source: C EQUIVALENCE OF COFLOW · pp. 14–15

D UNBIASEDNESS OF COFLOW

CoFlow’s interpolant preserves the prior and data endpoints when its repulsive adjustment vanishes at the boundaries, yielding an unbiased distribution transformation. The proof identifies the marginal velocity through the continuity equation and shows that CoFlow training matches standard flow matching.

D UNBIASEDNESS OF COFLOW

  • Vanishing repulsive adjustment at both boundaries guarantees an unbiased transformation from prior noise to target data.The paper focuses on this regime for a fair comparison with standard FM.Source: D UNBIASEDNESS OF COFLOW · p. 15
  • The interpolant starts at the data sample and ends at the prior sample when the adjustment vanishes at both endpoints.The endpoint conditions are ˆb0(z) = 0 and ˆb1(z) = 0.Source: D UNBIASEDNESS OF COFLOW · p. 15
  • A test-function derivative matches the continuity equation and identifies the marginal velocity as the field that achieves the distribution transformation.The field averages over possible sample pairings that pass through the current state.Source: D UNBIASEDNESS OF COFLOW · pp. 15–16
  • Training CoFlow uses the same mean-squared-error objective as standard flow matching.The network is trained against the conditional velocity target.Source: D UNBIASEDNESS OF COFLOW · p. 16

E TRAINING-FREE EQUIVALENCE OF DELTAFM

The optimal DeltaFM velocity is an affine transformation of the optimal standard FM velocity, enabling inference-time velocity centering without additional model training. In the reported comparison, FM with velocity centering performs comparably to models trained with the DeltaFM objective.

E TRAINING-FREE EQUIVALENCE OF DELTAFM

  • The DeltaFM objective uses an independently sampled data-noise pair as a negative alongside the positive pair.The positive input is constructed as x_t = α_t x̂ + σ_t ε; the negative pair is independently sampled from the batch.Source: E TRAINING-FREE EQUIVALENCE OF DELTAFM · p. 16
  • The optimal DeltaFM velocity field is an affine transformation of the optimal standard FM field.Applying the transformation to a learned standard FM velocity gives an approximation to the optimal DeltaFM field.Source: E TRAINING-FREE EQUIVALENCE OF DELTAFM · p. 17
  • Velocity centering applies the transformation at inference using an estimated data mean and requires no additional model training.The data mean is estimated by averaging SD-VAE latents across the training set.Source: E TRAINING-FREE EQUIVALENCE OF DELTAFM · p. 17
  • FM+VC achieves performance comparable to models trained with the DeltaFM objective.Table 5 compares models trained with REPA and uses NEF=50 for inference; VC is applied at inference time.Source: E TRAINING-FREE EQUIVALENCE OF DELTAFM · p. 17

F LIPSCHITZ CONSTANTS AND EULER DISCRETIZATION ERROR

Under bounded-velocity and spatial- and temporal-Lipschitz assumptions, the single-step Euler error has a bound that tightens when the spatial Lipschitz constant decreases, with the other constants fixed.

F LIPSCHITZ CONSTANTS AND EULER DISCRETIZATION ERROR

Loading 2610.01408v1…