权重衰减提升语言模型可塑性 Weight Decay Improves Language Model Plasticity
预训练时使用较大的weight decay能显著提升模型在下游微调任务中的表现
前置知识
Weight Decay（权重衰减）
一种正则化技术，在AdamW优化器中通过在每次更新时对参数施加一个与参数值成比例的惩罚项来防止过拟合。具体实现为 $\theta_{t+1} = \theta'_t - \gamma_t \lambda \theta_t$，其中 $\lambda$ 是权重衰减系数，$\gamma_t$ 是学习率。这会驱使模型参数趋向于零，从而控制模型复杂度。在大语言模型预训练中，$\lambda = 0.1$ 已成为一种事实上的默认设置，被广泛用于Llama、OLMo等主流模型的训练中。
本文的核心就是研究这个超参数对模型下游适应能力的影响，理解其工作机制是读懂论文的关键
Model Plasticity（模型可塑性）
指一个已经训练好的模型在进一步训练时适应新数据和新任务的能力。高可塑性的模型能够在不重新初始化的情况下有效学习新知识。衡量方法是：对预训练模型进行微调，然后评估其在下游任务上的表现——表现越好，说明可塑性越高。这与持续学习中的稳定性-可塑性困境密切相关：模型需要在保留旧知识（稳定性）和学习新知识（可塑性）之间取得平衡。
本文正是从可塑性这个新视角来重新审视weight decay的作用，这是理解全文的切入点
Tokens Per Parameter（TPP）
训练数据量与模型参数量的比值，用于衡量训练强度。20 TPP被认为是Chinchilla最优比例（Hoffmann et al., 2022），即每个参数对应约20个训练token。140 TPP则代表过度训练（overtrained）regime，训练数据量是计算最优的7倍。例如，一个1B参数的模型在20 TPP下使用20B tokens训练，在140 TPP下使用210B tokens训练。
实验结果表明weight decay的最优值在不同TPP regime下有显著差异，这是论文的重要发现之一
AdamW优化器
Adam优化器的一个改进版本，将权重衰减与梯度更新解耦。标准的Adam中，L2正则化会与自适应学习率产生交互，导致正则化效果不稳定。AdamW通过将权重衰减步骤与梯度更新步骤分离，使得正则化更加稳定和可控。目前是大语言模型预训练的标准优化器。
论文研究的是AdamW中weight decay超参数的作用，理解其解耦机制是前提
注意力矩阵的低秩性
注意力机制中的Query-Key矩阵的秩（rank）描述了注意力头表达能力的维度。由于该矩阵的秩上限是注意力头维度（通常远小于嵌入维度），它天然就是低秩的。weight decay可以进一步降低这个秩，使注意力矩阵更加压缩，从而学到更简洁、更泛化的模式。
这是论文解释weight decay提升可塑性的三个机制之一，低秩结构被认为有助于模型捕捉高层抽象模式
研究动机
在当前的大语言模型开发流程中，预训练和后训练（post-training）通常是分离进行的：一个团队负责预训练，另一个团队负责监督微调（SFT）、对齐和强化学习。预训练阶段选择超参数（如weight decay）的标准做法是基于验证集上的交叉熵损失——选择让损失最低的超参数值。然而，这种方法隐含了一个未经验证的假设：更低的预训练损失意味着更好的下游模型。具体来说，weight decay = 0.1 已成为预训练的事实标准（被Llama、OLMo等主流模型采用），但这个值是否在所有情况下都是最优的，尤其是从下游任务表现的角度来看，尚未被系统研究过。更关键的是，现有的scaling laws和超参数优化研究几乎完全聚焦于预训练损失，而忽略了模型在后续适应阶段的表现。
本文的目标是本文旨在从模型可塑性（model plasticity）的视角重新审视预训练超参数优化问题。具体目标包括：（1）系统性地研究weight decay对预训练模型可塑性的影响，即在微调阶段学习新任务的能力；（2）比较基于预训练损失选择的最优weight decay与基于下游任务表现选择的最优weight decay之间的差异；（3）揭示weight decay影响模型行为的机制，包括对模型表征、注意力矩阵结构和过拟合程度的影响。
与已有工作不同的是，本文的独特切入角度在于打破了预训练和后训练之间的解耦状态，提出了一种端到端（end-to-end）的超参数优化视角。与现有工作（如Bergsma et al., 2025研究weight decay scaling laws以最小化预训练损失）不同，本文关注的是最终目标——下游任务表现——而非中间目标——预训练损失。这种视角转换在概念上类似强化学习中的reward hacking问题：优化一个代理目标（预训练损失）并不总是优化真实目标（下游表现）的可靠方式。此外，本文还借鉴了计算机视觉和持续学习中关于模型可塑性的研究（Dohare et al., 2024），将其引入大语言模型领域，填补了该领域在语言模型可塑性研究上的空白。
核心方法
本文采用了一种系统性的实验设计，通过改变预训练阶段的weight decay值，然后在多种下游任务上微调并评估模型表现，来研究weight decay对模型可塑性的影响。整体思路是：首先在不同模型家族（Llama-2、OLMo-2）、不同尺寸（0.5B到4B）和不同训练强度（20 TPP和140 TPP）下预训练一系列模型，每个模型使用不同的weight decay值；然后在12个下游任务（涵盖CoT推理、语言理解、常识推理和安全对齐）上对这些预训练模型进行监督微调（SFT）；最后使用6个评估指标（包括Greedy/Pass@1、Maj@16、RM@16等）来评估微调后的模型表现，从而量化weight decay对模型可塑性的影响。
本文的核心创新点在于揭示了一个反直觉的现象：在某些训练regime下，预训练损失更高的模型在微调后反而表现更好。这与当前主流做法（选择预训练损失最低的超参数）形成了鲜明对比。具体来说，论文发现weight decay的最优值取决于优化目标：如果目标是最低预训练损失，最优weight decay可能是0.5或0.6（取决于模型和TPP）；但如果目标是最佳下游表现，最优weight decay可能需要进一步提高到1.0。这种现象在多个模型组中都观察到：例如，OLMo-2-1B-140x在weight decay=0.3或1.0时预训练损失略高于weight decay=0.1（分别为2.6208和2.7064 vs 2.6088），但在微调后的下游任务表现却显著更好。
方法步骤详情
实验分为三个主要阶段。第一阶段是预训练：在FineWeb-Edu（Llama-2）和OLMo-Mix-1124（OLMo-2）数据集上训练模型，sweep weight decay值从0到10（Llama-2-0.5B-20x和Llama-2-1B-20x使用10个值；OLMo-2-1B-20x使用11个值），记录每个模型的验证集交叉熵损失。第二阶段是微调：使用SFT在6个CoT推理任务（MetaMathQA、MedMCQA、PubMedQA、MMLUProCoT、RACE、SimpleScaling）、5个语言理解任务（HellaSwag、Winogrande、PiQA、Arc-Easy、Arc-Challenge）和1个安全对齐任务上进行微调，使用统一的超参数设置（batch size=64，学习率=1e-5或7.5e-6，3个epoch）。第三阶段是评估：对CoT任务使用6个指标（Greedy、Maj@16、Pass@16、Correct Ratio@16、RM@16、ORM Score），对语言理解任务使用cloze-style准确率，对安全任务使用harmfulness reward model评分。此外，还进行了机制分析：使用线性探针（linear probe）分析模型表征的线性可分离性，计算注意力矩阵的伪秩（pseudo-rank），以及计算训练-验证损失差距（train-val gap）来衡量过拟合程度。
技术新颖性
本文的技术新颖性体现在多个方面。首先，这是首次系统性地从模型可塑性视角研究预训练超参数选择问题，此前的工作要么关注预训练损失优化（Bergsma et al., 2025），要么在视觉模型或强化学习设置中研究可塑性（Dohare et al., 2024）。其次，论文提供了预训练损失不等于下游表现的严格实证证据，通过在5个模型组中展示相关性分析（图3）和留一法稳定性分析（图10），证明了这种现象的鲁棒性。第三，论文提供了三个机制来解释weight decay如何提升可塑性：（1）鼓励线性可分离表征——通过线性探针实验（图4）展示了weight decay与探针准确率的正相关；（2）降低注意力矩阵秩——通过伪秩分析（图5）展示了weight decay对注意力矩阵的不同影响（Query-Key矩阵对weight decay更敏感，秩降低约2倍）；（3）减少过拟合——通过train-val gap分析（图6）展示了weight decay与过拟合程度的单调递减关系。
实验结果
本文的核心发现可以归纳为三个主要结论。首先，较大的weight decay在预训练阶段可以提升模型可塑性，使模型在微调后获得更好的下游表现。在20 TPP的Chinchilla最优设置下，Llama-2-0.5B、Llama-2-1B和OLMo-2-1B的最优weight decay均为1.0，而Llama-2-4B的最优weight decay也是1.0。在140 TPP的过度训练设置下，OLMo-2-1B的最优weight decay为0.3。这些最优值都高于或等于当前默认的0.1。其次，预训练损失与下游表现之间没有完美的相关性。在5个模型组的相关性分析中（图3），Pearson相关系数r的范围从-0.96到+0.38不等，且留一法稳定性分析（图10）显示相关系数在去除单个数据点后可能在符号和大小上发生显著变化。具体例子包括：OLMo-2-1B-140x在weight decay=0.3或1.0时预训练损失分别比weight decay=0.1高约0.012和0.098，但在微调后的下游准确率却更高。第三，weight decay通过三种机制影响模型行为：在所有5个模型组中，更高的weight decay导致线性探针准确率更高（从约0.75提升到约0.85），注意力矩阵的秩在weight decay=1.0时降低约2倍，而train-val gap随weight decay单调递减（图6）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MetaMathQA (数学推理) | Greedy Accuracy | Llama-2-0.5B-20x: ~0.12, Llama-2-1B-20x: ~0.15 (weight decay=1.0) | Llama-2-0.5B-20x: ~0.08, Llama-2-1B-20x: ~0.10 (weight decay=0.1) | 相对提升约40-50% |
| MedMCQA (医学推理) | Greedy Accuracy | Llama-2-0.5B-20x: ~0.33, Llama-2-1B-20x: ~0.35 (weight decay=1.0) | Llama-2-0.5B-20x: ~0.30, Llama-2-1B-20x: ~0.32 (weight decay=0.1) | 相对提升约10% |
| PubMedQA (生物医学问答) | Greedy Accuracy | Llama-2-0.5B-20x: ~0.68, Llama-2-1B-20x: ~0.70 (weight decay=1.0) | Llama-2-0.5B-20x: ~0.65, Llama-2-1B-20x: ~0.68 (weight decay=0.1) | 相对提升约3-5% |
| RACE (阅读理解) | Greedy Accuracy | Llama-2-0.5B-20x: ~0.63, Llama-2-1B-20x: ~0.65 (weight decay=1.0) | Llama-2-0.5B-20x: ~0.55, Llama-2-1B-20x: ~0.58 (weight decay=0.1) | 相对提升约12-15% |
| CoT推理平均 | Pass@16 | Llama-2-0.5B-20x: ~0.53, Llama-2-1B-20x: ~0.55 (weight decay=1.0) | Llama-2-0.5B-20x: ~0.49, Llama-2-1B-20x: ~0.51 (weight decay=0.1) | 绝对提升约3-4个百分点 |
局限与改进
论文明确指出了几个重要局限。首先，机制分析是相关性的而非因果性的。虽然weight decay与线性探针准确率、注意力矩阵秩、train-val gap之间存在强相关关系，但难以分离这些机制与其他协变量的独立影响（例如，在保持模型表征不变的情况下单独改变注意力矩阵秩）。不过，weight decay改善可塑性这一效果本身是因果性的（实验设计只改变weight decay，其他变量保持恒定）。其次，实验范围有明确限制：只覆盖了Llama-2和OLMo-2两个模型家族，最大模型规模为4B参数，TPP比例最高为140x，下游任务主要是CoT推理、语言理解和安全对齐。在更大模型、更高TPP比例或其他下游任务类型上，结果可能会有所不同。第三，部分模型组（Llama-2-4B-20x和OLMo-2-1B-140x）的预训练模型数量较少（n<=10），可能影响相关性分析的统计可靠性。第四，论文只研究了SFT微调，没有覆盖RLHF、DPO等更复杂的后训练方法。此外，作者自己也指出，在极度过度训练的场景下，更低预训练损失带来的收益可能会超过可塑性的收益，此时默认的weight decay=0.1可能仍然更优。
独立分析的弱点
从独立分析的角度来看，本文存在几个值得关注的弱点。第一，实验的模型规模上限为4B参数，这在当今大模型时代（GPT-4约1.8T参数、Llama-3最大405B）显得相对较小，结果的外推性存疑——更大模型可能有不同的最优weight decay规律。建议未来在70B+规模上验证这些发现。第二，微调阶段使用了相对简单的SFT设置（3个epoch，固定的超参数），没有覆盖RLHF、DPO等现代对齐方法，而这些方法对预训练模型的敏感性可能不同。建议将实验扩展到包含奖励模型训练和偏好优化的完整后训练流程。第三，机制分析虽然提供了三个可能的解释，但缺乏因果验证实验——例如，能否通过后训练压缩注意力矩阵秩来获得类似效果？建议设计干预实验来建立因果关系。第四，论文没有讨论weight decay过大（如3.0或10.0）导致模型退化时，可塑性是否仍然提高——这些值在预训练损失分析中被排除，但在实践中可能遇到。
未来方向
基于本文的发现，有几个有前景的研究方向。首先，在更大规模模型（70B、405B甚至更大）上验证weight decay对可塑性的影响，并研究最优weight decay是否随模型规模变化（可能遵循某种scaling law）。其次，将研究扩展到其他后训练方法（RLHF、DPO、GRPO等），探索weight decay对这些方法效果的影响。第三，研究weight decay在视觉、多模态基础模型中的作用，验证这些发现是否是语言模型特有的。第四，深入理解weight decay如何影响遗忘——论文指出weight decay减少过拟合（减少对训练数据的记忆），这与可塑性提升密切相关，但具体的遗忘-学习平衡机制需要进一步探索。第五，开发新的超参数选择标准，同时考虑预训练损失和下游可塑性，例如在预训练过程中监控下游代理任务的表现，或者开发基于模型表征质量（如线性探针准确率）的early stopping指标。第六，研究不同weight decay值对模型在长上下文推理、代码生成、多轮对话等更复杂任务上的影响。
复现评估
论文的可复现性较好。作者公开了代码（GitHub链接在论文中），使用了公开可获取的数据集（FineWeb-Edu、OLMo-Mix-1124、MetaMathQA、MedMCQA等）和模型（Llama-2和OLMo-2架构）。预训练的算力需求相对适中：OLMo-2-1B-20x模型在单个H100节点上训练2天，OLMo-2-1B-140x模型训练2周。对于Llama-2-4B-20x，使用最多8个A100或16个H100 GPU。论文详细列出了所有超参数设置（表3-6），包括学习率、beta值、warmup比例等。微调阶段的算力需求也相对可控（batch size=64，3个epoch）。总体而言，对于有适当算力资源的研究团队来说，复现本文的核心实验结果是可行的。
论文图表
