Smoother Flow Matching via Contrastive Trajectory Repulsion
Flow Matching (FM) suffers from trajectory crossings that create high-frequency velocity targets and amplify numerical integration errors during few-step sampling. This paper proposes CoFlow, a framework that injects a repulsive contrastive drift into the forward SDE to push training trajectories away from negative samples, thereby smoothing the learned velocity field. The authors derive a simulation-free stochastic interpolant, prove that the transformation preserves the endpoint distributions when boundary terms vanish, and show consistent improvements on ImageNet 256x256 at 20 sampling steps without extra training cost.
The paper offers a conceptually elegant and theoretically grounded approach to a well-known practical bottleneck in flow matching. The SDE formulation via Doob's h-transform intuition and the derived closed-form interpolant are rigorous, and the ImageNet experiments demonstrate clear few-step improvements. However, the key claim that CoFlow reduces the local Lipschitz constant of the learned velocity field is directly validated only in a 2D toy experiment, leaving its verification in the high-dimensional image regime to indirect FID evidence.
The strongest aspects are the rigorous derivation of the interpolant and the practical simulation-free reparameterization. The authors prove unbiasedness when repulsion vanishes at the boundaries, provide a clean training algorithm, and show broad compatibility with REPA, OT-CFM, and DeltaFM. The finding that DeltaFM reduces to an inference-time affine transformation is a sharp theoretical insight that turns an expensive training modification into a zero-cost post-processing step.
First, the central mechanistic hypothesis that the method lowers the local Lipschitz constant of the learned velocity field is substantiated only in a 2D Gaussian-to-two-moons example (Fig. 3). On ImageNet, the paper reports no direct Lipschitz estimates or truncation-error measurements, so the empirical improvements, while real, are not firmly tied to the motivating theoretical quantity.
Second, the method introduces a hyperparameter $\lambda$ whose optimal value appears sensitive to model scale: Table 4 shows SiT-S/2 reaches its best FID at $\lambda=0.5$, whereas the larger B/2 and XL/2 models use $\lambda=0.2$. This scale-dependent tuning, combined with a restricted design space imposed by the boundary-vanishing requirement and time-independent negative sampling, suggests the method may need careful calibration in new settings.
The ImageNet ablations support the claim that state-residual repulsion helps few-step generation, but they do not isolate Lipschitz reduction as the cause. The comparison with DeltaFM is handled fairly: the authors both reproduce DeltaFM training and derive an equivalent inference-time velocity-centering transformation, the latter confirmed in Appendix Table 5. CoFlow consistently outperforms standard FM at low NFE and matches it at high NFE (Fig. 5), which is consistent with a numerical-discretization story, yet direct evidence linking the FID gap to lower $L$ would strengthen the argument considerably.
Algorithm 1 and Section 3.4 provide exact training steps, and the authors release an anonymous code repository. Reproduction requires standard ImageNet 256x256, a pre-trained SD-VAE latent encoder, and SiT backbones trained for 400K steps at batch size 256. Because the method is a lightweight modification of the interpolant and target velocity, it can be plugged into existing FM implementations with minimal effort. The main barrier is computational: reproducing the SiT-XL/2 run demands substantial GPU resources.
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
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.