仅用前向传播微调语言模型
机器学习
2024-01-12 v3 计算与语言
摘要
微调语言模型（LMs）在多种下游任务上取得了成功，但随着模型规模的增长，反向传播所需的内存量大得令人望而却步。零阶（ZO）方法原则上可仅用两次前向传播估计梯度，但理论上被认为在优化大型模型时极其缓慢。本文提出一种内存高效的零阶优化器（MeZO），将经典的 ZO-SGD 方法改造为原地操作，从而以与推理相同的内存占用微调语言模型。例如，在单张 A100 80GB GPU 上，MeZO 可训练一个 300 亿参数模型，而使用反向传播微调在相同预算下只能训练 2.7B 的 LM。我们在多种模型类型（掩码与自回归 LMs）、模型规模（最高 66B）和下游任务（分类、多选与生成）上进行了全面实验。结果表明：(1) MeZO 显著优于上下文学习与线性探测；(2) 在多个任务上 MeZO 达到与反向传播微调相当的性能，在我们的实现中内存最多减少 12 倍、GPU 小时最多减少 2 倍；(3) MeZO 兼容全参数与参数高效微调技术，如 LoRA 与前缀微调（prefix tuning）；(4) MeZO 可有效优化不可微目标（如最大化准确率或 F1）。我们以理论见解支撑实证发现，阐明充分的预训练与任务提示如何使 MeZO 能够微调巨型模型，尽管经典 ZO 分析另有结论。
引用
@article{arxiv.2305.17333,
title = {Fine-Tuning Language Models with Just Forward Passes},
author = {Sadhika Malladi and Tianyu Gao and Eshaan Nichani and Alex Damian and Jason D. Lee and Danqi Chen and Sanjeev Arora},
journal= {arXiv preprint arXiv:2305.17333},
year = {2024}
}
备注
Accepted by NeurIPS 2023 (oral). Code available at https://github.com/princeton-nlp/MeZO