ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

王林
发布: 2024-07-16 03:17:30
原创
904 人浏览过
ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存
AIxiv专栏是本站发布学术、技术内容的栏目。过去数年,本站AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。如果您有优秀的工作想要分享,欢迎投稿或者联系报道。投稿邮箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

本文共同第一作者简介:张逸骅:密歇根州立大学计算机系博士三年级学生,师从Sijia Liu教授,主要研究方向是大模型的安全、隐私和效率问题。李平治:本科毕业于中国科学技术大学,将于2024 Fall博士入学北卡罗来纳大学教堂山分校,师从陈天龙教授,主要研究兴趣集中在高效机器学习和AI4Science领域。洪骏远:德州大学奥斯汀分校博后,导师是Zhangyang Wang教授。博士毕业于密歇根州立大学,师从Jiayu Zhou教授,目前主要的研究方向是可信大语言模型和人工智能的医疗应用。李佳翔:明尼苏达大学博士后,目前在洪明毅教授和张树中教授指导下做数值优化理论、机器学习理论和大规模机器学习优化问题的研究。

开源大语言模型(LLM)百花齐放,为了让它们适应各种下游任务,微调(fine-tuning)是最广泛采用的基本方法。基于自动微分技术(auto-differentiation)的一阶优化器(SGD、Adam 等)虽然在模型微调中占据主流,然而在模型越来越大的今天,却带来越来越大的显存压力。因此,如何高效地在微调中降低显存使得单卡可以满足微调需求已经成为一个热门研究问题。值得注意的是,虽然反向传播是这些一阶优化器的基石,被用于计算神经网络每个权重的梯度,同时却也是显存杀手,其中庞大计算图的保存所带来的开销也在大模型时代被凸显得尤为突出。与此同时,零阶优化器(Zeroth-Order Optimization)则完全无需保存计算图,转而使用有限差分来近似计算网络的梯度,通过完全避免反向传播(back-propagation; BP)来大大减少神经网络更新中的内存开销。

类似于一阶优化器中随机梯度下降的各式变种,零阶优化器也有着各种此前无人探索的改进算法。近日,来自密歇根州立大学、北卡罗来纳大学教堂山分校、德克萨斯大学奥斯汀分校、明尼苏达大学双城分校、IBM 研究中心、普林斯顿大学、以及阿里巴巴达摩院的众多研究者联合推出全面评测(benchmark)文章:Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark 。这篇文章覆盖六种无需反向传播(BP-free)的优化器、五类大模型、三种复杂度的各项任务、四类微调方案,以及三项增强零阶优化器的全新算法。目前,相关论文已被 ICML 2024 高分接收,代码已开源;详情如下。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

  • 论文地址:https://arxiv.org/abs/2402.11592
  • 代码地址:https://github.com/ZO-Bench/ZO-LLM
  • 零阶优化讲义地址 (AAAI 2024 Tutorial):https://sites.google.com/view/zo-tutorial-aaai-2024/ 

零阶优化器是什么?为何如此重要?

零阶优化器(Zeroth-Order Optimization)仅仅依靠神经网络的输出进行梯度估计,以完全不需要计算反向传播和极少的内训消耗而闻名。尽管在零阶优化器领域也存在不同的梯度估计方法,本文特指基于随机梯度估计器(Random Gradient Estimator, RGE)的一系列算法。简单来说,就是通过从高斯分布中抽取的随机扰动来计算有限差分,并将其作为梯度的近似估计,RGE 数学公式如下所示。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

在此之前,零階優化已經大量應用於機器學習問題中,例如對抗樣本生成和防禦,黑箱模型解釋,強化學習和自動機器學習;詳細的演算法和應用介紹請參看 [1]。在大模型領域,MeZO [2] 率先提出將零階隨機梯度下降法(ZO-SGD)用作大模型微調,並展示了零階優化器的無限潛力。同時,ZO-SGD 是最簡單、基本的 BP-free 優化器,它的許多更高級的變種 [3] 能否在大模型微調領域帶給我們更多驚喜,是一個亟待研究的課題。本文系統性地評測了以下無需反向傳播(BP-free)的最佳化演算法在大模型微調任務上的效能、效率和相容性,目的是向社群展示零階優化器在多種大模型任務上的廣泛潛力:

  • ZO-SGD:零階隨機梯度下降[4]
  • ZO-SGD-Sign:基於符號的(sign-based)零階隨機梯度下降[5]
  • ZO-SGD-MMT:具有動量(momentum)的零階隨機梯度下降[6]
  • ZO-SGD-Cons:保守(conservative)梯度更新的零階隨機梯度下降[7]
  • ZO-Adam:零階Adam 優化器[8]

本研究還包括了Forward-Grad [9] 方法,該方法基於沿隨機方向向量的方向導數來無偏估計梯度。值得注意的是,Forward-Grad 雖然不直接使用梯度回傳,但仍然使用了自動微分演算法,因此它是一階的 BP-free 演算法。

綜上所述,本文的評測包含了上述 5 種零階優化器以及 Forward-Grad 方法,同時對比一階優化器中最常用的 FO-SGD 和 FO-Adam。在具體微調形式上,評測全面涵蓋了5 種LLM 架構(RoBERTa, OPT, LLaMA, Vicuna, Mistral),3 種不同複雜度的任務(SST2, COPA, WinoGrande),以及4 種微調方案(full-tuning , LoRA, prompt tuning, prefix tuning)。

大模型微調準確度評測

作者指出,為了有效地利用零階優化器對大型模型在下游任務上進行微調,必須合理地運用輸入模板,以便將下游任務模板與預訓練任務進行對齊。例如對 SST2 來說,使用範本 “
SENTENCE. It was [terrible|great].” 能夠在 ZO-SGD 上帶來約 10% 的效能提升。然而對於一階優化器(如 FO-SGD),無論是否使用模板效能差異都不大,都襯托出零階優化器的獨特性。

SST2 作為一個較為基礎的任務,其實驗結果可以支持以下結論:

  • ZO-Adam 似乎是最有效的零階最佳化器
  • ZO-Adam 似乎是最有效的零階優化器:在8 個微調設定中的4 個微調中表現最佳。
  • Forward-grad 是一種競爭力強但以前被忽視的方法,特別是在全模型微調 (full fine-tuning) 中。
  • ZO-SGD-Cons 和 ZO-SGD-MMT 也展示了強大的性能,而 ZO-SGD-Sign作為最簡單的零階優化器,往往是最弱的方法。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

進一步,研究使用更大的模型OPT-13B,在更複雜、更困難的任務(COPA 和WinoGrande)上進行實驗,得出以下結論:
  • 在更複雜的任務中,不同優化器的效能差異被進一步放大。
  • ZO-Adam 和 ZO-SGD-MMT 在各種實驗下展示了非常好的穩定性,這可能歸因於減少了變異數的設計。
  • LoRA 微調對於零階演算法始終表現出強大的穩健性,在各種實驗環境中穩定且可靠。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

🎜
大模型微調記憶體開銷評測與詳解

以 OPT-13B 模型在 MultiRC 資料集上微調為例,作者也進一步比較分析了不同零階與一階最佳化器的記憶體與時間成本最佳化。如下表所示:首先,從記憶體效率的角度來看,ZO-SGD、ZO-SGD-Cons 和ZO-SGD-Sign 顯示了類似的高記憶體效率,只需要一個A100 GPU 來進行大型語言模型的微調。這並不令人驚訝,因為這些零階優化器採用相對簡單的最佳化步驟,主要依賴零階梯度估計器 RGE 的利用。其次,Forward-Grad 似乎是零階最佳化方法在記憶體效率方面超過一階方法的臨界點(例如與 ZO-Adam 相比)。最後,與一階方法相比,零階優化每次迭代的運行時間成本降低了約 41.9%(以 ZO-SGD 與 FO-SGD 為例)。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

作者進一步深入比較了 ZO-SGD 與 FO-SGD 在不同序列長度下的記憶體效率。可以看到,ZO-SGD 的記憶體消耗保持一致,因為其峰值記憶體消耗僅由模型參數大小決定,相比之下,隨著序列長度的增加,FO-SGD 的峰值記憶體消耗先保持不變,然後開始增加。因此,在長上下文長度的設定中,ZO-SGD 將展示出更好的記憶體效率優勢。具體的記憶體理論值和實驗值可參考原論文。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

三種改進演算法用以增強零階優化器

零階優化器在 LLM 時收斂效率受限,主要是因為它們在梯度估計上的變異數較大。為了進一步增強零階優化器,作者從降低梯度估計方差的角度入手,提出了三種進階演算法,包括:分塊零階微調(block-wise ZO fine-tuning)、零階和一階混合微調(hybrid ZO and FO fine-tuning)、引入稀疏性的零階梯度估計(sparsity-induced ZO gradient estimation)。

分塊零階微調(Block-wise ZO fine-tuning)此方法的主要出發點在於,如果零階優化器在估計梯度時,對LLM 中參數分塊分別進行擾動,透過降低問題規模的方式來見效每次對梯度估計的方差,從而改進最佳化性能。這種方法的優點體現在能夠對模型梯度進行更準確的估計,但是完成一次梯度估計所需的前向傳播的次數會增加。例如,OPT-1.3B 可以分成 26 個參數區塊(24 個 Transformers 層、嵌入層和 LM 分類頭),那麼零階優化器每次計算模型梯度時就會計算 26 次前向傳播。為了公平比較ZO-SGD 和ZO-SGD-Block,作者還比較了另一種ZO-SGD 變體的性能,該變體每次對完整的模型進行參數擾動,並將多次擾動後的梯度估計求平均值(例如OPT-1.3B 的26 次),以此來確保比較時的前向傳播次數相同。 OPT-1.3B 上實驗結果表明,ZO-SGD-Block 大幅超越了兩種 ZO-SGD。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

零階和一階混合訓練(Hybrid ZO and FO fine-tuning)反向傳播(BP)從深層至淺層神經網路依序計算權重梯度。由於零階優化器在記憶體佔用上有遠超傳統一階優化器的優勢,但一階優化器的性能往往更好。因此,採用零階和一階優化器的結合將達到一種記憶體使用和效能之間的平衡(trade-off)。具體而言,對於較深層網絡,可以利用一階優化器透過反向傳播精確計算梯度;對於淺層網絡,則可以透過零階優化器進行梯度估算。實驗結果表明,在淺層部分(例如OPT-1.3B 的前8/24 層)採用零階優化器,而在剩餘的深層部分使用一階優化器,可以在節約大約三分之一的顯存的同時,達到與完全使用一階優化器相同的性能等級。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

使用稀疏梯度的零階優化器(ZO with gradient pruning)在一階優化器中,梯度剪裁通常用於加速訓練過程;而在零階優化器中,透過梯度剪裁引入的稀疏梯度可以進一步降低梯度估計的方差,從而提高性能。本文研究了在零階優化器中應用基於幅值的剪裁策略來獲取每一層的稀疏率,然後根據這些稀疏率生成隨機的稀疏梯度掩碼(mask),並將其應用於隨機梯度估計的擾動上。實驗結果顯示,適度的梯度稀疏性(約 20% 左右)能為零階優化器帶來一定程度的效能提升。

ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存

結語

在本文中,我們展示了零階優化器在大型語言模型微調中的有效應用。透過利用損失差分來近似梯度,零階優化方法避免了反向傳播和激活儲存的需求,並大大節省了記憶體資源。我們透過擴大現有的研究範圍,將不同的零階優化方法、任務類型及評估指標容納到了本次評測中,進行了首次系統的零階優化技術基準研究。我們的研究不僅揭示了這些方法在精確度和效率方面的表現,還深入探討了任務對齊和前向梯度的關鍵作用。利用這些實驗分析,我們提出了分塊優化、零階與一階混合訓練、梯度稀疏化等技術,以進一步增強基於零階優化的大模型微調。這些改進技術旨在保持記憶體效率的同時,提高微調的準確性。

我們堅信,這些發現和技術的應用可以大力降低大模型研究對硬體資源的要求,使得大模型微調在低端GPU 也成為可能,從而進一步推動學術研究並在工業界產生實際而有價值的影響。我們鼓勵廣大研究人員和技術開發者關注我們的研究成果,並探索更多利用 ZO 優化的可能性。未來的研究將繼續探索這一領域的深層問題,以解鎖 LLM 微調中的更多潛力。

了解更多內容請參考論文與 GitHub 倉庫,以取得更多資訊與資源。
? 2020): 43-54.

[2] Malladi, et al., “Fine-Tuning Language Models with Just Forward Passes.” NeurIPS' 2023.
[3] Liu, et al., “A Primer on Zeroth-Order Optimization in Signal Processing and Machine Learning.” IEEE Signal Processing Magazine.
[4] Ghadimi, et al., “Stochastic First- and Zerothd Methods for Stochastic Programming.” 
[5] Liu, et al., “signSGD via Zeroth-Order Oracle. ” ICLR' 2019.ated
, ated
er -Order and First-Order Momentum Methods from Mini to Minimax Optimization.” JMLR' 2022. 
[7] Kim, et al., “Curvature-Aware Derivative-Free Opt[7] Kim, et al., “Curvature-Aware Derivative-Free Optoization.” [8] Chen, et al., “ZO-AdaMM: Zeroth-Order Adaptive Momentum Method for Black-Box Optimization.”
[9] Baydin, et al., “Gradients without Backpropagation.”

以上是ICML 2024高分论文 | 零阶优化器微调大模型,大幅降低内存的详细内容。更多信息请关注PHP中文网其他相关文章!

相关标签:
来源:jiqizhixin.com
本站声明
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系admin@php.cn
最新问题
热门教程
更多>
最新下载
更多>
网站特效
网站源码
网站素材
前端模板
关于我们 免责声明 Sitemap
PHP中文网:公益在线PHP培训,帮助PHP学习者快速成长!