MuonIO：面向嵌入表和语言模型输出头的有原则的范数感知下降

arXiv:2610.02705v1 公告类型：新 摘要：Muon 优化器通过求解一个由谱范数惩罚的损失局部线性化问题，为隐藏线性层推导其更新规则；这一方法的动机来自针对稠密线性层的 RMS 稳定性论证。然而，标准的 Muon 实现将输入层（嵌入表）和输出层（语言模型输出头）排除在这一有理论依据的处理之外，而是对它们使用 AdamW。我们提出 MuonIO，为这两类层提供一种统一的 Muon 风格更新。对于语言模型输出头 $\mathbf{L} \in \mathbb{R}^{V \times d}$，鉴于 softmax 输出几何结构的 Lipschitz 连续性，我们提出使用 $2\to\infty$ 算子范数；对于嵌入表 $\mathbf{E} \in \mathbb{R}^{d \times V}$，则基于 Bernstein & Newhouse (2025) 所确定的 one-hot 输入几何，采用 $1 \to 2$ 算子范数。恒等式 $\lVert\mathbf{L}\rVert_{2\to\infty}=\lVert\mathbf{L}^\top\rVert_{1\to2}$ 由此将两个矩阵置于相同的面向词汇表的几何结构中：MuonIO 应用统一的归一化向量规则，对 $\mathbf{E}$ 表现为列归一化，对 $\mathbf{L}$ 表现为行归一化。实证评估表明，我们的方法有效：在 C4 上进行 1B LLaMA 预训练时，与 Muon 相比，MuonIO 将 I/O 优化器状态内存减少了 50%，将 I/O 更新 FLOPs 减少了约 46%，同时还改善了验证集困惑度。