从Muon到梯度裁剪：探讨QK稳定性问题

背景

本文来自 Hacker News 上的一篇技术博客，标题为“From Muon to Gradient Clipping: Some Thoughts on QK Stability”。作者从深度学习优化器 Muon 的优雅理论出发，探讨了其在 Transformer 的 Query（Q）和 Key（K）权重矩阵上直接应用时引发训练不稳定甚至崩溃的问题，并给出了基于第一性原理的分析与改进思路。文章融合了理论推导、工程瓶颈与启发式方案，是一篇深入的技术思考记录。

核心内容

1. 理论回顾：为什么标准 Muon 不适合 QK

Muon 的推导与本质

优化器是深度学习工具箱中的“引擎”。大多数研究者习惯使用 Adam 及其变体，但基于不同哲学构建的 Muon 优化器逐渐引起关注。Keller Jordan 和 Jeremy Bernstein 在 [1] 中首次提出，其深度在于不再停留在熟悉的参数空间，而是从函数空间的角度出发构建更新规则。该想法近期在 Kimi K2 [7] 等工作中被推进并成功实践，显示出巨大潜力。然而 K2 的公开材料也指出一个问题：如果直接将原始 Muon 优化器应用于 Transformer 的 Query（Q）和 Key（K）权重矩阵，训练会变得极端不稳定，甚至崩溃。

要理解这种不匹配，首先需要理解 Muon 的初衷及其与 Adam 等优化器的根本区别。

对于权重矩阵 $W$，假设已计算其关于损失 $L$ 的梯度 $G$。我们希望找到更新 $\Delta W$ 以最大程度减小损失。一阶泰勒展开给出： [ \Delta L \approx \langle G, \Delta W \rangle = \operatorname{Tr}(G^T \Delta W). ] 为了最小化 $L$，我们希望 $\Delta W$ 尽可能指向 $G$ 的反方向。这可以形式化为一个约束优化问题： [ \min_{\Delta W} \operatorname{Tr}(G^T \Delta W) \quad \text{subject to a size constraint on } \Delta W. ] 若不加约束，$\Delta W$ 的元素会趋于无穷大，无意义。每个优化器都必须施加某种约束，限制更新步的“大小”，区别在于如何定义这个大小。

传统约束在参数空间上，如 SGD、Adam 等，直接施加在 $\Delta W$ 上，最常用的是 Frobenius 范数： [ |\Delta W|_F \le \eta. ] 其几何意义直观：参数空间中的欧几里得距离不应太大。

Muon 的约束在函数空间，完全不同。它认为我们不应主要关心参数 $W$ 移动了多少，而应关心这种移动对模型函数的影响有多大。对于线性层，函数将输入 $x$ 映射为 $Wx$，因此参数更新带来的函数变化是 $\Delta W x$。Muon 的核心思想是：这种函数变化应该被界住，且对于每一个可能的输入都成立。为消除 $x$ 幅度的影响，只考虑单位范数输入。因此 Muon 的约束变为： [ \sup_{|x|_2 = 1} |\Delta W x|_2 \le \eta. ] 这意味着：对于任何长度为 1 的输入向量，经过更新矩阵 $\Delta W$ 变换后，输出向量的长度不应超过 $\eta$。熟悉矩阵理论的人会认出，这个上确界正是谱范数 $|\Delta W|_2$ 的定义，等于矩阵的最大奇异值。

现在 Muon 的优化问题清晰了： [ \min_{\Delta W} \operatorname{Tr}(G^T \Delta W) \quad \text{subject to} \quad |\Delta W|_2 \le \eta. ] 这个谱范数约束下的优化问题有一个优美的解析解：特征矩阵符号函数， [ \Delta W^* = -\eta , \operatorname{msign}(G). ] $\operatorname{msign}$ 函数保留 $G$ 的奇异向量，但将所有奇异值改为 1。这是一种极端的定向控制，将更新“能量”均匀分布在所有奇异方向上。

QK 更新中的几何冲突

理解了 Muon 的机制，我们可以重新审视其应用于 QK 更新时的失败。正如苏剑林在博客 [2] 中分析的，问题在于注意力机制的双线性性质。

一个较浅层的直觉是函数耦合：Q 和 K 的功能是成对出现的，它们只有通过乘积 $QK^T$ 才能产生有意义的注意力分数。独立约束 $W_Q$ 和 $W_K$ 的函数变化，并不能保证它们耦合乘积的函数变化也受到控制。

更深层的几何解释，按照苏剑林的猜想 [2]，与 Muon 的全秩更新行为有关。实践中观察到的“MaxLogit 爆炸”通常意味着 $W_Q$ 或 $W_K$ 的谱范数正在爆炸。那么问题变成：为什么 Muon 更容易导致谱范数爆炸？由于谱范数就是最大奇异值，这又等价于：为什么 Muon 倾向于增大奇异值？

Muon 与 Adam 的区别是什么？Muon 的更新经过 $\operatorname{msign}$ 处理，因此更新的所有奇异值相等，从这个意义上说其有效秩是满的。而一般的矩阵，其奇异值通常大小不一，主要由前几个主导，从有效秩角度看是低秩的。我们可以假设 Adam 的更新也是低秩的，高阶 muP 就做了类似假设。

用公式表示，令参数 $W$ 的 SVD 为 $U \Sigma V^T$，Muon 更新的 SVD 为 $U_m \Sigma_m V_m^T$，Adam 更新的 SVD 为 $U_a \Sigma_a V_a^T$。更新后的权重可写为： [ W_{\text{new}} = W + \Delta W. ] 如果 $W$ 的某个奇异向量对与更新的奇异向量对接近，那么对应的奇异值会直接相加，从而增大 $W$ 在该方向上的奇异值。

由于 Muon 的更新是满秩的，其能量均匀分布在所有奇异方向上，因此与 $W$ 的奇异向量“碰撞”的概率远高于低秩的 Adam 更新。Muon 自然更容易增大参数的奇异值。

这种效应在注意力机制中被放大。注意力分数的核心是双线性形式 $QK^T$。如果 Muon 同时增大 $W_Q$ 和 $W_K$ 的谱范数，那么最终点积计算时这种增大会被乘起来，爆炸风险急剧上升。形成“坏者更坏”的反馈循环，最终可能导致训练崩溃。

至此，我们有了一个理论解释：为何如此优雅设计的优化器变得不适合 Transformer 的 QK 权重。根本问题在于必须约束乘积项的变化，而不是两个独立的变化。这一洞察是下一步的起点。

2. 原则性尝试：为 QK 定制 Muon

第一部分给出了关键洞察：分别约束 $W_Q$ 和 $W_K$ 的更新是行不通的。在注意力中，两者是功能耦合的。最终重要的是它们的乘积。独立约束可能允许两个向量的范数变化很小，但方向变得对齐，导致它们的点积——即注意力分数——爆炸。

因此，一个原则性的方法必须考虑乘积的约束。作者尝试沿着 Muon 的第一性原理，提出一种适合 QK 的修改更新思路。文章接下来的部分（由于原文在此处截断，未完整呈现，但核心方向已明确）将记录从严格理论推导，到工程瓶颈，再到可行启发