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

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

随着大模型模型规模的增大，反向传播（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处的梯度： 这里，是从标准高斯分布