摘要
Muon及相关矩阵符号优化器在大型语言模型的预训练中越来越受到重视，但它们对各个权重矩阵内部几何的影响尚不明确。本报告提出了一个统一框架，基于一个理想化假设——在权重重缩放下损失的精确尺度不变性，这在重归一化网络中大致成立。
在这一假设下，普通的随机梯度下降（SGD）对更新大小施加了内置的 $1/||W||$ 刹车，而Muon的矩阵符号步骤移除了这一刹车，使得Frobenius范数和谱范数的漂移速度加快（分别为 $t^{1/2}$ 和 $t^{1/4}$）。我们进一步观察到，谱范数扰动具有非负的二阶项。这意味着一个轻量级的“谱限”可以控制输出协方差 $W K_X W^T$，而不会冻结训练：权重通过非顶方向、顶方向旋转和顶切换继续学习。
我们将此谱限与奇异值谱的最小熵（$H_{\infty}$）联系起来。我们研究了三个使用Muon训练的系统：一个nanoGPT前馈投影、一个64个专家的混合专家路由器，以及一个bf16 FlashAttention块的查询/键投影。在每种情况下，谱限增加了各向同性，并在边际上防止了具体的失败（如路由器崩溃为单一专家，以及一个注意头的近乎发散），同时验证损失基本保持不变。
我们强调，尺度不变性假设是强的，这些小规模的结果是初步的；欢迎评论。
博主点评： 该研究为Muon优化器在大型语言模型中的应用提供了新的视角，特别是在控制权重矩阵的几何结构方面。通过引入谱限的概念，研究者不仅展示了如何提升各向同性，还为未来的优化算法设计提供了重要参考。值得注意的是，研究中的假设需要进一步验证，以确保在更广泛的应用场景中有效。