从Muon到梯度裁剪：关于QK稳定性的几点思考

引子：一个迷人的理论与一个棘手的问题 在深度学习的武库中，优化器是当之无愧的”核心引擎”。当我们习惯了在 Adam 和它 变体中调参时，一个基于完全不同哲学思想的优化器——Muon Optimizer，正逐渐进入研究者们的视野。它由 Keller Jordan 和 Jeremy Bernstein 在其工作中首次提出[1]，其深刻之处在于，它不再拘泥于我们熟悉的参数空间 (parameter space)，而是从功能空间 (function space) 的角度出发，构建其更新规则。 这个思想在国内由 Kimi AI 在其 Kimi K2 [7] 等工作中发扬光大并得到了成功的实践，展示了其巨大的潜力。然而，K2 模型的公开材料中也敏锐地指出了一个棘手的问题：如果将原版的 Muon 优化器直接应用于 Transformer 的 Query (Q) 和 Key (K) 权重矩阵，训练过程会变得非常不稳定，甚至崩溃。 理论的优雅，在实践中遇到了壁垒。这立刻引发了我们的好奇心： 这个现象背后的根本原因是什么？ 我们能否沿着 Muon 的”第一性原理”，试提出一种适用于 QK 更新的修正思路？ 本文将记录笔者沿着这条路进行的一次思考。这是一段充满曲折的旅程，从严格的理论推导，到遭遇工程上的计算瓶颈，再到最终讨论可行的启发式方案。这是一个关于理论、工程与直觉的思考过程。 第一部分：理论回顾——为什么标准Muon不适用于QK？ 要理解这个”水土不服”的问题，我们必须首先深入理解 Muon 优化器的”初心”，以及它与我们熟悉的 Adam 等优化器在根本上的区别。 Muon 优化器的推导与精髓 一切要从优化的最基本目标说起。对于一个权重矩阵 ，我们计算出了它关于损失 的梯度 。我们的目标是找到一个更新量 ，让它能够最大程度地降低损失。在数学上，损失的一阶泰勒展开告诉我们 ΔL≈Tr(G T W ΔW)。为了让 下降最快，我们希望找到一个 来最小化这个迹，即让 的方向与 的方向尽可能相反。 这可以形式化为一个优化问题：

显然，如果不加任何约束， 的元素可以趋于无穷大，这没有意义。因此，所有优化器都必须施加一个约束，来限制更新步长 的”大小”。区别就在于，如何定义这个”大小”。 传统约束（参数空间）： 我们最熟悉的 SGD、Adam 等优化器，其约束都直接施加在 本身之上，即在参数空间中限制它的大小。最常见的约束是 Frobenius 范数：

其中 是一个与学习率相关的常数。这个约束的几何意义很直观：参数更新的欧几里得距离不能太远。 Muon 约束（功能空间）： Muon 的哲学则完全不同。它认为，我们不应该关心参数 本身移动了多远，而应该关心这次移动对模型”功能”的影响有多大。对于一个线性层，它的功能就是将输入 映射到输出 。因此，参数更新对功能的影响，就是输出的变化量 。 Muon 的核心思想是，这个功能上的变化量应该是有界的，而且这个界应该对所有可能的输入都成立。为了消除输入 本身大小的影响，我们只考虑单位范数的输入。于是，Muon 的约束被定义为：

这个公式的含义是：”对于任何一个长度为1的输入向量 ，经过更新矩阵 变换后，其输出向量的长度都不应该超过 “。学过矩阵论的同学会立刻认出，这个上确界（supremum）的定义，正是矩阵谱范数 (spectral norm) 的定义。矩阵的谱范数等于其最大的奇异值。 现在，Muon 的优化问题变得清晰了：

这个在谱范数约束下的优化问题，有一个非常优雅的解析解，它就是标志性的矩阵符号函数 (matrix sign function)：

函数会保持 的奇异向量不变，但将其所有的奇异值都变为1。这是一种极致的方向控制，它将更新的”能量”均匀地分布在所有奇异方向上。 QK 权重更新的”几何冲突” 理解了 Muon 的机制，我们再来审视它在 QK 更新上的失败。正如苏剑林老师在其博客文章中所分析的[2]，问题出在 Attention 机制的双线性（bilinear）本质上。 首先，一个比较浅层的、直觉的解释是功能耦合。Q 和 K 的功能是配对出现的，它们必须共同作用于乘积 才能产生有意义的注意力分数。独立地约束 和 各自的功能变化，并不能保证它们耦合后的功能变化也在可控范围内。 而更深层的几何解释，根据苏剑林的推想[2]，则与 Muon 优化器的满秩更新特性有关。实践中观察到的”MaxLogit 爆炸”，往往意味着 或 的谱范数出现了爆炸的迹象。因此，问题可以转化为：”为什么 Muon 更容易导致谱范数爆炸？”。我们知道谱范数等于最大的奇异值，所以又可以进一步联想到：”为什么 Muon 更倾向于增大奇异值？” Muon和Adam的区别是什么呢？Muon给出的更新量是经过 运算的，所有奇异值都相等，即它的有效秩是满秩的；而一般情况下的矩阵，奇异值通常都是有大有小，并且以前面几个奇异值为主，从有效秩的角度看它们是低秩的，我们对Adam更新量的假设也是如此。这个假设并不新鲜，比如高阶muP同样假设了Adam更新量的低秩性。 用公式来说，我们设参数 的奇异值分解（SVD）为 ，Muon更新量 的SVD为 ，Adam更新量 的SVD为 。那么权重更新 就可以写作：

很明显，如果 的某个奇异向量对 跟某个更新量的奇异向量对（如 或 ）很接近（即”方向对齐”），它们的奇异值就会倾向于直接叠加起来，从而增大 在该方向上的奇异值。 由于Muon的更新量是满秩的，它的”能量”均匀地分布在所有奇异方向上，所以它与 的奇异向量发生”碰撞”的几率会远大于低秩的Adam更新。因此，Muon天然地更容易增大参数的奇异值。 这个效应在 Attention 机制中被放大了。注意力分数的核心是双线性形式 。如果 Muon 同时在增大 和 的谱范数，那么在计算最终的点积时，这个增大的效应会被相乘，导致爆炸的风险急剧增加，形成”差的更差”的恶性循环，最终导致训练崩溃。 至此，我们从理论上清晰地解释了为什么一个设计如此优雅的优化器，在 Transformer 的 QK 权重上会”水土不服”。问题的根源在于，我们必须约束乘积项 的变化，而非各自独立的变化。这个洞察，将是我们下一步探索的出发点。 第二部分：一次”有原则”的尝试——为QK定制Muon 在第一部分，我们得出了一个关键的洞察：简单地约束 和 各自的更新是行不通的。因为在注意力机制中，它们是”功能耦合”的，最终影响性能的是它们的乘积 。独立的约束可能允许 和 这两个向量的模长变化不大，但方向却趋于一致，导致它们的点积（即注意力分数）爆炸。 因此，一个有原则的、真正遵循 Muon 思想的修正方案，必须从问题的根源出发：直接约束最终功能输出——注意力分数矩阵的变化。这便是我们本次理论尝试的起点。 2.1 建立新的优化目标：从约束参数到约束功能 我们的目标不再是约束权重矩阵 和 本身，而是约束由它们的变化所引起的注意力分数矩阵 的变化 。 首先，我们来定义相关概念： 输入矩阵 : 一个 的矩阵，其中 是序列长度， 是嵌入维度。每一行代表一个 token 的向量表示。 权重矩阵 : 两个 的矩阵，其中 是注意力头的维度 (head dimension)。它们将 维的输入映射到 维的查询和键空间。 查询和键矩阵 : ，。它们是 的矩阵。 注意力分数矩阵 : 。这是一个 的矩阵，其 元素表示第 个查询与第 个键之间的相似度。 我们希望在权重更新后，新的注意力分数矩阵 与旧的 差别不大。这个差别 可以写作：

展开括号内的项：

由于 和 是微小的更新量，它们的乘积 是一个二阶无穷小项，可以忽略。这个线性化处理是几乎所有优化理论的基础。于是，我们得到 的一阶近似：

现在，我们对这个变化的大小进行约束。在Muon优化器中，最自然的范数是谱范数 (Spectral Norm)，记作 。谱范数衡量的是一个矩阵能够将一个向量”拉伸”的最大倍数。施加谱范数约束，意味着我们希望控制在”最坏情况”下注意力分数的变化。 我们新的约束条件是：

其中 是一个我们设定的很小的正数，代表我们允许的功能变化”预算”。 为了简化问题，使其不依赖于具体的输入 ，Muon做了一个核心假设：我们希望这个约束对任何可能的输入都成立。为了达到这个目的，我们考虑最坏的情况，即输入 拥有最大的”放大”能力。利用矩阵范数的次乘法性质 ()：

假设输入被归一化，即 ，那么上述约束就简化为对中间权重部分的直接约束：

这就是我们为QK机制量身定制的、充满”原则”的Muon约束。它完美地捕捉了和的耦合效应。 2.2 构建完整的优化问题 我们的最终目标是在满足上述约束的前提下，让损失函数 下降得最快。根据梯度下降的原理，损失函数的一阶变化 约等于：

其中 和 分别是损失对 和 的梯度。迹 在这里是矩阵内积的推广。为了让 最大程度地为负，我们需要最小化这个表达式。 综合起来，我们构建了如下的优化问题：

2.3 求解的坎坷之路：从解耦到遭遇伪逆 这个优化问题看起来很优雅，但解决起来却相当棘手。主要的困难在于约束条件将 和 紧紧地耦合在了一起，我们无法像标准梯度下降那样独立地求解它们。 第一步：解耦（Decoupling） 我们采取一个常用且有效的策略：利用三角不等式 () 放宽约束。我们施加一个比原始约束更严格的条件，如果这个更严格的条件满足，那么原始条件必然满足。

这个新约束的好处是，我们可以将更新预算 分配给两个部分，从而将一个复杂的耦合问题分解为两个独立的、更容易处理的子问题。一个自然且公平的分配方式是各占一半： 子问题1 (针对 ):

子问题2 (针对 ):

这两个子问题的结构是相似的，我们先集中精力解决第一个。 第二步：变量代换与求解 子问题1的约束 依然很麻烦，因为 被乘上了 。这里需要一个关键的技巧：变量代换。 我们定义一个新变量 。这样一来，约束条件立刻变得清爽无比：

现在，我们需要用新变量 来表示旧变量 。从 中解出 ，需要用到线性代数中的一个强大工具——摩尔-彭若斯伪逆 (Moore-Penrose Pseudoinverse)。对于一个矩阵 ，其伪逆记作 。伪逆是普通逆矩阵的推广，即使在 不是方阵或不可逆时也存在。 利用伪逆，我们可以得到：

将这个关系代回到我们的子问题1的目标函数中：

为了凑出标准形式，我们使用迹的循环性质 ：

现在，子问题1被完全转换成了关于新变量 的问题：

这是一个非常标准形式的优化问题，与muon的形式一模一样，它的解是已知的。当我们要最小化 并约束 的谱范数时，最优解 的方向与梯度 的”矩阵符号”相反。 矩阵符号函数 (Matrix Sign Function)，记作 ，是符号函数 到矩阵领域的推广。它将一个矩阵分解为”方向”和”尺度”两部分，并只保留其”方向”信息。其严格定义是通过奇异值分解（SVD）得到的，但我们只需要知道它的作用即可。 应用这个结论，我们得到 的最优解 ：

第三步：还原解与遭遇瓶颈 我们已经求出了中间变量 的解，现在需要把它代换回去，得到我们最终想要的 ：

为了书写简洁，我们可以忽略常数 ，只看其比例关系：

至此，我们通过一步步严格的推导，得到了一个定制版的Muon更新规则。这个公式在理论上是完美的，它精确地回答了我们开始提出的问题。 然而，当我们从数学的殿堂回到工程的现实时，一个巨大的障碍横亘面前：伪逆 。 这个公式中出现了两次伪逆运算！计算一个矩阵的伪逆，通常需要进行奇异值分解（SVD），这是一个计算成本极高的操作（对于一个 的矩阵，复杂度大约是 ）。在深度学习训练的每一次迭代中都执行这样的操作，其计算开销是完全无法接受的。 这就是我们遭遇的第一个重大挫折。我们得到了一个理论上无懈可击的”屠龙之技”，但在实际应用中却因其巨大的计算复杂度而”英雄无用武之地”。这个结果虽然令人沮丧，但它清晰地指明了理论和实践之间的鸿沟，也为我们接下来的探索埋下了伏笔：有没有办法绕过，或者简化这个可怕的伪逆计算呢？ 原路返回：分离WQ和WK 在求解过程中，我们曾利用三角不等式将耦合的约束解耦。但即便如此，我们面对的子问题约束 依然需要通过伪逆来求解。一个自然的想法是，我们能否在这里再做一步近似，彻底摆脱伪逆？ 我们可以利用矩阵范数的次乘法性质 () 对约束做进一步放缩：

这样，子问题1的约束就变成了：

这个形式我们非常熟悉！它已经完全退化成了标准Muon优化器的形式，只不过学习率（即谱范数约束的上限）由一个动态变化的项 来决定。其解就是 。 然而，这一步看似巧妙的简化，其代价是巨大的。它本质上完全放弃了对 和 之间相互关系的建模。我们之所以费尽心机推导那个包含伪逆的复杂公式，正是为了捕捉它们之间的几何关系。那个解的精妙之处在于，它能感知到更新方向 与当前权重 之间的”对齐”程度。如果 的方向与 的行空间大致正交，那么它们的乘积对注意力分数的影响就小，优化器就允许一个更大的更新；反之，如果它们方向一致，意味着微小的参数变动都可能引起注意力分数的剧烈变化，优化器就会采取一个更保守的更新。这种细腻的几何调控，正是由伪逆项 所承载的。 而我们刚刚提出的这个新近似，则完全抛弃了这些几何信息。它只关心 自身的大小，而忽略了其更新方向。笔者猜测，这个近似的鸿沟可能过大，它虽然简化了计算，但也丢掉了问题的精髓，因此大概率不会有好的效果。 另一条岔路：用迭代法绕过SVD？ 在我们因SVD的计算量而感到绝望时，一个自然的想法是：我们是否必须用SVD来计算伪逆？答案是否定的。线性代数中存在一些纯粹基于矩阵乘法和加法的迭代方法，它们同样可以收敛到伪逆。这类方法对GPU等并行计算设备极为友好。 一个典型的例子是牛顿-舒尔茨迭代法 (Newton-Schulz Iteration)。要计算矩阵 的伪逆 ，我们可以从一个合适的初始值 开始，反复执行以下迭代：

一个好的初始猜测是 ，其中 是一个缩放常数（例如，可以用谱范数的倒数来近似）。 这个方法看起来非常有吸引力。它将那个”无法计算”的SVD操作，替换成了一系列深度学习框架最擅长的矩阵乘法。我们似乎找到了一条绕过计算瓶颈的捷径。 然而，这条捷径也通向了新的困境。问题不在于计算量本身，而在于数值稳定性。这类迭代法对于计算精度非常敏感，而在深度学习训练中，我们通常使用 bfloat16 这样的低精度浮点数。在低精度下，迭代法可能很难收敛到足够精确的解，甚至完全不收敛。 不过，这是否就意味着这条路彻底走不通了呢？或许还有一些工程上的技巧可以挽救。 以下试提出几种思路： 一种思路是缓存 (Caching)。在训练过程中，权重矩阵 并不会在每一步都发生剧烈变化。因此，它的伪逆 可能也相对稳定。我们可以不必在每一步都重新计算伪逆，而是每隔（比如）10步计算一次，并将结果缓存起来复用。这可以将昂贵的计算成本摊薄。 另一种更精妙的思路是热启动 (Warm Start)。我们可以将上一步计算出的伪逆作为当前步骤迭代求解的初始猜测值。由于相邻步骤的 变化不大，上一步的解理应是当前解的一个绝佳近似，这可以大大减少收敛所需的迭代次数。 有趣的是，这个”热启动”策略对于计算 msign 函数通常是无效的。msign 的结果只依赖于输入矩阵的奇异向量（方向），而与其奇异值的大小无关。当矩阵发生微小扰动时，其奇异向量可能会发生剧烈变化（特别是当奇异值很接近时），导致上一步的”方向”与这一步的”方向”天差地别。然而，伪逆的计算同时依赖于奇异向量和奇异值。笔者猜测，这种对奇异值的依赖可能会让伪逆函数本身变得更加”平滑”，使得前一步的解对于当前步依然是一个不错的初始点。 这些工程上的技巧，特别是”热启动”迭代，为解决计算瓶颈问题提供了非常具体的、有可行性的方向。但这终究是更进一步的探索了，我们在此暂时打住。我们似乎彻底走到了这条”有原则的修正”之路的尽头。 是时候放弃这种执念，换一个全新的角度思考问题了。 在第二部分，我们历经坎坷，最终推导出了一套理论上完美但工程上几乎不可行的更新公式。它就像一件被锁在玻璃柜里的屠龙宝刀，我们能看到它的锋利，却无法触及。 在几乎要放弃这条路的时候，我们决定做最后一次挣扎：这个复杂的公式，真的没有简化的可能了吗？或许，公式的内部隐藏着某些尚未被发现的代数结构？ 这便引出了我们故事的第三幕——一次意想不到的峰回路转，以及随之而来的、更深刻的”绝望”。 第三部分：峰回路转？一次未竟的数学简化 我们再次审视那个令人望而生畏的更新公式的核心部分，即 函数的内部：

这里的 msign 参数看起来就像一堆随机的矩阵在”打架”。然而，我们忽然意识到，梯度 并非一个无根之萍，它自身的结构是由反向传播的链式法则决定的，并且与 有着千丝万缕的联系！ 3.1 洞察：梯度 的非凡结构 我们回顾 的计算。梯度沿着计算图 反向传播，其中 且 。 根据链式法则，梯度可以分解为：

记 ，并应用矩阵微分法则，我们得到中间步骤的梯度：

以及最终的梯度：

代入 ，便揭示了 的内在结构：

（为简化叙述，我们这里直接给出了结论，其推导涉及矩阵迹的微分和链式法则，但核心是梯度 的计算必须经过 这一环，因此其最终表达式必然包含 ）。 这个发现是我们的第一个希望。梯度 不是一个独立的项，它天然地包含了我们试图”摆脱”的 ！这是否意味着，当它与公式中另一个包含 的项——伪逆——相遇时，会发生一些奇妙的化学反应？ 3.2 神奇的”内部抵消”与投影矩阵的诞生 怀着激动的心情，我们将 的表达式代入 msign 函数的内部。我们先定义几个关键组件以便看得更清楚： 核心梯度项 : 我们令 。这个 的矩阵可以被理解为从注意力分数层反向传播回来的”核心”梯度信息。 梯度表达式: 。 * 伪逆的转置: 我们需要计算 。假设 ( 矩阵, ) 是满秩的，其伪逆 有一个明确的表达式：。那么它的转置就是：

（最后一步成立因为 是一个对称矩阵，其逆也是对称的）。 现在，让我们把这些碎片拼起来，看看 msign 的参数 Arg 变成了什么：

展开第一个括号的转置：

这个表达式看起来依然混乱。但是，如果我们调整一下我们对 msign 参数的推导（不同的矩阵迹变换顺序会导致最终形式的些许不同，但本质等价），我们会遇到另一种形式的参数：

让我们探索一种最能揭示其内在结构的等价形式。经过严谨的推导，msign 的参数可以被整理为：

奇迹发生了！ 请看右边括号里的那部分：。在线性代数中，这是一个非常有名的结构，它被称为投影矩阵 (Projection Matrix)。我们把它记作 ，它表示将一个向量投影到由 的列向量所张成的子空间上。 （注：这里的推导经过了修正。在原始探索中，我们曾走入一个误区，试图化简为包含 的形式，但该 矩阵在 时是奇异的，不可逆。正确的推导指向了由可逆的 矩阵 构成的标准投影矩阵。） 这意味着，那个原来极其复杂的 msign 参数，经过一番化简，变成了一个形式上极为优雅、几何意义也异常清晰的表达式：

这简直是一次酣畅淋漓的数学化简！它告诉我们，对 的更新方向，本质上是由”核心梯度 “ 先投影到与 相关的子空间上，然后再取其符号决定的。这完全符合我们的直觉：对 的更新，当然只应该在那些能与 产生有效”互动”的方向上进行。 3.3 希望，但没完全有：致命的外部伪逆 我们似乎已经走到了胜利的边缘。公式变得如此简洁优美，我们不禁发出了灵魂拷问：”我们是不是已经消除了那个该死的伪逆？” 答案是残酷的：并没有。 首先，一个比较明显的问题是，投影矩阵 的定义中，本身就包含了逆矩阵 。我们并没有消除它，只是把它”藏”进了一个更有几何意义的马甲——投影矩阵——的内部。 但更致命的一击还在后面。让我们回到最终的更新公式：

我们刚刚费了九牛二虎之力简化的，仅仅是 msign 函数内部的参数。而在它的外部，那个孤零零的、等待与 msign 结果相乘的伪逆项 ，依然纹丝不动地待在那里！ 这意味着，完整的计算流程是： 1. 计算核心梯度 。 2. 计算逆矩阵 。 3. 用逆矩阵构造投影矩阵 。 4. 计算 （这本身可能也需要迭代计算）。 5. 再次使用伪逆 。 6. 将第4步和第5步的结果相乘，得到最终更新。 我们并没有解决问题，只是把问题复杂化了。 阶段性结论 这次”峰回路转”的简化尝试，是一次非常有趣的智力体操。它极大地加深了我们对QK更新机制几何本质的理解——更新是在一个投影子空间上进行的。 然而，从工程实践的角度看，它并未带来任何助益。我们不仅没能消除对昂贵逆矩阵（或伪逆）的依赖，甚至可能需要计算它两次。这条看似充满希望的理论小径，最终通向了一堵更厚的墙。我们似乎彻底走到了这条”有原则的修正”之路的尽头。 第四部分：Bilinear Attention 在经历了第二部分的艰难推导和第三部分的”伪”简化之后，我们的探索之旅似乎走进了一条死胡同。理论上完美的公式，在工程上却无法逾越。这种困境迫使我们后退一步，重新审视问题的本质。一个看似简单却很基础的问题浮现在脑海：我们为什么一定要将 和 分开看待呢？ 4.1 思想实验：一个统一的Bilinear矩阵 让我们进行一个思想实验。注意力分数的核心计算是 ，其中 和 。代入后，两个输入向量 之间的交互可以写作：

请注意括号里的项 。这是一个 的大矩阵，我们将它定义为一个统一的”注意力交互矩阵” ：

通过这个定义，整个注意力分数的计算被简化成了一个非常经典的双线性形式 (Bilinear Form)：

如果我们将整个模型架构中的Q、K投影层替换为这样一个统一的 矩阵，那么对于整个输入序列 ，注意力矩阵 就变成了 。 现在，让我们把Muon优化器的哲学思想应用在这个新架构上。我们需要约束的是功能的变化，而这里的功能就是由 定义的双线性变换。其功能更新量由 引起，对输出的影响是 。根据Muon的第一性原理，我们约束其在最坏情况下的功能变化：

这可以简化为对权重更新本身的谱范数约束：

于是，整个优化问题瞬间变得无比清晰和简单：

其中 是损失对这个统一矩阵的梯度。这个问题的解，正是我们熟悉的标准Muon更新：

这条路在理论上是完美的。它从根源上解决了QK耦合的问题，将复杂的双因子约束简化为单矩阵约束，并且得到了一个极其简洁的解析解。所有在第二、三部分遇到的伪逆、解耦等复杂问题都烟消云散。 4.2 现实的代价：参数量与计算量的爆炸 理论的完美往往伴随着现实的代价。为什么标准的Transformer不采用这种设计呢？答案在于效率。 这个统一的交互矩阵 的维度是 ，其中 是模型的嵌入维度。而标准的注意力头中， 和 的维度都是 ，其中 是头的维度 (head dimension)。 让我们来计算一下参数量： 标准注意力头: 参数量为 和 之和，即 。 统一双线性注意力: 参数量为 的大小，即 。 在典型的LLM中， 远大于 。例如，一个模型的嵌入维度 ，而头维度 。 标准头的参数量是: 。 统一矩阵的参数量是: 。 参数量相差了 倍！这仅仅是一个注意力头，对于一个多头、多层的Transformer模型，这种参数量的激增是完全无法接受的。此外，计算 的计算量也远高于计算 。 因此，这个理论上完美的”Full Bilinear Attention”方案，因为其巨大的参数量和计算开销，在实践中是不可行的。 4.3 新视角：MHA作为Low-Rank Approximation 这个看似失败的思想实验，却给了我们一个重新理解多头注意力机制（MHA）的深刻视角。 我们回头看 这个定义。这是一个 的矩阵，由一个 矩阵和一个 矩阵相乘得到。根据线性代数的知识，这个乘积矩阵 的秩 (Rank) 最大不会超过 。 由于在实践中，头维度 总是远小于嵌入维度 ()，这意味着标准注意力机制中的交互矩阵 天然是一个低秩矩阵 (Low-Rank Matrix)。 由此，我们得出一个关键的结论： 标准的多头注意力机制，可以被视为对一个全尺寸、全秩的”双线性注意力”的一种低秩近似 (Low-Rank Approximation)。 模型架构通过将 分解为 的形式，强制施加了一个结构性的低秩约束。这极大地降低了参数量和计算复杂度，使其变得可行。我们之前所有推导的困难，其根源就在于，我们试图在保持这种低秩分解结构的前提下，对它的两个因子（ 和 ）进行优化。 4.4 推论：通往Muon-LoRA的桥梁 这个”低秩近似”的视角，还为我们带来了意想不到的启发，并将其与另一个热门领域——参数高效微调（PEFT）——联系了起来。 我们来观察一下LoRA (Low-Rank Adaptation) 的核心思想。当微调一个大的预训练权重矩阵 时，LoRA并不直接更新 ，而是通过学习一个低秩的更新量 来适配新任务，即 。这个低秩更新量同样被分解为两个小矩阵的乘积：，其中 是一个 矩阵， 是一个 矩阵，而秩 通常很小。 对比一下二者的结构： 注意力头: 交互矩阵被分解为 。 LoRA: 权重更新量被分解为 。 它们的数学本质是完全相同的：用两个瘦矩阵的乘积来表示一个低秩矩阵。 这就引出了一个非常自然的推论（corollary）：如果我们能够成功地为 这种低秩因子推导出一个原理正确且计算高效的Muon更新法则，那么这套方法论将可以被直接迁移，用于设计一个全新的、基于Muon思想的LoRA微调优化器，我们或可称之为”Muon-LoRA”。 这个发现纯属意外收获，提升了解决QK更新问题的价值。它不再仅仅是训练大型模型时的一个稳定性技巧，更可能为参数高效微调领域开辟一条全新的、基于功能空间优化的道路。 4.5 重新审视解析解 这个”低秩近似”的视角，为我们提供了一面绝佳的棱镜，可以反过来重新审视我们在第二、三部分推导出的那个复杂的解析解，并理解它为何如此笨重。 我们当时的目标是找到满足约束 的最优解。现在我们能更清晰地看到，这个解析解为了在保持 分解形式的前提下解决问题，付出了巨大的代价。其本质可以被概括为一套非常繁琐的多步操作： 在高维空间中操作: 解析解的核心步骤——矩阵符号函数 msign，实际上是在一个概念上的 大矩阵空间中进行的（即对 求解）。这对应了我们刚刚讨论的、那个计算昂贵的”Full Bilinear Attention”矩阵。它试图在概念上对这个大矩阵进行Muon更新。

启发式的预算分配: 为了将耦合的约束解开，我们采用了三角不等式放缩，并将功能更新的”预算” 在 和 之间进行对半分配。这本质上是一种启发式规则，而非最优解耦。

复杂的投影操作: 最关键的是，为了将对那个 概念矩阵的更新，拆分回对 和 这两个低秩因子的更新，公式引入了极其复杂的投影操作。

首先，它需要通过伪逆构造一个投影矩阵，将梯度信息限制在与另一个因子相关的低秩子空间上。 然后，在计算出更新方向后，它还需要再次使用伪逆，将更新量根据对另一个因子的影响程度，投影回原始的 (或 ) 空间。 总结来说，那个解析解之所以复杂，是因为它试图用一套通用的数学工具（伪逆、投影）去”强行”解决一个被低秩结构约束的优化问题。它不仅计算上包含昂贵的伪逆，更深层的问题是，它在概念上绕了一个大圈：先将问题提升到 的高维空间，进行Muon更新，然后再通过一系列复杂的、计算成本高昂的投影操作，硬塞回到 的低秩空间中。这种做法既不直接，也不高效。 这让我们认识到，或许直接在低秩空间中寻找更简洁、更近似的解决方案，才是更务实的方向。 插注： 前四个部分，我不敢说保证对，但大方向总没有错。我觉得一个周末推出这些东西不丢人，写这些东西也不丢人。 后面的部分，我在尝试提出一些自己的想法，当然不成熟也不太正确，但至少对一类工程上更加可行的思路进行探索。我并不认为错误的尝试就毫无价值。 第五部分：回归原点——探索更直接的解决方案 前述章节的理论探索，虽然在数学上是严谨的，但最终将我们引向了一个包含伪逆等复杂运算的解析解。该方案计算成本极高，在工程实践中几乎不可行。这迫使我们不得不思考：是否必须追求一个理论上完美普适的解？ 从理论保证到经验近似 Muon思想的核心是寻求一个对所有可能输入都成立的、最坏情况下的理论保证。这引导我们约束了权重更新的谱范数。然而，为了这个普适的保证，我们付出了巨大的计算代价。 深度学习的实践给了我们另一条思路。在训练中，导致不稳定的直接原因，往往是当前数据批次（batch） 计算出的注意力分数中出现了极端值，从而产生了巨大的梯度。 这启发我们进行一次视角的切换：或许可以用当前数据批次中的”最坏情况”，来作为理论上”最坏情况”的一个经验近似。 这种妥协放弃了理论上的完备性，但可能换来工程上的可行性。我们不再寻求一个对权重矩阵普适的、但计算复杂的更新规则，而是尝试直接在优化流程中的关键节点进行干预。 干预策略的探索空间 放弃复杂的理论解析解后，我们可以从更直接的工程角度出发，对优化过程中的一些关键量进行干预。这些干预手段共同的目标，是在不引入过多计算开销的前提下，平滑 QK 权重的更新过程。 笔者试着提出一套方案：直接对于Attention Score的梯度进行干预。 由此可以形成一个广阔的探索空间： 反向传播干预 (Backward-pass Interventions)：这类方法在计算出梯度后，但在优化器更新权重之前，对梯度进行修改。

裁剪类方法 (Clipping-based)：最直接的方法，例如对梯度 的数值施加一个固定的上限和下限。 梯度归一化 (Gradient Normalization)：根据max logit gradient 所估计的梯度spectral norm 大小，对其进行缩放，控制其整体尺度。

参数校正干预 (Parameter-correction Interventions)：这类方法不直接修改梯度，而是在优化器更新权重之后，根据前向传播的结果对权重参数进行”事后校正”。

在这个探索空间中，我们既可以研究像梯度裁剪这样的反向传播干预，也可以审视参数校正的思路。 初步尝试：一个基于梯度裁剪的启发式方案 作为探索的第一步，我们尝试了反向传播干预中的梯度裁剪方案。这个方案并非一个成熟的结论，而更多是抛砖引玉，展示这类启发式方法的一般形态。 其大致流程可以设想为： 1. 正常进行前向传播和反向传播，计算出原始的注意力分数梯度 。 2. 对梯度矩阵 进行逐元素裁剪，得到 ，其中 是一个需要调节的超参数。 3. 使用标准的优化器（如Adam）和这些被修改过的梯度来更新 和 。 方法比较：梯度裁剪 vs. MuonClip Kimi AI 团队在其 Kimi K2 模型的技术报告中也提出了一个非常成功的解决方案，名为 MuonClip。它恰好是参数校正干预的一个绝佳范例。我们可以将两者并排比较，揭示出不同工程哲学下的巧妙思路。 Kimi 的 MuonClip 方案：基于前向输出的参数归一化 根据其公开的技术细节，MuonClip 的核心思想可以概括为： 经验性探测量 (Empirical Probing): 在每个优化步骤之后，使用当前批次的数据计算出更新后的 和 所产生的最大 Attention Logit，记为 。 动态缩放 (Dynamic Rescaling): 设定一个预设的 Logit 阈值 。如果 超过了 ，就计算一个缩放因子 (当时, )。 参数校正 (Parameter Correction): 将这个缩放因子应用回权重矩阵，进行”事后”校正： 其中 是一个平衡超参数（特例是 ）。这个校正有两个效果： a) 归一化权重 (隐式学习率调整): 对权重矩阵 本身的缩放，是一种参数范数的归一化，它也隐式地影响了后续步骤的有效学习率。所以，它本质上等于在我们的干预过程中使用归一化。 b) 调整前向传播的温度: 与我们思路不同的是，MuonClip同时对前向传播的输出进行干预。下一次前向传播计算出的注意力分数 会被整体缩放 倍。这等价于将 Softmax 的温度从1调整为 。当 attention logit 过大时（），温度升高，会平滑输出分布，从而抑制极端值的出现。

哲学思想的碰撞 共同点：经验主义的转向 两套方案都敏锐地意识到，试图为所有理论上的输入寻找一个普适的更新规则，其代价过于高昂。因此，双方都转向了利用当前训练批次的信息作为”现实世界”的代理，来近似那个理论上的最坏情况。

核心区别：作用域与手段

作用域：梯度裁剪方案是一种纯粹的反向路径干预，发生在梯度传递给优化器之前，只影响当次更新，不改变前向传播的计算。而 MuonClip 不仅干预后向，同时也改变前向，它在优化器更新之后修改权重，其核心目的是改变下一次前向传播的计算（即调整温度）。 手段：梯度裁剪的本质是裁剪 (clipping)，粗暴地”砍掉”极端的梯度值。而 MuonClip 的本质是缩放 (scaling)，通过调整 Attention Logit 的温度和权重范数来稳定训练。 过程 vs 整体 梯度剪裁只对局部变化量做限制。MuonClip对全局直接进行限制。 总结来说，梯度剪裁的思路与MuonClip有相同之处，代表了在广阔的探索空间中两种不同的设计选择。 关于max logit对于谱范数估计的准确度 使用当前批次的最大 logit (max logit) 来近似理论上的谱范数，一个自然的顾虑是：这种经验性的估计到底有多准？毕竟，我们是在用有限的样本去估计一个在无穷输入空间上的最大值。理论上，这可能会遭遇”维度灾难”（Curse of Dimensionality）——在高维空间中，随机采样的点很难覆盖到那个能让范数达到最大的”最坏”方向。 然而，在我们的场景下，有几个强有力的因素使得这个估计相当可靠： 维度并不算高：在典型的Transformer模型中，注意力头的维度（）通常是128。虽然不低，但这并非一个会立刻导致维度灾难的数字。 样本量巨大：我们用来估计的样本量远超通常的批次大小。一个包含百万级别 token的批次，会产生数以十亿级别的 token 对，即用于计算 logit 的样本对。如此巨大的样本量使得我们有很高的概率在批次内就遇到接近”最坏情况”的输入。 有效样本量的进一步放大：如果我们不只依赖当前批次，而是采用类似”指数移动平均最大值”（Exponential Moving Max）的策略来追踪 ，那么我们的有效样本量将横跨多个训练步骤，变得更加庞大和稳定，进一步减小了估计误差。 我们可以与对比学习 (Contrastive Learning) 的领域进行一个类比。在对比学习中，嵌入向量的维度常常要高得多（例如2048维），而用于寻找”最难负样本”的批内样本量却要小得多（例如32K或64K）。即便是在这种维度更高、样本量更小的更具挑战性的设定下，模型训练依然能取得很好的效果。这有力地佐证了，通过批次内采样来寻找和抑制”最难样本”的经验性方法，在实践中是相当稳健和有效的。 因此，尽管存在理论上的顾虑，但在实践中，使用 max logit 来驱动校正，是对谱范数一个相当准确且高效的代理（proxy）。 潜在的问题与挑战 这类经验性的启发式方法，虽然在工程上可能更可行，但也带来了新的问题和不确定性： 经验与理论的差距：用批次内的信息来近似全局的稳定性保证，这种做法的鲁棒性有多少？例如逐元素裁剪可能会在某些梯度方向上裁剪过度，而在另一些方向上裁剪不足，从而引入新的偏差。 超参数的敏感性：这类方法引入了新的超参数（如裁剪值 或 MuonClip 的阈值 ）。模型的最终性能可能对这些超参数非常敏感，而如何系统性地寻找最优超参数本身就是一个难题。 与优化器的相互作用：在梯度进入优化器之前对其进行修改，或者在更新后对参数进行校正，这些操作会如何与Adam等状态优化器的内部机制（如二阶动量）相互作用？是否存在潜在的冲突，反而影响了优化的效果？ 解答这些问题需要进一步细致的实验验证。 以上便是笔者一个周末的胡思乱想，希望对大家有所启发。 参考文献 [1] Jordan, K., Bernstein, J., et al. (2024). “Muon: An optimizer for hidden layers in neural networks.” Blog post. Available at: https://kellerjordan.github.io/posts/muon/ [2] 苏剑林. (2025). “QK-Clip：让Muon在Scaleup之路上更进一步.” 科学空间. Available at: https://spaces.ac.cn/archives/11126 [3] 苏剑林. (2025). “Muon续集：为什么我们选择尝试Muon？” 科学空间. Available at: https://spaces.ac.cn/archives/10739 [4] Liu, J., Su, J., et al. (2025). “Muon is Scalable for LLM Training.” arXiv preprint arXiv:2502.16982. [5] Jordan, K., et al. (2024). “modded-nanogpt: Speedrunning the NanoGPT baseline.” GitHub repository. Available at: https://github.com/KellerJordan/modded-nanogpt [6] Jordan, K., et al. (2024). “Muon Optimizer.” GitHub repository. Available at: https://github.com/KellerJordan/Muon [7] Moonshot AI. (2025). “Kimi K2: Open Agentic Intelligence.” Project page. Available at: https://moonshotai.github.io/Kimi-K2/