本文是LLM系列文章,针对《Memory-Efficient LLM Training with Online Subspace Descent》的翻译。
基于在线子空间下降的记忆高效LLM训练
摘要
最近,各种内存高效的LLM训练算法得到了广泛的应用。这些方法利用梯度的低秩结构,使用奇异值分解(SVD)得到的投影矩阵将优化器状态投影到子空间中。然而,这些算法的收敛性在很大程度上取决于其投影矩阵的更新规则。在这项工作中,我们为投影矩阵的任意更新规则提供了第一个收敛保证。这个保证通常适用于可以用哈密顿下降分析的优化器,包括最常见的优化器,如LION、Adam。受我们理论理解的启发,我们提出了在线子空间下降,这是一个新的无SVD子空间下降优化器家族。在线子空间下降不是用特征向量更新投影矩阵,而是用在线PCA更新投影矩阵。在线子空间下降是灵活的,只会给训练带来最小的开销。我们发现,对于在C4数据集上预训练60M至7B参数的LLaMA模型的任务,在线子空间下降在不同设置下比最先进的低秩训练方法具有更低的困惑度和更好的下游任务性能,并缩小了与全秩基线的差距。