AdaLomo: Low-memory Optimization with Adaptive Learning Rate

📄 arXiv: 2310.10195v3 📥 PDF

作者: Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, Xipeng Qiu

分类: cs.LG, cs.CL

发布日期: 2023-10-16 (更新: 2024-06-06)

备注: ACL 2024 camera ready version

🔗 代码/项目: GITHUB

💡 一句话要点

提出AdaLomo以解决大语言模型训练中的低内存优化问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 低内存优化 自适应学习率 大语言模型 优化算法 非负矩阵分解 收敛稳定性 机器学习

📋 核心要点

  1. 现有的低内存优化方法在超参数敏感性和收敛效果上存在不足，无法与AdamW优化器的性能相匹配。
  2. 论文提出的AdaLomo方法通过为每个参数提供自适应学习率，并采用非负矩阵分解来优化内存使用。
  3. 实验结果显示，AdaLomo在性能上与AdamW相当，同时显著降低了内存需求，降低了训练门槛。

📝 摘要（中文）

大型语言模型取得了显著成功，但其庞大的参数规模需要大量内存进行训练，设置了较高的门槛。尽管最近提出的低内存优化（LOMO）减少了内存占用，但其优化技术对超参数敏感，收敛效果不佳，无法与主流优化器AdamW相媲美。通过对Adam优化器的实证分析，我们发现自适应学习率在缩小差距方面比动量更为关键。基于此，我们提出了具有自适应学习率的低内存优化方法（AdaLomo），为每个参数提供自适应学习率。为了保持内存效率，我们在优化器状态中采用非负矩阵分解进行二阶矩估计。此外，我们建议使用分组更新归一化以稳定收敛。我们的实验表明，AdaLomo在指令调优和进一步预训练中取得了与AdamW相当的结果，同时显著降低了内存需求，从而降低了训练大型语言模型的硬件门槛。

🔬 方法详解

问题定义：论文旨在解决大型语言模型训练中的内存占用问题，现有的低内存优化方法（LOMO）在超参数设置和收敛速度上存在不足，无法达到主流优化器AdamW的性能。

核心思路：提出AdaLomo方法，通过为每个参数提供自适应学习率，利用非负矩阵分解进行二阶矩估计，以降低内存需求并提高收敛速度。

技术框架：AdaLomo的整体架构包括自适应学习率计算模块、非负矩阵分解模块和分组更新归一化模块，确保在内存效率和收敛稳定性之间取得平衡。

关键创新：最重要的创新点在于引入了自适应学习率机制，并结合非负矩阵分解技术，显著提升了优化过程的效率和效果。与传统的动量方法相比，AdaLomo在性能上更具优势。

关键设计：在参数设置上，AdaLomo采用了动态调整的学习率策略，损失函数设计上则注重于稳定收敛，网络结构上结合了分组更新归一化以减少内存占用。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AdaLomo在指令调优和进一步预训练中表现出色，其性能与AdamW相当，同时内存需求减少了显著的比例，降低了训练大型语言模型的硬件门槛，提升了模型训练的可及性。

🎯 应用场景

该研究具有广泛的应用潜力，特别是在需要训练大型语言模型的场景中，如自然语言处理、对话系统和智能助手等。通过降低内存需求，AdaLomo可以使更多的研究者和开发者能够在普通硬件上进行大规模模型训练，从而推动相关领域的发展。

📄 摘要（原文）

Large language models have achieved remarkable success, but their extensive parameter size necessitates substantial memory for training, thereby setting a high threshold. While the recently proposed low-memory optimization (LOMO) reduces memory footprint, its optimization technique, akin to stochastic gradient descent, is sensitive to hyper-parameters and exhibits suboptimal convergence, failing to match the performance of the prevailing optimizer for large language models, AdamW. Through empirical analysis of the Adam optimizer, we found that, compared to momentum, the adaptive learning rate is more critical for bridging the gap. Building on this insight, we introduce the low-memory optimization with adaptive learning rate (AdaLomo), which offers an adaptive learning rate for each parameter. To maintain memory efficiency, we employ non-negative matrix factorization for the second-order moment estimation in the optimizer state. Additionally, we suggest the use of a grouped update normalization to stabilize convergence. Our experiments with instruction-tuning and further pre-training demonstrate that AdaLomo achieves results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models. The code is accessible at https://github.com/OpenLMLab/LOMO.