Gradient Variance Reveals Failure Modes in Flow-Based Generative Models

TL;DR

This work reveals that gradient variance can diagnose fundamental failure modes in flow-based generative models, showing that deterministic training of Rectified Flows can cause memorization of training pairings even when interpolants intersect, thereby hindering true transport. The authors develop a Gaussian-to-Gaussian analytic framework and extend to finite data, proving that the minimizer can memorize exact data pairings (Prop. det_rect) and that stochastic interpolation or noisy pairings break memorization, restoring generalization. Empirical validation on Mixtures of Gaussians and CelebA demonstrates that deterministic interpolants induce memorization, while injecting noise improves OT alignment and sample quality; adversarial and simulated 1-ReFlow experiments further illustrate the perils of determinism and the efficacy of noise. These findings challenge the reliability of purely deterministic Rectified Flows and emphasize the need for stochasticity to guide models toward proper transport, with implications for OT-based learning and fast inference in flow-based generative modeling.

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.

Figures (12)

• Figure 1: Intuition behind the main results.
• Figure 2: Schematic representing a hypothetical loss ($L_\text{MC}$) landscape. The gradient variance of the loss acts as an indicator of solution quality. The schematic illustrates how the variance of the loss gradient reveals information about the optimality of the vector field under different interpolant types.
• Figure 3: Gradient variance over time for a vector field rotating by $120^\circ$, under various pairing types including rotated and random couplings. In the top right, the two Gaussians are shown along with sample trajectories, color-coded by integration time. Optimal pairings exhibit the highest variance, while random pairings do not yield the maximum variance.
• Figure 4: Gradient variance (with 95% confidence intervals) for Gaussian transport paired $\left(X_0, R_{30^\circ} X_0 + \mu\right)$ under different noise levels ($\sigma$ of the two interpolants mentioned in the title of the figures). We compare vector fields inducing $0^\circ$ (OT, blue), $30^\circ$ (pair-optimal, orange), and $60^\circ$ (non-OT, non-pair-optimal, green) rotations. Shaded regions show confidence intervals calculated over 100 bootstrap samples. As noise increases ($\sigma=0\to4$), the optimal transport (OT) field exhibits significantly reduced variance ($p<0.01$, paired t-test) while maintaining lower transport error.
• Figure 5: Gradient variance over time for several pairing types: $\left(X_0, {\bm{M}}_d^{1/2} {\bm{R}}_{k^\circ} X_0 + \mu\right)$ with $k \in \{0, 30, 60, 120, 180\}$, and $(X_0, X_1)$ with $X_1 \sim \mathcal{N}(\mu, {\bm{M}}_d)$, over the OT vector field. Shaded regions indicate 95% confidence intervals computed via 100 bootstrap samples. Notably, variance does not increase in regions where interpolant trajectories come closer together: for example, with $180^\circ$ rotation, all interpolants intersect at $t=1/2$, yet variance is lowest there. Random pairings exhibit lower variance under the OT field than structured pairings with $120^\circ$ (straight) or $180^\circ$ (not-straight) rotations, highlighting that variance is not simply a function of interpolant density.

Theorems & Definitions (15)

• Lemma 1
• Proposition 1: Informal
• Lemma 2: Idempotence of Rectified Flows
• Proposition 2: The Minimizer Memorizes
• Remark 1: Condition for Recovering the Same Couplings
• Remark 2: Noise Breaks Proof Assumptions
• proof
• proof
• proof
• proof