Adaptive gradient methods such as Adam and AdamW fix the first-order momentum coefficient β 1 (typically 0.9) for all timesteps and all parameters, regardless of gradient dynamics. This causes overshooting in high-variance regimes and misses faster-convergence opportunities near stationarity. We propose Stable Training with Adaptive Momentum (STAM), which adapts β 1 based on a per-tensor gradient variance proxy derived from momentum residuals. High variance reduces β 1 to damp oscillations; low variance preserves or increases β 1 to accelerate convergence. We further introduce STAMLITE, a memory-efficient variant with only O(1) extra state per parameter-half the memory of full STAM and the same footprint as AdamW. Across 16 benchmark phases spanning synthetic tasks, image classification, language modeling, robustness tests, and hyperparameter sweeps, STAM/STAMLITE achieve top-3 performance on 10 of 12 scored phases (83%). Notably, STAMLITE wins outright on hyperparameter robustness benchmarks, demonstrating that adaptive β 1 makes optimization more forgiving to suboptimal hyperparameters. Both variants are implemented as drop-in Optax optimizers and available on PyPI (stam-optimizer).

Congrats OP