Learnable Multipliers: Freeing the Scale of Language Model Matrix Layers

📄 arXiv: 2601.04890v1 📥 PDF

作者: Maksim Velikanov, Ilyas Chahed, Jingwei Zuo, Dhia Eddine Rhaiem, Younes Belkada, Hakim Hacid

分类: cs.LG

发布日期: 2026-01-08

💡 一句话要点

提出可学习乘子，解除语言模型矩阵层权重的尺度限制，提升模型性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 权重衰减 可学习乘子 muP 模型优化

📋 核心要点

  1. 现有大语言模型训练中，权重衰减与梯度噪声形成的权重范数平衡可能是次优的，限制了模型性能。
  2. 论文提出可学习乘子，通过学习权重矩阵的尺度，自适应地调整权重，打破了固定的权重范数平衡。
  3. 实验表明，该方法优于muP基线，降低了计算开销，并且在Adam和Muon优化器下均能提升下游任务性能。

📝 摘要（中文）

在大语言模型预训练中，对矩阵层应用权重衰减(WD)是一种标准做法。先前的工作表明，随机梯度噪声会导致权重矩阵W的布朗运动式扩展，这种增长会被WD抵消，从而达到一个具有特定权重范数||W||的WD-噪声平衡。本文将平衡范数视为训练过程中一种有害的人为因素，并通过引入可学习的乘子来学习最佳尺度来解决这个问题。首先，我们将一个可学习的标量乘子附加到W上，并证实WD-噪声平衡范数是次优的：学习到的尺度适应数据并提高性能。然后，我们认为单个行和列范数也受到类似的约束，并通过引入可学习的每行和每列乘子来释放它们的尺度。我们的方法可以看作是muP乘子的可学习的、更具表现力的泛化。它优于经过良好调整的muP基线，减少了乘子调整的计算开销，并提出了诸如前向传递对称性和学习到的乘子的宽度缩放等实际问题。最后，我们使用Adam和Muon优化器验证了可学习的乘子，结果表明，下游评估的改进与从Adam切换到Muon的改进相匹配。

🔬 方法详解

问题定义：在大语言模型训练中，权重衰减（WD）被广泛用于防止过拟合。然而，WD与随机梯度噪声之间的相互作用会导致权重矩阵的范数达到一个平衡状态，这个平衡状态可能并非最优，限制了模型的表达能力和泛化能力。现有方法，如muP，试图通过预定义的缩放规则来缓解这个问题，但这些规则可能不够灵活，无法适应不同的数据和模型结构。

核心思路：论文的核心思路是引入可学习的乘子，取代固定的权重缩放规则。通过让模型自身学习权重矩阵的最佳尺度，可以打破WD-噪声平衡的限制，使权重能够更好地适应数据，从而提高模型性能。这种方法可以看作是对muP的一种泛化，使其更具适应性和表达能力。

技术框架：该方法的核心是在权重矩阵的每一层引入可学习的乘子。具体来说，可以引入标量乘子、按行乘子或按列乘子。这些乘子与权重矩阵相乘，然后进行正常的模型训练。乘子本身也通过梯度下降进行优化，目标是最小化损失函数。整体训练流程与标准的大语言模型训练流程类似，只是在权重更新之前，需要先更新乘子。

关键创新：该方法最重要的创新在于将权重矩阵的尺度学习问题转化为一个可学习的优化问题。与传统的固定缩放规则相比，可学习的乘子能够更好地适应不同的数据和模型结构，从而获得更好的性能。此外，该方法还探讨了前向传递对称性和学习到的乘子的宽度缩放等实际问题，为进一步研究提供了方向。

关键设计：论文中，乘子被实现为可学习的标量、行向量或列向量。这些乘子在每次前向传播时与权重矩阵相乘。乘子的初始化策略和学习率是重要的超参数，需要仔细调整。损失函数与标准的语言模型训练损失函数相同，但梯度会通过乘子反向传播，从而更新乘子的值。论文还研究了不同的优化器（如Adam和Muon）对乘子学习的影响。

📊 实验亮点

实验结果表明，可学习乘子方法优于经过良好调整的muP基线。例如，在下游任务评估中，该方法取得了显著的性能提升，并且减少了乘子调整的计算开销。更重要的是，使用可学习乘子后，模型性能的提升与从Adam优化器切换到Muon优化器所带来的提升相当，这表明可学习乘子具有很大的潜力。

🎯 应用场景

该研究成果可广泛应用于大语言模型的预训练和微调，提升模型在各种下游任务上的性能，例如文本生成、机器翻译、问答系统等。通过自适应地调整权重尺度，可以提高模型的泛化能力和鲁棒性，使其更好地适应不同的数据分布和任务需求。此外，该方法还可以应用于其他类型的神经网络，例如卷积神经网络和图神经网络。

📄 摘要（原文）

Applying weight decay (WD) to matrix layers is standard practice in large-language-model pretraining. Prior work suggests that stochastic gradient noise induces a Brownian-like expansion of the weight matrices W, whose growth is counteracted by WD, leading to a WD-noise equilibrium with a certain weight norm ||W||. In this work, we view the equilibrium norm as a harmful artifact of the training procedure, and address it by introducing learnable multipliers to learn the optimal scale. First, we attach a learnable scalar multiplier to W and confirm that the WD-noise equilibrium norm is suboptimal: the learned scale adapts to data and improves performance. We then argue that individual row and column norms are similarly constrained, and free their scale by introducing learnable per-row and per-column multipliers. Our method can be viewed as a learnable, more expressive generalization of muP multipliers. It outperforms a well-tuned muP baseline, reduces the computational overhead of multiplier tuning, and surfaces practical questions such as forward-pass symmetries and the width-scaling of the learned multipliers. Finally, we validate learnable multipliers with both Adam and Muon optimizers, where it shows improvement in downstream evaluations matching the improvement of the switching from Adam to Muon.