A Physics-Inspired Optimizer: Velocity Regularized Adam

TL;DR

The paper addresses instabilities and oscillatory dynamics observed when training deep networks with adaptive optimizers like Adam by introducing VRAdam, a physics-inspired variant that penalizes high velocity through a quartic kinetic term. This leads to a dynamic learning-rate gate that automatically slows updates in high-velocity regimes, improving stability and convergence near the adaptive edge of stability. The authors provide theoretical results on global exponential stability and a non-convex convergence bound with rate $\mathcal{O}(\ln(N)/\sqrt{N})$, alongside extensive empirical benchmarks across vision, language, and generative tasks showing improvements over AdamW and related optimizers. The work suggests that physics-inspired regularization of velocity can yield robust, efficient optimizers with broad applicability in modern deep learning pipelines.

Abstract

We introduce Velocity-Regularized Adam (VRAdam), a physics-inspired optimizer for training deep neural networks that draws on ideas from quartic terms for kinetic energy with its stabilizing effects on various system dynamics. Previous algorithms, including the ubiquitous Adam, operate at the so-called adaptive edge of stability regime during training, leading to rapid oscillations and slowed convergence of loss. However, VRAdam adds a higher order penalty on the learning rate based on the velocity such that the algorithm automatically slows down whenever weight updates become large. In practice, we observe that the effective dynamic learning rate shrinks in high-velocity regimes, and damping oscillations. By combining this velocity-based regularizer for global damping with per-parameter scaling of Adam, we create a powerful hybrid optimizer. For this optimizer, we provide rigorous theoretical analysis of operation at the edge of stability from a physical and control perspective for the momentum. Furthermore, we derive convergence bounds with the rate $\mathcal{O}(\ln(N)/\sqrt{N})$ for a stochastic non convex objective under mild assumptions. We demonstrate that VRAdam exceeds the performance against standard optimizers including AdamW. We benchmark various tasks such as image classification, language modeling, and generative modeling using diverse architectures and training methodologies including Convolutional Neural Networks (CNNs), Transformers, and GFlowNets.

Figures (10)

• Figure 1: VRAdam optimizer. $f(\theta)$: objective function; $\beta_1,\beta_2\in[0,1)$; $v_t$: velocity estimate; $m_t$: second-moment estimate; $\eta_t$: dynamic learning rate at step $t$; $\alpha_0$: maximal learning rate; $\alpha_0/(1+\alpha_1)$: minimal learning rate; $\beta_3$: velocity penalizer
• Figure 2: (a) Training loss curves for VRAdam, Adam, and SAM foret2021sharpnessaware of ResNet 32 on CIFAR-10 (b) training accuracy curves (c) plot of maximal eigenvalues of the loss Hessian d effective learning rate during training. Hyperparameters for these plots are provided in Appendix \ref{['edge-of-stability']}.
• Figure 3: Train (left) and validation (right) loss curves with error envelopes calculated using different run values for language modeling using AdamW.
• Figure 4: Train (left) and validation (right) loss curves with error envelopes calculated using different run values for language modeling using VRAdam.
• Figure 5: Train (left) and validation (right) loss curves calculated using different run values for language modeling using SGD Nesterov with momentum. The dots on the top indicate NaN values.

Theorems & Definitions (2)

• Theorem 4.1: Uniform exponential stability of VRMomentum
• Theorem 4.2: Convergence of VRAdam