可学习乘子：释放语言模型矩阵层的尺度自由度 Learnable Multipliers: Freeing the Scale of Language Model Matrix Layers
通过引入可学习乘子打破权重衰减-噪声平衡，提升LLM预训练性能
前置知识
权重衰减（Weight Decay）
权重衰减是AdamW等优化器中的正则化技术，在参数更新中加入 θ_t ← θ_t − ηλθ_t 项，其中 η 是学习率，λ 是权重衰减超参数。它通过惩罚大的权重值来防止过拟合，同时对稳定大规模训练至关重要。在LLM预训练中，权重衰减是标准实践，几乎所有现代大模型训练都采用AdamW而非原始Adam。
本文的核心论点是权重衰减与梯度噪声共同形成的平衡态限制了矩阵层学习最优尺度的能力，理解权重衰减机制是理解本文动机的关键。
噪声-权重衰减平衡（Noise-WD Equilibrium）
梯度噪声会导致权重矩阵产生类似布朗运动的扩张，而权重衰减则对抗这种扩张。两者达到动态平衡时，权重范数满足 ∥W∥ ∝ √(η/λ)，其中 η 是学习率，λ 是权重衰减系数。这个平衡态的范数由优化超参数决定，而非从数据中学习得到，因此可能是次优的。
这是本文提出可学习乘子的理论基础，理解噪声-权重衰减平衡态的存在及其限制是理解本文核心创新的前提。
最大更新参数化（µP, Maximal Update Parametrization）
µP是一种参数化方法，通过引入非学习的缩放因子来控制前向传播和学习率的缩放，使得模型在不同宽度下保持稳定的激活和更新幅度。具体来说，µP使用固定的乘子 s 配合关于模型维度的缩放规则，而不是让 s 可学习。许多预训练语言模型还会在小规模上调整µP乘子，然后迁移到目标规模。
本文的可学习乘子可以看作µP乘子的可学习泛化版本，理解µP的工作原理有助于理解本文方法与现有技术的关系和区别。
Adam Brownian Motion (ABM) 模拟
为了验证噪声-平衡态的理论预测，作者设计了一种模拟方法：生成零均值独立同分布的梯度序列 g_t ~ N(0, 1)，然后用与实际训练相同的 λ 和 η 调度在AdamW优化器中更新参数。这个模拟产生纯噪声驱动的轨迹，用于与实际训练中的权重范数轨迹对比。
ABM模拟是本文实验验证的关键工具，通过对比实验轨迹和模拟轨迹，作者确认了矩阵层确实遵循噪声-权重衰减平衡态，而可学习乘子则不受此约束。
参数空间对称性（Parameter Space Symmetries）
模型参数空间中存在不改变模型输出的变换。本文重点讨论两种：乘法对称性，当两个可学习因子 a 和 b 仅通过乘积 ab 出现时，变换 (a,b) → (sa, s⁻¹b) 不改变前向映射；归一化对称性，当模型仅使用归一化后的量 c/∥c∥_rms 时，缩放 c → sc 不改变输出。这些对称性会导致参数漂移和数值不稳定。
对称性是可学习乘子训练中的关键挑战，需要通过轻量级权重衰减来抑制对称性导致的不稳定，这是本文方法实际可用的重要实践细节。
研究动机
在大语言模型预训练中，权重衰减与梯度噪声的相互作用会导致矩阵层权重范数被困在一个由优化超参数决定的平衡态中，即 ∥W∥ ∝ √(η/λ)。这意味着权重矩阵的尺度并非从训练数据中学习得到，而是被优化过程的固有特性所限制。具体而言，当梯度噪声较大时，权重会产生布朗运动式的扩张，而权重衰减则将这种扩张限制在一个固定的平衡值。这个平衡值取决于学习率和权重衰减的比值，而非数据本身的特性。论文通过实验验证了这一限制：在投影器实验中，当强制改变平衡态范数 S(η_P, λ_P) 时，没有乘子的模型性能明显下降，而配备乘子的模型保持稳定。在MLP实验中，冻结所有RMSNorm权重后，MLP矩阵层在较大范数尺度下出现损失退化。这些实验表明，矩阵层无法自主逃离噪声-权重衰减平衡态来学习数据最优的尺度。
本文的目标是本文的具体目标是通过引入可学习乘子来打破噪声-权重衰减平衡态对矩阵层尺度的限制，使模型能够学习数据依赖的最优尺度。作者希望验证以下假设：矩阵层的范数被平衡态固定是次优的，通过附加可学习的标量或向量乘子可以让模型自由学习最优尺度，从而提升预训练性能。同时，作者希望探索可学习乘子与µP参数化的关系，分析乘子的宽度缩放行为，以及解决乘子训练中的对称性和梯度裁剪等实际问题。最终目标是在真实的长周期预训练设置中验证可学习乘子的有效性，证明其是一种通用的、无需额外推理开销的性能提升方法。
与已有工作不同的是，本文的独特切入角度在于将µP中的固定乘子重新解释为噪声-权重衰减平衡态的产物，并提出用可学习乘子替代固定乘子。与以往关注µP缩放规则或超参数迁移的工作不同，本文直接质疑平衡态范数的最优性。这一视角解释了为什么RMSNorm等标量/向量权重不需要权重衰减也能稳定训练：因为它们不受噪声-平衡态约束，可以自由学习最优尺度。与已有的重参数化方法如权重归一化、残差缩放不同，本文的动机来自于噪声-平衡态理论，这为乘子的放置位置和训练策略提供了理论指导。此外，本文发现可学习乘子可以自动学习正确的µP宽度缩放行为，减少了µP工作流中昂贵的乘子调优需求，这是对现有µP实践的重要简化。
核心方法
本文的方法基于一个直觉性观察：在标准LLM训练中，标量和向量型参数（如RMSNorm权重）不应用权重衰减也能稳定训练并自由学习最优尺度，而矩阵层却被噪声-权重衰减平衡态所束缚。作者将这一现象推广，提出在矩阵层上附加可学习乘子来释放尺度自由度。技术路线上，首先定义标量乘子 W_ij = s·W̃_ij 和向量乘子 W_ij = r_i·W̃_ij·c_j 两种重参数化形式。然后分析乘子的梯度特性，发现行/列乘子会累积对应行/列的梯度，标量乘子累积整个矩阵的梯度，这种额外平均降低了乘子的梯度噪声水平，解释了为什么乘子不需要权重衰减来对抗布朗扩张。在模型架构层面，作者推荐在gated MLP、attention和mamba2块中的特定位置放置乘子，避免冗余使用导致的对称性问题。训练时对乘子施加轻量级权重衰减 λ_lrm = 2×10⁻³ 来抑制对称性引起的不稳定，并将乘子梯度排除在全局梯度裁剪之外。
本文的核心创新在于将可学习乘子从µP的固定缩放工具重新定义为打破噪声-权重衰减平衡态的学习机制。与µP使用非学习的缩放因子 s 不同，本文让 s 可学习，使其能够适应数据分布。这一改变有三个关键区别：首先，可学习乘子不再强制执行µP的维度缩放规则，而是自动从数据中学习正确的尺度；其次，可学习乘子消除了µP工作流中对前向乘子和权重衰减乘子的调优需求，只需保留学习率乘子调优；第三，可学习乘子在推理时可以合并到权重矩阵中，不引入额外的内存或延迟开销。与已有的重参数化方法相比，本文的方法有明确的理论动机：乘子的存在是为了逃离噪声-平衡态陷阱，而非仅仅增加表达能力或改善优化景观。
方法步骤详情
方法的具体步骤如下。第一步是选择乘子类型和放置位置。对于每个矩阵层，可以选择标量乘子 s∈R 或向量乘子 r_i∈R^{d_out}, c_j∈R^{d_in}。放置时需避免冗余，例如RMSNorm权重已经是第一个线性层的列乘子，不应再添加列乘子；MLP上投影的行乘子和下投影的列乘子也是冗余的。作者为gated MLP、attention和mamba2块提供了推荐配置。第二步是处理对称性问题。乘法对称性出现在Q/K乘子等场景中，归一化对称性出现在残差输出等场景中。解决方案是对乘子施加轻量级权重衰减 λ_lrm = 2×10⁻³，这足以抑制漂移而不影响性能。第三步是处理梯度裁剪。实验发现乘子在训练初期会产生较大的梯度范数，导致过度激进的裁剪。解决方案是将乘子梯度排除在全局梯度范数计算之外，仅对矩阵权重进行裁剪。第四步是乘子初始化。可以使用µP调优的值作为初始化，但实验表明这对性能影响不大，只要学习率乘子调优正确即可。第五步是推理时合并。乘子可以预先合并到权重矩阵中：W = s·W̃ 或 W = diag(r)·W̃·diag(c)，不引入推理开销。
技术新颖性
本文的技术新颖性体现在多个层面。首先，理论视角新颖：将µP乘子重新解释为噪声-权重衰减平衡态的产物，并提出用可学习乘子打破这一平衡态，这是对µP理论的重要补充。其次，梯度分析新颖：通过链式法则推导出乘子梯度的表达式，证明乘子梯度相对于矩阵梯度有更低的噪声水平（通过跨行/列的平均），解释了为什么乘子不需要权重衰减。第三，实验设计新颖：通过ABM模拟将纯噪声驱动的轨迹与实际训练轨迹对比，清晰地展示了矩阵层遵循平衡态而乘子不遵循的现象。第四，宽度缩放分析新颖：发现可学习乘子自动学习正确的µP缩放行为（如投影器乘子∝d⁻¹，QK乘子∝d⁻²），但缩放略慢于理论预测，暗示特征-权重对齐随宽度衰减。第五，实践贡献新颖：发现乘子梯度应排除在梯度裁剪之外、对称性需要轻量级权重衰减等实用技巧，这些对实际部署至关重要。
实验结果
本文的核心发现可以分为几个层面。在验证性实验中，投影器实验证明了在强制改变平衡态范数 S(η_P, λ_P) 时，无乘子的配置（FPN）在极端值处性能明显下降（最终损失从约0.76升至0.78以上），而标量乘子（SPN）和向量乘子（VPN）配置保持稳定。ABM模拟确认了矩阵范数完全遵循平衡态缩放，而乘子轨迹与ABM版本毫无相似之处。MLP实验在Adam和Muon优化器上都观察到相同的行为，证明噪声-平衡态陷阱是跨优化器的普遍现象。在特征尺度多样性方面，添加标量乘子后，残差块输出范数在模型深度方向呈现整体递增趋势，注意力层在模型后半部分有显著不同的尺度差异，SSM的dt投影显示跨层变化显著。添加向量乘子后，行范数分布从非常狭窄变得显著宽阔，表明内部特征尺度多样性对模型有益。在宽度缩放实验中，矩阵层范数几乎不随宽度变化，确认平衡态范数不随模型规模缩放；投影器和内部激活的范数保持稳定，说明乘子自动学习了正确的宽度缩放。在长周期预训练验证中，200GT训练的结果显示：Adam基线平均分数30.80%，Adam+LRM为32.01%（+1.21%）；Muon基线31.88%，Muon+LRM为32.98%（+1.10%）。推理相关基准提升更大：BBH上Adam+LRM提升+2.33%，MATH lvl5上Adam+LRM提升+1.55%，GSM8K上Muon+LRM提升+2.61%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Hellaswag | 准确率 | Adam+LRM: 49.89%, Muon+LRM: 50.56% | Adam: 48.91%, Muon: 50.31% | Adam: +0.98%, Muon: +0.25% |
| ARC-C | 准确率 | Adam+LRM: 38.73%, Muon+LRM: 39.39% | Adam: 38.70%, Muon: 39.04% | Adam: +0.03%, Muon: +0.35% |
| MMLU | 准确率 | Adam+LRM: 45.33%, Muon+LRM: 47.96% | Adam: 44.18%, Muon: 47.98% | Adam: +1.15%, Muon: -0.02% |
| MMLU-PRO | 准确率 | Adam+LRM: 19.92%, Muon+LRM: 19.32% | Adam: 18.26%, Muon: 18.96% | Adam: +1.66%, Muon: +0.36% |
| BBH | 准确率 | Adam+LRM: 12.03%, Muon+LRM: 13.52% | Adam: 9.70%, Muon: 10.13% | Adam: +2.33%, Muon: +3.39% |
| GSM8K | 准确率 | Adam+LRM: 49.10%, Muon+LRM: 50.63% | Adam: 48.35%, Muon: 48.02% | Adam: +0.75%, Muon: +2.61% |
| MATH lvl5 | 准确率 | Adam+LRM: 9.07%, Muon+LRM: 9.49% | Adam: 7.52%, Muon: 8.69% | Adam: +1.55%, Muon: +0.80% |
| 平均 | 准确率 | Adam+LRM: 32.01%, Muon+LRM: 32.98% | Adam: 30.80%, Muon: 31.88% | Adam: +1.21%, Muon: +1.10% |
局限与改进
本文的局限性主要体现在以下几个方面。首先，实验规模有限：虽然200GT的训练时长约为Chinchilla计算最优时长的20倍，但模型规模仅为0.5B参数（Falcon-H1-0.5B），更大规模模型上的验证尚未进行，乘子的效果是否随模型规模增长而变化尚不清楚。其次，宽度缩放分析发现学习乘子的缩放略慢于µP理论预测的 d⁻¹ 和 d⁻² 趋势，作者将其归因于特征-权重对齐 α 随宽度衰减，但这与标准µP假设矛盾，需要进一步研究。第三，对称性处理仍需架构特定的工程：虽然轻量级权重衰减是通用解决方案，但最优的权重衰减值可能因架构而异，且需要避免过度正则化影响性能。第四，投影器乘子的添加需要谨慎：作者发现添加行乘子到LM头会导致性能退化（因为直接学习logit创造了绕过特征学习的捷径），而移除列乘子并扫描固定标量乘子也没有改善，说明乘子设计仍有探索空间。第五，可学习乘子打破了µP的维度缩放规则，这意味着需要重新分析模型规模缩放行为，增加了超参数迁移的复杂性。最后，论文承认还有许多开放问题：信号-噪声比在参数张量上形成连续谱，何时参数获得尺度适应能力的精确准则尚未建立。
独立分析的弱点
本文在几个方面存在可以改进的弱点。首先，实验设置的局限性：所有实验都在0.5B参数规模的Falcon-H1上进行，缺乏更大规模（如7B、13B、70B）的验证，读者无法确认乘子效果是否随规模增长而持续或变化。改进方向是在多个模型规模上进行系统性验证，并分析乘子值和性能增益的缩放规律。其次，乘子放置策略的探索不够充分：虽然作者为gated MLP、attention和mamba2提供了推荐配置，但这些配置是基于理论推理而非大规模消融实验得出的。对于其他架构（如混合专家MoE、稀疏注意力等），乘子的最优放置策略未知。改进方向是进行架构无关的自动乘子放置搜索。第三，µP工作流的整合不够清晰：可学习乘子打破了µP的维度缩放规则，但作者没有提供新的缩放规则来替代µP，这使得从µP迁移到可学习乘子的工作流不够明确。改进方向是建立新的宽度缩放理论。第四，对称性处理的理论分析不足：作者通过实验发现轻量级权重衰减可以抑制对称性漂移，但没有从理论上分析为什么这个特定值有效、是否存在最优值、以及如何根据不同架构自动确定。第五，计算开销分析不完整：虽然作者提到可以通过在优化器级别手动处理乘子动态来减少吞吐量损失，但没有给出具体的吞吐量测量数据，读者无法量化实际开销。
未来方向
基于本文的发现和作者的讨论，未来研究方向包括以下几个方面。首先，建立参数张量信号-噪声比的连续谱理论：作者指出矩阵和标量/向量动力学的区别可能源于梯度中信号-噪声比的差异，未来需要找到一个可经验测量的指标来判断参数张量是否获得尺度适应能力，并建立展示两种训练机制的最小数学模型。其次，开发包含可学习乘子的完整缩放规则：需要研究是否应随模型规模缩放乘子的学习率和权重衰减，以及可学习乘子是否自动确保无限宽度极限下的最大特征学习而无需手动缩放规则。第三，研究乘子对不同模型能力的差异化影响：作者发现推理相关基准（BBH、MATH、GSM8K）的提升远大于知识相关基准（ARC-C、MMLU），需要确认这是否意味着乘子只增强特定类型的电路。第四，探索乘子与其他优化技术的交互：包括与学习率调度、梯度累积、混合精度训练等的协同效应。第五，将可学习乘子推广到其他架构：包括稀疏注意力、混合专家、多模态模型等，以及微调和强化学习阶段。第六，研究乘子初始化策略的理论基础：虽然实验表明µP调优值的初始化不是必需的，但理论上最优的初始化策略仍有待探索。
复现评估
本文的复现性较好。首先，论文提供了详细的实验配置，包括Falcon-H1-0.5B架构、200GT训练时长、Adam和Muon优化器的使用、µP乘子的调优值来源（来自Zuo et al., 2025）等。其次，论文的方法实现简单：可学习乘子只需在矩阵层上附加标量或向量参数，使用标准自动微分和优化器实现即可，不涉及复杂的自定义操作。第三，论文提供了具体的超参数建议：乘子权重衰减 λ_lrm = 2×10⁻³、梯度裁剪值1、排除乘子梯度等。然而，也存在一些复现挑战：论文没有提供代码仓库或具体的模型配置文件，µP乘子调优值来自另一篇论文需要额外获取，Falcon-H1架构的完整实现需要参考其他文档。此外，200GT的训练需要显著的计算资源，小规模团队可能难以复现完整的长周期训练验证。总体而言，方法本身易于实现，但完整的实验复现需要中等以上的计算资源。
论文图表
