Adaptive Momentum and Nonlinear Damping for Neural Network Training

AI summary

Authors: Aikaterini Karoni (University of Bristol), Rajit Rajpal and Benedict Leimkuhler (University of Edinburgh), Gabriel Stoltz (CERMICS, CNRS, École des Ponts, Institut Polytechnique de Paris; Inria Paris). Preprint, arXiv:2602.00334v2 [cs.LG].

Overview

Research area: Optimization for deep learning, viewed through the lens of continuous-time dynamical systems (dissipative Hamiltonian dynamics, damping and friction models) and applied to vision and language transformer training.

Technical level: Intermediate. The paper is written for readers comfortable with gradient-based optimizers (mSGD, Adam) and basic ordinary differential equations; the proofs are relegated to appendices, so the main text is readable without a heavy analysis background.

One-sentence scope: The paper proposes per-parameter adaptive momentum (friction) and cubic damping as alternatives to Adam's per-parameter learning rates, proves exponential convergence for two of the resulting schemes, and benchmarks them on TinyViT, DistilBERT, and GPT2 models.

What This Paper Is About

Momentum SGD applies one fixed momentum coefficient to every parameter, which ignores the fact that modern loss landscapes are highly anisotropic and coordinate-dependent. The authors recast momentum as a friction force in a continuous-time dynamical system and argue that friction should instead adapt per parameter, based on each parameter's kinetic energy. They show this adaptive friction is mathematically connected to a cubic damping term, and they build three optimizers out of these ideas, testing whether they can close the well-known performance gap between Adam and mSGD on transformer tasks without using per-parameter learning rates.

Key Contributions

1. Individual Kinetic Friction Adaptive Descent (iKFAD): a continuous-time optimizer in which each parameter has its own adaptive friction coefficient ξ, evolved as an exponentially weighted average of past squared momenta (kinetic energies), unlike the earlier single global adaptive friction of FAD.

2. A link between adaptive friction and cubic damping: the authors show that in a near-equilibrium regime (ξ̇ ≈ 0), iKFAD's momentum equation reduces to a cubically damped momentum method, ṗ ≈ −∇f(x) − (αρ)⁻¹[p]³.

3. Two cubic-damping optimizers: Cubically Damped mSGD (CD), which augments LDHD's momentum equation with a −c[p]³ term, and Cubically Damped Adam (CADAM), which augments the continuous-time Adam dynamics the same way.

4. Theory and benchmarks: exponential convergence proofs for the continuous-time dynamics of iKFAD and CD and for their discrete-time schemes, plus an empirical study on TinyViT, ResNet-18, DistilBERT (66M), GPT2-Nano (0.85M), GPT2-XS (45M), and GPT2-S (123M).

Main Findings

• Per-parameter friction matters more than a global one: in a direct comparison, the single-friction KFAD (from prior work) improved on mSGD but was clearly worse than iKFAD. Best test losses: TinyViT (CIFAR-10) iKFAD 0.612 ± 0.014 versus KFAD 0.6424 ± 0.014; GPT2-Nano (SPC) iKFAD 1.641 ± 0.006 versus KFAD 1.7363 ± 0.018. On TinyViT the best losses were close, but iKFAD reached its loss value substantially faster.

• The Adam–mSGD gap narrows without per-parameter learning rates: on the transformer tasks considered, iKFAD and CD substantially reduce the gap between Adam and mSGD despite not using per-parameter adaptive learning rates (Figure 1 and Figure 4).

• Benchmark results (best test loss, mean ± standard deviation, ten runs): TinyViT (CIFAR-10): Adam 0.613 ± 0.018, mSGD 0.647 ± 0.012, CD 0.619 ± 0.016, iKFAD 0.612 ± 0.014. DistilBERT (SST-2): Adam 0.299 ± 0.012, mSGD 0.345 ± 0.022, CD 0.268 ± 0.009, iKFAD 0.276 ± 0.011. DistilBERT (QNLI): Adam 0.340 ± 0.012, mSGD 0.361 ± 0.016, CD 0.328 ± 0.010, iKFAD 0.329 ± 0.013. GPT2-Nano (SPC): Adam 1.647 ± 0.010, mSGD 1.784 ± 0.011, CD 1.664 ± 0.008, iKFAD 1.641 ± 0.006. GPT2-XS (OWT): Adam 3.710 ± 0.011, mSGD 3.797 ± 0.008, CD 3.691 ± 0.016, iKFAD 3.693 ± 0.018. GPT2-S (OWT): Adam 3.3008, mSGD 3.4047, CD 3.2736, iKFAD 3.2633 (single seed).

• Linear damping can be switched off in practice: sweeping with γ = 0 versus γ ∈ [10⁻⁶, 10] showed little sensitivity to γ for CD and iKFAD, so γ can be set to zero, reducing CD and iKFAD to two and three hyperparameters respectively (matching mSGD and Adam). Optimal γ values after tuning in the γ > 0 setting were typically near the low end of the search range.

• Robustness across learning rate and friction: on γ–δt grids for NanoGPT (SPC), DistilBERT (SST2), and TinyViT (CIFAR-10), iKFAD and CD maintained substantially higher robustness and accuracy than mSGD and LDHD. mSGD and LDHD share the same continuous dynamics but show different stability regions under their respective discretizations.

• CADAM adds little over Adam: CADAM did not show significant improvement over Adam across the benchmarks considered; the authors suggest cubic damping provides limited additional benefit when per-parameter learning rates are already in use. CADAM also relied more heavily on linear damping, with optimal γ values significantly higher than for iKFAD and CD, and it performed poorly without linear friction.

• Low-dimensional behavior: on a two-dimensional Rosenbrock function (γ = 1, h = 0.005, α = ρ = c = 1, minimum at (1,1)), CD and iKFAD produced smooth, direct trajectories toward the minimum while momentum gradient descent showed large oscillations and overshooting. On a 200-dimensional anisotropic quadratic with eigenvalues from 1 to 10⁴, both methods converged faster than a theoretically optimally tuned mGD (h = 2/√M, γ = 2√m, m = 1, M = 10⁴), and suppressed the spectral peak mGD showed along the highest-eigenvalue direction.

• Resource cost: CD keeps the same 2N optimizer-state memory footprint as mSGD, with only negligible extra cost from element-wise cubing of the momentum vector. iKFAD and CADAM match Adam's 3N memory, each keeping one extra auxiliary vector (ξ for iKFAD, ζ for CADAM).

• Theory: under bounds 0 < m ≤ ∇²f(x) ≤ M < +∞ (or the more general condition in Equation 9 / 11), Theorems 1–4 establish exponential convergence of the form f(x(t)) − f(x*) + ‖p(t)‖ + ‖ξ(t)‖ ≤ C e^(−κt) for iKFAD and f(x(t)) − f(x*) + ‖p(t)‖² ≤ C e^(−κt) for CD, in both continuous and discrete time for sufficiently small δt. Convergence for CADAM is left for future work.

Methodology in Plain English

The authors start from the observation that discrete momentum SGD corresponds, in the zero-learning-rate limit, to Linearly Dissipative Hamiltonian Dynamics: a particle moving in a potential well given by the loss function, with a friction term −γp. This gives an explicit equivalence between the momentum coefficient and the friction coefficient, μ = 1 − γ√(δt). A fixed momentum coefficient therefore means uniform damping on every coordinate, which is a poor fit for anisotropic landscapes.

They replace this with a coordinate-wise friction governed by an auxiliary variable ξ that tracks an exponentially weighted average of each parameter's squared momentum. High recent momenta raise ξ, adding damping and cooling the system; low momenta lower ξ, allowing progress through flat regions. Solving the ξ equation exactly shows it is precisely an exponential moving average of past kinetic energies. When ξ changes slowly, the scheme reduces to a cubic damping term, which motivated the simpler CD optimizer that damps directly with −c[p]³. Cubic damping is gentler than linear damping at low momenta and more aggressive at high momenta. CADAM applies the same cubic term to the continuous-time Adam dynamics. The continuous dynamics are integrated numerically with an operator-splitting scheme (the CDBA scheme for iKFAD, Euler for CD), and convergence is proved by Lyapunov-style arguments following the earlier FAD analysis.

Experiments use CIFAR-10 with ResNet-18 and TinyViT, DistilBERT (66M) fine-tuned on SST-2 and QNLI, and GPT2-Nano (0.85M) on Shakespeare (SPC) plus a custom GPT2-XS (45M) and GPT2-S (123M) on OpenWebText. Language modeling used batch size 16; TinyViT and ResNet-18 used batch size 128. Hyperparameters were tuned with Optuna Bayesian search at a fixed budget of 80 trials per experiment, with two sweeps for the proposed optimizers (γ > 0 and γ = 0), taking the better configuration. Regularization and acceleration tricks such as learning rate schedules, weight decay, and gradient accumulation were deliberately omitted for a clean comparison. ResNet-18 results appear in Appendix D because all optimizers performed similarly there.

Why This Matters

Impact on research. The paper offers a unified dynamical-systems explanation for why adaptive methods help, reframing Adam's advantage as per-parameter momentum control rather than only per-parameter learning rates. It also gives provable exponential convergence for two nonlinearly damped first-order methods, extending the convergence toolkit beyond classic linear friction. The observation that a nonlinear, kinetic-energy-dependent damping coefficient can substitute for coordinate-wise gradient rescaling is a distinct conceptual claim within recent work on the Adam–mSGD gap.

Real-world applications (implied by the tasks and properties studied):

• Pretraining and fine-tuning transformer language models, where the Adam–mSGD gap has been a persistent practical obstacle (the paper evaluates GPT2-Nano, GPT2-XS, and GPT2-S on Shakespeare and OpenWebText).
• Vision transformers for image classification, demonstrated on CIFAR-10 with TinyViT and ResNet-18.
• Natural language understanding fine-tuning, demonstrated on DistilBERT with SST-2 and QNLI.
• Memory-constrained training: CD (with γ = 0) claims the same memory footprint and the same number of hyperparameters as mSGD while remaining competitive with Adam on several transformer benchmarks.

Industry relevance. The methods are drop-in modifications of existing momentum updates with negligible computational overhead, no per-parameter adaptive learning rates required, and (for CD with γ = 0) the same memory and hyperparameter budget as standard mSGD. That combination is attractive where scaling memory or tuning budgets is costly.

Future Directions

• Convergence theory for CADAM: the authors have not obtained exponential convergence guarantees for CADAM analogous to those proved for iKFAD and CD, and explicitly leave this open.
• Scaling to large language models: the authors state their benchmarks were limited to models under 100M parameters and that evaluating these schemes on large-scale LLMs remains future work (the largest model reported, GPT2-S, has 123M parameters).
• Hyperparameter robustness: tuning was capped at 80 trials per experiment, and the tuned ρ for iKFAD and c for CD varied widely in magnitude across tasks, indicating strong dependence on gradient scale. The authors suggest gradient normalization between training iterations as a possible fix, which they did not explore.
• Explaining and exploiting the iKFAD–CD relationship: since the two methods behave almost identically on the benchmarks, the heuristic correspondence c ≈ 1/(αρ) and the conditions under which iKFAD stays near steady state (ξ̇ ≈ 0) warrant further study, as does whether CADAM's weaker results stem from overlap with Adam's per-parameter learning rates.

Target Audience

Machine learning researchers and graduate students working on optimization algorithms, especially those interested in continuous-time or dynamical-systems perspectives on training; practitioners who train transformer models and want Adam-competitive alternatives with mSGD-like memory and hyperparameter budgets; and numerical analysts or applied mathematicians interested in damping, operator splitting, and convergence proofs for nonlinear dissipation. Readers without any background in ODEs or optimizer internals will need to work through the momentum-friction equivalence and the splitting schemes to follow the derivations.

Authors’ abstract

Momentum Stochastic Gradient Descent (mSGD) relies on a fixed momentum coefficient shared across all parameters, failing to account for the heterogeneous structure of modern loss landscapes. In this work, we adopt a continuous-time formulation to introduce individual, adaptive momentum coefficients regulated by the kinetic energy of each model parameter. This mechanism automatically adjusts to evolving training dynamics to maintain stability without sacrificing convergence speed. We demonstrate that this adaptive friction is inextricably linked to cubic damping, a suppression mechanism from structural dynamics. We additionally introduce two 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.