An Isotropy-Preserving Spectral Cap for Muon: Theory and Three Case Studies
作者: Jiachun Li
分类: cs.LG, cs.AI
发布日期: 2026-07-22
备注: Preliminary Report; Larger scale experiments ongoing; Comments are welcome
💡 一句话要点
提出一种保持各向同性的谱帽以优化Muon训练效果
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Muon优化器 谱帽 尺度不变性 大型语言模型 各向同性 深度学习 权重更新
📋 核心要点
- 现有的Muon优化器对权重矩阵的几何影响尚不清晰，导致训练效果不稳定。
- 提出了一种基于尺度不变性假设的谱帽，通过控制谱方向的增长来优化训练过程。
- 在nanoGPT、64专家混合路由器和FlashAttention模块的实验中，谱帽显著提高了各向同性，且验证损失保持不变。
📝 摘要（中文）
Muon及相关的矩阵符号优化器在大型语言模型的预训练中越来越受到重视，但其对单个权重矩阵内部几何结构的影响尚不明确。本文提出了一个统一框架，基于权重重缩放下损失的精确尺度不变性假设，适用于重归一化网络。在此假设下，普通的SGD更新受限于1/||W||的制约，而Muon的矩阵符号步骤则去除了这一制约，使得Frobenius和谱范数的漂移速度加快。我们进一步观察到谱范数扰动具有非负的二阶项，这意味着一种轻量级的“谱帽”可以在不冻结训练的情况下控制输出协方差。通过在三个系统上进行实验，结果表明该谱帽提高了各向同性，并在边缘情况下防止了具体失败，同时验证损失基本保持不变。
🔬 方法详解
问题定义：本文旨在解决Muon优化器对权重矩阵内部几何结构影响不明确的问题。现有方法在训练过程中可能导致不稳定的几何变化，影响模型性能。
核心思路：论文提出了一种轻量级的“谱帽”，通过控制权重更新中的谱方向增长，保持训练的灵活性和稳定性。该方法基于尺度不变性假设，确保在权重重缩放时损失保持不变。
技术框架：整体架构包括权重更新、谱帽应用和输出协方差控制三个主要模块。首先进行权重更新，然后应用谱帽以限制谱方向的增长，最后控制输出协方差以保持模型的稳定性。
关键创新：最重要的技术创新在于提出了谱帽的概念，它能够在不冻结训练的情况下，控制权重更新的方向，避免了传统方法中的几何不稳定性。
关键设计：在设计中，谱帽仅投影出单个主奇异方向的第一阶增长，确保其他方向的学习不受影响。损失函数和网络结构的设计也围绕着保持尺度不变性进行优化。
🖼️ 关键图片
📊 实验亮点
实验结果显示，谱帽在三个不同系统中均显著提高了各向同性，且在边缘情况下有效防止了模型崩溃。验证损失几乎保持不变，表明该方法在优化训练过程中的有效性和稳定性。
🎯 应用场景
该研究的潜在应用领域包括大型语言模型的预训练和优化，尤其是在需要高效训练和稳定性的场景中。通过提高各向同性，模型在处理复杂任务时的表现将更加可靠，未来可能对自然语言处理和其他深度学习领域产生深远影响。
📄 摘要（原文）
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.