Stopgrads clarify flow map learning with improved convergence theory
How I learned to stop worrying and love StopGrads: Stationarity, Convergence, and a case study on Flow Map Learning
Machine Learning
Summary
Training machine learning models often uses a trick called stopgrads that changes how gradients behave, making it unclear if the training is valid. The paper identifies a general pattern for objectives using stopgrads that guarantees where the training aims to settle, ensuring it matches the true solution. They show that using stopgrads for learning flow maps leads to stable and correct results with faster memory-efficient training methods. This helps understand why stopgrad techniques work and how to apply them reliably in complex machine learning tasks.
What this means in practice
- •For machine learning engineers: Ensure training of models with stopgrad operations converges correctly for flow map tasks, improving stability and memory usage.
- •For reinforcement learning practitioners: Use stopgrad regression principles to design algorithms with guaranteed stationary points matching true policy maps.
A theory result. No direct application yet.
Authors
Max W. Shen, Mark Goldstein, Zichu Wang, Aahlad Puli, Rajesh Ranganath
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.