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

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