带残差校正的内存高效子空间优化：从动量 SGD 到 Muon

内存高效的子空间优化已成为一种实用方法，能够在保持优化性能的同时降低大语言模型训练的内存开销。以往的算法通常依赖随机梯度的主奇异子空间，但在训练后期，梯度噪声可能主导该子空间，并阻碍收敛。为克服这一局限，我们采用残差校正框架，将优化器状态保留在子空间中，并将正交残差梯度加入参数更新。我们首先将这一方案融入动量 SGD，提出 SMRC。在固定正交归一基下，SMRC 达到 $O(T^{-1/4})$ 的收敛速率，与全空间动量 SGD 的收敛速率一致。我们的分析引入了子空间虚拟迭代点，以解耦动量与残差校正。此外，我们将 SMRC 扩展为多阶段变体，通过周期性地重新计算基来跟踪不断变化的主导子空间，并保持相同的收敛速率。由于每次基的刷新都会破坏虚拟迭代点的一致性，我们设计了一种边界校正方法，使其在不同阶段之间保持对齐。除动量 SGD 外，我们还将该框架融入 Muon，提出 PMRC，对投影后的动量应用矩阵符号函数。通过考虑残差梯度的加权动量跟踪分析，我们证明 PMRC 具有 $O(T^{-1/4})$ 的收敛速率，与全空间 Muon 的收敛速率一致。我们还设计了周期性刷新基的多阶段 PMRC，其收敛速率保持同阶。大量实验结果证实了我们算法的有效性和效率。