Why Adam Works Better with $\beta_1 = \beta_2$: The Missing Gradient Scale Invariance Principle

Technical Summary: Why Adam Works Better with $\beta_1 = \beta_2$

Problem Statement
Despite Adam being the default optimizer for large-scale neural network training for nearly a decade, a persistent empirical observation remains unexplained: training stability and validation performance often improve when the momentum parameters are set equal ($\beta_1 = \beta_2$). While recent studies (e.g., Orvieto and Gower [2025]) have documented this phenomenon across various architectures, there has been no rigorous theoretical explanation for why this specific configuration yields better behavior or how it relates to the internal dynamics of the optimizer. Furthermore, a parallel trend in modern optimizer design (e.g., Lion, Muon, Scion) involves explicitly enforcing scale-robust updates, yet the connection between these new methods and the "balanced" regime of Adam has remained unclear.

Methodology
The authors introduce a structural property termed gradient scale invariance to analyze Adam's dynamics. They formalize this concept as follows: an optimizer is gradient scale invariant if rescaling the current gradient by a positive factor $\lambda$ (while keeping the internal state fixed) leaves the update vector essentially unchanged.

To investigate this, the authors employ a continuous-time limit analysis of the Adam algorithm. They model Adam as the time discretization of a dynamical system governed by Ordinary Differential Equations (ODEs). By introducing relaxation times $\tau_1$ and $\tau_2$ corresponding to the discrete parameters $\beta_1$ and $\beta_2$ (via $\beta_i = e^{-\Delta t/\tau_i}$), they derive a continuous-time flow for the first and second moments ($m(t)$ and $v(t)$).

The core of the analysis involves decomposing the normalized update vector $R(t) = m(t)/\sqrt{v(t)}$ into components dependent on the logarithmic scale drift of the gradient, defined as $\delta(t) = \frac{d}{dt} \log |g(t)|$. The authors prove that the update can be decomposed into a scale-invariant sign component, an explicit "scale-lag" term proportional to $(\tau_2 - \tau_1)\delta(t)$, and higher-order remainder terms.

Key Contributions

1. Formalization of Gradient Scale Invariance: The paper defines first-order gradient scale invariance, where the explicit linear response to changes in gradient scale vanishes.
2. Theoretical Characterization of Adam: The authors prove that Adam achieves first-order gradient scale invariance if and only if $\beta_1 = \beta_2$ (or equivalently $\tau_1 = \tau_2$).
   • When $\beta_1 \neq \beta_2$, the update contains a first-order term proportional to the gradient's scale drift, making the update sensitive to fluctuations in gradient magnitude.
   • When $\beta_1 = \beta_2$, this first-order term cancels exactly, pushing the sensitivity to second order. Consequently, the update behaves more like a sign-based method, which is inherently scale-invariant.
3. Synthetic and Empirical Validation:
   • Synthetic Experiments: Controlled experiments with rescaled gradient signals demonstrate that when $\beta_1 = \beta_2$, the norm of the update vector remains stable under scale changes, whereas $\beta_1 \neq \beta_2$ results in significant first-order deviations.
   • Real-World Training: Extensive experiments across vision (EfficientNet, ResNet, ViT) and language (NanoGPT, T5) tasks show that the oscillation of the update norm ($\|R_k\|$) is minimized when $\beta_1 = \beta_2$. The authors define an oscillation metric $\omega$ to quantify this stability, finding that diagonal configurations consistently yield smoother update trajectories.
   • Statistical Significance: A binomial test across multiple model-dataset pairs confirms that the minimizing $\beta_2$ for update stability coincides with $\beta_1$ with high statistical significance.

Results
The experiments confirm that the theoretical cancellation of the scale-lag term translates directly to practical training dynamics.

• Update Norm Stability: In all tested architectures (NanoGPT, EfficientNet-B0, ResNet18, ViT-B16, T5), the update norm $\|R_k\|$ exhibits significantly lower oscillation when $\beta_1 = \beta_2$ compared to off-diagonal settings.
• Decomposition Analysis: A diagnostic replay on fixed-data probe gradients reveals that when $\beta_1 \neq \beta_2$, the explicit scale-lag term accounts for approximately 50% of the absolute budget of the normalized update. When $\beta_1 = \beta_2$, this term vanishes, and the update becomes dominated by the scale-invariant sign component.
• Robustness: The findings hold across different random seeds, learning rate schedules, and architectural families, suggesting the phenomenon is a fundamental property of the optimizer's structure rather than an artifact of specific hyperparameters.

Significance and Claims
The paper claims to provide a coherent theoretical explanation for the empirical success of the balanced regime ($\beta_1 = \beta_2$) in Adam. It posits that this configuration is not merely a convenient heuristic but a structural alignment that grants Adam first-order gradient scale invariance.

The authors argue that this property places the balanced Adam regime in the same conceptual class as recent state-of-the-art optimizers (like Lion and Muon) that explicitly design for scale robustness. By identifying that this desirable property is latent within Adam's standard formulation, the work bridges the gap between traditional adaptive methods and modern scale-invariant designs.

The paper concludes that reducing sensitivity to gradient scale is a structural principle for stable adaptive optimization. While the work does not propose a new optimizer, it offers a guiding principle for future optimizer design and suggests that setting $\beta_1 = \beta_2$ is a principled choice for improving training stability, particularly in regimes where gradient magnitude fluctuates. The authors remain modest, noting that their analysis is local and coordinate-wise, and that the optimal value of the common $\beta$ along the diagonal remains an open question dependent on specific training dynamics.