Gradient Variance Reveals Failure Modes in Flow-Based Generative Models

TL;DR

This paper reveals that flow-based generative models with straight-path objectives can memorize training pairs due to low gradient variance, leading to failure modes, which can be mitigated by adding noise.

Contribution

It analyzes how low gradient variance causes memorization in flow models and demonstrates that noise injection can restore proper generalization.

Findings

Deterministic training leads to memorization of training pairs.

Adding small noise during training improves generalization.

Memorization occurs even when interpolant lines intersect.

Abstract

Rectified Flows learn ODE vector fields whose trajectories are straight between source and target distributions, enabling near one-step inference. We show that this straight-path objective conceals fundamental failure modes: under deterministic training, low gradient variance drives memorization of arbitrary training pairings, even when interpolant lines between pairs intersect. To analyze this mechanism, we study Gaussian-to-Gaussian transport and use the loss gradient variance across stochastic and deterministic regimes to characterize which vector fields optimization favors in each setting. We then show that, in a setting where all interpolating lines intersect, applying Rectified Flow yields the same specific pairings at inference as during training. More generally, we prove that a memorizing vector field exists even when training interpolants intersect, and that optimizing the straight-path objective converges to this ill-defined field. At inference, deterministic integration reproduces the exact training pairings. We validate our findings empirically on the CelebA dataset, confirming that deterministic interpolants induce memorization, while the injection of small noise restores generalization.