Source-linked AI summary
How I learned to stop worrying and love StopGrads: Stationarity, Convergence, and a case study on Flow Map Learning
Max W. Shen, Mark Goldstein, Zichu Wang, Aahlad Puli, Rajesh Ranganath
TL;DR
Stopgrad objectives can have semi-gradients that differ from the gradients of their stated losses, leaving their stationary points and convergence without general theoretical grounding. The paper introduces a stopgrad regression principle and applies it to flow-map objectives, establishing true-flow-map stationarity and positive convergence results. It also derives slim stopgrad placements that reduce training memory by 2×.
Problem
Stopgrad objectives use semi-gradients that need not be gradients of the original objective, and general theoretical grounding for their stationary points and convergence is lacking.
Source: §1 · pp. 1–2Method
The paper introduces a stopgrad regression principle that characterizes stationary points as target-operator fixed points and applies it to flow-map objectives.
Source: §1 · p. 2, §3 · p. 5, §4.1 · p. 7Results
The true flow map is the unique stationary point for MeanFlow, and positive convergence results are shown for MeanFlow, improved MeanFlow, slim ESD, and slim LSD.
Source: §1 · p. 2Takeaways & Limitations
Slim ESD and slim LSD preserve the true flow map's unique stationarity while reducing training memory by 2×.
Source: §1 · p. 2, §8 · p. 12The analysis assumes infinite capacity and continuous optimization, so finite-width networks and stochastic gradient descent may behave differently; convergence results for IMF, slim ESD, and slim LSD assume perfect flow-matching pretraining.
Source: §8 · p. 12
Abstract
from arXiv · showhide
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.
1 Introduction
Stopgrad updates can differ from gradients of the stated objective, leaving their stationary points and convergence unclear. The paper develops a general analysis and applies it to flow-map objectives, deriving stationary-point and convergence results.
1 Introduction
- Stopgrad semi-gradients need not be gradients of the original objective, and their non-conservative update fields may loop without progress.These concerns motivate analyzing stationary points and convergence for stopgrad optimization.Source: §1 · p. 1
- The stopgrad regression principle characterizes stationary points as fixed points of a target operator and covers flow maps, reinforcement learning, and diffusion samplers.It also gives a sufficient condition for a unique attainable minimizer of the ordinary loss to be the unique stopgrad stationary point.Source: §1 · p. 2
- The true flow map is the unique stationary point for MeanFlow and improved MeanFlow.Removing MeanFlow's stopgrad removes the true flow map as a stationary point [8].Source: §1 · p. 2
- Slim ESD and slim LSD preserve the true flow map's unique stationarity while reducing training memory by 2×.They match or improve on FID on CIFAR-10 and IMAGENET.Source: §1 · p. 2
- Under functional semi-gradient flow, the learned flow map is a closed-form composition of its initialization and the true flow map, with exponential convergence under standard assumptions.The hand-off time varies with optimization time.Source: §1 · p. 2, Figure 2 · p. 3
2 Stopgrads in ML: from PyTorch to Proofs
The paper formalizes stopgrad differentiation from automatic-differentiation rules to function-space derivatives. This yields a semi-gradient update rule and a way to characterize stationarity.
2 Stopgrads in ML: from PyTorch to Proofs
- Stopgrad acts as the identity in the forward pass but returns a zero gradient in the backward pass, matching the practice in autodifferentiation frameworks.The formalization proceeds from library operations to directional and then functional derivatives.Source: §2 · p. 3
- Function-space analysis represents first variations through an inner product and makes them independent of model parameterization.The analysis uses a pairing space H and a model class M.Source: §2 · pp. 3–4
- Writing a loss as O[f, SG{f}] means differentiating only its first argument, so its functional gradient can differ from the ordinary gradient.The resulting update direction is called a semi-gradient.Source: §2 · p. 4
- The semi-gradient flow follows the negative semi-gradient, and updates stop precisely where the semi-gradient is zero.Stationarity is defined by a vanishing first variation for every admissible perturbation.Source: §2 · p. 4
3 Stopgrad Regression Principle
The stopgrad regression principle links stationary points to fixed points of a target operator. Its conditions also yield a uniqueness criterion that applies across several stopgrad objectives.
3 Stopgrad Regression Principle
- The stopgrad regression principle characterizes stationary points as fixed points of the target operator when the learned function appears once outside stopgrads and the operator is affine in labels.The theorem assumes input-distribution support and regularity conditions for the candidate stationary function.Source: §3 · pp. 4–5
- The semi-gradient is the residual between the learned function and the target operator, so stationarity is equivalent to that residual vanishing on the support of the input distribution.The principle unifies objectives for flow maps, reinforcement learning, and diffusion samplers.Source: §3 · p. 5
- Additive label noise holds for slim LSD, slim ESD, IMF, and adjoint RAM, but not for MeanFlow or TD(0)/Q-learning.Uniqueness can still be established directly for MeanFlow and the RL objectives when noise is not additive.Source: §3 · p. 5
- With additive label noise, a function is stationary exactly when its loss reaches the conditional-variance floor; attaining that floor at one function makes it the unique stationary point.Additive noise requires the label coefficient Q[f] to be independent of f.Source: §3 · p. 5
- A unit-step semi-gradient update for stopgrad regression is fixed-point iteration, and the RL value-function fixed point is unique when the Bellman operator is a contraction.The perspective also applies to PDE-based flow-map loss functionals.Source: §3 · pp. 5–6
4 Stationary Points of Stopgrad Flow Map Objectives
The section analyzes stationary points of stopgrad flow-map objectives, finding the true flow map as the unique stationary point for MeanFlow, improved MeanFlow, and slim LSD/ESD objectives. The slim LSD and ESD placements also avoid expensive second-order derivatives.
4 Stationary Points of Stopgrad Flow Map Objectives
- Flow maps approximate probability-flow ODE integration for few-step sampling, and the analyses assume regularity conditions including Lipschitz marginal velocities and square-integrable conditional velocities.The time sampler satisfies t ≤ u, and the training distribution has full support on the interior of the time triangle.Source: §4 · pp. 6–7
4.1 MeanFlow
- MeanFlow’s stopgrad objective has the true flow map as its unique stationary point; uniqueness follows from its transport equation, not the additive-label-noise condition.Removing the stopgrad makes the true flow map cease to be a stationary point of the simplified objective.Source: §4.1 · p. 7
4.2 Improved MeanFlow
- Improved MeanFlow’s unique stationary point is the true flow map, which also uniquely attains the conditional-variance floor under additive label noise.Its stationarity condition at u = t gives f(t, t, ·) = v(t, ·), recovering the MeanFlow equation for u ≠ t.Source: §4.2 · pp. 7–8
4.3 Lagrangian Self-Distillation
- Slim LSD avoids second-order derivatives; combined with flow matching, stationarity enforces the model-velocity condition and Lagrangian identity, uniquely selecting the true flow map.The original LSD stopgrad placement does not fit the stopgrad regression principle, leaving its stationary points unclear.Source: §4.3 · p. 8
4.4 Eulerian Self-Distillation
- Slim ESD avoids second-order derivatives; combined with flow matching, stationarity enforces model-velocity matching and the transport equation, uniquely selecting the true flow map.The original ESD placement does not fit the stopgrad regression principle, while the slim placement clarifies stationary points.Source: §4.4 · pp. 8–9
5 Convergence of Functional Semi-Gradient Flow for Flow Map Objectives
Although the studied semi-gradients are not gradients of any scalar objective, functional semi-gradient flow converges to the true flow map for MeanFlow, improved MeanFlow, slim ESD, and slim LSD. The flow has a closed-form expression composing the initial and true flow maps.
5 Convergence of Functional Semi-Gradient Flow for Flow Map Objectives
- The non-gradient nature of these update fields leaves convergence uncertain under standard optimization theory, motivating a separate convergence analysis.The section studies MeanFlow from scratch, while improved MeanFlow and slim ESD/LSD use distillation training after perfect flow-matching pretraining.Source: §5 · p. 9, §5.2 · p. 10
5.1 When is a semi-gradient the gradient of an objective?
- MeanFlow, improved MeanFlow, slim LSD, and slim ESD have semi-gradients that are not gradients of any twice continuously differentiable scalar functional.Theorem 2 gives a necessary symmetry condition: when the derivative of T is not symmetric, the semi-gradient has curl and is not a gradient of a scalar functional.Source: §5.1 · p. 9
5.2 Convergence of functional semi-gradient flow for MeanFlow, improved MeanFlow, slim ESD & LSD.
- The learned flow map is a closed-form composition of the initial and true flow maps, governed by a hand-off time that varies with optimization time.Under additional mild conditions, convergence is exponentially fast, with error bounded by c e^-τ(u − t).Source: §5.2 · p. 10
- Functional semi-gradient flow converges to the true flow map for MeanFlow, improved MeanFlow, slim ESD, and slim LSD.For improved MeanFlow, slim ESD, and slim LSD, the convergence result assumes perfect flow-matching pretraining.Source: §5.2 · p. 10
6 Experiments
The numerical experiments support the predicted evolution of the learned flow map. On CIFAR-10 and ImageNet, the proposed slim variants reduce training memory while matching or improving on original stopgrad variants in the reported settings.
6.1 Numerical semi-gradient flow experiments
- Simulated semi-gradient flow agrees with the predicted flow-map evolution, with Pearson correlation 1.000000 over 220M points and matching theoretical and simulated pointwise-error curves.The norm comparison uses four spatial dimensions; the pointwise-error comparison uses one.Source: §6.1 · p. 11, Figure 3 · p. 11
6.2 Slim-ESD and LSD: Stopgrad variants of Lagrangian and Eulerian self-distillation
- At CIFAR-10 single-step sampling, LSD SLIM reaches FID-1 47.6 versus 25.3 for LSD ORIG SG; ESD SLIM improves on ESD ORIG SG at every sampling budget.SLIM variants reduce memory by 2× over non-stopgrad ESD/LSD and 1.5× over ORIG SG variants.Source: §6.2 · p. 11, Table 2 · p. 11
- On ImageNet, SLIM LSD is competitive with LSD ORIG, and SLIM ESD remains the best among ESD variants.These are shorter, ablation-style runs reporting one-step FID after 50k flow-matching pretraining steps and 200k training steps.Source: §6.2 · p. 11, Table 3 · p. 12
7 Related work
The paper places its flow-map analysis alongside work on stopgrad training in flow maps, consistency models, reinforcement learning, and differentiable games. It distinguishes its contribution through stationary-point and convergence results for specific flow-map objectives.
7 Related work
- For MeanFlow, improved MeanFlow, and slim ESD/LSD, the paper proves the true flow map is the unique stationary point and proves convergence of functional semi-gradient flow.Related SGFlow work proves stationary-point results but still differentiates through forward-mode derivatives.Source: §7 · p. 12
- Semi-gradients are studied in reinforcement learning, while non-conservative update fields also arise in differentiable games such as GANs.The cited RL work describes TD-style updates that are not gradients of a fixed loss but can converge under conditions tied to Bellman operators and fixed points.Source: §7 · p. 12
8 Discussion
The paper formalizes stopgrads in the calculus of variations and uses the resulting principle to analyze flow-map objectives and devise memory-saving placements. The analysis is idealized, and some convergence results rely on perfect pretraining.
8 Discussion
- Slim stopgrad placements for LSD and ESD avoid differentiating through derivatives of f and reduce training memory by 2×.The paper also suggests that the stopgrad regression principle may help propose placements for physics-informed neural network losses.Source: §8 · p. 12
- The analysis assumes infinite capacity and continuous optimization, so finite-width networks and stochastic gradient descent may behave differently.Source: §8 · p. 12
- Convergence results for improved MeanFlow, slim ESD, and slim LSD assume perfect flow-matching pretraining.Source: §8 · p. 12
Glossary of symbols
The glossary defines the general regression setting, then the flow-map parameterization and its associated distributions and operators.
General setting
- The general setting uses inputs x, labels y, input distribution µ, model class M, perturbation class D, and the function space H = L2(µ).The target operator T[f, y] is affine in y, and g[f] denotes the residual against its conditional-mean target.Source: Glossary of symbols · p. 16
Flow maps
- A flow map is parameterized as F(t, u, x) = x + (u − t)f(t, u, x), with F* denoting the true flow map.The flow-map input distribution µ may include diagonal mass; ν is the full-support distribution of (t, x) used for diagonal flow matching.Source: Glossary of symbols · p. 16
- For flow maps, gSD is a stopgrad self-distillation semi-gradient, while gFM[f](t, x) = f(t, t, x) − v(t, x) is the flow-matching residual.Ev and Lv denote the Eulerian and Lagrangian residuals, respectively.Source: Glossary of symbols · p. 16
A Proof of Theorem 1 (Stopgrad Regression Principle)
The proof converts stationarity into a pointwise fixed-point condition, then uses a loss decomposition to characterize stationary points under additive label noise; the flow-map results rely on stated regularity and sampling assumptions.
A Proof of Theorem 1 (Stopgrad Regression Principle)
- Under the theorem’s continuity and full-support conditions, stationarity holds exactly when f equals its target operator evaluated at the conditional-mean label.The proof uses compactly supported test functions and the fundamental lemma of the calculus of variations to pass from vanishing first variations to pointwise vanishing.Source: A Proof of Theorem 1 (Stopgrad… · p. 17
- The loss decomposes into half the squared semi-gradient norm plus a conditional-variance floor; with additive label noise, stationary points are exactly the functions attaining that floor.If the loss attains the floor at exactly one model, that model is the unique stationary point; if its infimum is above the floor, no stationary point exists.Source: A Proof of Theorem 1 (Stopgrad… · pp. 17–18
B Standing Regularity Conditions for Flow Maps
- Flow-map results assume C1 models and velocity fields, finite relevant moments, full support on the open time triangle and diagonal distribution, and inclusion of the true flow map in M.The sampling distribution may place mass on the diagonal, and the true flow map is assumed admissible in the model class.Source: B Standing Regularity Conditions for… · p. 18
C Uniqueness of stationary points for flow map objectives
The section establishes uniqueness results for flow-map objectives and applies the stopgrad regression principle to reinforcement-learning and adjoint-matching examples.
C Uniqueness of stationary points for flow map objectives
- For a combined self-distillation and flow-matching objective, stationarity requires both the self-distillation residual and the diagonal flow-matching residual to vanish.The splitting result rules out cancellation between the two terms’ first variations.Source: C Uniqueness of stationary points for… · p. 20
- The true flow map is the unique stationary point for MeanFlow, improved MeanFlow, and the combined slim LSD and slim ESD objectives.The MeanFlow and improved MeanFlow results concern their respective objectives; the slim LSD and slim ESD results concern each objective combined with flow matching.Source: C Uniqueness of stationary points for… · pp. 20–21
- The combined slim LSD and slim ESD losses also have the true flow map as their unique global minimizer.For both losses, subtracting the flow-matching variance floor leaves the corresponding stopgrad-free objective with marginal velocity v in place of the sampled velocity.Source: C Uniqueness of stationary points for… · p. 21
D Unifying stopgrad objectives under the stopgrad regression principle
- The stopgrad regression principle applies beyond flow maps to continuous-state reinforcement-learning and diffusion-sampler settings, including measure-valued targets represented by transition kernels.For a random next-state Dirac measure, its conditional expectation is the transition kernel.Source: D Unifying stopgrad objectives under… · p. 22
D.1 Temporal difference learning: TD(0)
- For semi-gradient TD(0), stationary points are exactly the solutions of the Bellman evaluation equation; when γ < 1, its contraction gives the unique true value function.The uniqueness result is classical and is not a consequence of the stopgrad regression principle.Source: §D.1 · p. 22
D.2 Q-learning
- For Q-learning, stationary points are exactly the solutions of the Bellman optimality equation, whose operator has a unique fixed point when γ < 1.The unique fixed point is the optimal action-value function.Source: §D.2 · p. 23
D.3 Adjoint matching: the inner RAM objective
- In the inner RAM step for adjoint matching, stationarity requires the control field to equal the negative noise scale times the conditional expected terminal-objective gradient.The condition is u*(x, t) = −σ(t) E[∇g(X1) | Xt = x].Source: §D.3 · p. 24
E The update direction may not be the gradient of any scalar functional
This section analyzes functional semi-gradient flows for flow-map methods, showing that the semi-gradients of four methods are not gradients of twice continuously differentiable scalar functionals and stating assumptions used to express the flows through Eulerian and Lagrangian residuals. It also specifies the flow direction and notes that its first two right-hand-side terms can be combined using ∂uF.
E The update direction may not be the gradient of any scalar functional
- LMF, LiMF, LLSD-slim, and LESD-slim semi-gradients are not gradients of any twice continuously differentiable scalar functional in the stated inner product.A necessary condition for a semi-gradient to be a scalar gradient is symmetry of the derivative of its associated T; the paper exhibits asymmetric witnesses for each method.Source: E The update direction may not be the… · pp. 24–26
F MeanFlow without stopgrad loses stationarity at the true flow map
- Without stopgrads, MeanFlow’s nonlinear dependence on the sampled velocity drops a conditional-covariance correction, so its true flow map need not be stationary.For the stated 1D Gaussian interpolation with λ > 0, the true flow map is not stationary; applying stopgrads removes the interaction and restores it as the unique stationary point.Source: F MeanFlow without stopgrad loses… · pp. 27–28
G Proofs for convergence of functional semi-gradient flow
- The convergence proofs assume the semi-gradient maps the chosen function space into itself, using sufficient conditions that cover broad model classes but are stricter than necessary.The stated conditions include smooth functions with polynomial growth, a globally Lipschitz velocity, and a data measure with all polynomial moments.Source: G Proofs for convergence of functional… · p. 28
G.1 Form of functional semi-gradient flows
- Slim ESD and LSD preserve their diagonal predictions during post-training; with perfect flow-matching pretraining, these remain the true marginal velocity and yield Eulerian and Lagrangian residual flow forms.Their diagonal residuals vanish for every function, making diagonal values invariant along semi-gradient flow; MeanFlow needs no pretraining assumption for its residual form.Source: §G.1 · p. 29
- The functional semi-gradient flow direction is gsLSD[f](t, u, x) = f(t, u, x) + (u − t) ∂u f(t, u, x) − v.Source: §G.1 · p. 30
- Using ∂uF, the first two terms on the right-hand side can be combined.Source: §G.1 · p. 30
G.2 Proof for proposition 1
The proofs track combined learned and true flow maps to characterize semi-gradient evolution, then establish exponential convergence under regularity conditions. The section also gives Improved MeanFlow’s stationary-point and convergence results and describes numerical simulations and time-variable conventions.
G.2 Proof for proposition 1
- Combined maps for the Eulerian and Lagrangian cases obey transport equations, which the proofs use to characterize semi-gradient evolution.The Eulerian construction is double-forward; the Lagrangian construction backtracks through the true flow.Source: §G.2 · pp. 30, 32
G.3 Exponential convergence rates under bounded derivatives.
- Under a globally L-Lipschitz true velocity field, the Eulerian and Lagrangian learned flow maps converge exponentially.For the Eulerian case, the stated error bound is M_a e^(-τ(u − t)).Source: §G.3 · pp. 33–34
H Stationary points and convergence of Improved MeanFlow
- Improved MeanFlow has the true flow map as its unique stationary point, and under perfect flow-matching pretraining its learned map follows the Eulerian flow and converges to the true map.The true flow map is also the unique point attaining the conditional-variance floor.Source: H Stationary points and convergence of… · p. 35
I Numerical experiments on convergence
- The numerical experiments simulate functional semi-gradient flow on discretized grids in 1–5 spatial dimensions, enforcing the diagonal identity at each step.They compare the simulated learned map with its theoretical value and use fourth-order Runge–Kutta integration.Source: I Numerical experiments on convergence · p. 35
J Differentiating w.r.t. t versus u
- The forward-time and reverse-time formulations both yield backward transport in the time variable carrying the spatial argument, with the other time index held as a parameter.Source: J Differentiating w.r.t. t versus u · p. 36
K.1 CIFAR-10 Experiments
The appendices specify CIFAR-10 and ImageNet experiment settings and summarize the flow-map objectives compared in the paper. The supplied passages describe the comparison dimensions, but not experimental outcomes or table entries.
K.1 CIFAR-10 Experiments
- CIFAR-10 experiments use 200k gradient steps, evaluate FID at 1, 10, 50, and 100 function evaluations, and train each method on eight A100 GPUs.The effective batch size is 512.Source: §K.1 · p. 36
K.2 ImageNet Experiments
- ImageNet experiments use the 90M-parameter B4 diffusion transformer, with 10 epochs of flow-matching warmup and 40 total training epochs.Each run uses four B200 GPUs and takes about eight hours.Source: §K.2 · p. 37
L Table of all Flow Map Objectives
- Table 4 compares flow-map objectives by derivative cost, stationary-point characterization, and whether each objective was introduced in prior work or this work.Its caption specifies F(t, u, x) = x + (u − t)f(t, u, x) and the convention that t = 1 is the data endpoint.Source: L Table of all Flow Map Objectives · p. 37, Table 4 · p. 37