English

On Variance Reduction in Learning Mean Flows

Machine Learning 2026-05-12 v1 Artificial Intelligence Machine Learning

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.

Keywords

Cite

@article{arxiv.2605.09235,
  title  = {On Variance Reduction in Learning Mean Flows},
  author = {Juanwu Lu and Ziran Wang},
  journal= {arXiv preprint arXiv:2605.09235},
  year   = {2026}
}

Comments

25 pages, 7 figures, 6 tables

R2 v1 2026-07-01T13:01:00.483Z