论文
MuonIO:嵌入表与语言模型头的范数感知下降
MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads
摘要
Muon优化器通过求解谱范数惩罚下损失的局部线性化导出隐藏线性层的更新规则,其动机是稠密线性层的RMS稳定性论证。但标准Muon实现把输入(嵌入表)与输出(语言模型头)层排除在这种有原则的处理之外,代之以AdamW。我们提出MuonIO,对这两层使用单一的Muon式更新:对语言模型头L(V×d)论证使用2→∞算子范数(源于softmax输出几何的Lipschitz连续性),对嵌入表E(d×V)使用1→2算子范数(基于Bernstein & Newhouse 2025识别的one-hot输入几何)。恒等式‖L‖2→∞=‖Lᵀ‖1→2把两个矩阵放进同一面向词表的几何:MuonIO施加单一的归一化向量规则,对E表现为列归一化、对L表现为行归一化。经验评估证明方法有效:C4上1B LLaMA预训练中,MuonIO较Muon减少50%的I/O优化器状态内存与约46%的I/O更新FLOPs,同时改善验证困惑度。