面向内存受限的逆向友好优化：在近似梯度下训练大型语言模型
机器学习
2025-10-28 v1 人工智能
摘要
大型语言模型 (LLM) 的完整微调因依赖缓存激活值而导致内存消耗极大，常规优化器如 SGD 或 Adam 均假设可获得从缓存激活值派生的精确梯度。现有解决方案要么修改模型结构（如可逆网络），要么以计算换内存（如激活检查点），但优化器本身 remains 未被触动。在本工作中，我们引入 GradLite，一种逆向友好优化器，通过放宽对精确梯度的要求，使即使在中间激活被激进丢弃或近似的情况下也能高效训练。GradLite 利用两项关键技术：(i) 低秩雅可比近似，降低反向传播误差信号的维度；(ii) 误差反馈校正，通过跨迭代累积和补偿近似误差来保持收敛保证。我们提供理论分析表明，GradLite 保持无偏梯度估计且方差有界，确保收敛率可与 Adam 相当。实验上，GradLite 在优化器状态和激活内存消耗上降低最高 50%，且在推理 (MMLU, GSM8K)、多语言和对话基准测试中实现与检查点和优化器中心基线 (LoMo, GaLore) 相当或更好的下游性能。
引用
@article{arxiv.2510.22467,
title = {Backward-Friendly Optimization: Training Large Language Models with Approximate Gradients under Memory Constraints},
author = {Jing Yang and Kaitong Cai and Yijia Fan and Yufeng Yang and Keze Wang},
journal= {arXiv preprint arXiv:2510.22467},
year = {2025}
}