Kimi Team (Moonshot AI), 164+ co-authors | 2025-07 | arXiv:2507.20534 Category: algorithm | Tags: optimizer, moe, agentic-data, reinforcement-learning, open-source
MuonClip(Muon + per-head QK-Clip)解决 Muon 在 1T MoE 上的 attention logit 爆炸,15.5T tokens 零 spike 训练;三阶段 agentic 数据合成 + joint RLVR/self-critique RL 产出开源最强 non-thinking 模型(Tau2 66.1, SWE-bench 65.8, AIME'25 49.5)。
Pre-training: Muon optimizer 的 token efficiency 显著优于 AdamW(Moonlight 已证明),但在大规模 MoE 训练中 attention logits 爆炸(9B/53B MoE 规模即出现 >1000 的 logit spikes)。现有缓解方案均不适用:logit soft-cap 在 pre-cap 阶段无法阻止 dot product 增长;QK-Norm 不兼容 MLA(inference 时 key 矩阵未完全 materialized)。
Post-training: Agentic capabilities(multi-step tool use、long-horizon planning)在自然数据中极其稀缺,人工标注成本不可承受。需要在 scale 上合成高质量 agentic trajectories,但合成数据需要 realistic tool execution environment 和 multi-turn trajectory verification。
Problem formulation:
Objective 1(Pre-training): 最小化 autoregressive loss $\mathcal{L} = -\sum_t \log p(x_t \mid x_{ Objective 2(RL): 最大化 $\mathbb{E}_{x \sim \pi}[R(x)]$,其中 verifiable tasks $R(x) \in \{0, 1\}$(RLVR),subjective tasks $R(x) = \text{critic}(x, \text{rubric})$。 Assumptions: pre-training i.i.d. batch sampling;RL trajectories 符合 MDP 结构;critic 与 actor 共享参数(closed-loop)。 Inputs/outputs per MuonClip step: 消耗 gradient $\mathbf{G}_t$ 和当前权重 $\mathbf{W}_{t-1}$,Newton-Schulz 正交化后 per-head QK rescaling → $\mathbf{W}_t$。 Core mechanism — MuonClip: Muon 的 Newton-Schulz 正交化产生 equal singular values 的 weight updates(对比 Adam 的 skewed spectrum),在 Q/K projection weights 上跨步累加。由于 attention logits 是双线性形式 $\mathbf{Q}^h \mathbf{K}^{h\top}$,spectral norm 增长被平方放大。QK-Clip 的解决方案:per-head post-update rescaling——当 $S_{\max}^h > \tau$ 时,计算 $\gamma_h = \tau / S_{\max}^h$,然后: $$\mathbf{W}_{qc}^h \leftarrow \sqrt{\gamma_h} \cdot \mathbf{W}_{qc}^h, \quad \mathbf{W}_{kc}^h \leftarrow \sqrt{\gamma_h} \cdot \mathbf{W}_{kc}^h, \quad \mathbf{W}_{qr}^h \leftarrow \gamma_h \cdot \mathbf{W}_{qr}^h$$ Shared rotary key $\mathbf{W}_{kr}$ 不 rescale(防 cross-head 污染)。关键性质:QK-Clip 是 self-deactivating 的——K2 训练中约 30% 后所有 logits 自然降至 $\tau$ 以下,不再触发。 Secondary — Agentic Data Synthesis: 三阶段 pipeline:(1) Tool spec generation(3,000+ 真实 MCP 工具 + 20,000+ LLM 合成工具,t-SNE 验证分布对齐);(2) Agent/task generation(rubric-based quality constraints);(3) Multi-turn trajectory generation(LLM-simulated user personas + stateful tool execution world model + LLM-judge filtering)。 Tertiary — Joint RL: RLVR(verifiable binary rewards for math/code/safety)+ self-critique rubric reward(actor 生成 → critic pairwise 评估 → closed-loop refinement)。Prescriptive rubric 包含 "No Initial Praise" 和 "No Explicit Justification" 以防 reward hacking。 核心技术壁垒: 理解 Muon 为何导致 attention logit 爆炸。 Paper's Figure 2, verbatim (caption: "Left: During a mid-scale training run, attention logits rapidly exceed 1000 ... Right: Maximum logits for Kimi K2 with MuonClip and τ = 100 over the entire training run."). 左图展示 vanilla Muon 在 9B activated / 53B total MoE 上 attention logits 失控超过 1000,多处 spike 表明随时可能 diverge。右图 K2 全量训练使用 MuonClip($\tau=100$):初期仅 12.7% heads 触发,max logit 被控制在 ~100,约 30% 训练后所有 logits 自然降至阈值以下——QK-Clip 不再触发。这种 self-deactivating 行为是 MuonClip 最显著的性质:它不是 permanent fix,而是 temporary training stabilizer。 Paper's Figure 8, verbatim (caption: "Large-scale agentic data synthesis pipeline for tool use."). Pipeline 分两部分:(a) 左侧从 3,000+ 真实 MCP 工具出发,通过 LLM 扩展到 20,000+ 合成工具(t-SNE 验证分布对齐),再向上生成 agents(distinct system prompts + 工具组合)和 rubric-based tasks;(b) 右侧 trajectory 生成流程:LLM-simulated user persona 提供多轮输入 → model 调用工具 → stateful tool simulator(维护状态、引入受控随机性)→ LLM-judge multi-agent review → 仅保留通过的 trajectories。 MuonClip 完整算法(Algorithm 1): Step 1 — Muon optimizer update: $$\mathbf{M}_t = \mu \mathbf{M}_{t-1} + \mathbf{G}_t$$ $$\mathbf{O}_t = \text{Newton-Schulz}(\mathbf{M}_t) \cdot \sqrt{\max(n,m)} \cdot 0.2$$ $$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta(\mathbf{O}_t + \lambda \mathbf{W}_{t-1})$$ Step 2 — QK-Clip(per-head): For each head $h$, if $S_{\max}^h > \tau$, set $\gamma_h = \tau / S_{\max}^h$: $$\mathbf{W}_{qc}^h \leftarrow \sqrt{\gamma_h} \cdot \mathbf{W}_{qc}^h, \quad \mathbf{W}_{kc}^h \leftarrow \sqrt{\gamma_h} \cdot \mathbf{W}_{kc}^h, \quad \mathbf{W}_{qr}^h \leftarrow \gamma_h \cdot \mathbf{W}_{qr}^h$$ 其中 $S_{\max}^h = \frac{1}{\sqrt{d}} \max_{\mathbf{X} \in B} \max_{i,j} \mathbf{Q}_i^h \mathbf{K}_j^{h\top}$,shared rotary key $\mathbf{W}_{kr}$ 不 rescale 以防 cross-head 污染。 Architecture summary: 1.04T total / 32.6B activated MoE Transformer, 61 layers, MLA (64 heads), 384 routed + 1 shared expert (top-8, sparsity 48), hidden dim 7168, expert hidden dim 2048, 1 dense layer, no expert grouping. Position: RoPE + YaRN (128k). Compared to DeepSeek-V3: +50% experts, −50% heads, −13% activated params, +54% total params. 无形式化作者证明 — 仅实证。 论文未提供 MuonClip 收敛保证、spectral norm 增长率 bound、或 QK-Clip 对 optimization trajectory 的形式化影响分析。 Paper's Figure 3, verbatim (caption: "Per-step training loss curve of Kimi K2, without smoothing or sub-sampling. It shows no spikes throughout the entire training process."). Loss 曲线完全平滑,两阶段可辨:前 10T tokens 的 constant LR 阶段 loss 平稳下降;后 5.5T tokens 的 cosine decay 阶段 loss 下降加速。这是在 1T MoE 上使用 Muon 系优化器的首次公开零 spike 训练记录。论文未披露是否存在 checkpoint rollback 或 data skip 等应急措施。 Paper's Figure 5, verbatim (caption: "Sparsity Scaling Law. Increasing sparsity leads to improved model performance."). 四条 scaling curves 表明 sparsity 48(384 experts / 8 active)在所有 compute budgets 下均达到最低 validation loss。定量:在 val loss 1.5 处,sparsity 48 相比 sparsity 8/16/32 分别节省 1.69×/1.39×/1.15× FLOPs。这驱动了 K2 从 DeepSeek-V3 的 256 experts (sparsity 32) 扩展到 384 experts (sparsity 48)。Trade-off:更多 experts → 更多 EP all-to-all communication 和 serving memory。 Paper's Figure 1, verbatim. K2 non-thinking 设定下全面 benchmark 对比。关键观察:(1) K2 在 LiveCodeBench v6 (53.7)、OJBench (27.1)、AIME 2025 (49.5)、HMMT 2025 (38.8)、ZebraLogic (89.0)、Multi-Challenge (54.1)、IFEval (89.8)、LiveBench (76.4) 上超越所有模型含 proprietary;(2) SWE-Bench Verified agentic 65.8 vs Claude S4 72.7——显著缩小差距但仍落后;(3) Tau2-telecom 65.8 远超 Claude S4 (45.2),表明 agentic data synthesis 对 complex tool orchestration 尤其有效。 Rephrasing ablation (Table 1): 10 rephrasings × 1 epoch (SimpleQA 28.94) > 1 rephrasing × 10 epochs (27.39) > raw × 10 epochs (23.76)。Diverse surface forms strictly better than repetition。 Architecture heads ablation: 64 vs 128 attention heads 仅 0.5–1.2% val loss 差距,但 128k sequence length 下减少 83% inference FLOPs。 Pre-training: 15.5T tokens across Web Text, Code, Mathematics, Knowledge。Exact per-domain percentages [论文未披露]。Quality filters follow Kimi K1.5。Contamination check [论文未披露]。 Synthetic data (rephrasing): Knowledge data via WRAP-inspired style/perspective-diverse prompts + chunk-wise autoregressive generation + fidelity verification。Math data via SwallowMath learning-note style + cross-lingual translation。Each corpus rephrased at most twice。 Agentic synthetic data: Generator = in-house models (K1.5 + domain experts)。Seed = 3,000+ real MCP tools (GitHub)。Expansion to 20,000+ LLM-synthesized tools (t-SNE validated)。Validation = multi-agent LLM review, only passing trajectories retained。Volume = "tens of thousands" diverse training examples。Annotation [无人工标注,全合成]。 [实现未公开] — 训练代码未开源。 MuonClip 的实现难点不在于 QK-Clip 本身(Algorithm 1 清晰可实现),而在于三个层次:Q2 · 方法 #
Aspect Before (Muon) After (MuonClip) Weight update Newton-Schulz 正交化 + weight decay 同左 Logit constraint 无(logits 爆炸 >1000) Per-head QK-Clip,threshold $\tau$ MLA 兼容性 N/A Compressed Q/K 用 $\sqrt{\gamma}$,rotary Q 用 $\gamma$,shared rotary K 不动 自我调节 N/A Self-deactivating(~30% training 后停止触发) 性能影响 N/A 无退化(Appendix D ablation) msign 操作产生 equal singular values 的 weight updates,增加了 update singular vectors 与 existing weight singular vectors 的 cross-step alignment 概率。在 $\mathbf{Q}^h \mathbf{K}^{h\top}$ bilinear form 中此 spectral norm 增长被平方放大。Per-head post-update rescaling(而非 pre-softmax capping)是正确的 intervention point 这一 insight,以及 self-deactivation 性质,要求对 Muon 在 scale 下的优化动力学有深入理解。Q3 · 结果 #
§3 架构 / 方法图 #


§4 作者证明 #
Notation table #
Symbol Definition Domain $\mathbf{X}$ Transformer layer input representation $\mathbb{R}^{n \times d_{model}}$ $\mathbf{W}_q^h, \mathbf{W}_k^h, \mathbf{W}_v^h$ Per-head Q/K/V projection weights $\mathbb{R}^{d_{model} \times d_h}$ $d$ Per-head dimension $\mathbb{Z}^+$ $S_{\max}^h$ Max pre-softmax attention logit for head $h$ over batch $\mathbb{R}^+$ $\tau$ QK-Clip threshold(K2 使用 100) $\mathbb{R}^+$ $\gamma_h$ Per-head clipping ratio $\min(1, \tau / S_{\max}^h)$ $(0, 1]$ $\alpha$ Naïve clipping balance parameter(0.5) $[0, 1]$ $\mu$ Momentum coefficient $[0, 1)$ $\eta$ Learning rate $\mathbb{R}^+$ $\lambda$ Weight decay coefficient(0.1) $\mathbb{R}^+$ $\mathbf{W}_{qc}^h, \mathbf{W}_{kc}^h$ MLA compressed Q/K components per head $\mathbf{W}_{qr}^h$ Per-head rotary Q component $\mathbf{W}_{kr}$ Shared rotary K component(不被 QK-Clip 触及) 方程物理意义 #
6 minimum checks #
# Check Status 1 QK-Clip 不改变当前 step 的 forward/backward(仅 post-update rescaling) ✓ 论文明确声明 2 Per-head clipping 仅影响 logit 超阈值的 heads ✓ 设计意图,实证验证 3 Shared rotary key 不被 rescale 以避免 cross-head 污染 ✓ MLA 特有考量 4 $\tau=100$ 对 loss 无显著影响 ✓ Appendix D ablation(含 $\tau=30$) 5 Self-deactivation:~30% 训练后不再触发 ✓ Fig. 2 right panel 6 15.5T tokens 零 loss spike ✓ Fig. 3(unsmoothed per-step loss) 期望但缺失的形式化保证 #
§5 实验与数据 #
Training Recipe & Scale #
Stage Purpose Data Tokens LR Batch Special Pre-train (constant) 主训练 Web + Code + Math + Knowledge 10T 2e-4 (500-step warmup) 67M tokens MuonClip, $\tau=100$ Pre-train (decay) 精调 Same 5.5T 2e-4 → 2e-5 cosine 67M tokens WSD schedule Anneal (4k ctx) Quality boost High-quality mix 400B 2e-5 → 7e-6 67M tokens — Anneal (32k ctx) Long context Long-context data 60B Same 67M tokens — Context extension 128k — — — — YaRN SFT Instruction tuning Diverse + agentic synthetic [论文未披露] — — Muon optimizer RL (joint) Alignment RLVR tasks + subjective prompts [论文未披露] — — RLVR + self-critique, budget control, PTX loss
Convergence & Stability #

Sparsity Scaling Law #

Key Results #

Dataset Analysis #
§6 论证链 #
Step Claim Evidence Validity 1 Muon 比 AdamW token-efficient 但在 scale 上 attention logits 爆炸 Moonlight [47] (efficiency);Fig. 2 left (9B/53B MoE logits >1000) Strong — 独立工作已验证 efficiency;instability 有实验展示 2 QK-Norm 和 logit soft-cap 对 MLA 不适用 逻辑论证:QK-Norm 需要 materialized K(MLA 不满足);soft-cap 不阻止 pre-cap 增长 Moderate — 逻辑成立但未实验 ablate 这些替代方案 3 Per-head QK-Clip 约束 logits 且不退化性能 Fig. 2 right (logits ≤ $\tau$);Appendix D ($\tau=30$ 也有效;loss 无退化) Strong — 多 $\tau$ 值 ablation 4 MuonClip 在 1T MoE / 15.5T tokens 上实现零 spike Fig. 3 (unsmoothed per-step loss curve) Strong — 直接视觉证据,论文最强 claim 5 Ultra-sparse MoE (sparsity 48) 在固定 activated params 下提升 compute efficiency Fig. 5 (scaling law: 4 sparsity levels,1.69× FLOPs saving) Strong — controlled 实验 6 Agentic data synthesis 产出 SOTA tool-use capabilities Tau2-Bench 66.1, ACEBench 76.5 (Table 3) Moderate — 无 ablation 隔离 agentic data 贡献 vs RL 贡献;仅有 end-to-end 结果 7 Joint RLVR + self-critique 优于 SFT alone Table 3 全面结果 Weak — 无 SFT-only baseline;无 RLVR vs self-critique vs joint ablation §7 实现 cross-reference #
核心技术壁垒详解 #
关键实现细节 #
Reproducibility & Ecosystem #