Kimi K2は何をしているのか
Muon の学習不安定性と QK-Clip、MuonClip、Self-Critique RL のポイント
背景
大規模言語モデルの性能は、どれだけ効率よくトークンを使って学習できるかに左右される。近年は高品質なデータの枯渇が問題になり、**token efficiency（同じデータ量でどれだけ賢くなるか）**が重要になった。
この文脈で Muon optimizer が注目されている。Muonは、勾配を Newton–Schulz 近似で正規化し、勾配スケールを均一化しつつ少ないデータで高精度に到達する。Kimi K2 もこのMuonを採用している。
しかし、Muonには致命的な問題がある。
Attention の logit が爆発しやすいのだ。Muonで大規模モデルを学習すると、Query-Key 内積が際限なく大きくなり、softmax前の値（logit）が暴走する。
学習ログから、Muon使用時に attention logits が「1000」以上へ急上昇したことが確認されている。
従来対策では限界がある。
- Logit soft-cap：logitを直接クリップするが、クリップ前に値が膨れ上がる
- QK-Norm：Multi-head Latent Attention（MLA）では使用不可
この問題を解決するため、Kimi K2 は QK-Clip を導入した。
さらに、Muon・Weight Decay・RMS Matching・QK-Clip を統合した MuonClip optimizer を作り、安定的な大規模学習を実現した。
加えて、モデル自身の出力を自身で評価する Self-Critique RL を導入し、推論品質を強化した。
提案手法
全体像
Kimi K2 の学習基盤は、次の3点で成り立つ：
- Muon optimizer
- Attention logit explosion を抑える QK-Clip
- すべてを統合した MuonClip optimizer
- Self-Critique RL による自己評価型強化学習
この記事では、特に Muon → 不安定化 → QK-Clip → MuonClip の流れを中心に解説する。
1. Muon とは何か
Muonは、重み更新に Newton–Schulz 近似を使い、勾配方向を正規化する。
更新式には「勾配の寄与を均一化して rank を高める」特徴がある。
図として適切なのは Algorithm 1: MuonClip Optimizer（p.4）で、その前半がMuonの更新ステップに対応する。
図1：MuonClip（Muon + QK-Clip）のアルゴリズム
引用：Figure 2 付近の Algorithm 1（p.4）
図が示すポイント：
- 行 1〜7 が Muon の更新ループ
- Newton–Schulz による正規化（L25）
- Adam RMS に合わせたスケール調整（0.2 × sqrt(max(n,m))）
- 更新後に weight decay を適用
なぜ重要か：
Muonの特徴である「高ランクな更新」が、後述する logit explosion を引き起こす原因となるため。
読者が注目すべき点：
Muon は token-efficient だが unstable、という構造的問題を理解できる。
2. Muon の問題：Attention Logit Explosion
Muonは効率的であるにも関わらず、Attention の Query-Key 内積が異常に大きくなる現象が発生する。
論文の説明（p.3–4）：
- AdamWでは起きにくいが、Muonでは顕著
- QK 内積が急上昇し、softmax前の値（logit）が1000を超える
- これが loss spike や divergence を誘発する
図2：Muon使用時の Attention Logit の暴走
引用：Figure 2 Left（p.4）
図の説明：
- 左図：Muonで学習すると max logits が1000超
- 学習ステップが増えるほど発散
- 典型的な loss spike につながる
ポイント：
問題は損失ではなく、重みスペクトルの成長にある。
Muonは全特異値を拡大する「full-rank 更新」を生成するため、
Wq / Wk の spectral norm（固有値最大値）が巨大化しやすい。
この “構造的に発散しやすい性質” が、通常の optimizer にないMuon固有の問題。
3. QK-Clip：Attention Head ごとの爆発だけを抑える
Kimi K2 は、logit の直接クリップではなく、重み（Wq, Wk）をスケールするという設計を採用した。
仕組みの要点（p.3–4）
- 各 attention head h ごとに、
そのバッチでの最大 logit Sh_max を計算 - Sh_max > τ（閾値）なら
Wq, Wk を γ = τ / Sh_max だけ縮小 - 他の head には干渉しない
- forward/backward には影響しない（次ステップから効く）
図3：QK-Clip あり／なしの logit の挙動
引用：Figure 2 Right（p.4）
図が示すこと：
- QK-Clip を入れると logit は τ=100 に一気に収束
- 30%のステップ付近から安定
- 発散が完全に抑えられている
なぜ重要か：
- Muon そのものを改善したのではなく、
Muon の欠点だけを局所的に抑制できる点が独創的 - 特に「爆発している head だけ」をスケールするため、
precision を落とさず安定性だけを得られる
読者が注目すべき点：
「logit をクリップする」のではなく「重みのスペクトルを抑える」アプローチであること。
4. MuonClip：Muon + Weight Decay + RMS Matching + QK-Clip の統合
Kimi K2 では、Muon と QK-Clip を統合し、
MuonClip optimizer として正式に採用している。
Algorithm 1 の後半（p.4）がそれに対応。
MuonClip の特徴
-
まず Muon の更新を行う（Newton–Schulz）
-
その後に各 head の logit を確認し、必要なら QK-Clip
-
MLA（Multi-head Latent Attention）向けに特別な処理
- rotary成分は共有のため、k_R はスケールしない
効果
- Muon の高速学習能力を完全に保持
- logit explosion を完全に抑制
- 小規模実験でも性能劣化なし
特に重要なのは Self-deactivation：
QK-Clip は学習初期だけ 12.7% の head が発火し、
7万ステップ以降は自然に消滅する。
これは、
「クリップは一生効き続けるのではなく、初期安定化のためだけに必要」
という示唆になる。
5. Self-Critique RL：自分の出力を自分で採点する
論文後半（RLセクション）では、Kimi K2 の推論強化のために Self-Critique RL を導入した。
仕組みの概要：
- モデルがまず「回答」を生成
- その回答について「自己評価（critique）」を生成
- Critique のスコアを reward signal として学習
- Critique は人間に似た採点基準を持つよう最適化されている
重要な点：
- 外部ラベルなしで RL が回る
- hallucination 抑制
- reasoning の質向上
まとめ
Kimi K2 で最も重要な要素は次の4つに整理できる。
-
Muon optimizer
- token-efficient
- しかし full-rank 更新により logit explosion が発生
-
QK-Clip
- attention head 単位で Wq / Wk を縮小
- 発散する head のみを抑制
- forward/backward への影響なし
-
MuonClip optimizer
- Muon + weight decay + RMS matching + QK-Clip
- 大規模学習の安定化を実現
- 性能劣化なし、小規模実験でも変動なし
-
Self-Critique RL
- モデルが自身の出力を自己採点
- hallucination を抑制し reasoning 性能を向上
特に QK-Clip の“局所的な重みスケーリング” は、
大規模LLMの安定化問題に対して非常に実用的かつ新しいアプローチであり、
MuonClip 全体の肝と言える。