AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient Large Language Models Fine-Tuning

📄 arXiv: 2406.18060v3 📥 PDF

作者: Yifan Yang, Kai Zhen, Ershad Banijamal, Athanasios Mouchtaris, Zheng Zhang

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-06-26 (更新: 2024-12-02)

备注: Accepted for publication in EMNLP 2024

💡 一句话要点

提出AdaZeta框架，提升MeZO方法在大语言模型微调中的性能和收敛性

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

关键词: 大语言模型微调 零阶优化 张量分解 自适应学习率 内存效率 模型适配 收敛性优化

📋 核心要点

  1. 现有MeZO方法虽然内存效率高，但存在性能下降和发散风险，限制了其广泛应用。
  2. AdaZeta框架通过引入张量化适配器和自适应查询数调度策略，提升ZO估计精度并保证收敛。
  3. 实验结果表明，AdaZeta在精度、内存效率和收敛速度上优于现有MeZO方法，适用于Roberta-Large和Llama-2-7B等模型。

📝 摘要（中文）

本文提出了一种自适应零阶张量列适配(AdaZeta)框架，旨在提高零阶(ZO)方法在大语言模型(LLM)微调中的性能和收敛性。现有的Memory-efficient Zeroth-order (MeZO)方法仅使用前向传递来微调LLM，避免了反向传播图的需求，但存在显著的性能下降和发散风险。AdaZeta通过引入快速前向、低参数的张量化适配器来增强维度相关的ZO估计精度。此外，为了解决大规模ZO微调任务中频繁出现的发散问题，提出了一种自适应查询数调度策略，以保证收敛。在Roberta-Large和Llama-2-7B模型上的实验结果表明，AdaZeta框架在精度、内存效率和收敛速度方面均表现出有效性。

🔬 方法详解

问题定义：现有MeZO方法在微调大型语言模型时，虽然避免了反向传播，降低了内存需求，但由于零阶估计的固有误差，导致性能显著下降，并且在大规模任务中容易出现训练发散的问题。因此，如何提高零阶估计的精度，并保证训练的稳定性是需要解决的关键问题。

核心思路：AdaZeta的核心思路是通过引入一个低参数的张量化适配器来提升零阶估计的精度，该适配器能够捕捉不同维度之间的依赖关系，从而更准确地估计梯度。同时，采用自适应查询数调度策略，根据训练过程中的梯度变化动态调整查询次数，以保证训练的收敛性。

技术框架：AdaZeta框架主要包含两个核心模块：张量化适配器和自适应查询数调度器。首先，输入数据通过一个快速前向的张量化适配器，该适配器将原始特征映射到低维空间，并利用张量分解技术捕捉维度间的关系。然后，利用零阶优化方法，基于适配器输出的特征估计梯度。最后，自适应查询数调度器根据梯度估计的方差动态调整查询次数，以保证训练的稳定性。

关键创新：AdaZeta的关键创新在于：1) 提出了一个快速前向的张量化适配器，能够以较低的参数量提升零阶估计的精度；2) 引入了自适应查询数调度策略，能够根据训练过程中的梯度变化动态调整查询次数，有效避免训练发散的问题。与现有MeZO方法相比，AdaZeta在精度和稳定性上都有显著提升。

关键设计：张量化适配器采用Tensor-Train (TT)分解，将高维参数分解为一系列低维张量的乘积，从而显著降低参数量。自适应查询数调度器根据梯度估计的方差动态调整查询次数，当梯度方差较大时，增加查询次数以提高估计精度；当梯度方差较小时，减少查询次数以降低计算成本。损失函数采用标准的交叉熵损失函数，并结合L2正则化防止过拟合。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AdaZeta在Roberta-Large和Llama-2-7B模型上均取得了显著的性能提升。例如，在Llama-2-7B模型上，AdaZeta相比于基线MeZO方法，在多个NLP任务上的平均准确率提升了3%以上，并且收敛速度更快，有效解决了MeZO方法容易发散的问题。

🎯 应用场景

AdaZeta框架可应用于各种需要高效微调大型语言模型的场景，例如资源受限的边缘设备、移动设备等。该方法能够以较低的计算成本和内存需求，实现对LLM的快速定制化，从而加速LLM在各个领域的应用，例如智能客服、文本生成、机器翻译等。

📄 摘要（原文）

Fine-tuning large language models (LLMs) has achieved remarkable performance across various natural language processing tasks, yet it demands more and more memory as model sizes keep growing. To address this issue, the recently proposed Memory-efficient Zeroth-order (MeZO) methods attempt to fine-tune LLMs using only forward passes, thereby avoiding the need for a backpropagation graph. However, significant performance drops and a high risk of divergence have limited their widespread adoption. In this paper, we propose the Adaptive Zeroth-order Tensor-Train Adaption (AdaZeta) framework, specifically designed to improve the performance and convergence of the ZO methods. To enhance dimension-dependent ZO estimation accuracy, we introduce a fast-forward, low-parameter tensorized adapter. To tackle the frequently observed divergence issue in large-scale ZO fine-tuning tasks, we propose an adaptive query number schedule that guarantees convergence. Detailed theoretical analysis and extensive experimental results on Roberta-Large and Llama-2-7B models substantiate the efficacy of our AdaZeta framework in terms of accuracy, memory efficiency, and convergence speed.