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

2026-09-14

AI Summary

本文解决了StopGrads在机器学习模型训练中影响收敛性的问题，通过引入一种回归原则来统一并优化StopGrad目标，保证了理论基础和训练效率。

Abstract

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.

Problem

• Stopgrads
• stationary points
• convergence guarantees

Innovation

• stopgrad regression principle
• stationary points
• convergence guarantees
• flow map learning
• training memory reduction