On the Stability and Convergence of Flow Map Training

1 INTRODUCTION

• Flow maps learn the ODE solution between any two times, offering a one-step alternative to sampling methods that require hundreds of network evaluations.The high inference cost motivates few-step generative models.Source: §1 · pp. 1–2
• Flow-map training uses stop gradients and other stabilization tricks, but their theoretical soundness and effects on learning dynamics remain poorly understood.Stop gradients alter gradient directions, and their ability to recover the true flow map had not been shown.Source: §1 · p. 2
• Adaptive normalization bounds updates and, with specific weighting, guarantees convergence toward a solution neighborhood; stable-update reweighting also lowers FID on CIFAR-10 and FFHQ 64×64.The paper also reports that synthetic-data training dynamics match the functional gradient-flow analysis.Source: §1 · p. 2
• The gradient flow of the MeanFlow and Lagrangian objectives converges to the desired flow map under a specific functional setting, relying on careful stop-gradient placement.The paper reports this result pointwise, in L∞, and in L2 for MeanFlow.Source: §1 · p. 2

2 BACKGROUND AND PRELIMINARIES

• Flow maps learn the ODE solution from any starting point at time s to its later-time state, enabling sampling with a single network evaluation.Numerical integration of a learned velocity field can require hundreds of evaluations, especially for high-dimensional data.Source: §2 · p. 3
• The MeanFlow objective substitutes the average-velocity parametrization into the Eulerian identity using a stop-gradient operator.The average-velocity parametrization enforces Φ(x, s, s) = x while learning; stop-gradient terms have their gradients set to zero.Source: §2 · pp. 3–4
• Differentiating the flow-map definition with respect to t yields the Lagrangian identity, which gives a Lagrangian objective equivalent to AYF-LMD up to reweighting.The cited work is Sabour et al. (2025).Source: §2 · p. 4
• Stop-gradient losses avoid double backpropagation and have flow maps as their zero-loss minimizers, but the operator may prevent optimization from decreasing the loss.The passage notes that MeanFlow loss has been observed not to converge to zero over iterations, and that its update cannot be the gradient of any scalar loss.Source: §2 · p. 4
• The paper situates its analysis as the first functional convergence result for MeanFlow and Lagrangian flow-map training dynamics.This claim is made in conjunction with concurrent work by Shen et al. (2026).Source: §2 · p. 4

3.1 SETTING

• With exact deterministic gradients and an L2(Rd × ∆) inner product, the weighted loss gradient is ρ(u − utgt).The sampling choices for x, s, and t induce density ρ; the analysis disregards optimization stochasticity.Source: §3.1 · pp. 4–5

3.2 CONVERGENCE RESULTS OF THE GRADIENT FLOW DYNAMICS

• Under regularity assumptions, both MeanFlow and Lagrangian gradient flows converge pointwise to the true flow map as optimization time grows.For Lagrangian flow, this result additionally requires continuity of Φ⋆ in (x, s).Source: §3.2 · p. 5
• The convergence proof uses characteristics: MeanFlow has a linear induced PDE, while the Lagrangian PDE is nonlinear but also admits this analysis.The sampling density ρ determines the speed of the characteristics.Source: §3.2 · p. 5
• MeanFlow advances transport to later times while Lagrangian training moves toward earlier initial times, replacing the initialization with true-flow transport from opposite directions.The L2 average-velocity convergence result requires an additional integrability condition on the evolving velocity error.Source: §3.2 · p. 6

3.3 INSTABILITY AND NORMALIZATION OF GRADIENT DESCENT

• Gradient descent can diverge for either objective, even with zero underlying velocity, because its differential operators can amplify arbitrarily small initial errors.The counterexample applies for every positive step size and arbitrarily small initial error.Source: §3.3 · p. 7
• Normalization strength p ≥ 0.5 gives uniformly bounded updates, while p = 1 keeps MeanFlow iterates in an O(η^1/2) neighborhood under the stated residual-growth condition.The neighborhood guarantee assumes the residual norm grows at most linearly with iteration.Source: §3.3 · p. 7

4 RELATED WORK

• Consistency models learn trajectory endpoints, and continuous-time variants use a MeanFlow-like loss with fixed t = 1 and also employ tangent normalization.The passage describes these methods as amenable to the same analysis.Source: §4 · p. 7
• Flow-map work spans LSD and ESD objectives [Boffi et al. (2025b)], MeanFlow [Geng et al. (2025)], and other distillation and reparameterization approaches.The passage also contrasts claims of spurious minimizers for ESD with its finding that stop-gradient MeanFlow optimization does not stall in a local minimum.Source: §4 · p. 7
• Concurrent work by Shen et al. (2026) derives the same functional gradient-flow closed form, but assumes ρ = 1; this paper also gives L2 rates and gradient-descent results.Shen et al. consider a broader class of problems, including reinforcement learning.Source: §4 · p. 8

5 EXPERIMENTS

• Learned synthetic-data average velocities match the theorized closed-form snapshots at each optimization time.The experiment pointwise-parameterizes average velocity on a discrete grid to approximate the functional setting.Source: §5 · p. 8
• Image experiments evaluate MeanFlow with adaptive normalization on CIFAR-10 and FFHQ 64×64 using pretrained velocity teachers.FID is reported after 200K training steps across normalization exponents; the table reports mean ± standard deviation over three runs.Source: §5 · pp. 8–9, Table 1 · p. 9
• FID results show p=0.5 to p=1 gives good performance across datasets, while p<0.5 has noticeably worse FIDs and corresponds to potentially unbounded updates.For a fixed seed, p≤0.25 and p=1.25 also yield the poorest generated images.Source: §5 · pp. 8–9, Table 1 · p. 9, Figure 3 · p. 9

6 CONCLUSION

• With the stated stop-gradient placement, gradient flow recovers the true flow map pointwise, in L∞, and in L2 for MeanFlow.The stop-gradient operators also prevent double backpropagation.Source: §6 · p. 9
• Gradient-descent iterations on MeanFlow and Lagrangian losses can diverge, while adaptive normalization provides a theoretical and practical solution and enables higher-fidelity learning on CIFAR-10 and FFHQ 64×64.Source: §6 · p. 10
• The convergence results cover only losses in Equations (5) and (7), an idealized functional gradient-flow setting, and do not establish convergence for neural-network parameterizations or stochastic training challenges.The authors leave convergence for other objectives and quantifying design-choice effects to future work.Source: §6 · p. 10

AI USE STATEMENT

• Generative AI assisted with brainstorming, method implementation, rephrasing, scientific visuals, mathematical proof ingredients, and proof checking.Source: AI USE STATEMENT · p. 10
• The authors exclude AI assistance from developing theoretical models, formulating claims, designing experiments, and interpreting results, among other listed tasks.Source: AI USE STATEMENT · p. 10
• The authors state that they reviewed AI-assisted work and take responsibility for the final content, including AI-assisted text, claims, and artifacts.Source: AI USE STATEMENT · p. 10

A.1 FRÉCHET DERIVATIVES

• The appendix differentiates the objective over an infinite-dimensional C1 function space, treating the stop-gradient target as constant, and fixes an L2 inner product for the gradient.The chosen geometry treats changes equally across space and time without spatial mixing; neural-network parametrizations do not naturally have this geometry.Source: §A.1 · p. 13

B.1 PROOFS OF PDE CONVERGENCE

• The closed-form solutions express each iterate using the true flow map and a ρ-dependent intermediate time that approaches the relevant endpoint as training time increases.For MeanFlow, tτ approaches t; for the Lagrangian objective, sτ approaches s.Source: §B.1 · pp. 13–14
• MeanFlow training converges to the true average velocity, while Lagrangian training converges pointwise to the true flow map when that map is continuous in (x, s).The Lagrangian result follows as its intermediate time sτ approaches s.Source: §B.1 · pp. 13–15

B.2 CONVERGENCE RATES

• MeanFlow convergence rates are bounded under a positive lower bound on ρ and bounded initialization and velocity; Lagrangian rate bounds require stronger assumptions, including map closeness and a Lipschitz bound.A positive lower bound on ρ cannot hold for a density over all of R^d, though the stated inequality applies on compact trajectory domains.Source: §B.2 · pp. 15, 17

B.3 SUFFICIENT CONDITIONS FOR CONVERGENCE

• The characteristic-based proofs require local Lipschitz regularity of v and ρ, and avoiding finite-time divergence of the integrating factor requires additional control on ρ.An upper bound on ρ is one stated sufficient condition for that control.Source: §B.3 · p. 17

B.4 GRADIENT DESCENT PROOFS

• Unnormalized gradient descent can diverge from arbitrarily small initial errors for MeanFlow and for the Lagrangian loss when v = 0.The authors attribute the need for increasingly small step sizes to stiffness inherent in both PDEs.Source: §B.4 · pp. 18–19
• For p ≥ 0.5, adaptive normalization gives uniformly bounded update norms; for MeanFlow, p = 1 yields convergence to a neighborhood when residual norms grow at most linearly.The neighborhood guarantee also assumes the stated regularity, density, and initialization conditions.Source: §B.4 · pp. 19, 21
• For p > 1, similar convergence guarantees are described as likely possible, but they would require stricter assumptions on residual norms..Source: §B.4 · p. 21

B.5 EQUIVALENCE TO ALIGN YOUR FLOW

• AYF-LMD and the Lagrangian loss induce the same gradient update, up to their choice of loss weightings.Source: §B.5 · p. 21

• The observed convergence rates match the derived trajectories, with error-floor convergence at or shortly after the predicted boundary-hitting time for a < 0.The tested weightings are ρ = (t −s)^a; the theoretical last diagonal hitting time is marked for a = −1 and a = −0.5.Source: §C.1 · p. 22, Figure 5 · p. 23
• For a = −0.5 and a = −1, varying the diagonal over optimization time yields higher flow-map error than a = 0.This matches the prediction that the interior tracks the boundary for a < 0; the imperfect recovery for a = 0 is attributed to lookup-table discretization error.Source: §C.1 · p. 22, Figure 6 · p. 23
• In the reported Figure 7 results, when η is too large compared with 1/k, the table is unable to learn, and the error grows.Source: §C.1 · p. 22

• MeanFlow achieves lower relative L2 error than Lagrangian across all tested λ values, and Lagrangian fails to learn the flow map above λ = 3.5.The comparison uses v(x, t) = λx and averages results over five seeds.Source: §C.2 · p. 23, Figure 8 · p. 24
• Figure 7 indicates that η affects whether gradient descent learns an error of the form (t − s)^k; as k increases, a smaller η is required to prevent error accumulation.Source: Figure 7 · p. 24
• Table 2 reports CIFAR-10 FID for the Lagrangian objective across normalization exponents p, with results summarized as mean ± standard deviation over three runs.Source: Table 2 · p. 25

• For Lagrangian-trained CIFAR-10 flow maps, the lowest FIDs occur at p ≥ 0.5, while p = 0 gives the worst performance.Source: D ADDITIONAL EXPERIMENTAL DETAILS FOR… · p. 25
• The adaptive-normalization weight is w = sg((∥u_k −u_tgt,k∥2^2 + c)^−p), with c = 10^−3; the experiments use OT-CFM or EDM teachers depending on dataset.The passage specifies OT-CFM for CIFAR-10 and EDM for FFHQ 64×64.Source: D ADDITIONAL EXPERIMENTAL DETAILS FOR… · p. 25
• Time pairs use (s, t −s) embeddings and specialized distributions, with separate sampling rules for off-diagonal and diagonal points.The off-diagonal rules sample s and t from LogitNormal distributions and order them; diagonal points sample s from LogitNormal(0.6, 1.62).Source: D ADDITIONAL EXPERIMENTAL DETAILS FOR… · p. 25