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
cs.LG
2026-09-15
Stopgrad regression makes the true flow map the unique stationary point of MeanFlow and slim ESD/LSD, with exponential convergence and 2× lower training memory.
A stopgrad is an identity in the forward pass and a hard zero in the backward pass. Self-distillation, temporal-difference learning, and flow-map training all rely on it. The update is a semi-gradient: only some occurrences of the model are differentiated, so the direction need not be the gradient of any scalar loss. A field with curl can loop. Figure 1 draws both outcomes. Some non-gradient fields still have a stable stationary point. Others do not: a small push leaves the minimum and runs into a limit cycle.
Few-step samplers learn a flow map, the map taking a particle from position x at time t to time u along the probability-flow ODE. MeanFlow puts a stopgrad in the loss and gives no stationary-point or convergence argument. Eulerian and Lagrangian self-distillation (ESD, LSD) show that the loss without a stopgrad has a unique minimizer. They leave the stopgrad versions uncharacterized. Training successfully and stopping at the true map are different claims.
Theorem 1 is a template you can check directly. The learned f appears once outside the stopgrad and regresses onto a target that is affine in the label. The semi-gradient is the pointwise residual: f minus the target at the conditional mean of the label. Zeros of that residual are stationary points, and they are the fixed points of the target operator. One semi-gradient step is one fixed-point iteration. TD(0), Q-learning, and the inner regression of adjoint matching all fit. Uniqueness of the Bellman fixed point for a discount below 1 is the classical contraction, not a new result.
If the coefficient of the label does not depend on f, the noise is additive. The loss is then a residual norm plus a conditional-variance floor independent of f, and the unique function that hits the floor is the unique stationary point. Slim ESD, slim LSD, and improved MeanFlow are in this class. MeanFlow is not: its coefficient contains a spatial derivative, so uniqueness is proved on its own.
The parameterization F = x + (u − t) f makes the diagonal condition F(u, u, x) = x automatic. The true flow map is the only solution of the transport equation. MeanFlow places the product of the spatial derivative and the sample velocity inside the stopgrad. Stationarity becomes that transport equation, whose only solution is the true flow map. Drop the stopgrad and the first variation at the true map is proportional to the posterior covariance Cov(ẋt | xt). When that covariance does not vanish, the true flow map is no longer stationary. Improved MeanFlow swaps in the model's instantaneous velocity f(t, t, x), the coefficient drops to zero, and the true flow map is both the unique stationary point and the unique point that attains the floor.
Original LSD stops the gradient only on the evaluation at the transported endpoint. Original ESD leaves a term that still contains f outside the stopgrad. Neither placement matches the template, and their stationary points are marked unknown. Slim placements keep f outside exactly once, so training no longer backpropagates through derivatives of f. With the flow-matching term added, a splitting lemma says the two first variations cannot cancel. The true flow map remains the unique stationary point.
All four semi-gradients have curl. The target operator fails a symmetry test, so none of them is the gradient of a twice continuously differentiable scalar functional. Functional semi-gradient flow still converges, and the solution is a closed-form composition. On Eulerian objectives the true flow carries the particle from t to a hand-off time, and the initial map finishes the trip to u. Lagrangian slim LSD runs in the opposite order, and only after the diagonal velocity already equals the true velocity. The hand-off time is a convex combination of t and u, weighted by e to the power −τ, with τ the optimization time. The error is at most c · e^{−τ} · (u − t). Improved MeanFlow and both slim objectives need perfect flow-matching pretraining for this argument; under that pretraining the improved MeanFlow flow coincides with MeanFlow. MeanFlow itself is proved from scratch.
On grids in one to five spatial dimensions, simulated semi-gradient flow and the closed-form composition agree with Pearson correlation 1.000000 across 220 million points. That checks the formula, not a neural net.
On CIFAR-10 each method runs 200k steps, checkpointed every 25k. For each sampling budget the FID is the best EMA checkpoint, over 3 seeds. Slim peak memory is about 2× lower than ESD and LSD without a stopgrad, and about 1.5× lower than the original stopgrad placement.
| Method | FID-1 | FID-100 | Peak GiB |
| ESD | 34.3±0.5 | 4.9±0.3 | 64.6 |
| ESD original stopgrad | 122.0±24.5 | 5.8±1.2 | 50.0 |
| ESD slim | 25.4±2.8 | 2.8±0.03 | 31.2 |
| LSD | 291.9±6.5 | 3.2±0.4 | 60.8 |
| LSD original stopgrad | 25.3±0.4 | 2.7±0.004 | 41.0 |
| LSD slim | 47.6±0.5 | 2.7±0.05 | 26.9 |
ESD slim beats the original stopgrad at 1, 10, 50, and 100 steps. FID-10 and FID-50 are 4.2±0.05 and 3.1±0.05. LSD slim matches the original stopgrad from 10 steps up (3.4±0.07 at 10, 3.0±0.06 at 50) and loses at one step, 47.6 against 25.3.
ImageNet uses SD-VAE latents. A 90M MeanFlow-B4 is flow-matching pretrained for 50k steps, where one-step FID is 315, then trained to 200k steps (40 epochs). ESD slim reaches 128.7±1.3, ahead of ESD at 287.6±2.0 and the original stopgrad at 195.1±49.2. LSD slim reaches 159.5±8.6, slightly behind the original stopgrad at 152.8±3.0. LSD with no stopgrad is 390.8±23.5, worse than the pretrained 315. These are short ablation runs.
The template is a placement rule: leave f outside the stopgrad once, keep the target affine in the label, and stationary points are fixed points. The MeanFlow stopgrad is load-bearing. Delete it and the true map leaves the stationary set. The loss is still defined. The update is not its gradient. In function space the flow still reaches the true map.
Slim training skips second-order backpropagation through derivatives of f. On CIFAR-10, peak memory falls from 60.8 and 64.6 GiB to 26.9 and 31.2 GiB, and ESD FID improves at the same time. Concurrent work wrote down equivalent objectives; this derivation starts from the stationarity condition and records that the original placement is uncharacterized. The same template covers TD learning and adjoint matching, and the discussion points it at PINN residual losses. There is no head-to-head with a fully trained published MeanFlow. The contribution is the stationary-point theory plus a cheaper placement.
Two limits are stated up front. The analysis is variational: infinite capacity and a continuous flow in function space, so finite-width networks and SGD can behave differently. Convergence for improved MeanFlow, slim ESD, and slim LSD assumes perfect flow-matching pretraining, with instantaneous velocity already equal to the true marginal velocity on the diagonal.
The original ESD and LSD stopgrads sit outside Theorem 1, and the table marks their stationary points unknown. Slim guarantees do not transfer. Curl is proved for twice continuously differentiable scalar functionals on function space. The grid experiment checks the closed form in dimensions 1 through 5. It does not check minibatch noise or EMA. CIFAR-10 FID is the best EMA across checkpoints, not the final iterate. Slim LSD is clearly worse at one step, so the memory cut is not free everywhere. On ImageNet, ESD slim still has one-step FID 128.7 and LSD slim 159.5, after only 200k steps. Equivalent slim objectives already appear in concurrent work. The addition is the unifying principle, uniqueness, and the closed-form flow.