Weight Decay Improves Language Model Plasticity

Paper: 2602.11137 Authors: Tessa Han, Sebastian Bordt, Hanlin Zhang, Sham Kakade Categories: cs.LG, cs.AI, cs.CL

Abstract

This paper challenges the standard approach to language model pretraining by examining model plasticity—the capacity to adapt effectively through fine-tuning—rather than just validation loss. The authors demonstrate that weight decay, a fundamental regularization hyperparameter, plays a crucial role in determining downstream adaptability. Through systematic experiments, they reveal that models pretrained with higher weight decay values exhibit superior plasticity, achieving better performance on downstream tasks even when their base pretraining loss is higher. This counterintuitive finding suggests that optimizing solely for pretraining loss may be suboptimal for real-world deployment scenarios where fine-tuning is essential.

Key Contributions

• Introduces model plasticity as a critical evaluation metric for pretraining, complementing traditional validation loss measurements
• Demonstrates empirically that higher weight decay during pretraining leads to more plastic models with better fine-tuning performance
• Reveals a counterintuitive trade-off: base models with worse pretraining loss can outperform lower-loss models after fine-tuning
• Provides mechanistic insights into how weight decay affects model behavior: promoting linearly separable representations, regularizing attention matrices, and reducing overfitting
• Challenges current hyperparameter optimization practices that focus exclusively on pretraining metrics

Methodology and Experimental Design

The research employs a systematic experimental framework to isolate the effects of weight decay on model plasticity. The authors pretrain language models with varying weight decay values while controlling for other hyperparameters, then evaluate both pretraining performance and downstream task adaptation through fine-tuning.

The experimental setup includes multiple downstream tasks spanning different domains and difficulty levels, ensuring that plasticity measurements are robust across diverse scenarios. By comparing models with identical architectures but different weight decay settings, the study establishes a clear causal relationship between this regularization parameter and downstream adaptability.

The authors also conduct mechanistic analyses to understand why weight decay improves plasticity. They examine learned representations through linear probing, analyze attention pattern distributions, and measure generalization gaps between training and validation performance. These investigations reveal that weight decay acts as more than simple L2 regularization—it fundamentally shapes how models encode and organize information.

Results and Findings

The experimental results demonstrate a consistent pattern: models trained with higher weight decay values show significantly larger performance improvements when fine-tuned on downstream tasks. This effect holds across various model sizes and task types, suggesting it’s a fundamental property rather than an artifact of specific experimental conditions.

Particularly striking is the discovery of performance inversions. In several cases, models with higher pretraining loss (due to stronger weight decay) ultimately achieve better downstream performance than models with lower pretraining loss. This finding directly contradicts the common assumption that lower pretraining loss always translates to better overall model quality.

The mechanistic analyses reveal three key effects of weight decay:

1. Linearly separable representations: Higher weight decay encourages the model to learn representations where different concepts and categories are more easily distinguished through linear boundaries, facilitating efficient adaptation to new tasks.

2. Attention regularization: Weight decay prevents attention matrices from becoming overly peaked or specialized, maintaining flexibility that enables the model to redirect attention patterns during fine-tuning.

3. Reduced overfitting: Models with higher weight decay show smaller gaps between training and validation performance, indicating better generalization that transfers to downstream scenarios.

Implications for LLM Development

This research has significant implications for how we approach language model development and hyperparameter optimization. The current paradigm of optimizing pretraining hyperparameters based solely on validation loss may be leaving performance on the table, particularly for models intended for fine-tuning applications.

The findings suggest that practitioners should consider multi-objective optimization strategies that balance pretraining performance with downstream plasticity. This might involve evaluating candidate hyperparameter configurations on a small set of representative downstream tasks during the pretraining phase, rather than waiting until after pretraining is complete.

For organizations deploying LLMs, these results indicate that selecting base models should involve more than comparing pretraining benchmarks. Models with slightly higher pretraining loss but better plasticity characteristics may ultimately deliver superior performance in production environments where task-specific fine-tuning is standard practice.

The work also raises important questions about scaling laws and their relationship to plasticity. If weight decay affects how models scale with compute and data, existing scaling law formulations may need revision to account for downstream adaptability alongside pretraining loss.

Takeaways

1. Model plasticity—the ability to adapt through fine-tuning—should be considered alongside pretraining loss when evaluating language models and optimizing hyperparameters.

2. Higher weight decay during pretraining consistently improves downstream task performance, even when it increases pretraining loss, revealing a fundamental trade-off in model development.

3. Weight decay’s benefits stem from multiple mechanisms: promoting linearly separable representations, regularizing attention patterns, and reducing overfitting on pretraining data.

4. Current hyperparameter optimization practices that focus exclusively on pretraining metrics may be suboptimal for real-world deployment scenarios involving fine-tuning.

5. The relationship between a single hyperparameter and model behavior is multifaceted, suggesting that simple interpretations of regularization effects may miss important dynamics in large-scale model training.