Weight Decay Improves Language Model Plasticity

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.

Modern language models are usually built in two stages: first, a general-purpose base model is created, and then it is taught to perform specific tasks. Current approaches often judge the quality of the base model by how well it performs right after it is built. However, this misses an important question: how well can the base model continue to learn new data for specific tasks? In this paper, we study this ability of the base model. We find that building the base model in a way that encourages it to learn simpler, higher-level patterns can make it more adaptable and thus learn new data better. Surprisingly, this means a base model that looks worse at first can end up being refined into a better final model. Through further analyses, we find that this way of building the base model helps the model organize information more clearly, simplifies its internal structure, and helps it learn more generalizable patterns. Our findings suggest that researchers and engineers should not judge the base model only by its immediate performance—its ability to learn new data is also crucial—and that the way the base model is built strongly shapes how useful the model becomes.