Understanding Optimization in Deep Learning with Central Flows
This provides a theoretical tool for understanding optimization in deep learning, addressing a foundational gap in traditional theories, though it is incremental as it builds on existing concepts of time-averaging.
The paper tackled the challenge of describing optimization dynamics in deep learning, particularly in the oscillatory 'edge of stability' regime, by developing a theory based on 'central flows' that predict long-term trajectories with high numerical accuracy for generic neural networks.
Traditional theories of optimization cannot describe the dynamics of optimization in deep learning, even in the simple setting of deterministic training. The challenge is that optimizers typically operate in a complex, oscillatory regime called the "edge of stability." In this paper, we develop theory that can describe the dynamics of optimization in this regime. Our key insight is that while the *exact* trajectory of an oscillatory optimizer may be challenging to analyze, the *time-averaged* (i.e. smoothed) trajectory is often much more tractable. To analyze an optimizer, we derive a differential equation called a "central flow" that characterizes this time-averaged trajectory. We empirically show that these central flows can predict long-term optimization trajectories for generic neural networks with a high degree of numerical accuracy. By interpreting these central flows, we are able to understand how gradient descent makes progress even as the loss sometimes goes up; how adaptive optimizers "adapt" to the local loss landscape; and how adaptive optimizers implicitly navigate towards regions where they can take larger steps. Our results suggest that central flows can be a valuable theoretical tool for reasoning about optimization in deep learning.