通过分组优化注意力头，提升模型训练效率
When and Why Grouping Attention Heads Accelerates Muon Optimization
- 将注意力头分组进行正交化更新，平衡加速与误差代价
- 在GPT-2 Small上，最优分组使验证损失低于全矩阵和逐头方法
- 适用于追求训练加速的Transformer模型优化场景
Muon通过正交化矩阵更新提升优化性能，但多头注意力天然以头为单位运作。这种粒度不匹配引发疑问：Muon应作用于完整投影、单个头，还是中间头组？我们通过一步下降比较全矩阵Muon与分组式Muon，发现存在权衡：分组正交化带来的增益与因替换全矩阵正交化而引入的额外更新范数代价之间存在平衡。受此启发，我们提出Group Muon，将头组大小与分组规则设为优化器超参数。在FineWeb上训练的GPT-2 Small实验表明，适当分组可使验证损失优于全QKV Muon和完全头级分离（MuonSplit）。
原文摘要 · Abstract (English)
Muon orthogonalizes matrix updates, but multi-head attention naturally operates at the level of heads. This granularity mismatch raises the question of whether Muon should be applied to the full attention projection, to individual heads, or to intermediate head groups. We study this question through a one-step descent comparison between full-matrix Muon and group-wise Muon. Our analysis reveals a trade-off between the \textbf{group-wise whitening gain} from group-wise updates and the \textbf{grouping-induced norm cost}, an additional update-norm cost caused by replacing full-matrix whitening with group-wise whitening. Motivated by this trade-off, we propose \textbf{Group Muon}, which treats head group size and grouping rule as optimizer hyperparameters. On GPT-2 Small trained on FineWeb, appropriate grouping improves validation loss over both full-QKV Muon and fully head-wise MuonSplit.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。