Muon优化器易导致权重矩阵失衡，提出谱范数约束可保持训练稳定性。
An Isotropy-Preserving Spectral Cap for Muon: Theory and Three Case Studies
- 引入谱范数截断机制，仅抑制最大奇异方向的一阶增长
- 三组实验中均提升权重分布各向同性，避免路由坍缩与注意力发散
- 保留学习能力的同时防止模型失效，适合大规模语言模型训练
Muon及其相关矩阵符号优化器在大语言模型预训练中日益普及，但其对单个权重矩阵内部几何结构的影响尚不明确。本报告基于一个理想化假设——损失函数在权重缩放下精确尺度不变（在归一化密集网络中近似成立），构建统一框架。在此假设下，普通SGD具有1/||W||的更新幅度制动，而Muon的矩阵符号步则移除了该制动，导致弗罗贝尼乌斯和谱范数以更快速度向外漂移（t^{1/2} vs t^{1/4}）。我们进一步发现谱范数扰动存在非负二阶项，暗示可设计轻量级“谱范数截断”：仅从每次更新中投影出单一最大奇异方向的一阶增长，即可控制输出协方差W K_X W^T而不冻结训练——权重仍可通过非主方向、主方向旋转及主方向切换持续学习。该截断与奇异值谱的极小熵（H-infinity）相关。随后研究了三个使用Muon训练的系统：nanoGPT前馈投影、64专家混合专家路由器、bf16 FlashAttention块的查询/键投影。在每种情况下，该截断均提升了各向同性，在边缘情形（路由器坍缩至单个专家，某注意力头接近发散）有效预防具体故障，同时验证损失基本不变。强调尺度不变性假设较强，且结果为初步小规模试验，欢迎讨论。
原文摘要 · Abstract (English)
Muon and related matrix-sign optimizers are increasingly used to pre-train large language models, but their effect on the internal geometry of individual weight matrices is not well understood. This preliminary report proposes a unified framework built on a single idealizing assumption -- exact scale invariance of the loss under weight rescaling, which holds approximately in normalization-heavy networks. Under this assumption, plain SGD carries a built-in 1/||W|| brake on its update size, whereas Muon's matrix-sign step removes that brake, so both the Frobenius and spectral norms drift outward faster (t^{1/2} versus t^{1/4}). We further observe that the spectral-norm perturbation has a non-negative second-order term. This implies that a lightweight "spectral cap" -- which projects out only the first-order growth of the single top singular direction from each update -- can control the output covariance W K_X W^T without freezing training: the weight keeps learning through non-top directions, top-direction rotation, and top switching. We relate this cap to the min-entropy (H-infinity) of the singular-value spectrum. We then study three systems trained with Muon: a nanoGPT feed-forward projection, a 64-expert mixture-of-experts router, and the query/key projections of a bf16 FlashAttention block. In each case the cap increases isotropy and, at the margins -- a router collapsing to a single expert, and the near-divergence of one attention head -- prevents a concrete failure, while leaving validation loss essentially unchanged. We emphasize that the scale-invariance assumption is strong and that these small-scale results are preliminary; comments are welcome.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。