Understanding Optimization in Deep Learning with Central Flows

TL;DR

This paper introduces the concept of central flows, differential equations that model the smoothed trajectories of optimizers in deep learning, providing new insights into their dynamics in the edge of stability regime.

Contribution

It develops a novel theoretical framework called central flows to analyze the long-term behavior of optimizers in deep learning, especially in oscillatory regimes.

Findings

Central flows accurately predict optimization trajectories.

They explain how gradient descent progresses despite loss oscillations.

They reveal how adaptive optimizers navigate the loss landscape.

Abstract

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.