Weight Decay › Weight Decay in Transformer Training

The practical application of weight decay in transformer-based language models involves several details that go beyond the basic recipe. Published training recipes show both common parameter choices and the reasoning behind them.

Examining how major language models have used weight decay provides useful reference points.

BERT (2018) used $\lambda = 0.01$ with AdamW for pretraining, with bias and LayerNorm parameters excluded. This configuration became one of the most influential templates for transformer pretraining. The relatively modest $\lambda = 0.01$ reflects the large training corpus (BookCorpus plus Wikipedia, roughly 3 billion tokens), where the data itself provides strong regularization against memorization.

GPT-3 (2020) used $\lambda = 0.1$ with AdamW. This stronger regularization was chosen despite the massive training corpus (300 billion tokens) because GPT-3 has 175 billion parameters, far exceeding the number of training tokens. The high model-to-data capacity ratio justified stronger weight decay.

LLaMA (2023) used $\lambda = 0.1$ with AdamW. LLaMA was trained on roughly 1 trillion tokens with models ranging from 7B to 65B parameters, a much more compute-efficient regime. The $\lambda = 0.1$ choice suggests that even with large data quantities, the training team preferred stronger regularization for large models.

PaLM (2022) used $\lambda = 0.1$ with Adafactor, a memory-efficient adaptive optimizer. This reinforces the trend toward $\lambda = 0.1$ for very large models, regardless of the specific optimizer variant.

The pattern across these configurations is clear: as model size grows relative to dataset size, stronger weight decay is preferred. The $0.01$ to $0.1$ range covers most practical scenarios, with $0.1$ being increasingly common for billion-parameter models.

A practical tip when starting a new project: rather than searching over a wide range, start with $\lambda = 0.01$ as your default. If you observe clear signs of overfitting (train loss much lower than validation loss), increase to $0.05$ or $0.1$. If you observe underfitting (both train and validation loss are high and adding more epochs does not help training loss), reduce to $0.001$ or remove weight decay entirely. This three-point search covers the most important range without expensive grid searches.

Weight decay and gradient clipping are both tools for keeping the optimization process well-behaved, and they interact in subtle ways. Gradient clipping (covered in Chapter 13) caps the gradient norm to prevent explosive updates. Weight decay prevents weights from growing large over time.

In practice, most transformer training recipes use both together. Gradient clipping handles sudden large gradient spikes that can destabilize training. Weight decay handles the gradual tendency for weights to grow large over many steps of stable training. Neither fully substitutes for the other.

When gradient clipping is active and clips a gradient, it scales down all gradient components proportionally. This means the task gradient term in the AdamW update is clipped, but the weight decay term $\eta \lambda \theta_i$ is not affected by clipping. This behavior is desirable: weight decay continues to apply its steady shrinkage pressure even during gradient spikes, preventing the weight magnitudes from growing large precisely when gradient instability might push them in that direction.

A natural question is: if gradient clipping prevents explosive gradient steps, and weight decay prevents large weights, do you really need both? The answer is yes, because they address different timescales. Gradient clipping is a reactive mechanism that fires when a gradient is too large at a specific step. Weight decay is a proactive mechanism that continuously prevents the slow accumulation of large weights over thousands of steps. A model trained with only gradient clipping can still develop large weights gradually; each individual step is clipped, but the cumulative effect of many moderate steps can still push weights to large magnitudes. Weight decay prevents this slow drift.

Weight decay is not the only regularization technique available, and understanding where it sits relative to alternatives helps you choose the right tool for different situations.

Dropout randomly zeros activations during the forward pass, which prevents neurons from co-adapting. It is most effective for fully-connected layers and is less useful for convolutional layers (where spatial correlations are expected) or attention layers (where the patterns should be consistent). Weight decay, by contrast, applies everywhere in the network.

Data augmentation increases the effective dataset size by transforming training examples. It addresses overfitting at the data level rather than the model level. Combining weight decay with data augmentation is common and effective: data augmentation increases training diversity, while weight decay constrains the model's capacity to memorize.

Batch normalization was once used as an implicit regularizer, since its stochastic nature (batch statistics vary per mini-batch) provides a form of regularization. However, modern transformers rely on layer normalization, which has less regularization effect. For these architectures, explicit weight decay is more important.

Early stopping monitors validation loss and stops training before overfitting worsens. This is compatible with weight decay but is a less principled approach: early stopping discards training steps that might still be beneficial, while weight decay allows full training by constraining the solution space.

For transformer-based language models, the typical recipe combines weight decay with dropout (especially in the attention and feed-forward layers), data preprocessing and augmentation at the corpus level, and no early stopping (since training to convergence is generally better when weight decay is calibrated correctly). This combination works reliably across many model scales and tasks.

A practical complication arises when combining weight decay with learning rate schedules. As discussed in earlier chapters, most large-scale training runs use a warmup period followed by cosine or linear decay of the learning rate. With AdamW, the effective weight decay per step is $\eta \lambda$. If $\eta$ varies by a factor of 1000 over training (e.g., from a peak of $3 \times 10^{-4}$ to a minimum of $3 \times 10^{-7}$), the effective regularization per step also varies by a factor of 1000.

This means the model experiences very different amounts of regularization at different phases of training. During the high-learning-rate phase, the combined effect of large gradient steps and large weight decay steps can provide strong regularization but also high noise. During the low-learning-rate phase, both the gradient updates and the weight decay become very small, and the model is essentially making tiny adjustments with minimal regularization.

Some research groups have proposed learning rate-independent weight decay, where $\lambda$ is adjusted to keep the effective decay rate constant relative to the parameter update magnitude. In practice, this adds complexity without consistent benefit for most training runs. The standard approach remains using a fixed $\lambda$ throughout training and accepting that the effective per-step regularization varies with the learning rate schedule.

A related question is whether to use a weight decay schedule: starting high early in training and reducing it later. Intuitively, you might want strong regularization early (when the model is still far from any good solution and vulnerable to memorizing noise) and weaker regularization late (when you want precise fine-tuning). Some training setups experiment with this, but it adds another hyperparameter dimension and is rarely used in published large-scale training. The consensus is that a fixed $\lambda$ combined with a learning rate schedule already implicitly provides this behavior: since $\eta$ decreases over training, the effective per-step weight decay $\eta \lambda$ also decreases, achieving something similar to a decaying weight decay schedule without explicitly coding it.