SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training

文章主要内容和创新点总结

主要内容

1. 研究背景
   • 现有低秩梯度优化方法（如GaLore）虽能减少LLM训练内存消耗，但依赖标准最速下降技术，忽略了各向异性损失景观下的收敛加速潜力。
   • 传统优化器（如Adam、Shampoo）在高维空间中计算成本高，近似正交化方法（如Newton-Schulz）存在误差累积问题，尤其在LLM训练的病态条件下表现不佳。
2. 核心方法：SUMO优化器
   • 子空间感知矩正交化：利用动态低维子空间（通过随机截断SVD更新），对一阶矩矩阵进行精确SVD正交化，避免Newton-Schulz的近似误差。
   • 理论分析：证明矩矩阵在训练中呈指数级低秩化趋势，推导Newton-Schulz误差上界与条件数的关系，证明SVD正交化可显著提升收敛速率。
   • 算法设计：结合低秩梯度投影、矩变换、SVD正交化和范数增长限制，在减少内存的同时保持高效优化。
3. 实验验证
   • 微调任务（GLUE）：SUMO在RoBERTa-Base上比GaLore、LoRA准确率更高，内存减少20%以上。
   • 预训练任务（LLaMA）：在C4数据集上，SUM