Per-Head Muon

Per-Head Muon is the optimizer variation described in 177: 详解Kimi K3：强到冲击Anthropic估值的模型什么样？ as part of [[KimiK3|Kimi K3]]’s training recipe. The source explains Muon as an optimizer that approximately orthogonalizes momentum before applying updates, reducing the chance that learning collapses into a few dominant directions.

K3’s source-specific extension is to apply that logic per attention head. Zeng Zhiyuan / 曾志远 says per-head treatment can keep different attention heads’ update scales more balanced, while Zhao Chenyang / 赵晨阳 points to implementation difficulty around fused QKV tensors, distributed optimizer state, communication, and orthogonalization pipelines.

Key Claims

• Optimizer design affects whether large training runs stay stable and make useful long-term parameter updates.
• Per-head orthogonalization can reduce imbalance among attention heads.
• The implementation is an infrastructure problem as much as an algorithm problem because model parallelism and optimizer state are distributed.
• Optimizer research is a promising AI-agent domain because experiments are structured, measurable, and easier to automate than open-ended discovery.

Connections

• Kimi K3, Attention Residues, Kimi Delta Attention / KDA, and Transformer Architecture — model architecture and optimization context.
• Model-Infra Co-Design, Frontier Model Scaling, and ML Coding — training-infrastructure context.
• Kernel Development Agents, Recursive Self-Improvement, and AI Coding Verification — AI-assisted research and verification branch.