On Variance Reduction in Learning Mean Flows

• URL: http://arxiv.org/abs/2605.09235v1
• Date: Sun, 10 May 2026 00:32:53 GMT
• Title: On Variance Reduction in Learning Mean Flows
• Abstract summary: MeanFlow training is notoriously unstable, with non-decreasing loss and unbounded gradient variance.<n>We establish a theory that attributes this pathology to a misuse of the conditional velocity field.<n>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.
• Score: 7.705995307512647
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Sobolev Regularized Score Difference Estimation in Diffusion Models [3.8546389874826947]
  Estimating the difference of two Stein's score functions is a fundamental problem in generative modeling.<n>We propose a statistically consistent and scalable estimator for score differences based on Sobolev regularization.<n>We demonstrate its effectiveness on real-world tasks, including transfer learning for ECG signal generation.
  arXiv  Detail & Related papers  (2026-08-18T18:23:30Z)
• Constrained Diffusion Models with Primal-Dual Inference [76.64466170898369]
  We develop constrained diffusion models with primal-dual inference (PDI)<n>PDI infers the optimal primal distribution and its parametrizing dual variable.<n>We evaluate PDI on constrained sampling from a mixture of Gaussians, wireless resource allocation, and portfolio management.
  arXiv  Detail & Related papers  (2026-06-15T18:32:33Z)
• Sobolev Regularized MMD Gradient Flow [22.929984194375805]
  We propose a regularized variant of the gradient-regularized MMD flow.<n>The proposed regularization mitigates the difference between the mean kernel mean embeddings.<n>A key highlight of the proposed flow is that it is applicable in both sampling (from an unnormalized target distribution)
  arXiv  Detail & Related papers  (2026-05-12T09:58:19Z)
• Temporal Pair Consistency for Variance-Reduced Flow Matching [13.328987133593154]
  Temporal Pair Consistency (TPC) is a lightweight variance-reduction principle that couples velocity predictions at paired timesteps along the same probability path.<n>Instantiated within flow matching, TPC improves sample quality and efficiency across CIFAR-10 and ImageNet at multiple resolutions.
  arXiv  Detail & Related papers  (2026-02-04T00:05:21Z)
• Variational Entropic Optimal Transport [67.76725267984578]
  We propose Variational Entropic Optimal Transport (VarEOT) for domain translation problems.<n>VarEOT is based on an exact variational reformulation of the log-partition $log mathbbE[exp(cdot)$ as a tractable generalization over an auxiliary positive normalizer.<n> Experiments on synthetic data and unpaired image-to-image translation demonstrate competitive or improved translation quality.
  arXiv  Detail & Related papers  (2026-02-02T15:48:44Z)
• Statistical Properties of Rectified Flow [0.7136933021609079]
  Rectified flow is a method for defining a transport map between two distributions.<n>We study structural properties of the rectified flow, including existence, uniqueness, and regularity.<n>We are able to establish convergence at faster rates than the ones for the usual nonparametric regression and density estimation.
  arXiv  Detail & Related papers  (2025-11-05T05:09:12Z)
• Malliavin Calculus with Weak Derivatives for Counterfactual Stochastic Optimization [16.179901221618156]
  We study counterfactual optimization of conditional loss functionals under misspecified noisy gradient information.<n>We show using Malliavin calculus that the conditional loss functional of a diffusion process admits an exact representation as a Skorohod integral, yielding variance comparable to classical Monte-Carlo.
  arXiv  Detail & Related papers  (2025-09-30T21:37:54Z)
• On the Wasserstein Convergence and Straightness of Rectified Flow [54.580605276017096]
  Rectified Flow (RF) is a generative model that aims to learn straight flow trajectories from noise to data.<n>We provide a theoretical analysis of the Wasserstein distance between the sampling distribution of RF and the target distribution.<n>We present general conditions guaranteeing uniqueness and straightness of 1-RF, which is in line with previous empirical findings.
  arXiv  Detail & Related papers  (2024-10-19T02:36:11Z)
• TIC-TAC: A Framework for Improved Covariance Estimation in Deep Heteroscedastic Regression [109.69084997173196]
  Deepscedastic regression involves jointly optimizing the mean and covariance of the predicted distribution using the negative log-likelihood. Recent works show that this may result in sub-optimal convergence due to the challenges associated with covariance estimation. We study two questions: (1) Does the predicted covariance truly capture the randomness of the predicted mean? Our results show that not only does TIC accurately learn the covariance, it additionally facilitates an improved convergence of the negative log-likelihood.
  arXiv  Detail & Related papers  (2023-10-29T09:54:03Z)
• Joint control variate for faster black-box variational inference [32.38477249925875]
  Black-box variational inference performance is sometimes hindered by the use of gradient estimators with high variance. This variance comes from two sources of randomness: Data subsampling and Monte Carlo sampling.
  arXiv  Detail & Related papers  (2022-10-13T18:36:20Z)
• Efficient CDF Approximations for Normalizing Flows [64.60846767084877]
  We build upon the diffeomorphic properties of normalizing flows to estimate the cumulative distribution function (CDF) over a closed region. Our experiments on popular flow architectures and UCI datasets show a marked improvement in sample efficiency as compared to traditional estimators.
  arXiv  Detail & Related papers  (2022-02-23T06:11:49Z)
• Comparing Probability Distributions with Conditional Transport [63.11403041984197]
  We propose conditional transport (CT) as a new divergence and approximate it with the amortized CT (ACT) cost. ACT amortizes the computation of its conditional transport plans and comes with unbiased sample gradients that are straightforward to compute. On a wide variety of benchmark datasets generative modeling, substituting the default statistical distance of an existing generative adversarial network with ACT is shown to consistently improve the performance.
  arXiv  Detail & Related papers  (2020-12-28T05:14:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.