AdamHD: Decoupled Huber Decay Regularization for Language Model Pre-Training

Fu-Ming Guo, Yingfang Fan

Keywords: Robust Optimization, Non-Convex Optimization, Foundation models, LLMs, Optimization for deep networks

TL;DR: AdamHuberDecay (AdamHD), a decoupled regularization scheme that replaces the163 quadratic L2 penalty in AdamW with a Huber-style decay acting directly in parameter space

Abstract:

Adaptive optimizers with decoupled weight decay, such as AdamW, are the de facto standard for pre-training large transformer-based generative models. Yet the quadratic nature of the $\ell_2$ penalty embedded in weight decay drives all parameters toward the origin at the same rate, making the update vulnerable to rare but extreme gradient directions and often over-penalizing well-conditioned coordinates. We propose AdamHuberDecay, a drop-in replacement for AdamW that substitutes the decoupled $\ell_2$ penalty with a decoupled smooth Huber regularizer. The resulting update decays parameters quadratically while their magnitude remains below a threshold $\delta$, and linearly ($\ell_1$-like) once they exceed $\delta$, yielding (i) bounded regularization gradients, (ii) invariance to per-coordinate second-moment rescaling, and (iii) stronger sparsity pressure on overgrown weights.

We derive the closed-form decoupled Huber decay step and show how to integrate it with any Adam-family optimizer at $O(1)$ extra cost. Extensive experiments on GPT-2 and GPT-3 pre-training demonstrate that AdamHuberDecay (a) converges 10-15\% faster in wall-clock time, (b) reduces validation perplexity by up to 4 points, (c) delivers performance improvements of 2.5-4.7\% across downstream tasks, and (d) yields visibly sparser weight histograms that translate into 20-30\% memory savings after magnitude pruning, without tuning the decay coefficient beyond the default grid used for AdamW. Ablations confirm robustness to outlier gradients and large-batch regimes, together with theoretical analyses that bound the expected parameter norm under noisy updates. AdamHuberDecay therefore provides a simple, principled path toward more efficient and resilient training of next-generation foundational generative transformers.

Submission Number: 27