🤖 machine learning

Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors

本文引入了幅度-方向（MD）解耦，这是一种优化器改进方法，它将权重矩阵分解为固定范数的方向和以不同速率更新的可学习幅度，从而在各种模型架构和优化器中稳定训练动态，并消除对权重衰减和预热的需求。

原作者： Alexander Hägele, Alejandro Hernández-Cano, Atli Kosson, Martin Jaggi

发布于 2026-06-25

📖 1 分钟阅读☕ 轻松阅读

原作者： Alexander Hägele, Alejandro Hernández-Cano, Atli Kosson, Martin Jaggi

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个巨大的、复杂的机器人（神经网络）说一种新的语言。机器人的大脑由数十亿个被称为**权重（weights）的小开关组成。为了教导这个机器人，你使用了一位名为优化器（optimizer）**的老师（比如 Adam 或 Muon），它通过轻推这些开关来让机器人变得更聪明。

长期以来，这些老师都将每个权重视为一个单一的、坚固的整体。它们会朝着特定方向推动整个整体。但本文的作者意识到这里存在一个问题：每个权重实际上有两个截然不同的部分：

方向（Direction）： 开关指向哪个方向（就像指南针上的指针）。
幅度（Magnitude）： 这个开关有多强或多响（就像音量旋钮）。

问题：“音量”与“方向”的纠缠

在标准训练中，老师试图同时移动整个整体。这造成了一种混乱的纠缠：

音量漂移： 当老师试图旋转开关的方向时，即使老师并非本意，其“音量”（幅度）也会产生副作用，变得更大或更小。
方向变得混乱： 如果音量太大，微小的推动很难改变方向；如果音量太小，同样的推动会让方向剧烈旋转。
补救工具箱： 由于这种混乱，工程师不得不使用复杂的“创可贴”，比如权重衰减（weight decay）（一条不断尝试缩小音量的规则）和预热（warmup）（通过从极慢的速度开始以避免混乱）来保持训练的稳定。

解决方案：幅度-方向（MD）解耦

作者提出了一种教导机器人的新方法。他们不再将权重视为一个坚固的整体，而是在优化器内部将其拆分为两个独立的部分：

指南针（方向）： 他们强制要求方向保持在一个固定的“球面”上（就像一个地球仪）。老师可以在这个地球仪上旋转指南针的指针，但指针的长度永远不会改变。
音量旋钮（幅度）： 他们为开关的每一行和每一列添加了独立的、可学习的“音量旋钮”（增益）。这些旋钮独立地控制着响度。

类比：
想象你正在驾驶一艘船。

旧方法： 你有一个巨大的杠杆，同时控制着舵（方向）和引擎功率（幅度）。如果你试图转动舵，引擎功率会意外地升高或降低，导致船只难以操控。你必须不断地与引擎作斗争，以保持航向稳定。
新方法（MD 解耦）： 你有一个专门控制舵的转向轮和一个独立的控制引擎节流阀。你可以随心所欲地转动转向轮，而引擎功率不会发生变化。你也可以独立地调节节流阀，让船行得更快或更慢。

使用这种方法会发生什么？

论文表明，这种简单的分离带来了一些令人惊喜的好处：

不再需要“创可贴”： 因为方向被锁定在固定大小，且音量被单独控制，所以机器人不再需要权重衰减或预热。训练过程自然地保持稳定。
可预测的缩放： 通常情况下，如果你把机器人做大（增加开关数量），你必须从头开始重新调整老师的设置。使用这种新方法，无论机器人是大是小，“转向灵敏度”（学习率）都保持不变。你可以先调优一个小模型，然后直接将完全相同的设置用于一个庞大的模型。
学习更快： 机器人学得更好、更快。在针对大型模型（混合专家模型，Mixture-of-Experts）的测试中，这种方法达到了与现有最佳方法相同的性能水平，但使用的计算能力仅为一半。

核心结论

本文认为，通过将神经网络权重的“方向”和“响度”视为两个可以独立控制的变量，我们可以更高效地训练 AI 模型，减少规则的使用，并获得更好的结果。这就像是意识到，要开好车，你需要分别控制转向和油门，而不是试图用一只手同时完成这两件事。

技术摘要：通过解耦权重向量的幅值与方向来改进神经网络训练

1. 问题陈述

现代神经网络训练依赖于如 Adam 和 Muon 之类的优化器，这些优化器将权重矩阵视为单一的、整体的对象。然而，每个权重矩阵 $W$ 本质上都携带两个截然不同的量：幅值 ( $\|W\|$ ) 和方向 ( $\hat{W} = W/\|W\|$ )。在标准优化中，这两个量是耦合在一起的，导致不稳定且不可预测的训练动态：

方向性变化的依赖性： 更新引起的角度变化与当前幅值成反比。因此，学习率（LR）并不能直接控制方向变化的速率，而这种变化会随层级和时间而变化。
非预期的幅值漂移： 由于尺度不变性或噪声，更新往往趋于垂直于当前权重，这导致即使在没有径向梯度的情况下，幅值也会增长。这种“失控增长”迫使从业者依赖**权重衰减（weight decay）等间接机制来稳定训练，并依赖预热（warmup）**来防止早期不稳定。
缩放限制： 当前的缩放训练方法（例如 $\mu$ P）通常在改变模型宽度或深度时需要重新调整超参数。此外，最近的一些方法虽然将权重约束在固定范数的球面上，但缺乏可学习的幅值参数，使得尺度控制完全由归一化层承担。

2. 方法论：幅值-方向 (MD) 解耦

作者提出了 幅值-方向 (MD) 解耦，这是一种优化器修改方案，它将每个权重矩阵分解为一个固定范数的方向和可学习的幅值增益，并以独立的学习率分别更新它们。

核心分解

每个权重矩阵 $W$ 被重新参数化为：
$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$
其中：

$\hat{W}$ 是方向，被约束在固定的超球面上（例如 Frobenius 范数）。
$\gamma_{\text{row}}$ 和 $\gamma_{\text{col}}$ 是可学习的幅值增益（分别作用于行和列的向量）。

优化程序

该方法完全在优化器步骤内运行，保持模型的前向传播不变（模型看到的是一个单一的融合张量 $W$ ）：

恢复： 从融合后的权重 $W$ 中恢复 $\hat{W}$ 和增益 $\gamma$ 。
梯度拆分： 将梯度 $G = \partial L / \partial W$ 拆分为方向梯度 ( $\partial L / \partial \hat{W}$ ) 和增益梯度（沿各自轴求和）。
方向更新： 使用基础优化器（如 Adam, Muon）配合学习率 $\eta_W$ 更新 $\hat{W}$ ，然后将其投影回球面以保持固定范数。
增益更新： 使用独立的学习率 $\eta_\gamma$ 更新增益 $\gamma$ （可选使用如 softplus 等平滑映射以确保正定性）。
重组： 为下一次前向传播重建 $W$ 。

关键特性

优化器无关性： 该方法作为任何基础优化器（Adam, Muon, Shampoo 等）的封装器运行。
直接的学习率控制： 通过将方向固定在球面上，学习率直接设定了相对的角度更新，消除了对当前权重幅值的依赖。
消除启发式方法： 该方法消除了对权重衰减的需求（因为权重已被范数约束）以及对预热的需求（因为不稳定的早期更新从一开始就得到了调节）。
细粒度的尺度控制： 不同于仅依赖归一化层增益的方法，MD 解耦提供了显式的、可学习的个体行与列的尺度控制，允许模型放大或抑制特定特征。

3. 主要贡献与发现

实验结果

作者在稠密模型（181M 至 1.29B 参数）和大型混合专家（MoE）模型（总参数量高达 6.7B）上验证了 MD 解耦。

性能提升： MD 解耦优于经过精调的基准模型（AdamW 和 Muon）。在大型 MoE 模型上，MuonMD（结合了 Muon 的 MD 解耦）达到 AdamW 的损失水平仅需约 2 倍少的计算量。
学习率迁移： 该方法实现了跨模型宽度的零样本学习率迁移。在小模型上找到的最优矩阵学习率，在显著更大的模型上依然保持最优，且无需重新调优。这一特性对于宽度和深度的缩放均有效。
无需预热的训练： 实验表明，取消预热不仅能保持稳定性，还能提高最终损失，因为早期的训练步骤不再浪费在降低的学习率上。
消融研究洞察：
- 归一化轴： 通过 Frobenius 范数（整个矩阵）约束方向被发现是最灵活的，其表现与行/列约束相当。
- 增益参数化： 结合了每行和每列的增益（ $\gamma_{\text{row}} + \gamma_{\text{col}}$ ）优于标量或单轴增益。使用 softplus 参数化增益比直接更新或指数更新具有微弱但一致的优势。
- 嵌入（Embeddings）： 将嵌入向量约束为单位 L2 范数比保持无约束状态表现略好。
与相关工作的比较：
论文将 MD 解耦置于关于球面训练和重参数化的研究领域中：
- 不同于 nGPT，后者将球面约束与特定的架构变更（如 L2 归一化、插值残差）捆绑在一起，MD 解耦是一种优化器侧的修改，可以应用于标准架构。当应用于 nGPT 架构时，MD 解作为进一步提升了性能。
- 不同于 Muown 或 HyperP，这些方法可能会固定范数但不具备可学习的幅值，或者将该方法与特定优化器绑定，MD 解耦则明确学习幅值且与优化器无关。
- 它解决了先前研究（如 Kosson et al., 2024a）中识别出的“幅值-方向干扰”问题，通过分离两者的动力学过程。

4. 重要性与主张

论文声称，将幅值和方向视为受分别控制的量，可以产生更可预测的训练动态，并为现代优化器提供了一个简单且具有广泛适用性的改进。

简洁性： 该方法不需要架构变更，且在优化器步骤中增加的计算开销极小（仅为逐元素操作）。
鲁棒性： 它通过移除对权重衰减和预热的需求，简化了训练流程，减少了超参数搜索空间。
可扩展性： 它为跨模型宽度的学习率迁移提供了一种原则性的机制，解决了大规模语言模型缩放中的一个关键瓶颈。

作者总结道，虽然在优化器步骤中引入额外的逻辑会带来轻微的开销，但 MD 解耦在训练稳定性、计算效率和超参数迁移方面的优势，使其成为大规模神经网络训练领域的一项重要进展。