Addax：利用零阶梯度提高 SGD 在微调语言模型方面的内存效率与性能
机器学习
2024-10-10 v1 计算与语言
摘要
使用 Adam 优化器微调语言模型(LMs)通常需要大量内存，限制了可及性。已提出 SGD 的"原地"版本(IP-SGD)和高效零阶优化器(MeZO)以解决此问题。然而，IP-SGD 仍需大量内存，MeZO 由于其零阶性质导致收敛缓慢且最终性能下降。本文引入 Addax，这是一种新方法，通过将其与 MeZO 集成来提高 IP-SGD 的内存效率和性能。具体而言，Addax 根据数据点在 minibatch 中的内存消耗，计算这些数据点的零阶或一阶梯度，并将这些梯度估计值组合以更新方向。通过对内存消耗更大的计算数据点的零阶梯度、对其他数据点计算一阶梯度，Addax 克服了 MeZO 的慢收敛问题和 IP-SGD 的过度内存需求。此外，零阶梯度可作为一阶梯度的正则器，进一步提高模型的最终性能。在理论上，我们在轻微假设下建立了 Addax 的收敛性，证明其比 MeZO 更快的收敛速度和更不严格的超参数选择。我们的实验使用多样化的 LMs 和任务表明，Addax 在准确率和收敛速度方面始终优于 MeZO，而内存占用相当。当使用一块 A100 GPU 对 OPT-13B 进行微调时，Addax 在准确率/F1 分数方面平均优于 MeZO 14%，运行速度快 15 倍，同时保持与 MeZO 相当的内存占用。在我们对更大型的 OPT-30B 模型的实验中，Addax 在准确率/F1 分数方面平均优于 MeZO >16，运行速度快 30 倍，仅用一块 H100 GPU。此外，Addax 在大多数任务上都优于标准微调方法(如 IP-SGD 和 Adam)，且显著降低了内存需求。
引用
@article{arxiv.2410.06441,
title = {Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models},
author = {Zeman Li and Xinwei Zhang and Peilin Zhong and Yuan Deng and Meisam Razaviyayn and Vahab Mirrokni},
journal= {arXiv preprint arXiv:2410.06441},
year = {2024}
}