Adaptive memory tuning improves Adam optimizer for AI training
Early Memory Selection for Balanced Adam
Machine LearningArtificial Intelligence
Summary
Training AI models requires careful choice of settings, like how much past information to remember when adjusting learning steps. The authors propose a way to pick this memory setting early on using short tests, then keep it fixed for the full training. Their method balances the randomness from sampling and the delay from averaging past updates. Tests on various computer vision and language tasks show that this approach reduces errors compared to common fixed settings.
What this means in practice
- •For machine learning engineers: Choose adaptive memory parameters early to improve deep model training across vision and language tasks.
- •For ai infrastructure teams: Implement pilot-based tuning to reduce validation errors and accelerate model optimization pipelines.
Authors
Alberto Fernández-Hernández, Cristian Pérez-Corral, Jose I. Mestre, Manuel F. Dolz, Enrique S. Quintana-Ortí
Abstract
We propose a method for choosing the shared memory parameter $β_1=β_2=β$ in Adam from a short pilot training. The selected $β$ remains fixed during the subsequent full training. A local model of Adam's normalized direction balances sampling variability against the delay introduced by averaging past gradients. This balance gives a cubic memory rule, whose two coefficients are estimated from gradient probes at a few pilot checkpoints. The estimator uses the numerator and denominator jointly, preserving their covariance. With a 200-update pilot and sixteen probe gradients at each of four checkpoints, a seed-matched retrospective evaluation on eleven vision and language workloads reduces mean relative validation gap by 40.7% and worst-quarter mean gap by 44.3% against the grid representative of shared $β=0.95$. The mean gap is also 32.3% lower than that of the best constant $β$ chosen across all eleven workloads.