Appendix E 为什么 Muon 更容易出现 Logit 爆炸

本文为论文《Kimi K2：面向智能体能力的开源大模型》「中文全文翻译」板块的对应章节，与「论文原文」板块同一章节一一对应，层级保持一致。

logit 爆炸发生在最大的 softmax 前注意力分数

| | $S_{\max}=\max_{i,j}\bigl(q_{i}^{\vphantom{\top}}!\cdot k_{j}\bigr)$ | | (1) |

在训练过程中无界增长时。由于

| | $|q_{i}!\cdot!k_{j}|\leq\|q_{i}\|\|k_{j}\|\leq\|x_{i}\|\|x_{j}\|\|\mathbf{W}{q}\|\|\mathbf{W}{k}\|,$ | | (2) |

且 RMS-Norm 保证了 $|x_{i}||x_{j}|$ 有界，该现象主要由 $\mathbf{W}{q}$ 或 $\mathbf{W}{k}$ 不断增长的谱范数驱动。我们在经验上发现，Muon 更容易出现 logit 爆炸。下面给出我们的假设。

更新中的结构差异

Muon 产生的权重更新来自 $\mathrm{msign}$ 操作；因此，更新矩阵的所有奇异值都相等——其有效秩是满的。 相比之下，Adam 产生的典型更新矩阵呈现偏斜的谱：少数较大的奇异值占主导，有效秩较低。对 Adam 的这一低秩假设并不新鲜；高阶 muP 也作出了同样的假设。

这一现象在 16B 的 Moonlight 模型上得到了验证，该模型显示使用 Muon 训练的权重比使用 Adam 训练的权重具有更高的奇异值熵（即更高的有效秩），印证了理论直觉。

SVD 表述

设第 $t-1$ 步的参数矩阵具有奇异值分解

| | $\mathbf{W}{t-1}=\sum{i}\sigma_{i}\,u_{i}v_{i}^{\top}$ | | (3) |

我们将更新矩阵写作

| | $\displaystyle\Delta\mathbf{W}{t}$ | $\displaystyle=\sum{j}\bar{\sigma}\,\bar{u}{j}\bar{v}{j}^{\top}$ | | (4) |

因此下一步的参数更新为

| | $\displaystyle\mathbf{W}{t}\leftarrow\sum{i}\sigma_{i}u_{i}v_{i}^{\top}+\sum_{j}\bar{\sigma}\,\bar{u}{j}\bar{v}{j}^{\top}$ | | (5) |

在 Muon 中，由于权重和更新的有效秩都高于 Adam，我们假设奇异向量对 $u_{i}v_{i}^{\top}$ 与 $\bar{u}{j}\bar{v}{j}^{\top}$ 对齐的概率更高。这可能导致 $\mathbf{W}_{t}$ 对应的奇异值以加法方式增加。

注意力特有的放大效应

注意力 logit 通过双线性形式计算

| | $q_{i}\cdot k_{j}=(x_{i}\mathbf{W}{q})\cdot(x{j}\mathbf{W}_{k}).$ | | (6) |

乘积 $\mathbf{W}{q}\mathbf{W}{k}^{\top}$ 会使谱范数平方，因此任一矩阵中奇异值的增加都会被放大。Muon 增大奇异值的倾向因此转化为更高的 logit 爆炸风险。