Technical Summary: Adaptive Memory Momentum via a Model-Based Framework for Deep Learning Optimization

1. Problem Statement

Modern deep learning models predominantly rely on momentum-based first-order optimizers (e.g., SGD with Heavy Ball, AdamW, Lion). These methods utilize a momentum coefficient, $\beta$, to govern the optimizer's "memory"—determining how much past gradients influence the current update direction. In practice, $\beta$ is almost universally set to a fixed value (typically $\beta = 0.9$) and remains constant throughout training.

The authors identify two critical limitations with this static approach:

1. Suboptimality: Empirical evidence suggests that a single fixed $\beta$ is rarely optimal across the entire training trajectory or across different datasets and model architectures.
2. Instability: The performance landscape around the optimal fixed $\beta$ is highly unstable; slight deviations from the optimal value can drastically degrade convergence.
3. Theoretical Gap: While momentum accelerates convergence in deterministic convex settings, theoretical guarantees for stochastic, non-convex deep learning settings often show that stochastic Heavy Ball achieves no better convergence rates than plain SGD, despite its empirical success.

The central question posed is: Can the momentum coefficient adapt over time to better match the optimization landscape?

2. Methodology

The paper proposes Adaptive Memory (AM), a framework that replaces the constant momentum coefficient with a dynamic, time-varying coefficient $\beta_t$ computed online during optimization.

Core Framework: Model-Based Approximation

The method is derived from proximal model-based optimization. Instead of minimizing the true loss function $f(x)$ directly, the authors approximate it using a surrogate model $f^m_t(x)$ constructed from two planes:

1. The Gradient Plane: Based on the current gradient $\nabla f(x_t)$.
2. The Momentum Plane: Based on the previous descent direction $\frac{1}{\eta}(x_{t-1} - x_t)$, which encodes the accumulated momentum.

The surrogate model is defined as:
$$f^m_t(x) = \max \left\{ f(x_t) + g_t^\top(x - x_t), \quad \hat{f}(x_t) + \frac{1}{\eta}(x_{t-1} - x_t)^\top(x - x_t) \right\}$$
where $g_t = \nabla f(x_t)$ and $\hat{f}(x_t)$ is a bias term representing the momentum plane.

Derivation of Adaptive $\beta_t$

The update step is formulated as a proximal minimization problem with an additional regularization term to control the alignment between the current and previous descent directions:
$$x_{t+1} \in \arg\min_x \left( f^m_t(x) + \frac{1}{2\eta}\|x - x_t\|^2 + \frac{\lambda}{\eta}\langle x_{t-1} - x_t, x - x_t \rangle \right)$$

Solving this optimization problem yields a closed-form expression for the adaptive momentum coefficient $\beta_t^*$. The solution involves a quadratic program that balances the current gradient and the accumulated momentum. The resulting update rule resembles the Heavy Ball method but with a dynamically computed $\beta_t$:
$$d_{t+1} = \frac{\beta_t^* + \lambda}{1 + \lambda} d_t + \frac{1 - \beta_t^*}{1 + \lambda} g_t$$
$$x_{t+1} = x_t - \eta d_{t+1}$$

Practical Implementation Details

To ensure stability and efficiency in stochastic settings, the authors introduce several practical modifications:

• Clipping: $\beta_t$ is clipped to the interval $[0, \beta_{\max}]$, where $\beta_{\max}$ is typically set to the standard fixed value (e.g., 0.9).
• Loss Approximation: Calculating the exact loss difference $\Delta f = f(x_{t-1}) - f(x_t)$ on the same batch is inefficient, while using different batches introduces noise. The authors approximate this using a first-order expansion: $\Delta f \approx \eta g_t^\top d_t$.
• Overestimation Strategy: Unlike traditional bundle methods that use underestimating planes to guarantee convergence, the authors find that using an overestimating plane ($\hat{f}(x_t) = f(x_{t-1})$) leads to faster convergence in practice.
• Extension to AdamW: The framework is extended to AdamW by incorporating diagonal preconditioners ($P_t$) and decoupled weight decay. The adaptive $\beta_t$ is computed per layer to account for layer heterogeneity in large language models (LLMs).

3. Key Contributions

1. Motivation: The paper demonstrates empirically that fixed momentum is suboptimal and that finding the optimal static $\beta$ is tedious and non-robust.
2. Novel Methodology: It introduces a proximal framework using a two-plane approximation of the loss function to derive a closed-form, adaptive momentum coefficient. This is the first application of such a proximal framework specifically for momentum-based optimization.
3. Algorithmic Variants: The authors implement Adaptive Memory variants for both SGD (AM-MGD) and AdamW (AM-AdamW).
4. Theoretical Guarantees: Convergence rates of $O(1/\sqrt{T})$ are established for both convex and non-convex settings under standard smoothness and bounded gradient assumptions, placing the method within the standard landscape of stochastic first-order methods.
5. Hyperparameter Efficiency: The method requires no additional hyperparameter tuning beyond the standard baseline settings (e.g., $\lambda$ and $\beta_{\max}$ are fixed across tasks).

4. Experimental Results

The authors evaluated AM across a spectrum of tasks:

• Convex Problems: On logistic regression tasks (LIBSVM datasets), AM-MGD consistently outperformed fixed-momentum baselines (including the optimal grid-searched $\beta^*$) without requiring tuning.
• Image Classification: On CIFAR-10/100 and ImageNet using VGG, ResNet, and Wide-ResNet architectures, AM-MGD achieved faster optimization and better generalization.
  • Dynamics: The adaptive $\beta_t$ naturally exhibited three phases: low momentum (SGD-like) during noisy initial training, high momentum during the acceleration phase, and frequent "soft restarts" (drops in $\beta_t$) near interpolation points where gradients weaken.
• Large Language Model Pretraining: On LLaMA models (20M to 1B parameters) trained on C4:
  • AM-AdamW outperformed standard AdamW across all scales.
  • Warm-up Elimination: Notably, AM-AdamW stabilized early training even without a learning rate warm-up schedule, a regime where standard AdamW often fails or requires careful tuning. This suggests AM can serve as a hyperparameter-free alternative to warm-up.
• Ablation Studies:
  • The method is robust to the choice of $\lambda$ (within $[0.01, 1]$).
  • Performance gains persist across varying batch sizes and learning rates, with AM allowing for convergence at higher learning rates where static baselines fail.
  • Computational overhead is negligible (approx. 0.2% time increase for LLMs) with no additional memory footprint.

5. Significance and Claims

The paper claims that Adaptive Memory offers a principled, tuning-free alternative to static momentum. Its significance lies in:

• Bridging Theory and Practice: It provides a mechanism that adapts to the local geometry of the loss landscape, addressing the disconnect between theoretical convergence rates and empirical performance.
• Simplifying Training Pipelines: By stabilizing early training and potentially eliminating the need for learning rate warm-up schedules, AM simplifies the training of large-scale models.
• Robustness: The method demonstrates robustness to high learning rates and varying batch sizes, reducing the sensitivity of deep learning training to hyperparameter selection.

The authors conclude that this work opens doors for new ways of inducing adaptivity in optimization, specifically moving beyond adaptive learning rates to adaptive momentum, and suggests future work in per-parameter momentum coefficients and further investigation into momentum restarts.