Adaptive Momentum and Nonlinear Damping for Neural Network Training

TL;DR

This work reframes momentum-based optimization as a coordinate-wise dissipative process by introducing per-parameter adaptive friction and cubic damping. The authors propose iKFAD, Cubically Damped mSGD (CD), and CADAM, with iKFAD incorporating an adaptive friction state $\xi$ that tracks kinetic energy and induces a cubically damped momentum regime when near equilibrium. They prove exponential convergence for the continuous-time dynamics of iKFAD and CD and show discrete-time convergence for small step sizes, then validate the methods on Vision and language models, demonstrating robustness to learning-rate choices and achieving performance competitive with or surpassing Adam while bridging the Adam–mSGD gap. The results suggest a practical, memory-light approach to stabilize and speed training of transformers and related architectures, with opportunities to extend theory for CADAM and scale to larger models.

Abstract

We propose a continuous-time scheme for large-scale optimization that introduces individual, adaptive momentum coefficients regulated by the kinetic energy of each model parameter. This approach automatically adjusts to local landscape curvature to maintain stability without sacrificing convergence speed. We demonstrate that our adaptive friction can be related to cubic damping, a suppression mechanism from structural dynamics. Furthermore, we introduce two specific optimization schemes by augmenting the continuous dynamics of mSGD and Adam with a cubic damping term. Empirically, our methods demonstrate robustness and match or outperform Adam on training ViT, BERT, and GPT2 tasks where mSGD typically struggles. We further provide theoretical results establishing the exponential convergence of the proposed schemes.

Figures (11)

• Figure 1: Adam-mSGD gap: Selected experiments demonstrating that CD and iKFAD can close the gap between Adam and mSGD without per-parameter adaptive learning rates on language modeling tasks using Transformers.
• Figure 2: Phase portraits for a two-dimensional Rosenbrock function: Each red grid point corresponds to a different initialization in the $(x,y)$ domain and each blue line corresponds to a different optimization trajectory. The minimum $(x_{\text{min}}, y_{\text{min}}) = (1,1)$, is denoted by a green star. We set $\gamma=1$, $h=0.005$, and $\alpha=\rho=c=1$. mSGD exhibits significant oscillations and overshooting of the minimum (note the different scale of the mSGD y-axis). In contrast, our per-parameter, adaptive friction method iKFAD and cubic damping method CD lead to smoother and much more direct trajectories towards the minimum, similarly to Adam.
• Figure 3: Effect of adaptive friction and cubic damping on a two-hundred dimensional anisotropic quadratic: $f(\mathbf{x}) = \frac{1}{2} \, \mathbf{x}^T \mathbf{A} \mathbf{x}$, with eigenvalues ranging between $1$ and $10^4$. For momentum gradient descent (mGD), we choose the theoretically optimal learning rate $h = 2/\sqrt{M}$ and friction $\gamma = 2 \sqrt{m}$, where $m=1$ and $M=10^4$ are the minimum and maximum eigenvalues of $f$ respectively. Left: Objective function value $f(x)$ across optimization steps. Note that both iKFAD and CD outperform mGD in this setting. Right: Fourier spectrum of trajectory in direction of highest eigenvalue.
• Figure 4: Training and test losses for iKFAD, CD, CADAM, Adam, and mSGD across ResNet-18, DistilBERT, and GPT2. All curves are averaged over 10 random seeds. Standard deviations are omitted for readability. All experiments showcase a substantial gap in performance between Adam and mSGD. Both iKFAD and CD are able to overcome this gap despite not using per-parameter learning rates.
• Figure 5: Corresponding training and test accuracies for the experiments in Figure \ref{['fig:losses_all']}. Standard deviations are included.

Theorems & Definitions (22)

• Theorem 1
• proof
• Theorem 2
• Theorem 3
• Theorem 4
• Lemma 1
• proof
• Proposition 1
• Theorem 4
• proof