Appendix E Why Muon is More Prone to Logit Explosion

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

Logit explosion occurs when the largest pre-softmax attention score

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

grows unboundedly during training. Since

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

and RMS-Norm keeps $|x_{i}||x_{j}|$ bounded, the phenomenon is primarily driven by the growing spectral-norm of $\mathbf{W}{q}$ or $\mathbf{W}{k}$. Empirically, we found that Muon is more susceptible to logit explosion. We give our hypothesis below.

Structural difference in updates

Muon produces a weight update coming from the $\mathrm{msign}$ operation; as a result, all singular values of the update matrix are equal — its effective rank is full. In contrast, a typical update matrix produced by Adam exhibits a skewed spectrum: a few large singular values dominate, and the effective rank is low. This low-rank assumption for Adam is not new; higher-order muP makes the same assumption.

Such phenomenon is verified on the 16 B Moonlight model, which shows weights trained with Muon exhibit higher singular-value entropy (i.e. higher effective rank) than those trained with Adam, corroborating the theoretical intuition.

SVD formulation

Let the parameter matrix at step $t-1$ have the singular value decomposition

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

We write the update matrices as

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

The next parameter update is therefore

| | $\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) |

In Muon, as both the weights and the updates have a higher effective rank than Adam, we hypothesize there is a higher probability for singular-vector pair $u_{i}v_{i}^{\top}$ to align with $\bar{u}{j}\bar{v}{j}^{\top}$. This could cause the corresponding singular value of $\mathbf{W}_{t}$ to increase additively.

Attention-specific amplification

Attention logits are computed via the bilinear form

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

The product $\mathbf{W}{q}\mathbf{W}{k}^{\top}$ squares the spectral norm, so any singular-value increase in either matrix is compounded. Muon’s tendency to enlarge singular values therefore translates into a higher risk of logit explosion.