QK-Clip: Taking Muon Further on the Scaleup Journey

TL;DR

Jianlin Su’s July-2025 kexue.fm post explaining QK-Clip and MuonClip, the stability mechanisms behind Kimi K2’s 15.5T-token zero-loss-spike pretraining run. After Moonlight (16B MoE) validated Muon at moderate scale via weight decay + Update-RMS alignment, scaling to hundreds of billions of parameters hit a new roadblock: MaxLogit explosion — attention logits drifting upward in some heads, leading to numerical instability. QK-Clip is a post-update weight rescaling that minimally shrinks $W_q$ and $W_k$ to keep per-head MaxLogit below a threshold $\tau$ . It directly modifies weights rather than logits, so it composes with QK-Norm and MLA. In Kimi K2 ( $\tau{=}100$ ), Muon and QK-Clip “battled” for ~70k of ~220k steps; afterward the model self-regulated and the clip became inactive. The article is in Chinese; an English translation exists at frontier.soket.ai. Note: the original kexue.fm URL returned 403 at filing time; this page synthesizes the post’s claims from the Soket AI translation, Su’s announcement tweet, and the Fireworks / Lakernewhouse / Kimi K2 tech report secondary writeups, not from a direct fetch.

Key claims

Failure mode at scale: Muon scaled past hundreds of billions of parameters hits MaxLogit explosion — per-head max attention logits drift up over training, eventually exceeding 1000 in a mid-scale run and risking numerical divergence [Soket translation §Background; echoed in Kimi K2 tech report Fig. — “left panel”].
Existing fixes are insufficient: logit soft-capping distorts the attention distribution; QK-Norm normalizes Q and K vectors but doesn’t address the root cause (weight norms) and is awkward to combine with MLA [Soket translation §Why not soft-cap / QK-Norm].
QK-Clip mechanism: after each Muon update, if the per-head MaxLogit $L > \tau$ , rescale $W_q \leftarrow \eta^{\alpha} W_q$ and $W_k \leftarrow \eta^{1-\alpha} W_k$ with $\eta = \tau/L$ and $\alpha \approx 0.5$ — a minimal post-hoc weight shrink that caps scores at $\tau$ without changing the attention pattern’s relative geometry [Soket translation §QK-Clip procedure; Fireworks deep-dive §How qk-clip works].
Compatible with MLA and QK-Norm: because QK-Clip operates on weights post-update rather than on logits at forward time, it composes with MLA’s compressed K-cache and with existing QK-Norm stacks where soft-capping does not [Soket translation §Compatibility].
Empirical trajectory in Kimi K2 ( $\tau{=}100$ , 220k steps): around step 7k some heads first exceeded $\tau$ ; QK-Clip and Muon remained in tension for ~70k steps (Muon pushing MaxLogit up, the clip pulling it down); after step ~70k MaxLogit decayed below $\tau$ on its own and QK-Clip became inactive for the remainder of training [Soket translation §Training behavior; Kimi K2 tech report Fig. “right panel”].
Self-regulation hypothesis: under weight decay, once early instability is navigated, the model naturally adapts to reduce MaxLogit — QK-Clip is therefore “performance-safe”: it only acts during the unstable early phase [Soket translation §Self-regulation].
Performance-safety check: in small-model ablations, aggressive clipping (down to $\tau{=}30$ ) shows no measurable performance degradation [Soket translation §Small-model ablation].
Muon-specific: in matched experiments Muon exhibits MaxLogit blow-up more readily than AdamW, so QK-Clip is best read as an update rule supplementing Muon at scale rather than as a generic stability fix [Soket translation §Closing].
Result claimed for Kimi K2: 15.5T-token pretraining with zero loss spikes, attributed to the MuonClip = Muon + weight decay + Update-RMS alignment + QK-Clip stack [Kimi K2 tech report Abstract; also reproduced in the Kimi K2 blog and on the Soket translation §Background].

Method

QK-Clip is a per-step post-update intervention. After a Muon step updates the model parameters, for each attention head with current projection matrices $W_q$ and $W_k$ :

Measure $L = \max_{x,y} \mathrm{logit}(x, y)$ over the most-recent batch (or a running estimate), where the logit is $\frac{(W_q x)^\top (W_k y)}{\sqrt{d}}$ or the MLA equivalent.
If $L > \tau$ , compute $\eta = \tau / L < 1$ .
Rescale $W_q \leftarrow \eta^{\alpha} W_q$ and $W_k \leftarrow \eta^{1-\alpha} W_k$ , with $\alpha \in [0,1]$ (typically 0.5).
Otherwise leave the weights untouched.

Because $\eta^{\alpha} \eta^{1-\alpha} = \eta$ , the product $W_q^\top W_k$ is scaled exactly by $\eta$ , which scales every logit in the head by $\eta$ and therefore brings MaxLogit to exactly $\tau$ . The split between $W_q$ and $W_k$ (via $\alpha$ ) is what makes the intervention “minimal” in spectral terms — both matrices shrink by a small amount rather than one of them shrinking sharply.

The article emphasizes that this is fundamentally a weight fix, not an activation fix. Soft-capping changes the attention distribution non-trivially (clips one logit harder than another); QK-Norm changes the forward path and is awkward to graft onto MLA. QK-Clip leaves the forward path untouched and only nudges weights — so the attention pattern’s relative geometry is preserved, and the intervention sits on top of any forward-path stack (MLA, QK-Norm, GQA, etc.).

In Kimi K2’s setting, MaxLogit is measured per attention head, and $\tau = 100$ was chosen as the cap. Su notes in the comment thread on the original post that the choice of $\tau$ matters less than one might expect — too-small $\tau$ slows training negligibly in small ablations (down to 30 was fine).

Results

The article is primarily methodological. The empirical numbers come from the Kimi K2 training run that it explains:

15.5T tokens, zero loss spikes [Kimi K2 tech report Abstract; Su’s article §Background].
MaxLogit trajectory: caps at $\tau = 100$ from ~step 7k through ~step 70k, then decays freely below 100 for the remainder of the 220k-step run; QK-Clip becomes a no-op past step ~70k [Soket translation §Training behavior; Kimi K2 tech report Fig. right panel].
Small-model ablation: aggressive clipping at $\tau = 30$ shows no measurable degradation in final loss [Soket translation §Small-model ablation].
Mid-scale failure baseline (no QK-Clip): attention logits exceed 1000, training diverges [Kimi K2 tech report Fig. left panel; Soket translation references this same figure].

No standalone QK-Clip-vs-no-QK-Clip ablation at K2 scale is reported in this post (the K2 run was done with the clip throughout); the performance-safety argument is built from small-model ablations plus the observation that the clip self-deactivates past step 70k.

Why it’s interesting

This is the source-of-record for one of the two “remove the tuning axis” tricks that Jianlin Su / Kimi has pushed into the open optimizer literature — the other being Quantile Balancing for MoE routing (Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method). Both share the recipe “identify the exact constraint that the hand-tuned heuristic was approximating, solve for it in closed form, ship without a coefficient to tune.” The QK-Clip post is the older and more famous of the two — it’s the technique credited with Kimi K2’s loss-spike-free 15.5T-token run.

The article also fits cleanly into the architecture-vs-optimizer fork that Training stability at scale is tracking. The unified-sinks work (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training) argues outlier-driven rescaling via softmax/RMSNorm is load-bearing for training stability — QK-Clip is the dual move: bound the outliers explicitly so the rescaling stays in a numerically safe regime. SSO (Controlled LLM Training on Spectral Sphere) takes the “do this at the optimizer level for all layers” generalization. The “Muon scales past hundreds of billions only with extra weight-norm control” framing in this post is also the empirical motivation behind all the recent torch.optim Muon discussion (Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819)) and behind Aurora’s leverage-aware variant (Aurora: A Leverage-Aware Optimizer for Rectangular Matrices).

For Luma: if Muon (or a Muon variant) ever displaces AdamW for large- scale training internally, the QK-Clip recipe is essentially free insurance — one extra clip per step, scales to the head count, and self-deactivates once the model self-regulates.