发表机构
School of Electronic Information and Communications; Huazhong University of Science and Technology（电子信息与通信学院; 华中科技大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文提出二次谱下降（QSD）方法，在保持Muon谱范数约束下用二次模型替代线性模型，通过Frank-Wolfe步骤求解，在GPT预训练中优于Muon并减少训练时间。
AI 中文摘要
Muon 可以被解释为在谱范数球上优化一个线性局部目标。这产生了一个矩阵符号更新，该更新保留梯度的奇异方向，并为所有活跃的奇异模式分配相同的幅度。我们询问当考虑局部曲率时，这两个性质是否仍然最优。为了回答这个问题，我们保持 Muon 的谱范数约束不变，并将线性局部模型替换为二次模型。我们将所得方法称为二次谱下降（QSD）。我们表明，曲率可以改变最优更新的奇异值和奇异方向。为了使 QSD 实用，我们用克罗内克因子统计量近似曲率，并通过少量 Frank-Wolfe 步骤求解约束二次问题，每一步都有一个闭式矩阵符号子问题。我们进一步提供了一个最优性证书，在相同二次代理下与 Muon 的比较，以及内层求解器的 $O(1/K)$ 收敛率。在 GPT 预训练上的实验表明，QSD 在验证损失上持续优于 Muon 和最近的 Muon 变体，并在匹配验证损失时将墙钟训练时间减少最多 $8.49\%$。
英文摘要
Muon emerges as a strong competitor of the AdamW for LLM pretraining, because the matrix-wise update it employs can potentially incur smaller second-order penalty than the once dominating AdamW, which performs coordinate-wise update. However, the spectral flattening procedure in Muon is quite debatable since it discards the spectral amplitude information totally. To seek for better spectral allocation (and the associated spectral subspace), we propose to solve the quadratic model of loss function under the spectral norm constraint \textit{directly} (i.e., in a genuinely Newtonian way) and thus obtaining the Quadratic Spectral Descent (QSD) algorithm. In contrast, many existing curvature-aware methods either exploit the second-order information in an \textit{implicit} way by changing the weight update geometry (such as Mousse, FISMO) or rely on strong assumptions (such as the weight displacement isotropy assumption in Newton-Muon). QSD's potential advantage over these methods is best illustrated in the isotropic curvature scenario, where Mousse, FISMO and Newton-Muon all reduce to Muon while the spectral allocation in QSD is still \textit{non-flat} (since the spectral allocation in QSD depends on the \textit{gradient to curvature ratio}). Meanwhile, to control the complexity of QSD, we employ inversion-free K-FAC and \textit{online} Frank-Wolfe update which is essentially a matrix sign operator. Overall, the complexity increase can be rather mild. Experiments on GPT pre-training show that QSD consistently improves validation loss over Muon and recent Muon variants, while achieving up to an $8.49\%$ wall-clock speedup at matched validation loss.
