何时何地对分组注意力头加速Muon优化？
机器学习
2026-05-12 v1
摘要
Muon对矩阵更新进行正交化，但多头注意力机制本质上是在注意力头的层级上进行运算。这种粒度不匹配引发了一个问题：Muon是应该作用于完整的注意力投影、单个注意力头，还是作用于中间的注意力头组？我们通过对全矩阵Muon与分组Muon的单步梯度下降比较，研究了这一问题。我们的分析揭示了一种权衡：即分组白化收益（group-wise whitening gain）与分组导致的范数成本（grouping-induced norm cost）之间的的权衡——后者是由用分组白化取代全矩阵白化所导致的额外更新范数成本。鼓舞着这一权衡，我们提出了Group Muon（分组Muon），将注意力头的组大小和分组规则视为优化器的超参数。在在FineWeb数据集上训练的GPT-2 Small模型上，合适的分组能够在全QKV Muon和完全分头的Muonsplit之间取得更好的验证损失。
引用
@article{arxiv.2605.08933,
title = {When and Why Grouping Attention Heads Accelerates Muon Optimization},
author = {Hongtao Zhang and Wenjie Zhou and Wei Chen and Xueqi Cheng},
journal= {arXiv preprint arXiv:2605.08933},
year = {2026}
}
备注
16 pages, 4 figures