MuonEq：在正交化之前进行轻量级均衡化
机器学习
2026-05-12 v2 机器学习
摘要
正交化更新优化器如Muon可改进矩阵型参数的训练，但现有扩展方法通常要么在正交化后对更新进行缩放，要么在正交化前使用更重的白化预条件器。我们引入{\method}，这是一套轻量级的正交化前均衡化方案，包含三种形式：双向行/列归一化（RC）、行归一化（R）和列归一化（C）。通过在有限步长Newton--Schulz正交化之前对动量矩阵进行再平衡，{\method}改善了正交化看到的几何结构。我们指出，有限步正交化由输入谱决定，尤其是稳定秩和条件数，而行/列归一化充当白化的零阶近似。对于隐藏矩阵权重，R是默认变体。理论上，{\method} (R)保留标准的型Muong非凸稳定性保证，同时实现解耦权重衰减和无时限递减学习率计划，并将其推广到有限步NS5，仅受显式不完美常数的限制。在LLaMA2 C4数据集上进行预训练时，{\method} (R)在130M、350M和1B模型上始终优于Muon，收敛更快且验证困惑度更低。代码已公开于\href{https://github.com/MaeChd/muon-eq}{MuonEq codebase}。
引用
@article{arxiv.2603.28254,
title = {MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration},
author = {Da Chang and Qiankun Shi and Lvgang Zhang and Yu Li and Ruijie Zhang and Yao Lu and Yongxiang Liu and Ganzhao Yuan},
journal= {arXiv preprint arXiv:2603.28254},
year = {2026}
}