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

TL;DR

AdamHD introduces a decoupled Huber regularizer for optimizer updates in language model pre-training, improving convergence speed, sparsity, and robustness over traditional weight decay methods like AdamW.

Contribution

It proposes AdamHuberDecay, a novel optimizer that replaces the $ ext{L}_2$ penalty with a smooth Huber regularizer, enhancing training efficiency and model sparsity.

Findings

Converges 10-15% faster in wall-clock time.

Reduces validation perplexity by up to 4 points.

Achieves 2.5-4.7% performance improvements on downstream tasks.

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 $\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.