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

AI summary

Overview

Research area: Optimization algorithms for large-scale language model pre-training (the ScaleOpt workshop, GPU-Accelerated and Scalable Optimization track, arXiv:2511.14721v1 [cs.LG]).

Technical level: Advanced. The paper is written for readers comfortable with adaptive optimizers (Adam/AdamW), proximal operators, and Transformer pre-training pipelines, though the core idea — capping how hard the regularizer pulls on large weights — is intuitive on its own.

Scope: The paper proposes AdamHuberDecay (AdamHD), a drop-in replacement for AdamW's decoupled ℓ₂ weight decay that replaces the quadratic penalty with a Huber penalty, and evaluates it on GPT-2 and GPT-3 pre-training across five model scales.

What This Paper Is About

AdamW and similar adaptive optimizers are the default for pre-training large Transformer models, but their decoupled ℓ₂ weight decay shrinks every parameter toward zero at a rate proportional to its own magnitude. This means large, information-bearing weights get pulled hardest, and late in training the model's capacity can be under-utilized — what the authors call the "over-decay problem."

The goal is to replace that quadratic penalty with a Huber-style regularizer that behaves like ordinary ℓ₂ decay for small weights but clips the decay force to a constant once a weight grows past a threshold δ, yielding bounded regularization gradients and ℓ₁-like shrinkage on overgrown coordinates — all at O(1) extra cost inside any Adam-family optimizer.

Key Contributions

1. A Huber-style decoupled weight penalty. The authors replace the ℓ₂ penalty with H_δ(a) = ½a² for |a| ≤ δ and δ(|a| − ½δ) for |a| > δ, giving a decay gradient equal to the weight itself when small and δ·sign(a) when large (a clipping operator).

2. A closed-form proximal update. They derive prox_{τH_δ}(y) = y/(1+τ) when |y| ≤ (1+τ)δ, and y − τδ·sign(y) when |y| > (1+τ)δ, which applies elementwise as a "scaled shrink" for moderate weights and an ℓ₁-like soft shift with capped slope for large ones.

3. Data-driven, per-tensor thresholds. Rather than a single global δ, the threshold δ_t^(l) for each parameter tensor is set as c times the mean absolute magnitude of that tensor, optionally smoothed with an exponential moving average (β₀).

4. Empirical validation plus stability analysis. Experiments on GPT-2 and GPT-3 pre-training, alongside analysis showing the proximal step is firmly nonexpansive and caps regularizer-induced shrink at α_t·λ·δ per coordinate.

Main Findings

• Faster convergence: AdamHD converges 10–15% faster in wall-clock time to a target validation loss, as illustrated by the paper's step/computation comparison (Figure 2).

• Lower validation perplexity: Validation perplexity on FineWeb is reduced by up to 4 points relative to baselines.

• Downstream gains of 2.5–4.7%: On GPT-2 Large 774M, AdamHD reaches an average score of 36.97 versus AdamW 34.47, Lion 34.90, and Sophia 34.52. Per-task (HellaSwag 10-shot, MMLU 5-shot, ARC Challenge 25-shot, GSM8k 5-shot, TruthfulQA 0-shot, WinoGrande 5-shot): AdamW 57.80 / 25.86 / 30.46 / 0.15 / 35.18 / 57.35; Lion 58.01 / 25.81 / 30.43 / 0.16 / 35.11 / 57.58; Sophia 57.83 / 25.78 / 30.37 / 0.14 / 34.98 / 56.95; AdamHD 60.11 / 27.12 / 33.23 / 0.83 / 39.34 / 61.21.

• Sparser weights and memory savings: AdamHD produces visibly sparser weight histograms, translating into 20–30% memory savings after magnitude pruning.

• No extra hyperparameter sweeps: The reported benefits hold without tuning the decay coefficient beyond the standard decay grid used for AdamW.

• Robustness: Ablations confirm robustness to outlier gradients and to large-batch training regimes.

• Bounded decay displacement: The analysis shows the per-step shrinkage from the regularizer saturates at α_t·λ·δ per coordinate, and the proximal map is 1-Lipschitz, so it does not amplify perturbations in the pre-regularization iterate — though the authors note that robustness to loss-gradient spikes still requires standard practices such as gradient clipping.

• Limit cases: As δ → ∞, the update reduces to the decoupled proximal-ℓ₂ step y/(1+τ); as δ → 0, it tends to the identity (no regularization).

Methodology in Plain English

The authors start from AdamW's "decoupled" structure, where weight decay is applied directly to the parameters rather than folded into the gradient. They swap the ℓ₂ penalty for the Huber penalty, which is quadratic near zero and linear far from zero.

To make the update implementable in one pass, they show the regularized step is the proximal operator of the Huber penalty and solve it in closed form. The result is a rule with two branches: if a parameter is not too large, divide it by (1 + learning-rate × decay); if it is large, subtract a fixed amount (learning-rate × decay × δ) in the direction of its sign. Because the penalty is separable across coordinates, this applies elementwise and costs O(1) extra work.

The threshold δ is not fixed — it is computed per parameter tensor as a constant times that tensor's average absolute weight magnitude, optionally tracked with an exponential moving average.

For evaluation, the team pre-trains GPT-2 models at 124M, 350M, 774M, and 1.558B parameters and a GPT-3 125M model on the FineWeb corpus (over 15 trillion tokens, GPT-2 BPE tokenizer, vocabulary 50,257 padded to 50,304, 16-bit token storage). Sequence length is 1024 for GPT-2 and 2048 for GPT-3 125M; learning rates range from 6e-4 for smaller models to 2.5e-4 for larger ones, with cosine decay and 700 warmup steps. Weight decay is 0.1 for embedding and output layers and 0.0 elsewhere; gradient clipping is at a global norm of 1.0; training uses bfloat16 computation with float32 master weights. Downstream evaluation uses the Eleuther Evaluation Harness with standardized few-shot protocols.

Hardware and engineering: Lambda Cloud nodes with 8 NVIDIA A100 80GB SXM4 GPUs with NVLink per node, custom CUDA C kernels, FlashAttention-2, NVCC optimizations, and 4D parallelism (data, tensor, pipeline, sequence) coordinated via MPI for multi-node communication.

Why This Matters

Impact on research. Weight decay has been treated as a fixed, quadratic shrink since AdamW. This work shows the shape of the penalty — not just its strength — is a tunable design axis, and it connects LLM optimizer design to the proximal-operator toolbox and to robust statistics. It also supplies a theoretical account (bounded expected parameter norm under noisy updates) for why late-stage training behaves better.

Real-world applications (as implications of the training regime the paper studies; the paper does not deploy these directly):

• Pre-training large generative language models where wall-clock time and GPU cost dominate budgets.
• Producing smaller deployable checkpoints, since magnitude pruning of AdamHD-trained weights yields 20–30% memory savings.
• Training runs exposed to gradient outliers or large batches, where the capped decay step is reported to be more stable.
• Post-training or fine-tuning pipelines that inherit a sparser, better-conditioned weight distribution from the pre-training stage.

Industry relevance. Optimizer changes are among the cheapest interventions available: AdamHD is a drop-in replacement requiring no new hyperparameter sweep beyond the standard AdamW decay grid, adds O(1) overhead, and slots into the Adam-family optimizers (AdamW, DeMo, Lion) that production training stacks already use. A 10–15% wall-clock speedup to target perplexity translates directly into reduced GPU-hours at the scales at which foundation models are trained.

Future Directions

• Fixing δ and c. The paper uses a threshold proportional to each tensor's mean absolute magnitude with a coefficient c, but does not report a systematic study of how c should scale with model size, depth, or training horizon — a natural next step is a principled schedule or scaling law for δ.

• Kernel-level integration. The paper notes that deep-learning kernels typically fuse ℓ₂ decay into a single CUDA update, and the per-parameter clip breaks that fast path. Making the Huber branch as fusion-friendly as ℓ₂ decay is an open engineering question.

• Extending beyond the Adam family. The proximal view is generic, but the experiments cover Adam/AdamW-style updates. Whether the same Huber decay helps momentum-decoupled (DeMo) or sign-based (Lion) optimizers, and how it interacts with Sophia, is untested here.

• Separating the mechanisms. The method claims three distinct properties — bounded regularization gradients, invariance to per-coordinate second-moment rescaling, and stronger sparsity pressure — and reports sparsity benefits from magnitude pruning; quantifying how much of the downstream gain comes from sparsity versus stabilization remains open.

Target Audience

Optimizer and systems researchers working on large-scale model training; ML engineers who own pre-training or large-scale fine-tuning pipelines and can swap an optimizer; and graduate students or practitioners with working knowledge of Adam/AdamW who want a concrete example of how proximal operators and robust-statistics ideas show up in real optimizer design. Readers looking for a fully worked comparison against every recent optimizer variant, or for exact perplexity values and full hyperparameter sweeps, will find those details are not reported in this paper's truncated content.

Authors’ 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 $δ$, and linearly ($\ell_1$-like) once they exceed $δ$, 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.