零阶优化器MeZO在大型语言模型微调中的内存高效应用

语言模型（LM）的微调在各类下游任务中取得了显著的成功，然而，随着LM规模的增加，反向传播所需的内存量也随之增大。传统的零阶（ZO）方法理论上可以通过仅使用两次前向传递来估计梯度，但在优化大型模型时速度极其缓慢。在本研究中，我们提出了一种内存高效的零阶优化器，即MeZO。MeZO采用经典的ZO-SGD方法进行就地操作，使得在微调LM时只需使用与推理相同的内存占用。我们进行了综合实验，涵盖了不同模型类型（屏蔽和自回归LM）、不同模型规模（高达66B）和不同下游任务（分类、多项选择和生成）。研究结果表明：（1）MeZO相比上下文学习和线性探测取得了显著的优势；（2）MeZO在性能上达到了与跨多个任务的反向传播微调相当的水平，最多减少了12倍的内存；（3）MeZO兼容全参数和参数高效调优技术，如LoRA和前缀调优；（4）MeZO能够有效地优化不可微分的目标，例如最大化精度或F1。我们的实验发现得到了理论支持，强调充分的预训练和任务提示是如何使MeZO能够成功微调大型模型的关键。