Weight Decay

Weight decay is a regularization technique used during model training that adds a penalty proportional to the squared magnitude of the weights to the loss function. This encourages the model to keep weights small, which reduces overfitting and improves generalization. In practice, weight decay is often implemented as L2 regularization, where the weight update subtracts a small fraction (the decay rate) of the weight value each step. Operators fine-tuning models with Hugging Face Transformers or using tools like Axolotl set a weight_decay hyperparameter (commonly 0.01–0.1). It matters because proper weight decay helps the fine-tuned model retain generalization on diverse inputs rather than memorizing the training set.

Weight decay is mathematically equivalent to L2 regularization when using standard SGD, but differs under adaptive optimizers like AdamW. In AdamW, weight decay is decoupled from the gradient update, applying the decay directly to the weights after the optimizer step. This decoupling is important because adaptive optimizers scale updates per-parameter, and coupling weight decay with the learning rate can lead to suboptimal regularization. For operators fine-tuning large language models, AdamW is the standard optimizer, and weight_decay is a key hyperparameter. Typical values range from 0.01 to 0.1. Too high a value can underfit; too low can overfit. During fine-tuning, operators often use a weight decay of 0.01 for most layers but may set it to 0 for bias terms and layer norms, as those parameters don't benefit from regularization. The effect is visible in validation loss curves: proper weight decay reduces the gap between training and validation loss.

When fine-tuning Llama 3.1 8B on a custom dataset using Axolotl, the config includes weight_decay: 0.01. This adds a penalty to the loss proportional to the squared weights. After training, the model's weights are slightly smaller in magnitude, which helps it avoid memorizing rare patterns in the training data. If weight decay were set to 0, the model might achieve lower training loss but higher validation loss, indicating overfitting.

In a fine-tuning workflow with Hugging Face Transformers, you set weight_decay in the TrainingArguments. For example: TrainingArguments(weight_decay=0.01). The Trainer then applies L2 regularization during each optimizer step. If using AdamW, the weight decay is applied separately from the learning rate. Operators can monitor the effect by comparing training and evaluation loss curves in TensorBoard. A large gap suggests adjusting weight decay upward; a flat validation loss suggests reducing it.