MuonIO 用 2→∞ 与 1→2 算子范数统一嵌入表与语言模型头的更新，在 1B LLaMA 预训练中把优化器状态内存减半、更新 FLOPs 降约 46% 并改善验证困惑度
相关研究与后续进展
核心概要
MuonIO 针对标准 Muon 排除的嵌入表 E 与语言模型头 L，分别以 1→2 与 2→∞ 算子范数给出动机，并利用 ‖L‖_{2→∞}=‖Lᵀ‖_{1→2} 把两者统一为同一词汇导向几何下的单一归一化向量规则（E 为列归一化、L 为行归一化），在 C4 上 1B LLaMA 预训练中相比 Muon 将 I/O 优化器状态内存降低 50%、I/O 更新 FLOPs 降低约 46%，同时改善验证困惑度。
Figure 1: Validation loss trajectories on C4 at the 130M scale for Plain Muon (baseline 1), Muon Tuned (baseline 2), Ember and MuonIO. All four configurations use standard Muon hidden updates. The inset highlights the 0.7–0.9B training token range.
arXiv深度剖析
MuonIO 为嵌入表和语言模型头这两类被标准 Muon 排除的层提供了统一的 Muon 式更新规则。 标准 Muon 实现把输入嵌入表与输出语言模型头交给 AdamW，而 MuonIO 让这两层也进入同一套基于算子范数的原则化处理。 摘要给出的是方法层面的推导与统一，配合在 C4 上 1B LLaMA 预训练的经验评估。
语言模型头 L 使用 2→∞ 算子范数，其动机来自 softmax 输出几何的 Lipschitz 连续性；嵌入表 E 使用 1→2 算子范数，其动机来自 Bernstein & Newhouse (2025) 指出的 one-hot 输入几何。 把两层各自的几何来源分别对应到不同算子范数，而不是沿用同一套稠密线性层的谱范数论证。 摘要以几何论证与引用作为动机依据，并以经验评估作为效果支撑。
恒等式 ‖L‖_{2→∞}=‖Lᵀ‖_{1→2} 使两个矩阵落入同一词汇导向几何，MuonIO 因此只需一条归一化向量规则：对 E 表现为列归一化，对 L 表现为行归一化。 把原本分属输入侧与输出侧的两套处理合并为单一规则，简化了实现层面的更新形式。 该统一由算子范数恒等式在数学上给出，摘要未展开证明细节。
在 C4 上 1B LLaMA 预训练中，MuonIO 相比 Muon 将 I/O 优化器状态内存降低 50%、I/O 更新 FLOPs 降低约 46%，同时改善验证困惑度。 在保持或改善验证困惑度的同时降低优化器状态内存与更新计算量，指向大规模预训练中的资源效率。 摘要报告的是单一设定（1B LLaMA、C4 预训练）下的经验结果，未给出多规模或多数据集的对照。
启示与展望
该工作面向使用 Muon 类优化器进行 Transformer 语言模型预训练的实践者，尤其是需要控制优化器状态内存与更新计算量的场景。摘要所述结果适用于 1B LLaMA 在 C4 上的预训练设定；方法本身以算子范数与几何论证为出发点，因此其适用范围被表述为嵌入表与语言模型头这两类层。对希望把 Muon 的原则化处理扩展到输入输出层的读者，MuonIO 提供了可直接对照的单一归一化规则形式。
摘要未给出验证困惑度的具体数值、训练步数、批大小或硬件配置，也未说明 50% 内存与约 46% FLOPs 的统计口径与测量方式。恒等式 ‖L‖_{2→∞}=‖Lᵀ‖_{1→2} 的证明细节、列归一化与行归一化的具体实现形式，以及该规则在更大规模或其他架构上的表现，均需查阅正文。由于本次仅基于摘要，上述细节无法从现有文本确认。