Second-Order Fine-Tuning without Pain for LLMs:A Hessian Informed Zeroth-Order Optimizer

📄 arXiv: 2402.15173v4 📥 PDF

作者: Yanjun Zhao, Sizhe Dang, Haishan Ye, Guang Dai, Yi Qian, Ivor W. Tsang

分类: cs.LG

发布日期: 2024-02-23 (更新: 2025-02-18)

💡 一句话要点

提出HiZOO优化器以解决大语言模型微调中的内存问题

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

关键词: 大语言模型 微调 零阶优化器 Hessian 模型收敛性 GPU内存优化 机器学习

📋 核心要点

  1. 现有的微调方法在使用一阶优化器时，由于反向传播导致GPU内存消耗过大，限制了模型的规模和训练效率。
  2. 本研究提出的HiZOO优化器通过利用对角Hessian信息，改进了零阶优化器的性能，减少了内存使用并提高了收敛速度。
  3. 实验结果表明，HiZOO在350M到66B参数的多种模型上显著提高了训练效率和模型准确性，减少了训练步骤。

📝 摘要（中文）

微调大型语言模型（LLMs）时，传统的一阶优化器由于反向传播过程会消耗大量GPU内存。近期研究转向零阶优化器进行微调，虽然通过两次前向传播显著节省了内存，但这些优化器在不同维度的参数曲率异质性方面存在问题。本研究提出了HiZOO，这是一种基于对角Hessian的零阶优化器，首次利用对角Hessian来增强零阶优化器在LLMs微调中的表现。此外，HiZOO避免了高昂的内存成本，每步仅增加一次前向传播。大量实验表明，HiZOO提高了模型收敛性，显著减少了训练步骤，并有效提升了模型准确性。我们还可视化了HiZOO在测试函数上的优化轨迹，展示其处理异质曲率的有效性。最后，我们提供了HiZOO的收敛性理论证明。代码可在https://anonymous.4open.science/r/HiZOO27F8获取。

🔬 方法详解

问题定义：本论文旨在解决在微调大型语言模型时，传统一阶优化器因反向传播导致的高内存消耗问题。现有的零阶优化器虽然节省了内存，但在处理不同维度的参数曲率异质性时效果不佳。

核心思路：论文提出的HiZOO优化器通过引入对角Hessian信息，增强了零阶优化器的性能，旨在提高模型的收敛性和准确性，同时减少内存使用。

技术框架：HiZOO的整体架构包括两个主要模块：首先是基于对角Hessian的参数更新策略，其次是优化过程中的前向传播计算。每一步仅需一次前向传播，从而降低了内存需求。

关键创新：HiZOO的主要创新在于首次将对角Hessian应用于零阶优化器，显著改善了其在处理异质曲率时的表现。这一设计使得优化过程更加高效，避免了传统方法的内存瓶颈。

关键设计：在HiZOO中，关键参数设置包括对角Hessian的计算方式和更新策略，损失函数的选择也经过精心设计，以确保优化过程的稳定性和收敛性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，HiZOO在350M到66B参数的多种模型上，训练步骤减少了显著比例，模型准确性也得到了有效提升。具体而言，HiZOO在多个基准测试中表现优于传统的零阶优化器，验证了其在处理异质曲率时的优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和其他需要大规模语言模型的任务。HiZOO优化器的设计使得在资源受限的环境中，仍能高效地微调大型模型，从而推动相关技术的实际应用和发展。

📄 摘要（原文）

Fine-tuning large language models (LLMs) with classic first-order optimizers entails prohibitive GPU memory due to the backpropagation process. Recent works have turned to zeroth-order optimizers for fine-tuning, which save substantial memory by using two forward passes. However, these optimizers are plagued by the heterogeneity of parameter curvatures across different dimensions. In this work, we propose HiZOO, a diagonal Hessian informed zeroth-order optimizer which is the first work to leverage the diagonal Hessian to enhance zeroth-order optimizer for fine-tuning LLMs. What's more, HiZOO avoids the expensive memory cost and only increases one forward pass per step. Extensive experiments on various models (350M~66B parameters) indicate that HiZOO improves model convergence, significantly reducing training steps and effectively enhancing model accuracy. Moreover, we visualize the optimization trajectories of HiZOO on test functions, illustrating its effectiveness in handling heterogeneous curvatures. Lastly, we provide theoretical proofs of convergence for HiZOO. Code is publicly available at https://anonymous.4open.science/r/HiZOO27F8.