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的改进相匹配。

🔬 方法详解

问题定义：在大语言模型训练中，权重衰减（Weight Decay, WD）被广泛应用于矩阵层，以防止过拟合。然而，WD与随机梯度噪声之间存在一种平衡，导致权重矩阵的范数被限制在一个特定的范围内。这种平衡状态可能并非最优，限制了模型的表达能力和性能。现有的muP方法虽然尝试解决这个问题，但需要手动调整乘子，计算开销大，且不够灵活。

核心思路：本文的核心思路是通过引入可学习的乘子（Learnable Multipliers）来动态调整权重矩阵的尺度。不再依赖固定的权重衰减来控制权重范数，而是让模型自己学习最佳的权重尺度。通过这种方式，可以解除权重范数的限制，使模型能够更好地适应数据，从而提高性能。

技术框架：该方法的核心在于在权重矩阵W上附加可学习的乘子。首先，作者尝试附加一个标量乘子，然后扩展到逐行和逐列的乘子。训练过程中，这些乘子与权重矩阵一起更新。整体框架与标准的Transformer架构兼容，可以与不同的优化器（如Adam和Muon）结合使用。该方法不需要修改模型的其他部分，易于集成。

关键创新：该方法最重要的创新在于将权重矩阵的尺度调整问题转化为一个可学习的过程。与传统的固定权重衰减或手动调整的muP方法相比，可学习乘子能够更灵活地适应不同的数据和任务。此外，逐行和逐列乘子的引入，进一步提升了模型的表达能力，允许模型对不同维度的特征进行差异化的缩放。

关键设计：可学习乘子被实现为可训练的标量，逐行向量或逐列向量，具体取决于所应用的变体。这些乘子在训练过程中与权重矩阵一起更新，使用标准的梯度下降算法。损失函数保持不变，仍然是标准的交叉熵损失或类似的损失函数。关键在于乘子的初始化和学习率的设置，需要仔细调整以获得最佳性能。作者还探讨了前向传递对称性和学习到的乘子的宽度缩放等问题，为未来的研究提供了方向。

📊 实验亮点

实验结果表明，可学习乘子方法优于经过良好调整的muP基线。使用Adam优化器时，该方法能够显著提高模型的性能。更重要的是，该方法在使用Muon优化器时，性能提升与从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.