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
回归原理证明，MeanFlow、改进版和 slim 目标的唯一驻点都是真 flow map，函数空间里的半梯度流指数收敛，改摆 stopgrad 后训练显存约降一半。
stopgrad 前向是恒等，反向梯度为 0。自蒸馏、时序差分和 flow map 都用它。更新方向是半梯度，只对模型的一部分出现求导，可以不是任何标量损失的梯度。场若有旋度，轨迹会绕圈。Fig. 1 里，有的非梯度场仍有稳定驻点，有的驻点一碰就散，轨迹进入极限环。
少步采样学 flow map：概率流 ODE 上，粒子从时刻 t 的位置 x 走到时刻 u。MeanFlow 带 stopgrad，但没有驻点（半梯度为零的函数）和收敛分析。ESD 与 LSD 只证明去掉 stopgrad 后损失有唯一极小点。训得动，和知道会停在真映射上，是两件事。
定理 1 可以照着检查。f 在 stopgrad 外只出现一次，回归一个对标签仿射的目标。半梯度逐点等于残差：f 减去目标在标签条件期望上的值。残差为零即驻点，也就是该算子的不动点；走一步就是做一次不动点迭代。TD(0)、Q-learning 和 adjoint matching 的内层回归都落在这个模板里。折扣小于 1 时 Bellman 不动点的唯一性是压缩映射的旧结论。
标签系数若不依赖 f，叫加性噪声。损失等于残差范数加上与 f 无关的方差下界，打到下界的唯一函数就是唯一驻点。slim ESD、slim LSD、改进 MeanFlow 属于这类。MeanFlow 的系数含空间导数，不是加性，唯一性单独证。
F = x + (u − t) f 使对角线 F(u, u, x) = x 自动成立，真 flow map 是传输方程的唯一解。MeanFlow 把空间导数与样本速度之积放进 stopgrad，驻点方程就是传输方程，只有真 flow map 这一解。拿掉 stopgrad，真映射处的一阶变分正比于 Cov(ẋt | xt)；协方差不消失，真 flow map 不再是驻点。改进 MeanFlow 改用瞬时速度 f(t, t, x)，系数变为 0，真 flow map 既是唯一驻点也是唯一的下界点。
原始 LSD 只截断走到终点再求值的项，原始 ESD 把含 f 的项留在 stopgrad 外，都对不上模板，表中已知驻点是问号。slim 让 f 仍只在外面一次，省掉对导数的二阶反传。加上 flow matching 后，分裂引理说明两段变分不会对消，唯一驻点仍是真 flow map。
四个半梯度都有旋度，不是任何二阶连续可微标量泛函的梯度。函数空间里的流仍收敛，解是闭式复合。Euler 型先用真流走到交接时刻，再交给初始映射；Lagrange 型相反，且要求对角线上的速度已经等于真速度。交接时刻是 t 与 u 的凸组合，权重是 e 的 −τ 次方。误差不超过 c·e^{−τ}·(u−t)。改进 MeanFlow 和两种 slim 的证明要求 flow matching 已经完美，此时改进 MeanFlow 的流与 MeanFlow 重合；MeanFlow 本身可以从头训。
1 到 5 维网格上积分半梯度流，模拟值与闭式复合在 2.2 亿个点上的 Pearson 相关为 1.000000。这是在核对公式。
CIFAR-10 每种方法 200k step，每 25k 存档，每个采样步数取各存档 EMA 的最好 FID，3 个种子。slim 相对无 stopgrad 的 ESD/LSD，峰值显存约降为 1/2；相对原始 stopgrad，约降为 1/1.5。
| 方法 | FID-1 | FID-100 | 峰值显存 GiB |
| ESD | 34.3±0.5 | 4.9±0.3 | 64.6 |
| ESD 原始 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 原始 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 在 1、10、50、100 步都好于原始 stopgrad，10 步与 50 步为 4.2±0.05 和 3.1±0.05。LSD slim 在 10 步以上与原始 stopgrad 打平，10 步 3.4±0.07，50 步 3.0±0.06，单步则从 25.3 退到 47.6。
ImageNet 用 SD-VAE 的 latent。90M 参数的 MeanFlow-B4 先做 50k step 的 flow matching，此时单步 FID 为 315，再训到 200k step（40 epoch）。ESD slim 为 128.7±1.3，好于 ESD 的 287.6±2.0 和原始 stopgrad 的 195.1±49.2。LSD slim 为 159.5±8.6，略差于原始 stopgrad 的 152.8±3.0；无 stopgrad 的 LSD 是 390.8±23.5，比 315 更差。这是短程消融。
摆 stopgrad 可以先对模板：外面只留一次 f，目标对标签仿射，驻点就是不动点。MeanFlow 里那个 stopgrad 不能删，删了真映射离开驻点集。损失还在，训练跟随的却不是这份损失的梯度；函数空间里仍然收到真 flow map。
slim 不再对 f 的导数做二阶反传。CIFAR-10 峰值显存从 60.8 与 64.6 GiB 降到 26.9 与 31.2 GiB，ESD 的 FID 同时更好。同期工作写过等价目标，这里是从驻点条件推出来的。同一模板覆盖 TD 和 adjoint matching。讨论把 PINN 的 PDE 残差回归视为同类结构。生成质量没有和已发表 MeanFlow 的完整训练比。这是理论补全加上省显存的改摆法。
局限写了两条。分析是无穷容量下的连续函数空间流，有限宽网络和随机梯度下降可以不同。改进 MeanFlow、slim ESD、slim LSD 的收敛还假设 flow matching 预训练完美，对角线上的瞬时速度已经等于真边缘速度。
原始 ESD/LSD 的 stopgrad 不在定理 1 里，表上已知驻点是问号，slim 的保证不能挪过去。旋度是对函数空间里二阶连续可微的标量泛函证的；网格实验只核对闭式复合，维度停在 1 到 5，没有看小批量和 EMA 会不会绕圈。CIFAR-10 的 FID 是各存档里最好的 EMA，不是最后一步。LSD 单步上 slim 明显更差，省显存不是处处免费。ImageNet 上 ESD slim 的单步 FID 仍是 128.7，LSD slim 是 159.5，而且只训到 200k step。并发工作已经给出等价的 slim 目标，增量主要在统一原理、唯一驻点和闭式收敛。