acceptodds
Under review as a conference paper at ICLR 2027

What Does One-Step Flow Matching Learn with Adaptive Weighting?

Abstract

One-step generative models replace iterative sampling with one network evaluation, but training them remains challenging. MeanFlow uses adaptive loss weighting to stabilize training, downweighting large residuals that can reflect valid target variation rather than prediction error. Through improved MeanFlow's velocity-regression formulation, we characterize the selected conditional law, giving necessary and sufficient conditions on radial weights for it to have no greater total variance than the original law at every stationary prediction, accounting for mean shifts. Population flows reveal changes in generated proportions even with convex regression and no neural fitting, while image experiments show that similar FID can conceal substantial shifts in class composition. Our prediction weighting removes target dependence, retaining adaptive scaling and the pointwise conditional mean while better matching training class proportions in our unconditional generation experiments. Comparisons on CIFAR-10 and ImageNet-LT distinguish target preservation from sample quality, establishing weighting as a generative modeling choice rather than merely a stabilization technique.

Then back it, or bet against it.

Related papers

Open the market on this paper to see 7 more related papers.