通过解耦权重向量的幅度和方向改进神经网络训练 | Alexander Hägele

# 通过解耦权重向量的幅值和方向来改进神经网络训练 来源：https://haeggee.github.io/posts/magnitude-direction-decoupling **作者：** Alexander Hägele, Atli Kosson, Alejandro Hernández\-Cano, Martin Jaggi > *Machine Learning and Optimization Lab, EPFL\.* *2026年6月15日* 联系方式: \\\(\\texttt\{alexander\.hagele@epfl\.ch\}\\\) **图1：***幅值-方向解耦*的主要结果。*左图：*学习率扫描。无论基础优化器如何，将权重固定到球面上都能改善最优损失；引入可学习的幅值（我们的工作，MD）则进一步提升。*中图：*使用稀疏MoE的扩展定律，在广泛的计算量范围内改进持续有效。*右图：*跨模型宽度的学习率迁移：通过球面直接控制相对权重更新，最优学习率自动迁移。所有细节详见实验部分 (https://haeggee.github.io/posts/magnitude-direction-decoupling#decoupling-magnitude-and-direction-the-details)。 ## TL;DR 在这篇文章中，我们介绍一种新的优化器改进，称之为***幅值-方向（MD）解耦***。我们将在未来几天发布更详细的预印本。本文旨在作为其更易读的配套材料：它建立了方法的直觉，并引导你了解我们当前的实验。 **核心思想。** 我们将每个权重矩阵 \\\(W\\\) 分解为一个方向 \\\(\\widehat\{W\}=W/\\\|W\\\|\\\)（具有固定范数）和一个幅值 \\\(\\gamma\\\)，并分别学习两者，每个都用各自的学习率（LR）以精心调控的速度进行更新。幅值不必是单个数字：它可以作用于每个矩阵、每行或每行和每列（我们在下面发现后者效果最好）。 **收益。** 在Adam和Muon上，解耦幅值和方向都优于精心调参的基线，无需重新调参即可将最优学习率迁移到不同模型宽度，并且在我们扩展到大型混合专家模型时依然有效。这些结果总结在图1中 (https://haeggee.github.io/posts/magnitude-direction-decoupling#figure-1)。 **目录** 1. 问题：幅值-方向干扰 (https://haeggee.github.io/posts/magnitude-direction-decoupling#the-problem-magnitude-direction-interference)——为什么标准优化会纠缠两者。 2. 解决方案：幅值-方向解耦 (https://haeggee.github.io/posts/magnitude-direction-decoupling#the-solution-magnitude-direction-decoupling)——解决方案。 3. 优化器内部 (https://haeggee.github.io/posts/magnitude-direction-decoupling#inside-the-optimizer)——更新规则和伪代码。 4. 细节 (https://haeggee.github.io/posts/magnitude-direction-decoupling#decoupling-magnitude-and-direction-the-details)——完整的实验结果。 5. 相关工作 (https://haeggee.github.io/posts/magnitude-direction-decoupling#related-work) 与讨论 (https://haeggee.github.io/posts/magnitude-direction-decoupling#discussion)。 如果你想跳过前文，可以直接跳转到方法解释 (https://haeggee.github.io/posts/magnitude-direction-decoupling#the-solution-magnitude-direction-decoupling) 和实验细节 (https://haeggee.github.io/posts/magnitude-direction-decoupling#decoupling-magnitude-and-direction-the-details)。 > *注意：* 除非另有说明，全文中的 \\\(\\\|\\cdot\\\|\\\) 表示 Frobenius 范数，“方向”/“幅值”可能指整个矩阵或其行和列，具体取决于变体。 --- ## 问题：幅值-方向干扰 我们从以下观察开始。一个权重矩阵 \\\(W\\\) 自然地分为两部分：一个**幅值** \\\(\\\|W\\\|\\\) 和一个**方向** \\\(\\widehat\{W\}=W/\\\|W\\\|\\\)。这类似于使用极坐标 (https://en.wikipedia.org/wiki/Polar_coordinate_system) 指定向量。像 Adam 和 Muon 这样的标准优化器对整个 \\\(W\\\) 进行更新，这两个量*最终会相互干扰*。图2通过一个玩具示例展示了这一点。 在一个尺度不变损失上的归一化梯度下降：相同的步长在幅值较小时对方向改变更大，而幅值即使在没有径向梯度的情况下也在增长。 **图2：** 标准优化器中的幅值-方向干扰，以玩具尺度不变损失为例说明，其中只有权重方向影响损失。*左图：*极坐标下的损失景观，从小的（红色）和大的（橙色）起始幅值出发，采用相同的归一化优化器步长。*中图：*相同的步长在小幅值时对方向（进而对损失）的改变远大于大幅值时。*右图：*即使损失没有径向梯度，该步长仍然增加了幅值。 损失（*左图*）是尺度不变的：只有方向影响输出，幅值不影响。这是深度学习中的常见情况，矩阵后面经常跟着归一化层。然而，幅值仍然影响单步更新的效果：在幅值较小时，相同的步长对方向改变很大；在幅值较大时，几乎没什么改变（*中图*）；尽管损失没有径向梯度，该步长仍然会增加幅值（*右图*）。学习率对这两个影响都无控制：方向变化由当前幅值决定，而幅值则作为方向变化的副产品而漂移。这就是为什么*标准优化器难以学习权重矩阵的幅值，并且需要权重衰减才能长期保持对方向的学习*。我们将在下面更仔细地观察这两种影响。 **方向变化取决于幅值。** 我们可以通过角度更新 \\\(\\angle\(W, W\+\\Delta W\)\\\) 来衡量优化器更新 \\\(\\Delta W\\\) 引起的方向变化，这近似于相对更新 \\\(\\\|\\Delta W\\\|/\\\|W\\\|\\\)。对于像 Adam 或 Muon 这样的归一化优化器，更新大小由学习率设定且独立于权重范数，因此方向变化大致与当前幅值 \\\(\\\|W\\\|\\\) 成反比（图2的中图 (https://haeggee.github.io/posts/magnitude-direction-decoupling#figure-2)）。因此，学习率并不能直接设定方向变化的速度，并且这个速度可能在不同层之间以及随时间变化而损害学习。我们之前关于旋转平衡 (https://arxiv.org/abs/2305.17212) 的工作展示了权重衰减如何通过在时间上调节相对更新并在各层之间平衡它们来部分解决这个问题。 **尽管没有径向梯度，幅值仍会增长。** 方向变化也会反馈到幅值上。暂不考虑动量细节，更新通常大致垂直于当前权重——这源于尺度不变性或噪声的性质——而垂直更新*总是*会增加幅值。即使没有任何力量将权重向外拉，这种情况也会发生：尺度不变函数没有径向梯度，但范数仍然会悄悄增加（图2，*右图*）。对于非尺度不变函数，（负的）径向信号必须足够强才能抵消它。在实践中，幅值趋于由学习率和权重衰减设定的平衡点，而非任何学习到的最优点（参见旋转平衡 (https://arxiv.org/abs/2305.17212)），这种不必要的增长可能需要像 Kimi 的 QK-clip (https://arxiv.org/abs/2507.20534) 这样的技巧来加以控制。 > *注意*：为简单起见，我们在此将幅值描述为单个标量，但同样的干扰也适用于每行或每列：一个不能很好地学习每矩阵尺度的优化器，在更细粒度上也表现不佳。 --- 解决方案是以一种*类似*极坐标的形式优化权重——一个方向和一个幅值——分别更新两者，使彼此互不干扰。具体来说，我们将每个权重分解为一个具有固定范数（因此位于固定超球面上）的方向 \\\(\\widehat\{W\}\\\) 和可学习的幅值增益： \\\[W = \\operatorname\{diag\}\(\\gamma\_\{\\text\{row\}\}\)\\,\\widehat\{W\}\\,\\operatorname\{diag\}\(\\gamma\_\{\\text\{col\}\}\), \\qquad \\widehat\{W\}\\ \\text\{位于球面上\},\\\] 其中 \\\(\\gamma\_\{\\text\{row\}\}\\in\\mathbb\{R\}^\{d\_\\text\{out\}\}\\\) 和 \\\(\\gamma\_\{\\text\{col\}\}\\in\\mathbb\{R\}^\{d\_\\text\{in\}\}\\\) 是可学习的增益（单个标量或单侧增益是特例）。两者以分别控制的速度学习；更新规则如下 (https://haeggee.github.io/posts/magnitude-direction-decoupling#inside-the-optimizer)。 **但我们不是刚刚说过幅值不重要吗？** 这只在一定程度成立：归一化层使损失对每个矩阵的*单个整体标量*不变，但对更细粒度的尺度则不然。模型仍然需要控制其激活的尺度，放大某些特征同时抑制其他特征，并混合处于不同尺度的激活。每行和每列的幅值会改变函数，因此能够学习它们很重要；这正是 RMSNorm 层中可学习增益起作用的原因。但是，标准 Transformer 中的此类增益远少于其矩阵的行数和列数（而且归一化层并非无处不在），因此它们不能提供同样细粒度的控制。我们的增益 \\\(\\gamma\_\{\\text\{row\}\}, \\gamma\_\{\\text\{col\}\}\\\) 使这种控制显式化，并以良好调控的速度学习，而不会受到上一节中容易干扰的动态影响。 ## 优化器内部 **更新方向。** 我们保持 \\\(\\widehat\{W\}\\\) 的更新量与其幅值成比例，然后将 \\\(\\widehat\{W\}\\\) 投影回球面，使幅值保持恒定（依赖于产生归一化更新的优化器）。这样，相对权重更新在每一步都由学习率直接决定。由于没有需要漂移的平衡点，也不依赖于初始化范数或训练长度，学习率调度直接设定相对更新。 **更新幅值。** \\\(W\\\) 的幅值由增益 \\\(\\gamma\\\) 决定，其更新方式类似于归一化层中常见的其他可学习增益。增益可以是标量、作用于每行或每列的向量，或者是同时缩放行和列的两个向量。我们注意到，这些幅值增益并没有提供超出原始矩阵的额外表示能力；它们只影响学习动态。 **融合权重。** 在实践中，我们不想将 \\\(\\gamma\\\) 和 \\\(\\widehat\{W\}\\\) 分开保存，并在前向和反向传播中重建权重。这会增加不必要的内存往返。相反，模型持有*融合后的权重张量* \\\(W\\\)，并像往常一样计算梯度 \\\(G = \\partial L/\\partial W\\\)。然后，在每一步，优化器恢复方向和增益，将梯度分拆给它们，更新各自，将方向投影回球面，并重新组装 \\\(W\\\)。 让我们以最简单的情况来说明优化器步骤，即单个标量 \\\(\\gamma\\\)，其中 \\\(W=\\gamma\\odot \\widehat\{W\}\\\)： \\\[\\begin\{aligned\} \\widehat\{W\} &\\leftarrow W / \\gamma && \\text\{恢复球面上的方向\}\\\\\[2pt\] g\_\\gamma &\\leftarrow \\mathrm\{reduce\}\\big\(\\widehat\{W\} \\odot G\\big\) && \\text\{增益梯度：在增益不跨越的轴上求和\}\\\\\[2pt\] G\_\{\\widehat\{W\}\} &\\leftarrow \\gamma \\odot G && \\text\{方向梯度 \} \\partial L/\\partial \\widehat\{W\}\\\\\[2pt\] \\widehat\{W\} &\\leftarrow \\mathrm\{OptStep\}\\big\(\\widehat\{W\},\\, G\_\{\\widehat\{W\}\},\\, \\eta\_W\\big\) && \\text\{任意（归一化的）矩阵优化器（Adam / Muon / ...）\}\\\\\[2pt\] \\widehat\{W\} &\\leftarrow \\widehat\{W\} \\,/\\, \\lVert \\widehat\{W\}\\rVert && \\text\{投影回球面\}\\\\\[2pt\] \\gamma &\\leftarrow \\mathrm\{AdamStep\}\\big\(\\gamma,\\, g\_\\gamma,\\, \\eta\_\\gamma\\big\) && \\text\{更新增益（使用其自身的学习率）\}\\\\\[2pt\] W &\\leftarrow \\gamma \\odot \\widehat\{W\} && \\text\{为下一次前向重新组装\} \\end\{aligned\}\\\] 两个梯度都遵循 \\\(W = \\gamma \\odot \\widehat\{W\}\\\) 上的链式法则：方向看到的是经增益缩放的 \\\(G\\)，增益看到的是 \\\(G\\) 在当前方向上的投影。我们将分拆和缩放操作保留在优化器内部，因此模型只看到一个“正常”的权重张量。 **重参数化的增益。** 直接更新 \\\(\\gamma\\\) 不稳定，因为靠近零时其梯度可能会激增。我们改为存储一个“原始”增益 \\\(\\widehat\{\\gamma\}\\\)，并通过一个平滑映射 \\\(\\gamma = \\varphi\(\\widehat\{\\gamma\}\)\\\)，例如 softplus，来保持 \\\(\\gamma \> 0\\\) 并限制梯度。我们将在结果 (https://haeggee.github.io/posts/magnitude-direction-decoupling#part-1-the-axes-of-normalization-and-magnitude) 中回到这个选择；最佳增益参数化仍然是一个开放问题。 ### 重要性质 - 整个想法*独立于优化器*：权重更新可以被视为黑盒，因此更好的优化器（AdEMAMix, Muon, Shampoo, ...）应该能够直接迁移。 - *我们不再需要权重衰减*，因为权重已经在球面上。这也避免了它与学习率调度的复杂交互，且有效步长现在就是学习率。 - 对于任何足够长的训练运行，*我们获得了跨宽度的学习率迁移*，因为我们直接控制了相对权重更新。 - 像 Muon 一样，*我们不再需要预热（warmup）*，因为预热旨在防止的早期大更新永远不会出现。 --- ## 解耦幅值和方向：细节 现在我们展示在寻找最佳方案过程中所做的消融实验。 **设置。** 对于第一部分，我们使用从 181M 到 1.29B 参数的密集 GPT 风格语言模型，每个模型都具有头维度 128、GQA、QK-norm 和 Sandwich Norm [1 (https://arxiv.org/abs/2105.13290), 2 (https://arxiv.org/pdf/2502.02732)]。我们在 RMSNorm 之后的块输出上应用固定缩放 \\\(\\alpha=\\frac\{1\}\{L\}\\\)（用于正确的深度缩放；更多信息见下文 (https://haeggee.github.io/posts/magnitude-direction-decoupling#part-2-learning-rate-transfer)）。矩阵参数使用标准差 \\\(\\frac\{1\}\{\\sqrt\{d\}\}\\\) 初始化，嵌入通过 \\\(\\sqrt\{d\}\\\) 放大，使得进入模型的 RMS 为 \\\(1\\\)。代码是 Megatron-LM 的一个分支。 我们的消融基础是 181M 模型（\\\(d=512, L=12\\\)），在 FineWeb Edu (https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) 子集的 250 亿个 token 上训练。这是刻意的*强过度训练*（Chinchilla 意义上的）：在约 50k 步、批量大小约 0.5M token（4096 序列长度）的情况下，更长期的训练动态变得可见，更接近真实的预训练运行（数百万步！）。 我们关注 AdamW 和 Muon 作为最流行的基础优化器。在所有实验中（包括图1 (https://haeggee.github.io/posts/magnitude-direction-decoupling#figure-1) 的扫描），我们*固定*每个 Adam 优化参数组的学习率（经验证处于良好范围；增益和输出层为 \\\(10^\{\-3\}\\\)，嵌入为 \\\(3\\cdot10^\{\-3\}\\\)），并分别针对每个优化器或设置变化扫描*矩阵学习率*。这意味着每种方法都使用*相同的预算*进行调参。标准方法（AdamW 和 Muon）使用权重衰减 \\\(0\.1\\\)；幅值-方向变体不使用权重衰减，因为权重已经在球面上。对于图1 (https://haeggee.github.io/posts/magnitude-direction-decoupling#figure-1) 中的 Muon，我们使用缩放因子 \\\(\\sqrt\{d\_\\text\{out\}/d\_\\text\{in\}\}\\\)（我们发现这在扫描时明显优于 RMS grafting）。除非另有说明。