我如何学会不再担忧并爱上StopGrads:平稳性、收敛性及流映射学习案例研究
How I learned to stop worrying and love StopGrads: Stationarity, Convergence, and a case study on Flow Map Learning
- Frontiers Research(前沿研究)
- Prescient Design, Genentech(基因泰克公司前瞻设计部门)
- Flatiron Institute(熨斗研究所)
- MIT(麻省理工学院)
- NYU(纽约大学)
机构由 AI 辅助整理,请以论文原文为准。
AI总结:
本研究提出stopgrad回归原理,统一流映射、强化学习和扩散采样器的stopgrad目标函数,证明其平稳点唯一性及收敛性,并提出改进放置方式使训练内存减半。
AI中文摘要:
Stopgrads被广泛用于训练机器学习模型,但stopgrads可能改变原始目标函数的梯度、平稳点和收敛保证,这使得stopgrad训练在理论上缺乏依据。我们引入了一个stopgrad回归原理,该原理为stopgrad目标函数提供了一个通用模板,并给出了其平稳点及其唯一性的闭式刻画,统一了流映射、强化学习和扩散采样器的stopgrad目标函数。我们通过证明其唯一平稳点即为真实流映射,并展示欧拉和拉格朗日目标函数(包括MeanFlow和改进的MeanFlow)的正向收敛结果,为优化stopgrad流映射目标函数提供了理论基础。值得注意的是,我们证明在函数半梯度流下,学习到的流映射具有一个闭式表达式,该表达式由初始流映射和真实流映射复合而成。此外,我们利用stopgrad回归原理,为流映射目标函数提出了改进的stopgrad放置方式,将训练内存减少了2倍。
英文摘要:
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.