1 episodes
Paper: Kimi K2 and MuonClip
Kimi K2 (arXiv 2507.20534, Jul 2025) is a 1.04T-total / 32B-active MoE pretrained on 15.5T tokens with no loss spikes using MuonClip: Muon with weight decay and RMS matching, plus QK-Clip, a per-head weight rescale that caps the max attention logit at tau = 100. The module derives the rule from bilinearity, reproduces logit growth under Muon in a synthetic toy with seeds, builds the max-logit dashboard, and summarizes K2's agentic data-synthesis pipeline.
Progress — loading…
Objectives
Language
- Explain why lookup-like attention heads are pushed toward ever-larger margins
- Describe K2's agentic data pipeline (tools, agents, rubrics, judges, sandboxes)
Math & data structures
- Bilinearity of attention logits; derive the exact cap of QK-Clip
- Singular values, Newton-Schulz orthogonalization and RMS matching
- Per-head max-logit state as an L x H array; weight layouts per head
Code
- Implement QK-Clip and Muon direction in NumPy with tests
- Run a 4-condition x 3-seed logit-explosion experiment with attn/max_logit dashboards
- Port to a registered PyTorch optimizer and attention layer
Episodes
- Kimi K2: MuonClip and QK-Clip17 min · 28 cues
Attention logit explosion, QK-Clip rule (τ=100), MuonClip; agentic data synthesis (arXiv 2507.20534).
Exercises
| Exercise | Tier | Ladder | Time |
|---|---|---|---|
| QK-Clip from scratch: per-head max logits, exact capping, and the Muon direction | browser | Predict · Scratch · Harden · Compose | 50 min |
| Induce max-logit growth under Muon, then tame it with QK-Clip (4 conditions x 3 seeds) | browser | Predict · Scratch · Break it · Probe a claim | 75 min |
| Port: logged attention + MuonClip in PyTorch, matched to the NumPy QK-Clip | gpu | Port · Compose | 45 min |
Labs
- Kimi K2: MuonClip and QK-Clip browser F28.01 · 80 min
Library contributions
Files in nlpcore this module teaches you to write.
Sources
- Kimi Team, Kimi K2: Open Agentic Intelligence. arXiv 2507.20534 (v1 28 Jul 2025; v2 3 Feb 2026). · arXiv:2507.20534 · §abstract, 2.1 MuonClip / Algorithm 1, architecture, pre-training recipe, agentic data synthesis, rephrasing
- Induce max-logit growth under Muon, then tame it with QK-Clip (4 conditions x 3 seeds) reference runs (this curriculum): synthetic copy-lookup attention toy, 4 conditions x 3 seeds, NumPy.
- Jordan, K. (2024). Muon: An optimizer for hidden layers in neural networks. kellerjordan.github.io/posts/muon · linkunverified
- Liu et al., Muon is Scalable for LLM Training. arXiv 2502.16982 (24 Feb 2025). · arXiv:2502.16982