新审视零阶优化在内存高效大模型微调中的应用

人工智能咨询培训老师叶梓 转载标明出处

随着大模型模型规模的增大，反向传播（BP）所需的内存开销也日益增加，这对内存效率提出了挑战。尤其是在设备上训练等内存效率至关重要的应用场景中，解决这一问题变得尤为迫切。

本文提出了一种新的方法，即在LLM微调过程中，采用零阶（Zeroth-Order, ZO）优化算法替代传统的需要反向传播的一阶（First-Order, FO）优化算法，以降低内存成本。这种BP-free的优化方法基于Malladi等人在2023年提出的初始概念。

论文地址：https://arxiv.org/abs/2402.11592

项目地址：https://github.com/ZO-Bench/ZO-LLM

方法

ZO优化作为一种不依赖梯度信息的替代方法，通过基于函数值的梯度估计来近似传统的一阶梯度。这种方法通常反映了其对应的一阶优化算法的算法框架，但它用ZO梯度估计替换了FO梯度作为下降方向。文中特别关注了随机梯度估计器（Randomized Gradient Estimator, RGE），这是一种依赖于沿随机选择的方向向量的函数值的有限差分的方法。

RGE的数学表达式，用来近似函数在点x处的梯度： 这里，是从标准高斯分布 中抽取的随机方向向量，q 是函数查询的数量，而 μ 是一个小的扰动步长。

前向梯度（Forward-Grad）是一种在LLM微调中可能被忽视的BP-free基线方法。与RGE不同，Forward-Grad利用前向模式自动微分（AD），但省去了在深度模型微调或训练中反向传播的需要。这意味着Forward-Grad不需要进行反向传播，因此可以作为提高LLM微调内存效率的另一种替代梯度估计方法。

ZO优化方法的范围包含： ZO-SGD是由Ghadimi和Lan在2013年提出的，它通过估计梯度来更新模型参数。Malladi等人在2023年的研究中采用了这种技术进行LLM的微调。基于符号的梯度估计的ZO-SGD（ZO-SGD-Sign），这种方法通过符号函数处理梯度估计，简化了计算过程。另外，带有动量（MMT）的ZO-SGD，它利用历史梯度信息来优化更新步骤，以及保守梯度更新的ZO-SGD（ZO-SGD-Cons），它在选择下降方向时采用一种基于启发式的方法。最后，ZO-Adam，这是Adam优化器的ZO变体，它结合了动量和自适应学习率。

通用的优化框架是将上述ZO优化方法统一起来，用于解决最小化问题 min_x f(x)。框架中的更新公式为： 这里，表示第t次迭代的解，是学习率，而h(·) 是对梯度估计进行后处理的下降方向操作。例如，ZO-SGD可以直接将估计的梯度作为下降方向，而ZO-SGD-Sign则使用梯度估计的符号作为下降方向。

任务对齐指的是将微调任务与预训练任务的格式对齐，这通常通过精心设计的输入提示来实现。例如，可以将文本分类任务转换为下一个词预测任务，从而使得微调任务与预训练任务更加一致。这种对齐有助于提高ZO优化的性能。

使用预训练的Roberta-Large模型在SST2和RTE数据集上进行微调，并比较了使用和不使用文本对齐时的性能差异。结果表明，没有进行文本对齐时，ZO-SGD和ZO-Adam的性能分别下降了约10%和8%，而FO-SGD的性能下降较小。这表明任务对齐对于提高ZO优化的微调性能至关重要。

想要掌握如何将大模型的力量发挥到极致吗？叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。9月22日晚，实战专家1小时讲解让您轻松上手，学习如何使用 Llama Factory 微调模型。

加助理微信提供直播链接：amliy007，29.9元即可参加线上直播分享，叶老师亲自指导，互动沟通，全面掌握Llama Factory，关注享粉丝福利，限时免费CSDN听直播后的录播讲解。
LLaMA Factory 支持多种预训练模型和微调算法。它提供灵活的运算精度和优化算法选择，以及丰富的实验监控工具。开源特性和社区支持使其易于使用，适合各类用户快速提升模型性能。

实验

研究者选择了四个具有不同复杂性的任务：最简单的二元分类任务SST2，问答任务COPA，常识推理任务WinoGrande，以及多句阅读理解任务MultiRC（仅用于效率评估）。这些任务覆盖了从简单分类到复杂推理的多种场景。同时，研究者探索了四种参数高效的微调（PEFT）方案：全参数微调（FT）、低秩适应（LoRA）、前缀调整（Prefix）和提示调整（Prompt）。

研究涵盖了几种具有代表性的LLM，包括Roberta-Large、OPT、LLaMA2、Vicuna和Mistral，这些模型在规模和架构上有所不同。

优化方法： 实验中使用了多种ZO优化方法，包括ZO-SGD、ZO-SGD-Sign、ZO-SGD-MMT、ZO-SGD-Cons和ZO-Adam。这些方法在处理梯度估计和参数更新时采用了不同的策略。例如，ZO-SGD直接使用估计的梯度进行更新，而ZO-SGD-Sign则使用梯度估计的符号。此外，还比较了基于前向模式自动微分的Forward-Grad方法，以及两种一阶（FO）优化器：SGD和Adam。

实验迭代： 实验中，ZO优化器和FO优化器分别运行了20,000和625次迭代。这反映了ZO优化通常需要更多的迭代次数才能收敛。

精度和效率评估： 通过测试准确性和效率（包括内存效率、查询效率和运行时效率）来评估不同优化方法的性能。

实验结果：

表2展示了在SST2数据集上不同BP-free和基于BP的方法的性能比较。结果显示，ZO-Adam在多数设置中表现最佳，但需要更多的内存。Forward-Grad在全参数微调（FT）设置中表现出色，表明它可能适合大规模问题。ZO方法的性能波动较大，这可能与随机梯度估计（RGE）的高方差有关。

当问题规模增大且变得更复杂时，不同ZO方法之间的性能差距增大。FO方法（FO-SGD和FO-Adam）在WinoGrande任务上的性能明显优于所有ZO方法，这揭示了ZO方法在处理大型模型和复杂任务时的可扩展性瓶颈。

内存效率深入分析：

表5提供了不同优化器在全参数微调和LoRA微调中的即时峰值内存消耗的比较。ZO方法通常比FO方法更有效，因为它们不需要存储中间结果，且可以逐层估计梯度，从而避免了存储整个模型的梯度。

实验结果与理论分析一致，表明ZO方法在内存效率方面具有显著优势。随着输入序列长度的增加，ZO方法的内存优势更加明显，因为它们的峰值内存消耗仅由模型参数大小决定，与中间结果的大小无关。