On Variance Reduction in Learning Mean Flows

Abstract: One-step generative modeling has emerged as a leading approach to amortize the inference cost of diffusion and flow-matching models. Among distillation-free methods, MeanFlow training is notoriously unstable, with non-decreasing loss and unbounded gradient variance. In this work, we establish a theory that attributes this pathology to a misuse of the conditional velocity field: it plays two distinct statistical roles in the loss, both as an unbiased regression target and as a Monte Carlo control variate inside a Jacobi-vector product, with the original loss assigning the wrong coefficient to the latter. We derive the optimal coefficient in closed form, and show that a family of fixes in concurrent works corresponds to different practical realizations of the same optimum. A controlled sweep of this coefficient on two-dimensional benchmarks and on a latent Diffusion Transformer recovers the predicted bias-variance ordering. The optimal coefficient yields up to a %54 improvement in sample quality on two-dimensional benchmarks and a monotone FID trend at every matched-step DiT checkpoint. Crucially, the same DiT measurement also reveals a quantitative FID-MSE landscape mismatch: although gradient variance is minimized at an interior coefficient value, the coefficient that minimizes FID prefers the direct use of conditional velocity.

• The paper introduces an optimal mixing coefficient to reduce gradient variance by adjusting the dual role of the conditional velocity field.
• It explains that misapplying the conventional MeanFlow loss leads to uncontrolled variance amplified by the Jacobian-vector product.
• Empirical analyses on synthetic datasets and DiT models demonstrate that tuning the tangent improves training stability and sample quality.

Summary of "On Variance Reduction in Learning Mean Flows"

Problem Context and Theoretical Contributions

The paper "On Variance Reduction in Learning Mean Flows" (2605.09235) addresses instability issues in MeanFlow-based one-step generative modeling, particularly the empirical observation that gradient variance grows unbounded during training. The authors formulate a statistical theory explaining this pathology, attributing it to a misapplied coefficient when using the conditional velocity field both as a regression target and as a tangent inside a Jacobian-vector product (JVP). Crucially, the conditional velocity acts as a Monte Carlo control variate, but the conventional MeanFlow loss uses a statistically suboptimal coefficient.

The authors derive the optimal coefficient in closed form, showing that a variety of concurrent remedies from recent literature are interpretable as different practical realizations of this optimal coefficient. Their theoretical developments demonstrate that the Jacobi amplification factor inside the JVP induces uncontrolled variance, which is hidden from the optimizer due to the stop-gradient operation.

Variance Decomposition and Control Variate Theory

The analysis reveals that the conditional velocity field serves two distinct statistical roles: as an unbiased regression target and as a tangent for the JVP, amplifying variance through the Jacobi factor. Leveraging Reynolds decomposition and a detailed analysis of MeanFlow's per-sample loss, the paper shows that:

• Unbiased Target Role: Injects zero-mean conditional noise linearly, which forms an irreducible noise floor.
• Variance Amplifier Role: Multiplies this noise with the spatial Jacobian, leading to quadratic amplification and empirical instability.

A critical theorem establishes that gradient variance scales with $\Tr(Cov[\nabla_{\theta} \ell_{\text{MF}}]) \sim \Tr(J^{\top} J \Sigma_{v'})$, with $J$ as the Jacobi factor and $\Sigma_{v'}$ the covariance of velocity fluctuations. The stop-gradient operator suppresses part of the variance-driven gradient, resulting in a semi-gradient gap and non-decreasing loss behavior.

Figure 1: Spatial distribution of $\sqrt{\Tr(\Sigma_{v'})}$ shows variance concentration in mode-mixing regions, with amplification toward the latent endpoint.

Bias-Variance Trade-off and Optimal Coefficient Derivation

By formally treating the tangent as a control variate, the authors introduce a mixing coefficient $\beta$ that interpolates between a stochastic conditional tangent and a deterministic proxy. The optimal $\beta^\ast$ is shown to be:

$$\beta^\ast = \frac{\kappa}{\kappa+1} \cdot \frac{\sigma^2 d}{\sigma^2 d + \|\Delta\|^2}$$

where $\kappa$ depends on the spectral norm of the Jacobi factor, $\sigma^2 d$ is the total variance, and $\|\Delta\|^2$ represents proxy bias. This closed-form solution captures both noise cancellation and James-Stein shrinkage effects.

The theory clarifies empirical remedies: deterministic tangents, alternative Jacobian constructions, or loss schedule manipulations all aim to control the amplification described above. The paper provides a unified theoretical framework for recent advances such as AlphaFlow, Improved MeanFlow, Terminal Velocity Matching, and others.

Empirical Validation and Experimental Analysis

Controlled sweeps of the tangent-mixing coefficient $J$0 are performed on both low-dimensional synthetic datasets and on training DiT-B/4 architectures with ImageNet-256 latents.

• Gradient Variance Reduction: Experiments confirm monotonic gradient variance reduction as $J$1 increases. On high-curvature datasets (e.g., checkerboard, eight_gaussians), deterministic tangents yield up to $J$2 reduction; the reduction magnitude correlates with geometric curvature and the amplification factor $J$3.

  Figure 2: Empirical total gradient variance $J$4 shows consistent declines across datasets with increasing $J$5.

• Sample Quality and Bias-Variance Optima: Sample quality, measured by sliced Wasserstein-1 distance ($J$6), trades off against variance. On high-curvature datasets, $J$7 is near the deterministic-tangent extreme ($J$8), yielding up to $J$9 improvement in SW$\Sigma_{v'}$0. On low-curvature datasets, the variance is less amplified, and $\Sigma_{v'}$1 shifts toward zero, consistent with the theory.

  Figure 3: Empirical sample-quality as $\Sigma_{v'}$2 illustrates dataset-dependent bias-variance interplay and the empirical optimal $\Sigma_{v'}$3.

• Scaling with Feature Dimension: Extensive sweeps over DGMM datasets with increasing dimension confirm that in high $\Sigma_{v'}$4 regimes, the quality gap widens, and the interior optimum shifts accordingly.
• Large-Scale DiT Experiments: When training DiT-B/4 on ImageNet, per-step loss variance with stochastic tangents grows by two orders of magnitude with interpolation time $\Sigma_{v'}$5, while deterministic tangents remain flat.

  Figure 4: Per-step loss variance in DiT-B/4: stochastic tangents amplify variance, deterministic tangents suppress it as predicted by the Jacobi-factor theory.

• FID-MSE Landscape Mismatch: The empirical FID ordering aligns with bias-variance predictions, but the FID-optimal coefficient is at $\Sigma_{v'}$6 (unbiased tangent) despite the gradient-MSE optimum sitting at $\Sigma_{v'}$7. The FID landscape penalizes bias super-linearly, revealing a quantitative mismatch between minimizing MSE and optimizing FID.

Qualitative Results

The paper provides qualitative visualizations of generated samples under different $\Sigma_{v'}$8 values at matched training steps. These highlight the perceptual and metric effects of bias amplification.

Figure 5: Class-conditioned DiT samples from $\Sigma_{v'}$9 baseline, demonstrating lowest FID and best perceptual quality.

Figure 6: DiT samples from $\sqrt{\Tr(\Sigma_{v'})}$0 checkpoint, with intermediate FID values.

Figure 7: DiT samples from $\sqrt{\Tr(\Sigma_{v'})}$1 corner (deterministic tangent), showing highest FID and visible degradation.

Implications and Future Directions

Theoretical implications center on formalizing variance reduction via control variates in self-supervised generative learning. The results call for metric-aware $\sqrt{\Tr(\Sigma_{v'})}$2 schedules that interpolate between variance minimization and sample-quality optimization. Practically, deterministic tangent proxies (e.g., EMA velocity anchors) can stabilize training but degrade FID unless bias is carefully managed.

Future work involves:

• Developing adaptive per-sample or global schedules for $\sqrt{\Tr(\Sigma_{v'})}$3, informed by curvature and bias estimates.
• Extending theory and practice to more complex manifolds or anisotropic data distributions.
• Designing FID-aware losses to bridge the observed mismatch with gradient-MSE minimization.

The diagnosis applies broadly to any generative training scheme using total derivative bootstrapping in a MeanFlow/JVP-style objective.

Conclusion

This paper rigorously establishes the statistical structure underlying gradient variance amplification in MeanFlow training. By isolating two roles of the conditional velocity and deriving the optimal mixing coefficient for variance reduction, it unifies the empirical fixes from concurrent literature. Experimental evidence confirms the predicted bias-variance ordering, demonstrates significant quantitative improvements in sample quality and stability, and uncovers crucial metric-dependent landscape mismatches relevant for future generative model development.