Never Saddle for Reparameterized Steepest Descent as Mirror Flow

TL;DR

This paper introduces steepest mirror flows to unify and analyze how optimization geometry influences learning dynamics, implicit bias, and sparsity, explaining why Adam variants often outperform SGD in fine-tuning.

Contribution

It provides a theoretical framework for steepest descent methods, revealing their advantages in saddle-point escape and feature learning, especially in the context of Adam and AdamW.

Findings

Steepest descent facilitates saddle-point escape and feature learning.

Gradient descent requires large learning rates to escape saddles, which is uncommon in fine-tuning.

Decoupled weight decay stabilizes feature learning by enforcing new balance equations.

Abstract

How does the choice of optimization algorithm shape a model's ability to learn features? To address this question for steepest descent methods --including sign descent, which is closely related to Adam --we introduce steepest mirror flows as a unifying theoretical framework. This framework reveals how optimization geometry governs learning dynamics, implicit bias, and sparsity and it provides two explanations for why Adam and AdamW often outperform SGD in fine-tuning. Focusing on diagonal linear networks and deep diagonal linear reparameterizations (a simplified proxy for attention), we show that steeper descent facilitates both saddle-point escape and feature learning. In contrast, gradient descent requires unrealistically large learning rates to escape saddles, an uncommon regime in fine-tuning. Empirically, we confirm that saddle-point escape is a central challenge in fine-tuning. Furthermore, we demonstrate that decoupled weight decay, as in AdamW, stabilizes feature learning by enforcing novel balance equations. Together, these results highlight two mechanisms how steepest descent can aid modern optimization.