Adaptive Momentum and Nonlinear Damping for Neural Network Training

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.

This paper introduces new methods for neural network training. Many commonly used training methods employ a momentum mechanism, which can speed up convergence, but may also cause instability if not appropriately tuned. Our proposed methods automatically tune the momentum coefficient by adjusting the amount of damping applied to each model parameter during training. We test our methods on several transformer-based benchmarks and find that they can perform competitively with Adam, a state-of-the-art optimization scheme, on tasks where momentum-enhanced stochastic gradient descent typically struggles.