💬 NLP

Weight Decay Improves Language Model Plasticity

This paper demonstrates that increasing weight decay during the pretraining of large language models enhances their downstream plasticity and fine-tuning performance, even if it results in higher pretraining validation loss, by promoting linearly separable representations and regularizing attention mechanisms.

Original authors: Tessa Han, Sebastian Bordt, Hanlin Zhang, Sham Kakade

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tessa Han, Sebastian Bordt, Hanlin Zhang, Sham Kakade

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Don’t Just Study for the Test; Study to Learn

Imagine you are training a student (the AI model) for a very important career. The training happens in two stages:

General Education (Pretraining): The student reads millions of books to learn how the world works.
Specialized Job Training (Fine-tuning): The student learns specific skills for their future job, like solving complex math problems or writing safe code.

For a long time, engineers have judged the "General Education" phase by looking at one thing: How well did the student do on the practice quizzes? If the student got a high score on the general quizzes, engineers assumed they would be a great employee later.

This paper challenges that assumption. It argues that a student who gets a perfect score on the general quizzes might actually be a worse employee because they memorized the answers too rigidly. Instead, we want a student who is "plastic"—meaning they are flexible, adaptable, and ready to learn new things.

The Secret Ingredient: "Weight Decay"

In the world of AI, there is a setting called Weight Decay. Think of this as a "Mental Tidy-Up" rule.

Low Weight Decay: The student is allowed to keep every single tiny detail, weird exception, and random fact they read. Their brain becomes cluttered with specific memories.
High Weight Decay: The student is forced to throw away the clutter. They are encouraged to keep only the broad, general patterns and big ideas. It’s like cleaning out your closet and keeping only the clothes that match everything, rather than keeping every single unique outfit you’ve ever owned.

What the Researchers Found

The researchers experimented with different levels of this "Mental Tidy-Up" (Weight Decay) during the General Education phase. Here is what they discovered:

1. The Counter-Intuitive Trade-off

If you want the best quiz scores: Use low Weight Decay. The student memorizes the training data perfectly and gets a high score on the general validation tests.
If you want the best job performance: Use higher Weight Decay. The student’s quiz scores might be slightly lower because they didn’t memorize every tiny detail. However, when they move to the Specialized Job Training, they learn much faster and perform better.

The Analogy: Imagine two students.

Student A memorized every page of the history textbook. They ace the history quiz. But when asked to write a creative essay about history, they struggle because they are stuck reciting facts.
Student B understood the themes of history but forgot some dates. They got a B+ on the quiz. But when asked to write the essay, they excel because they can adapt their understanding to new topics.
The Paper’s Claim: We usually pick Student A because of the quiz score, but Student B is actually the better hire.

2. Why Does "Tidy-Up" Help? (The Mechanisms)

The researchers looked under the hood to see why higher Weight Decay makes the AI more adaptable. They found three reasons:

Linearly Separable Representations (The "Organized Filing Cabinet"):
With higher Weight Decay, the AI organizes its knowledge in a very structured, logical way. It’s like having a filing cabinet where every document is in the exact right folder. When the AI needs to learn a new task, it doesn’t have to dig through a messy pile of papers; it can easily find and adjust the relevant information.
Regularizing Attention Matrices (The "Focus Filter"):
AI models use "attention" to decide which words in a sentence are important. Higher Weight Decay simplifies this process. It’s like wearing noise-canceling headphones. Instead of trying to listen to every sound in a crowded room (which is overwhelming and distracting), the AI focuses only on the most essential signals. This makes it easier to switch focus to new tasks later.
Reducing Overfitting (The "Letting Go" Effect):
If you memorize the training data too perfectly, you "overfit." You become stuck in the past. Higher Weight Decay forces the AI to "forget" some of the specific training details. This sounds bad, but it’s actually good for learning new things. It’s the difference between a musician who can only play the exact songs they practiced versus one who understands music theory and can improvise new songs. By forgetting the specific "songs" (training data), the AI remains flexible enough to learn new "songs" (downstream tasks).

The Conclusion

The paper concludes that we need to change how we tune AI models.

Currently, engineers optimize for Pretraining Loss (the quiz score). The paper shows that this is a flawed metric if the goal is to have a useful AI. Instead, we should optimize for Plasticity (the ability to adapt).

To do this, we should use higher Weight Decay than is currently standard. This creates a model that is slightly less perfect at the start but significantly better at learning new skills later. It’s the difference between hiring a robot that recites a manual and hiring a robot that understands how to think.

Technical Summary: Weight Decay Improves Language Model Plasticity

Problem Statement
Large language models (LLMs) are typically developed in two distinct phases: pretraining to create a base model, followed by post-training (fine-tuning, alignment, etc.) to improve downstream performance. Current practices largely treat these phases as decoupled, optimizing pretraining hyperparameters—such as weight decay—primarily based on the base model’s validation loss. This approach assumes that minimizing pretraining loss yields the most capable downstream model. However, this overlooks model plasticity, defined as the ability of a pretrained model to effectively adapt to new data and tasks during subsequent training. The paper investigates whether optimizing for pretraining performance aligns with optimizing for downstream adaptability, specifically focusing on the role of the weight decay hyperparameter.

Methodology
The authors conduct systematic experiments across two model families (Llama-2 and OLMo-2), multiple sizes (up to 4B parameters), and two training regimes: compute-optimal (20 tokens-per-parameter, TPP) and overtrained (140 TPP).

Pretraining: Models are pretrained on standard datasets (FineWeb-Edu for Llama-2, OLMo-Mix-1124 for OLMo-2) while varying the weight decay hyperparameter ( $\lambda$ ) in the AdamW optimizer.
Fine-tuning and Evaluation: Pretrained models are fine-tuned on diverse downstream tasks, including:
- Six Chain-of-Thought (CoT) reasoning tasks (e.g., MetaMathQA, MedMCQA).
- Five language understanding and commonsense reasoning benchmarks (e.g., HellaSwag, ARC).
- One safety alignment task.
- Performance is evaluated using metrics such as Greedy accuracy, Pass@16, Maj@16, RM@16, Correct Ratio, and ORM scores.
Mechanistic Analysis: To understand how weight decay affects plasticity, the authors analyze:
- Linear Probing: Training linear probes on last-token embeddings to assess the linear separability of internal representations (sentiment and topic classification).
- Attention Matrix Rank: Measuring the pseudo-rank of query-key ( $W_{QK}$ ) and value-projection ( $W_{VP}$ ) matrices to assess structural compression.
- Overfitting: Calculating the "Train-Val Gap" (Validation Loss $-$ Training Loss) to quantify the degree of overfitting to pretraining data.

Key Contributions

Weight Decay and Plasticity: The paper demonstrates that weight decay is a critical factor in shaping model plasticity. Higher weight decay during pretraining facilitates better adaptation to new tasks during fine-tuning.
Decoupling of Objectives: The authors provide empirical evidence that the weight decay value minimizing pretraining validation loss is not necessarily the value that maximizes downstream performance. In some regimes, models with higher pretraining loss (due to higher weight decay) achieve superior downstream results.
Mechanistic Insights: The paper identifies three mechanisms through which weight decay influences model behavior: encouraging linearly separable representations, regularizing attention matrices (reducing their rank), and reducing overfitting on training data.

Results

Optimal Weight Decay for Downstream Performance: Across various model families, sizes, and tasks, the optimal pretraining weight decay for downstream performance is consistently larger than the standard default of 0.1.
- In the compute-optimal regime (20 TPP), the optimal weight decay for downstream performance was found to be 1.0 for most models.
- In the overtrained regime (140 TPP), the optimal value was lower (0.3 for OLMo-2-1B-140x) but still higher than the value optimal for pretraining loss alone.
Pretraining Loss vs. Downstream Performance: The relationship between pretraining loss and downstream accuracy is not perfectly predictive. Models with similar pretraining losses can exhibit significantly different downstream performances. Furthermore, models with slightly worse pretraining losses (higher weight decay) often outperform those with better pretraining losses (lower weight decay) after fine-tuning.
Mechanistic Effects:
- Linear Separability: Higher weight decay leads to higher linear probing accuracy, indicating that representations become more structured and linearly separable. This suggests that fine-tuning can more easily refine existing representations rather than learning new ones from scratch.
- Attention Rank: Higher weight decay reduces the pseudo-rank of attention matrices, particularly the query-key matrix ( $W_{QK}$ ). This low-rank structure is hypothesized to prevent overfitting to high-dimensional noise, capturing more broadly applicable patterns.
- Reduced Overfitting: Increasing weight decay monotonically decreases the train-val gap, indicating less overfitting to the pretraining data. This "forgetting" of specific training details may enhance the model's capacity to learn new information (plasticity).

Significance
The paper highlights the importance of considering model plasticity in the pretrain-then-post-train development pipeline. It challenges the prevailing practice of using cross-entropy loss as the sole metric for hyperparameter optimization during pretraining. By demonstrating that a single hyperparameter (weight decay) plays a multifaceted role—balancing pretraining loss, optimization stability, and downstream adaptability—the work suggests that hyperparameter selection should be aligned with the ultimate objective of maximizing post-trained performance, rather than just pretraining metrics. The findings underscore the limits of decoupling pretraining and post-training optimization and point to the need for re-evaluating standard hyperparameter defaults to better account for downstream adaptability.