Why Adam Works Better with $β_1 = β_2$: The Missing Gradient Scale Invariance Principle

TL;DR

This work addresses why Adam performs better when $β_{1}=β_{2}$ by introducing gradient scale invariance as a structural principle. It derives a continuous-time Adam Flow with dynamics for the first and second moments and shows that first-order gradient scale invariance holds if and only if $β_{1}=β_{2}$, yielding a leading update form $R(t)=\operatorname{sign}(g(t))\bigl(1 + (τ_{2}-τ_{1})δ(t)\bigr)$ with $δ(t)=\frac{d}{dt}\log|g(t)|$. The authors provide theoretical proofs and corroborate them with extensive experiments across vision and language models, demonstrating smoother updates when the diagonal configuration is used. The findings unify Adam with contemporary scale-robust optimizers and offer a principled guideline for designing future adaptive methods, highlighting robustness to gradient rescaling as a core design principle.

Abstract

Adam has been at the core of large-scale training for almost a decade, yet a simple empirical fact remains unaccounted for: both validation scores and the qualitative behaviour of the training runs improve when the momentum parameters satisfy $β_{1}=β_{2}$. Some recent studies have reported this pattern, but there is still no explanation for why this choice helps. We show that this choice is closely tied to a structural property that we refer to as \textit{gradient scale invariance}. We formalize this notion and prove that Adam becomes gradient scale invariant of first order if and only if $β_{1}=β_{2}$. This perspective places the balanced regime of Adam in direct alignment with the design principles underlying several recent optimizers that explicitly enforce scale-robust updates. The theory is supported by experiments across vision and language tasks, and across different architectural families, in which rescaling the gradient has a markedly smoother effect on the update when $β_{1}=β_{2}$. Overall, our results offer a coherent explanation for an open question in the behavior of Adam and provide a simple principle that helps guide the design of future optimizers.

Figures (8)

• Figure 1: Evolution of $\|\mathbf{R}_k\|$ in Adam for $\beta_{1}=\beta_{2}$.
• Figure 2: Evolution of $\|\mathbf{R}_k\|$ in Adam for $\beta_{1}\neq\beta_{2}$.
• Figure 3: Training dynamics for NanoGPT on WikiText. Each panel shows the evolution of the training loss (left axis) and the norm of the update $\mathbf{R}_k$ (right axis) for a fixed pair $(\beta_{1},\beta_{2})$. Rows correspond to fixed $\beta_{1}$ and columns to varying $\beta_{2}$, with identical axis scales within each row. Curves show the seed-averaged dynamics after exponential smoothing (window of 200 steps), with shaded regions indicating one standard deviation across seeds; the oscillation metric $\omega$ is reported in the title. Configurations with $\beta_{1}=\beta_{2}$ exhibit systematically smoother update norms, while off-diagonal choices lead to increased oscillations.
• Figure 4: Training dynamics for EfficientNet-B0 on TinyImageNet.
• Figure 5: Training dynamics for T5 on SQuAD.

Theorems & Definitions (14)

• Definition 3.1
• Proposition 3.2: First-order behavior of $\mathbf{m}$ and $\mathbf{v}$
• Theorem 3.3: First-order expansion of the normalized update
• Remark 3.4
• Definition 3.5
• Corollary 3.6: First-order gradient scale invariance of Adam
• Lemma 1.1: Tracking a slowly varying signal
• proof
• Proposition 1.2: First-order behavior of $\mathbf{m}$ and $\mathbf{v}$
• proof