论 Muon 的收敛性及拓展

Muon 优化器在处理用于训练神经网络的矩阵结构参数方面展现出了卓越的实证成效。然而，其实际表现与理论理解之间仍存在显著差距。现有分析表明，在随机非凸设定下，Muon 变体仅能达到 $\mathcal{O}(T^{-1/4})$ 的次优遍历收敛率，其中 $T$ 表示迭代次数。为了研究 Muon 的理论极限，我们分析了两种基于动量的方差缩减变体：单批次的 Muon-MVR1 和双批次的 Muon-MVR2。我们首次严格证明了，在无期限学习率调度下，方差缩减使 Muon-MVR2 能够获得最优的随时收敛率 $\widetilde{\mathcal{O}}(T^{-1/3})$，这与该问题类的下界相匹配。在 Polyak--\L{}ojasiewicz (PL) 条件下，我们为 Muon-MVR1 和 Muon-MVR2 建立了随时收敛保证：它们对于期望平方根次优性分别获得了 $\widetilde{\mathcal{O}}(T^{-1/4})$ 和 $\widetilde{\mathcal{O}}(T^{-1/3})$ 的最佳迭代率，并且在迭代过程中给定一个额外的梯度一致上界后，它们分别实现了目标函数差值 $\mathcal{O}(T^{-1/4})$ 和 $\mathcal{O}(T^{-1/3})$ 的最终迭代率。在 CIFAR-10 和 C4 上的实验支持了所提出的方差缩减 Muon 变体的实际有效性。代码可在 \href{https://github.com/MaeChd/MUON-MVR}{Muon-MVR} 代码库获取。