Weight Decay Improves Language Model Plasticity

Tessa Han; Sebastian Bordt; Hanlin Zhang; Sham Kakade

Source-linked AI summary

Weight Decay Improves Language Model Plasticity

Tessa Han, Sebastian Bordt, Hanlin Zhang, Sham Kakade

arXiv:2602.11137v2Feb 11, 2026cs.LGcs.AIcs.CL

TL;DR

Pretraining is commonly optimized for validation loss, although downstream adaptability may differ across base models. This paper studies weight decay through downstream fine-tuning and finds that larger values generally improve plasticity, even when they worsen pretraining validation loss.

Problem
Pretraining hyperparameters are primarily evaluated by base-model validation loss, despite limited attention to how well models adapt during downstream training.
Method
The authors vary pretraining weight decay across model families, sizes, training regimes, and downstream tasks, measuring plasticity through fine-tuning performance.
Results
Larger weight decay generally improves pretrained-model plasticity and downstream fine-tuning performance, despite sometimes producing higher pretraining validation loss.
Takeaways & Limitations
Pretraining hyperparameter selection should account for downstream plasticity rather than relying solely on minimal pretraining validation loss.
Takeaways & Limitations
The benefits of increased plasticity must be weighed against effects that depend on model size, training duration, and other training parameters.

Abstract

from arXiv · show

Large language models are typically trained in two broad phases: pretraining to produce a base model, followed by further training to improve downstream performance. However, hyperparameter optimization and scaling laws are studied primarily from the perspective of the base model's validation loss, overlooking a crucial model property: downstream adaptability. In this work, we study pretraining from the perspective of model plasticity, that is, the ability of the base model to successfully adapt to downstream tasks upon additional training. We focus on the role of weight decay, a key regularization parameter during pretraining, and show through systematic experiments that larger weight decay increases the plasticity of the pretrained model, resulting in greater performance gains downstream after fine-tuning. This effect can lead to counterintuitive trade-offs where base models that perform worse after pretraining can perform better after further training. Further investigation of weight decay's mechanistic effects on model behavior reveals that it encourages linearly separable representations, regularizes attention matrices, and reduces overfitting on the training data. Together, these findings highlight the importance of pretrained model plasticity, the limits of using cross-entropy loss as the sole metric for hyperparameter optimization, and the multifaceted role that a single optimization hyperparameter plays in shaping model behavior.

1. Introduction

The introduction frames pretrained-model plasticity—the ability to adapt effectively during further training—as a missing objective in pretraining optimization. It shows that weight decay strongly shapes this plasticity, with larger values improving downstream fine-tuning despite sometimes worsening pretraining validation loss.

Motivation: Modern language models use pretraining followed by post-training, but these stages and their optimization objectives are often treated as decoupled.Pretraining hyperparameters and scaling laws are predominantly studied through the base model’s validation loss.
Plasticity: Model plasticity is the ability of a trained model to adapt effectively to new data and tasks during further training, and similar pretraining loss can conceal important plasticity differences.Optimizing pretraining hyperparameters for loss alone may therefore fail to produce the best post-trained model.
Experimental scope: Experiments vary pretraining weight decay across Llama-2 and OLMo-2 models up to 4B parameters, compute-optimal and overtrained regimes, and multiple downstream task families.Fine-tuning includes six Chain-of-Thought reasoning tasks and five language understanding tasks.
Main findings: Weight decay facilitates downstream adaptation across model families, sizes, training regimes, tasks, and metrics, with evidence favoring an optimal value larger than the standard default of 0.1.The findings motivate re-evaluating standard pretraining hyperparameter choices around downstream adaptability.
Main findings: Larger weight decay can produce higher pretraining validation loss yet better downstream performance after fine-tuning in some training regimes.This provides a counterexample to selecting pretraining hyperparameters solely by minimizing validation loss.
Mechanistic perspective: Mechanistically, weight decay encourages linearly separable representations, regularizes attention matrices, and reduces overfitting on training data.These effects potentially preserve the model’s ability to adapt during subsequent training and help sustain plasticity.

2. Related Work

Prior work has studied weight decay as a standard language-model training hyperparameter and model plasticity across continual, transfer, and reinforcement learning. This paper addresses the underdeveloped study of language-model plasticity by investigating weight decay’s role in it.

Weight decay in language model training: Weight decay is a standard language-model training hyperparameter, commonly implemented with adaptive optimizers such as AdamW and classically associated with regularization and generalization.The passage also identifies prior work examining additional effects of weight decay, but the supplied text truncates before specifying them.
Plasticity of deep learning models: Model plasticity has been studied in continual learning, transfer learning, and reinforcement learning, where models often undergo multiple rounds of training.Prior studies cited in the passage examine plasticity in these settings, including image models and other mechanisms affecting plasticity.
Language model plasticity: Research on language-model plasticity remains underdeveloped, motivating this paper’s investigation of how weight decay affects it.The paper contrasts its focus with previous work on language-model plasticity involving active forgetting and tokenization.

3. Background and Methods

The section defines language-model plasticity through downstream fine-tuning performance and investigates how pretraining weight decay affects it. Experiments vary weight decay across model families, sizes, training regimes, downstream tasks, evaluation metrics, and additional hyperparameter settings.

Language model plasticity: Plasticity is measured by fine-tuning a pretrained model on a task and evaluating its downstream performance.Higher downstream performance indicates greater ability to learn new data during fine-tuning.
Research Question: The study asks how weight decay during language-model pretraining affects the pretrained model’s ability to learn new knowledge during subsequent training.The investigation empirically varies weight decay during pretraining, then fine-tunes and evaluates the resulting models.
Pretraining: Pretraining covers Llama-2 and OLMo-2 models trained on FineWeb-Edu or OLMo-Mix-1124 at 20 TPP and 140 TPP ratios, with different weight-decay variants.The five groups are Llama-2-0.5B-20x, Llama-2-1B-20x, Llama-2-4B-20x, OLMo-2-1B-20x, and OLMo-2-1B-140x.
Fine-tuning: Supervised fine-tuning spans CoT reasoning, language understanding and commonsense reasoning, and safety alignment tasks.CoT experiments use six datasets spanning math, medical, biomedical, multidisciplinary, and reading-comprehension domains.
Hyperparameter sweeps: Additional sweeps test whether weight decay’s plasticity effect depends on pretraining and fine-tuning hyperparameter choices.These experiments use additional OLMo-2-1B-20x models and are described further in Appendix E.

4. Weight decay Improves Language Model Plasticity

Pretraining weight decay improves language-model plasticity and downstream performance, with the plasticity-optimal value generally exceeding the default 0.1. Consequently, minimizing pretraining validation loss is not sufficient for selecting the best downstream model.

4.1 Pretraining performance: Weight decay values below 0.1 have little effect on pretraining loss, whereas values from 0.1 to 3 can change loss in either direction and values such as 10 can substantially degrade performance.The direction of moderate-to-large weight decay effects depends on the training setting.
4.2 Weight decay and plasticity: The plasticity-optimal weight decay is 1.0 in the compute-optimal 20 TPP regime and 0.3 in the overtrained 140 TPP regime.The 1.0 optimum applies to four listed 20 TPP models, while 0.3 applies to OLMo-2-1B-140x.
4.2 Weight decay and plasticity: The weight decay that maximizes fine-tuned performance is higher or equivalent to the value minimizing pretraining validation loss, depending on the model.It is higher for Llama-2-0.5B-20x, Llama-2-1B-20x, OLMo-2-1B-20x, and OLMo-2-1B-140x, and equivalent for Llama-2-4B-20x.
4.2 Weight decay and plasticity: The plasticity benefit of higher weight decay persists across sweeps that jointly vary pretraining learning rate and fine-tuning weight decay, learning rate, and batch size.The experiments report consistent performance across a range of pretraining and fine-tuning hyperparameter settings.
4.2 Weight decay and plasticity: Higher pretraining weight decay improves plasticity and downstream performance among models with reasonable pretraining validation losses, with the optimal value above the default 0.1.This finding holds across diverse downstream tasks, including chain-of-thought reasoning, language understanding, commonsense reasoning, and safety alignment.
4.3 Pretraining performance and downstream performance: Pretraining performance is an imperfect proxy for downstream performance: models with similar losses can differ downstream, and models with lower losses can perform either better or worse.The reported correlations are suggestive rather than conclusive because each model group has n ≤10, but the main finding does not depend on them.

5. A Mechanistic Perspective on Weight Decay and Model Behavior

Weight decay affects model plasticity through three mechanisms: it promotes linearly separated representations, regularizes attention matrices toward lower rank, and reduces overfitting on pretraining data. These effects provide potential explanations for why stronger weight decay improves downstream adaptability.

Internal representations: Higher weight decay produces more linearly separated and structured internal representations across model layers.Linear probes achieve higher accuracy for models pretrained with stronger weight decay, although the relationship is not perfectly monotonic.
Attention matrices: Stronger weight decay monotonically reduces the ranks of both query-key and value-projection attention matrices.Default weight decay still yields near full-rank matrices, despite the overall rank reduction.
Attention matrices: Attention matrices respond differently to weight decay: a weight decay of 1.0 roughly halves WQK rank, while WVP remains close to full-rank.For Llama-2 models, WVP rank remains essentially stable through weight decay 1.0 before collapsing, alongside a significant performance drop.
Pretraining overfitting: Larger weight decay reduces overfitting on pretraining data, evidenced by a monotonically decreasing train-val gap.A larger train-val gap indicates more overfitting and less forgetting of training data.
Mechanistic synthesis: Together, weight decay encourages linearly separated representations, regularizes attention matrices, and reduces training-data overfitting.These diverse mechanistic effects are identified as potential explanations for improved model plasticity.

6. Discussion and Concluding Remarks · Appendix

The discussion frames weight decay as a multidimensional determinant of language-model plasticity, not merely capacity control or optimization stability. Larger weight decay can improve downstream fine-tuning despite worse pretraining validation loss, with trade-offs depending on model scale and training duration.

6. Discussion and Concluding Remarks: Weight decay has nuanced effects across the modern language-model training lifecycle beyond capacity control and optimization stability.The paper characterizes its effects multidimensionally rather than through traditional single-purpose perspectives.
6. Discussion and Concluding Remarks: Smaller weight decay achieves lower post-pretraining validation loss, especially in over-trained regimes, whereas larger weight decay improves downstream plasticity.This establishes a direct trade-off between pretraining loss and adaptability after additional training.
6. Discussion and Concluding Remarks: Larger-weight-decay models can perform best after downstream fine-tuning despite having worse validation loss immediately after pretraining.The result makes pretraining validation loss insufficient as the sole criterion for model selection.
6. Discussion and Concluding Remarks: Weight decay may increase plasticity by promoting linearly separable representations, regularizing attention matrices, and reducing training-data overfitting.These mechanisms provide a mechanistic account of how one hyperparameter can shape model behavior in multiple ways.
6. Discussion and Concluding Remarks: The value of increased plasticity must be weighed against effects that vary with model size, training duration, and other training parameters.In heavily overtrained settings or very large models trained for many steps, markedly lower pretraining validation loss may be more beneficial.
6. Discussion and Concluding Remarks: Future work should test stability–plasticity trade-offs in large and heavily overtrained models and examine applications beyond language.Suggested extensions include vision and multimodal foundation models and downstream goals beyond CoT reasoning, language understanding, commonsense reasoning, and safety alignment.

A. Additional discussion · B. Pre-training

The discussion situates the findings relative to prior work and emphasizes limitations concerning causal interpretation, experimental scope, and sample size. Results in smaller model groups remain consistent with broader trends, but conclusions may not generalize beyond the tested settings.

A. Additional discussion: The study replicates prior findings on attention-matrix rank reduction in a larger, modern setting and links rank reduction to better pretraining and fine-tuning performance up to weight decay 0.6 for OLMo models.It also reports novel phenomena beyond Kobayashi et al. (2024).
A. Additional discussion: Prior work finds that weight decay can improve image-representation transferability, although standard self-supervised-learning evaluations may not reveal the optimal value.The discussion also references work on OOD compositional generalization and reasoning-to-memorization transitions under complexity control.
A. Additional discussion: The proposed mechanisms linking weight decay to plasticity—including linearly separable representations, reduced attention-matrix rank, and reduced overfitting—are correlational rather than causal.Causal identification is difficult because mechanisms may be entangled with other covariates.
A. Additional discussion: The experiments cover Llama-2 and OLMo-2 models up to 4B parameters, TPP ratios of 20x and 140x, and twelve downstream tasks across reasoning, understanding, commonsense, and safety alignment.Results may differ for larger models, higher TPP ratios, other fine-tuning tasks, or post-training beyond fine-tuning.
A. Additional discussion: Some model groups use fewer pretrained models because of computational costs, limiting statistical confidence for Llama-2-4B-20x and OLMo-2-1B-140x.These groups nevertheless show trends consistent with more extensively studied models.
B. Pre-training: Across the smaller model groups, better pretraining validation loss does not guarantee better downstream performance, and the best downstream weight decay exceeds the standard 0.1 default.This consistency supports the paper’s broader plasticity-related observations despite the limited sample size.

B.1. Model architectures and training regimes

The study compares five pretrained model setups spanning Llama-2 and OLMo-2 architectures, model sizes up to 4B, and training regimes of 20 TPP or 140 TPP. Each setup includes variants with different weight decay values.

B.1. Model architectures and training regimes: Five setups span Llama-2 and OLMo-2 families, sizes up to 4B, and training regimes of 20 TPP or 140 TPP.The setups are Llama-2-0.5B-20x, Llama-2-1B-20x, Llama-2-4B-20x, OLMo-2-1B-20x, and OLMo-2-1B-140x.
B.1. Model architectures and training regimes: Each model setup is pretrained with variants using different weight decay values.

B.2. Training details · C. Fine-tuning

The paper varies weight decay as the independent pretraining variable across Llama-2 and OLMo-2 models, while determining training data size from the TPP ratio. Training uses specified GPU setups, optimizers, learning-rate schedules, and model-specific configurations, with validation loss assessed across weight-decay values.

B.2. Training details: Training data size for each model is determined by the TPP ratio.
B.2. Training details: Table 2 reports the model configurations and training data sizes.
B.2. Training details: Llama-2 pretraining uses up to 8 A100 or 16 H100 GPUs, while OLMo-2 pretraining uses 8xH100 GPUs.OLMo-2-1B-20x models train for 2 days and OLMo-2-1B-140x models for 2 weeks on a single H100 node.
B.2. Training details: All models use AdamW with a warmup-cosine learning-rate schedule, except OLMo-2-1B-140x models, which use warmup-stable-decay.
B.2. Training details: Each model is trained with variants spanning weight-decay values specified in Table 3.Additional training hyperparameters are provided in Tables 4 and 5.
B.2. Training details: For Llama-2 models, hyperparameters follow Qi et al. (2025) except weight decay, which is varied independently.
B.2. Training details: For OLMo-2 models, hyperparameters follow OLMo-2 defaults except weight decay, which is varied independently.
B.2. Training details: The weight-decay value minimizing pretraining validation cross-entropy may equal or exceed the standard default of 0.1, depending on the training regime.

C.1. Training details … E.1. Varying pretraining hyperparameters

The paper evaluates pretrained-model plasticity through fine-tuning on reasoning, language, commonsense, and safety tasks, alongside analyses of stability, upstream performance, and hyperparameter interactions. Across these evaluations, higher pretraining weight decay generally improves downstream adaptability, while pretraining loss does not always identify the best fine-tuned model.

C.1. Training details: Fine-tuning uses five OLMo-2-1B-20x pretrained models with weight decay values {0, 0.1, 0.3, 0.6, 1.0}.The values include the standard default 0.1 and 0.6, which produced the lowest pretraining validation loss.
C.2. Chain-of-thought reasoning: The evaluation covers six chain-of-thought datasets, five language-understanding and commonsense datasets, and safety alignment with 20,000 general instructions plus 300, 500, 1,000, or 2,000 safety instructions.Chain-of-thought models are evaluated with six metrics, language and commonsense models with cloze-style accuracy, and safety models on 100 harmful prompts using an HRM score from 0 to 4.
C.3. Language understanding and commonsense reasoning: Higher pretraining weight decay improves fine-tuning performance on chain-of-thought and language and general-knowledge tasks.Figures 8 and 11 report the downstream performance trends attributed to increased pretraining weight decay.
D. Trade-off analyses: Pretraining validation cross-entropy loss is not fully predictive of fine-tuning accuracy, and its full-set Pearson correlation is unstable under leave-one-out analysis.The leave-one-out correlation can change noticeably in magnitude and sign.
Stability-plasticity trade-off: Models with higher plasticity also tend to have higher upstream performance after fine-tuning, but greater plasticity sometimes reduces stability.The paper measures upstream performance with average accuracy across HellaSwag, PiQA, Winogrande, Arc-Easy, and Arc-Challenge.
Can weight decay hurt upstream performance while improving downstream performance?: Weight decay can produce more plastic pretrained models with higher downstream performance without apparently affecting pretrained models’ upstream performance.The reported upstream accuracies include Llama-2-0.5B-20x {0.40, 0.41, 0.41*}, Llama-2-1B-20x {0.44, 0.45, 0.45*}, OLMo-2-1B-20x {0.49, 0.49, 0.50*}, and OLMo-2-1B-140x {0.61, 0.60*, 0.56}.
E. Weight decay’s effect on model plasticity across hyperparameter settings: Across learning rates lr = [2e-4, 4e-4, 8e-4] and weight decays wd pt = [0.1, 0.6, 1.0], the study trains 9 pretrained and 54 fine-tuned OLMo-2-1B-20x models.Each pretrained model is fine-tuned on six chain-of-thought reasoning datasets.
E.1. Varying pretraining hyperparameters: For each fixed learning rate, higher weight decay leads to higher fine-tuning performance, showing that weight decay’s plasticity effect persists across pretraining hyperparameter changes.The result is reported in Figure 19 and supports findings from Section 4.2 and Appendices C.2, C.3, and C.4.

E.2. Varying fine-tuning hyperparameters … F.2.1. ATTENTION PSEUDO-RANK COMPUTATION

Across fine-tuning settings, higher pretraining weight decay consistently improves downstream performance, even though pretraining loss is not a perfect predictor of fine-tuning outcomes. The paper also quantifies attention-matrix effective dimensionality with a pseudo-rank metric based on singular values.

E.2. Varying fine-tuning hyperparameters: Fine-tuning analyses evaluate OLMo-2-1B-20x models pretrained with wd pt = [0.1, 0.3, 0.6, 1.0] and default learning rate 4e-4 on MetaMathQA and SimpleScaling.The setup jointly varies fine-tuning learning rate, weight decay, and batch size.
E.2. Varying fine-tuning hyperparameters: Higher pretraining weight decay improves fine-tuning performance across learning rates, datasets, and varied fine-tuning hyperparameters.Experiments use MetaMathQA and SimpleScaling with lr = [1e-5, 3e-5, 6e-5], wd ft = [0, 0.1, 1.0], and bs = [32, 64, 128].
E.2. Varying fine-tuning hyperparameters: Pretraining performance is not perfectly predictive of fine-tuning performance: the weight decay yielding the best pretraining loss does not always yield the best fine-tuning loss.Overall, the model with best PT loss is not the model with best FT loss.
E.2. Varying fine-tuning hyperparameters: The trend holds both when fine-tuning performance is averaged over hyperparameter combinations and when selecting the best combination for each pretrained model.Figure 23 reports average performance in Rows 1, 3, 5 and best performance in Rows 2, 4, 6.
F.2.1. ATTENTION PSEUDO-RANK COMPUTATION: The pseudo-rank of a matrix W is k/n, where k is the smallest number of singular values needed to satisfy the stated 95%-capture condition.Singular values are ordered σ1 ≥ σ2 ≥ · · · ≥ σn.
F.2.1. ATTENTION PSEUDO-RANK COMPUTATION: The pseudo-rank measures the fraction of largest singular values required to capture at least 95% of the sum of all singular values.This metric is used to quantify effective dimensionality and monitor low-rank structures during training.
F.2.1. ATTENTION PSEUDO-RANK COMPUTATION: The analysis applies pseudo-rank computation to the product of key-query matrices, WQK = W T K WQ, and to value-projection matrices, WVP = WP WV.These matrices are monitored for the emergence of low-rank structures during training.

F.2.2. ADDITIONAL ANALYSES ON ATTENTION MATRIX RANK · F.3. Norm of parameters

Additional analyses show that weight decay lowers attention-matrix rank and model weight norms, while training time alone does not reduce attention rank. Increasing weight decay also slightly shifts total parameter norm toward the language-model head and embeddings relative to intermediate layers.

F.2.2. ADDITIONAL ANALYSES ON ATTENTION MATRIX RANK: Weight decay reduces the pseudo-rank of query-key and value-projection attention matrices in fully trained Llama-2 models.The analysis covers layers 5 and 15 at 20 TPP.
F.2.2. ADDITIONAL ANALYSES ON ATTENTION MATRIX RANK: In OLMo models trained at 140 TPP, weight decay 0.1 generally produces lower attention rank than the 20 TPP and fully trained OLMo-2-1B-0425 models.The passage conjectures that differing learning-rate schedules explain this pattern: warmup-stable-decay for 140 TPP versus cosine for the 1x and 144x models.
F.2.2. ADDITIONAL ANALYSES ON ATTENTION MATRIX RANK: Training time does not reduce the rank of attention matrices.This result is summarized in Figure 28.
F.2.2. ADDITIONAL ANALYSES ON ATTENTION MATRIX RANK: Weight decay reduces model weight norms, except in the input layer, whose weights are not decayed.The observation concerns OLMo-2-1B models trained at 20 TPP.
F.3. Norm of parameters: As weight decay increases, intermediate-layer parameter norms slightly decline as a fraction of total model norm.The relative norms of the lm head and the non-weight-decayed embedding layer grow in OLMo models.
F.3. Norm of parameters: Figure 30 measures the effect of weight decay using per-layer L2-norm fractions for OLMo-2-1B-20x models.The layer-wise values sum all parameters within each layer.