On the Stability and Convergence of Flow Map Training

Ishan Nath; Tom Stanic; Avishek Joey Bose; Gauthier Gidel; Quentin Bertrand

stat.MLOct 14, 2026

On the Stability and Convergence of Flow Map Training

Authors: Ishan Nath, Tom Stanic, Avishek Joey Bose, Gauthier Gidel, Quentin Bertrand

Organizations: Inria, ENS de Lyon, CNRS, Université Claude Bernard Lyon 1, LIP UMR 5668, 69342 Lyon Cedex 07, France · Université de Montréal · Mila - Quebec AI Institute · Imperial College London ·

Abstract

Flow maps are a recent method of probabilistic generation that amortize diffusion trajectories into one-step samplers, unlocking efficient generation. While progress has been made in improving their empirical performance, there is still a lack of understanding behind the mechanisms underlying their success. Specifically, many flow map learning methods employ stop gradient operators as a crucial part of their training target. Adaptive weighting also drastically improves the generation capabilities of flow maps. Although many interpretations have been proposed to support the inclusion of these heuristics, none have analyzed the impact of these assumptions on the training dynamics of flow maps themselves. In this work, we address this issue by providing grounded theory explaining why such tricks are required. Analyzing the functional gradient flow induced by the MeanFlow and Lagrangian losses, we show their convergence to the true flow map, settling the validity of these objectives. For finite updates, we highlight a severe roadblock to the usage of these losses in practice: the possible unboundedness of the update. Adaptive normalization not only provides bounded updates, but, combined with a specific weighting, additionally provides convergence guarantees towards a neighborhood of the solution. On synthetic data, we show that our functional gradient flow analysis exactly describes the true training dynamics for flow maps. Experiments on standard image datasets demonstrate that a proper reweighting scheme is essential for strong empirical performance, which is only achieved for a specific range of reweighting strengths.

Figures & tables

Figure 2: Synthetic Gaussian-to-Gaussian dataset. Snapshots of grid function learned average velocities (

u_{k}

, solid blue/orange lines) versus their theoretically predicted counterparts

u_{\tau}

at four optimization times

\tau

. The two show strong agreement across optimization times

\tau

and for both the MeanFlow (top) and Lagrangian average velocity target (bottom), and both relax towards the ground truth

u^{\star}

Figure 3: FFHQ dataset. The images are generated from a fixed seed, over different numbers of function evaluations (NFE), with flow maps learned with different normalization exponents $p$ .

	Bounded [-1pt]updates	CIFAR-10			FFHQ 64 $\times$ 64
$p$	Bounded [-1pt]updates	NFE 1	NFE 2	NFE 4	NFE 1	NFE 2	NFE 4
0.00	✗	$10.62$ $\pm\,0.21$	$10.00$ $\pm\,0.69$	$9.60$ $\pm\,0.61$	$15.38$ $\pm\,0.84$	$12.67$ $\pm\,0.60$	$12.49$ $\pm\,0.43$
0.25	✗	$6.81$ $\pm\,0.23$	$6.70$ $\pm\,0.30$	$6.49$ $\pm\,0.18$	$10.92$ $\pm\,0.43$	$9.46$ $\pm\,0.97$	$9.72$ $\pm\,0.68$
0.50	✓	$5.51$ $\pm\,0.06$	$5.55$ $\pm\,0.36$	$4.93$ $\pm\,0.37$	$8.93$ $\pm\,0.21$	$7.89$ $\pm\,0.68$	$8.15$ $\pm\,0.53$
0.75	✓	$\mathbf{5.46}$ $\pm\,0.06$	$\mathbf{5.05}$ $\pm\,0.16$	$4.78$ $\pm\,0.04$	$\mathbf{8.67}$ $\pm\,0.30$	$6.86$ $\pm\,0.91$	$\mathbf{7.69}$ $\pm\,0.74$
1.00	✓	$5.61$ $\pm\,0.20$	$5.17$ $\pm\,0.05$	$4.69$ $\pm\,0.04$	$10.11$ $\pm\,0.21$	$\mathbf{5.64}$ $\pm\,1.11$	$7.85$ $\pm\,0.72$

Table 1: FID

\downarrow

for MeanFlow with different normalization exponents

p

on CIFAR-10 and on FFHQ 64

\times

64. Results are reported as mean

\pm

standard deviation over three runs.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Synthetic Gaussian-to-Gaussian dataset. Each line corresponds to an optimization time $\tau$ for the empirically learned average velocity (left) and the derived closed-form average velocity (right). As the optimization time increases, the average velocities converge towards the true one $u^{\star}$ . The two average velocities, empirically learned and the closed-form, show strong agreement across optimization times $\tau$ and for both the MeanFlow (top) and Lagrangian target (bottom).

Figure 5: A comparison of the convergence rates of flow maps trained with different loss weightings

\rho=(t-s)^{a}

, versus their theorized trajectories. For

a=-1

and

a=-0.5

, the theoretical last diagonal hitting time is marked.

Figure 6: Instead of keeping the diagonal

u(x,s,s)

fixed, we vary it with optimization time, keeping it centered at the true velocity. The learned flow maps for

a<0

have a higher MSE against the true flow map, compared to

a=0

. We plot the errors achieved when the diagonal is fixed for comparison.

Figure 7: The choice of

\eta

impacts whether we learn an error of the form

(t-s)^{k}

. As

k

increases, a smaller

\eta

must be taken to ensure that gradient descent does not accumulate errors.

Figure 8: For the linear flow

v(x,t)=\lambda x

, we report the average relative

L^{2}(p_{s})

error over 5 seeds for average velocity networks against the ground truth

u^{\star}=(e^{\lambda(t-s)}-1)x/(t-s)

. Lagrangian is worse than MeanFlow for every

\lambda

tested, and is unstable at high

\lambda

$p$	NFE 1	NFE 2	NFE 4
0.00	$74.01\pm 8.71$	$25.82\pm 3.51$	$9.75\pm 1.16$
0.25	$56.61\pm 1.46$	$21.51\pm 0.85$	$8.14\pm 0.71$
0.50	$51.99\pm 8.56$	$19.24\pm 1.30$	$6.80\pm 0.15$
0.75	$42.54\pm 1.87$	$15.25\pm 1.06$	$\mathbf{5.94\pm 0.24}$
1.00	$\mathbf{39.58\pm 5.96}$	$\mathbf{13.94\pm 2.17}$	$6.35\pm 0.67$
1.25	$54.33\pm 7.59$	$21.00\pm 3.00$	$8.06\pm 1.17$

Appendix

Table 2: FID on CIFAR-10 for the Lagrangian objective with different normalization exponents

p

. Results are reported as mean

\pm

standard deviation over three runs.

	CIFAR distilled	FFHQ distilled
Teacher	CFM (EMA)	EDM VE (EMA)
Architecture	EDM2-style U-Net	EDM2-style U-Net
Parameters	54.7M	59.9M
Model Channels	128	128
Channel Mults	[2, 2, 2]	[1, 2, 2, 2]
Attention Dims	16, 8	16, 8

Appendix

Table 3: Training configuration for the CIFAR-10 and FFHQ experiments.

Flow Map Learning

Flow map learning trains neural models to predict finite-time state or distribution transitions directly, enabling fast simulation, control, and generation.

See all 49 papers

Sep 14, 2026

How I learned to stop worrying and love StopGrads: Stationarity, Convergence, and a case study on Flow Map Learning

Stopgrads are widely used in training machine learning models, but stopgrads can alter the gradient, stationary points and convergence guarantees of the original objective, which can make stopgrad training theoretically ungrounded. We introduce a stopgrad regression principle, which identifies a general template for stopgrad objectives with a closed-form characterization of stationary points and their uniqueness, unifying stopgrad objectives for flow maps, reinforcement learning, and diffusion samplers. We provide theoretical grounding for optimizing stopgrad flow map objectives by showing their unique stationary point is the true flow map, and showing positive convergence results for Eulerian and Lagrangian objectives, including MeanFlow and improved MeanFlow. Remarkably, we show that under functional semi-gradient flow, the learned flow map has a closed-form expression composing the initial flow map and the true flow map. We additionally use our stopgrad regression principle to propose modified stopgrad placements for flow map objectives which reduce training memory by 2x.

Max W. Shen, Mark Goldstein, Zichu Wang +2

Flatiron Institute · NYU

Flow Map Learning Stochastic Optimization Convergence Gradient Descent Dynamics

Ask PDF

Oct 21, 2026

Improved Training of Strong Stochastic Flow Maps

Stochastic flow maps promise few-step sampling of stochastic differential equations (SDEs) by learning to directly jump between times along the solution path, yet training these models to learn this solution map remains difficult. The primary obstacle is the Brownian path itself: it is an infinite-dimensional object which needs to be passed to the neural network and is nowhere differentiable, which makes the implementation of derivative-based losses that power the state-of-the-art deterministic flow maps just out of reach. We identify two core design principles which address these issues: a) represent the Brownian path with a basis which carries higher-order local information, and b) train with an objective that forces the network to use the full representation of the path, not just the first order. To address this, we leverage the shifted Legendre polynomial expansion with a Wong-Zakai styled objective, unlocking derivative-based losses. Finally, we propose a recipe for fine-tuning large pre-trained flow models into strong stochastic flow maps. We ablate through this design space on several toy problems and use our findings to transform a large flow model into a strong stochastic flow map on ImageNet, achieving roughly a three-fold improvement in FD-DINOv2 over SSFM at four NFE.

Zander W. Blasingame, Antonio Franca, Raul Minan +3

AITHYRA · Technische Universität Wien · Technische Universität München +1

Flow Map Learning Stochastic Differential Equations

Ask PDF

Jul 29, 2026

Flow Map Learning via Nongradient Vector Flow

Diffusion and flow-based models benefit from simple regression losses, but inference incurs significant overhead because sampling requires integration. Consistency models address this by directly learning the flow maps along the ODE trajectory, opening a design space between one-step and many-step approaches. However, existing methods face computational challenges such as requiring model inverses or backpropagation through iterated model calls, and do not always prove that the desired ODE flow map is a solution to the loss. We introduce SGFlow, an approach for learning flow maps that bypasses explicit invertibility constraints and expensive differentiation through model iteration. SGFlow trains a model to compute both the ODE solutions and the implied velocity from scratch by following non-conservative dynamics with a stationary point at the desired flow map. On the CIFAR image benchmark, no single method attains the best FID at every step count: SGFlow attains the best FID at 10 sampling steps and remains competitive with flow matching, Meanflow, and Lagrangian map matching at other step counts, while being the only one with a proven stationary-point guarantee for its stopgrad-based dynamics.

Mark Goldstein, Anshuk Uppal, Raghav Singhal +2

Center for Computational Mathematics, Flatiron Institute. · Department of Applied Mathematics and Computer Science, Technical University of Denmark. · Courant Institute, New York University.

Diffusion Sampling Flow Map Learning Flow-Based Generative Modeling Diffusion Model Acceleration

Ask PDF

Abstract

Flow Map Learning

Explore similar work

How I learned to stop worrying and love StopGrads: Stationarity, Convergence, and a case study on Flow Map Learning

Improved Training of Strong Stochastic Flow Maps

Flow Map Learning via Nongradient Vector Flow