ICML 2024高分論文 | 零階優(yōu)化器微調(diào)大模型，大幅降低內(nèi)存 精華
本文共同第一作者簡介：張逸驊：密歇根州立大學(xué)計算機系博士三年級學(xué)生，師從Sijia Liu教授，主要研究方向是大模型的安全、隱私和效率問題。李平治：本科畢業(yè)于中國科學(xué)技術(shù)大學(xué)，將于2024 Fall博士入學(xué)北卡羅來納大學(xué)教堂山分校，師從陳天龍教授，主要研究興趣集中在高效機器學(xué)習(xí)和AI4Science領(lǐng)域。洪駿遠(yuǎn)：德州大學(xué)奧斯汀分校博后，導(dǎo)師是Zhangyang Wang教授。博士畢業(yè)于密歇根州立大學(xué)，師從Jiayu Zhou教授，目前主要的研究方向是可信大語言模型和人工智能的醫(yī)療應(yīng)用。李佳翔：明尼蘇達(dá)大學(xué)博士后，目前在洪明毅教授和張樹中教授指導(dǎo)下做數(shù)值優(yōu)化理論、機器學(xué)習(xí)理論和大規(guī)模機器學(xué)習(xí)優(yōu)化問題的研究。
開源大語言模型（LLM）百花齊放，為了讓它們適應(yīng)各種下游任務(wù)，微調(diào)（fine-tuning）是最廣泛采用的基本方法?；谧詣游⒎旨夹g(shù)（auto-differentiation）的一階優(yōu)化器（SGD、Adam 等）雖然在模型微調(diào)中占據(jù)主流，然而在模型越來越大的今天，卻帶來越來越大的顯存壓力。因此，如何高效地在微調(diào)中降低顯存使得單卡可以滿足微調(diào)需求已經(jīng)成為一個熱門研究問題。值得注意的是，雖然反向傳播是這些一階優(yōu)化器的基石，被用于計算神經(jīng)網(wǎng)絡(luò)每個權(quán)重的梯度，同時卻也是顯存殺手，其中龐大計算圖的保存所帶來的開銷也在大模型時代被凸顯得尤為突出。與此同時，零階優(yōu)化器（Zeroth-Order Optimization）則完全無需保存計算圖，轉(zhuǎn)而使用有限差分來近似計算網(wǎng)絡(luò)的梯度，通過完全避免反向傳播（back-propagation; BP）來大大減少神經(jīng)網(wǎng)絡(luò)更新中的內(nèi)存開銷。
類似于一階優(yōu)化器中隨機梯度下降的各式變種，零階優(yōu)化器也有著各種此前無人探索的改進(jìn)算法。近日，來自密歇根州立大學(xué)、北卡羅來納大學(xué)教堂山分校、德克薩斯大學(xué)奧斯汀分校、明尼蘇達(dá)大學(xué)雙城分校、IBM 研究中心、普林斯頓大學(xué)、以及阿里巴巴達(dá)摩院的眾多研究者聯(lián)合推出全面評測（benchmark）文章：Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark 。這篇文章覆蓋六種無需反向傳播（BP-free）的優(yōu)化器、五類大模型、三種復(fù)雜度的各項任務(wù)、四類微調(diào)方案，以及三項增強零階優(yōu)化器的全新算法。目前，相關(guān)論文已被 ICML 2024 高分接收，代碼已開源；詳情如下。
- 論文地址：https://arxiv.org/abs/2402.11592
- 代碼地址：https://github.com/ZO-Bench/ZO-LLM
- 零階優(yōu)化講義地址 (AAAI 2024 Tutorial)：https://sites.google.com/view/zo-tutorial-aaai-2024/
零階優(yōu)化器是什么？為何如此重要？
零階優(yōu)化器（Zeroth-Order Optimization）僅僅依靠神經(jīng)網(wǎng)絡(luò)的輸出進(jìn)行梯度估計，以完全不需要計算反向傳播和極少的內(nèi)訓(xùn)消耗而聞名。盡管在零階優(yōu)化器領(lǐng)域也存在不同的梯度估計方法，本文特指基于隨機梯度估計器（Random Gradient Estimator, RGE）的一系列算法。簡單來說，就是通過從高斯分布中抽取的隨機擾動來計算有限差分，并將其作為梯度的近似估計，RGE 數(shù)學(xué)公式如下所示。
在此之前，零階優(yōu)化已經(jīng)大量應(yīng)用于機器學(xué)習(xí)問題中，比如對抗樣本生成和防御，黑箱模型解釋，強化學(xué)習(xí)和自動機器學(xué)習(xí)；詳細(xì)的算法和應(yīng)用介紹請參看 [1]。在大模型領(lǐng)域，MeZO [2] 率先提出將零階隨機梯度下降法（ZO-SGD）用作大模型微調(diào)，并展示了零階優(yōu)化器的無限潛力。于此同時，ZO-SGD 是最簡單、基本的 BP-free 優(yōu)化器，它的許多更高級的變種 [3] 能否在大模型微調(diào)領(lǐng)域帶給我們更多驚喜，是一個亟待研究的課題。本文系統(tǒng)地評測了以下無需反向傳播（BP-free）的優(yōu)化算法在大模型微調(diào)任務(wù)上的性能、效率和兼容性，目的是向社區(qū)展示零階優(yōu)化器在多種大模型任務(wù)上的廣泛潛力：
- ZO-SGD：零階隨機梯度下降 [4]
- ZO-SGD-Sign：基于符號的（sign-based）零階隨機梯度下降 [5]
- ZO-SGD-MMT：帶有動量（momentum）的零階隨機梯度下降 [6]
- ZO-SGD-Cons：保守（conservative）梯度更新的零階隨機梯度下降 [7]
- ZO-Adam：零階 Adam 優(yōu)化器 [8]
本研究還包括了 Forward-Grad [9] 方法，該方法基于沿隨機方向向量的方向?qū)?shù)來無偏估計梯度。值得注意的是，F(xiàn)orward-Grad 雖然不直接使用梯度回傳，但是卻仍然用到了自動微分算法，因此它是一種一階的 BP-free 算法。
綜上所述，本文的評測包含了上述 5 種零階優(yōu)化器以及 Forward-Grad 方法，同時對比一階優(yōu)化器中最常用的 FO-SGD 和 FO-Adam。在具體微調(diào)形式上，評測全面覆蓋了 5 種 LLM 架構(gòu)（RoBERTa, OPT, LLaMA, Vicuna, Mistral），3 種不同復(fù)雜度的任務(wù)（SST2, COPA, WinoGrande），以及 4 種微調(diào)方案（full-tuning, LoRA, prompt tuning, prefix tuning）。
大模型微調(diào)準(zhǔn)確性評測
作者指出，為了有效地利用零階優(yōu)化器對大型模型在下游任務(wù)上進(jìn)行微調(diào)，必須合理地運用輸入模板，以便將下游任務(wù)與預(yù)訓(xùn)練任務(wù)進(jìn)行對齊。例如對于 SST2 來說，使用模板 “<CLS>SENTENCE. It was [terrible|great].<SEP>” 能夠在 ZO-SGD 上帶來約 10% 的性能提升。然而對于一階優(yōu)化器（如 FO-SGD），無論是否使用模板性能差異都不大，襯托出零階優(yōu)化器的獨特性。
SST2 作為一個較為基礎(chǔ)的任務(wù)，其實驗結(jié)果可以支持以下結(jié)論：
- ZO-Adam 似乎是最有效的零階優(yōu)化器：在 8 個微調(diào)設(shè)置中的 4 個中表現(xiàn)最佳。
- Forward-grad 是一種競爭力強但以前被忽視的方法，特別是在全模型微調(diào) (full fine-tuning) 中。
- ZO-SGD-Cons 和 ZO-SGD-MMT 也展示了強大的性能，而 ZO-SGD-Sign作為最簡單的零階優(yōu)化器，往往是最弱的方法。
進(jìn)一步，該研究使用更大的模型 OPT-13B，在更復(fù)雜、更困難的任務(wù)（COPA 和 WinoGrande）上進(jìn)行實驗，得出以下結(jié)論：
- 在更復(fù)雜的任務(wù)中，不同優(yōu)化器的性能差異被進(jìn)一步放大。
- ZO-Adam 和 ZO-SGD-MMT 在各種實驗下展示了非常好的穩(wěn)定性，這可能歸因于減少了方差的設(shè)計。
- LoRA 微調(diào)對于零階算法始終表現(xiàn)出強大的魯棒性，在各種實驗環(huán)境中穩(wěn)定且可靠。
大模型微調(diào)內(nèi)存開銷評測與詳解
以 OPT-13B 模型在 MultiRC 數(shù)據(jù)集上微調(diào)為例，作者還進(jìn)一步對比分析了不同零階和一階優(yōu)化器的內(nèi)存與時間成本。如下表所示：首先，從內(nèi)存效率的角度看，ZO-SGD、ZO-SGD-Cons 和 ZO-SGD-Sign 顯示出了類似的高內(nèi)存效率，只需要一個 A100 GPU 來進(jìn)行大型語言模型的微調(diào)。這并不令人驚訝，因為這些零階優(yōu)化器采用相對簡單的優(yōu)化步驟，主要依賴于零階梯度估計器 RGE 的利用。其次，F(xiàn)orward-Grad 似乎是零階優(yōu)化方法在內(nèi)存效率方面超過一階方法的臨界點（例如與 ZO-Adam 相比）。最后，與一階方法相比，零階優(yōu)化每次迭代的運行時間成本降低了約 41.9%（以 ZO-SGD 與 FO-SGD 為例）。
作者進(jìn)一步深入比較了 ZO-SGD 與 FO-SGD 在不同序列長度下的內(nèi)存效率?？梢钥吹剑琙O-SGD 的內(nèi)存消耗保持一致，因為其峰值內(nèi)存消耗僅由模型參數(shù)大小決定，相比之下，隨著序列長度的增加，F(xiàn)O-SGD 的峰值內(nèi)存消耗先保持不變，然后開始增加。因此，在長上下文長度的設(shè)置中，ZO-SGD 將展示出更好的內(nèi)存效率優(yōu)勢。具體的內(nèi)存理論值和實驗值可參見原論文。
三種改進(jìn)算法用以增強零階優(yōu)化器
零階優(yōu)化器在應(yīng)用于 LLM 時收斂效率受限，主要是因為它們在梯度估計上的方差較大。為了進(jìn)一步增強零階優(yōu)化器，作者從降低梯度估計方差的角度入手，提出了三種進(jìn)階算法，包括：分塊零階微調(diào)（block-wise ZO fine-tuning）、零階和一階混合微調(diào)（hybrid ZO and FO fine-tuning）、引入稀疏性的零階梯度估計（sparsity-induced ZO gradient estimation）。
分塊零階微調(diào)（Block-wise ZO fine-tuning）此方法的主要出發(fā)點在于，如果零階優(yōu)化器在估計梯度時，對 LLM 中參數(shù)分塊分別進(jìn)行擾動，通過降低問題規(guī)模的方式來見效每次對梯度估計的方差，從而改進(jìn)優(yōu)化性能。這種方法的優(yōu)點體現(xiàn)在能夠?qū)δＰ吞荻冗M(jìn)行更準(zhǔn)確的估計，但是完成一次梯度估計所需要的前向傳播的次數(shù)會增加。例如，OPT-1.3B 可以分成 26 個參數(shù)塊（24 個 Transformers 層、嵌入層和 LM 分類頭），那么零階優(yōu)化器每次計算模型梯度時就會計算 26 次前向傳播。為了公平比較 ZO-SGD 和 ZO-SGD-Block，作者還比較了另一種 ZO-SGD 變體的性能，該變體每次對完整的模型進(jìn)行參數(shù)擾動，并將多次擾動后的梯度估計求平均（例如 OPT-1.3B 的 26 次），以此來保證比較時的前向傳播次數(shù)相同。OPT-1.3B 上實驗結(jié)果表明，ZO-SGD-Block 大幅超越了兩種 ZO-SGD。
零階和一階混合訓(xùn)練（Hybrid ZO and FO fine-tuning）反向傳播（BP）從深層至淺層神經(jīng)網(wǎng)絡(luò)依次計算權(quán)重梯度。由于零階優(yōu)化器在內(nèi)存占用上有遠(yuǎn)超傳統(tǒng)一階優(yōu)化器的優(yōu)勢，但一階優(yōu)化器的性能往往更好。因此，采用零階和一階優(yōu)化器的結(jié)合將達(dá)到一種內(nèi)存使用和性能之間的平衡（trade-off）。具體而言，對于較深層次網(wǎng)絡(luò)，可以利用一階優(yōu)化器通過反向傳播精確計算梯度；對于淺層網(wǎng)絡(luò)，則可以通過零階優(yōu)化器進(jìn)行梯度估算。實驗結(jié)果表明，在淺層部分（例如 OPT-1.3B 的前 8/24 層）采用零階優(yōu)化器，而在剩余的深層部分使用一階優(yōu)化器，可以在節(jié)約大約三分之一的顯存的同時，達(dá)到與完全使用一階優(yōu)化器相同的性能水平。
使用稀疏梯度的零階優(yōu)化器（ZO with gradient pruning）在一階優(yōu)化器中，梯度剪裁通常用于加速訓(xùn)練過程；而在零階優(yōu)化器中，通過梯度剪裁引入的稀疏梯度可以進(jìn)一步降低梯度估計的方差，從而提高性能。本文研究了在零階優(yōu)化器中應(yīng)用基于幅值的剪裁策略來獲取每一層的稀疏率，然后根據(jù)這些稀疏率生成隨機的稀疏梯度掩碼（mask），并將其應(yīng)用于隨機梯度估計的擾動上。實驗結(jié)果顯示，適度的梯度稀疏性（約 20% 左右）能給零階優(yōu)化器帶來一定程度的性能提升。
結(jié)語
在本文中，我們展示了零階優(yōu)化器在大型語言模型微調(diào)中的有效應(yīng)用。通過利用損失差分來近似梯度，零階優(yōu)化方法避免了反向傳播和激活存儲的需求，極大地節(jié)省了內(nèi)存資源。我們通過擴大已有的研究范圍，將不同的零階優(yōu)化方法、任務(wù)類型及評估指標(biāo)容納到了本次評測中，進(jìn)行了首次系統(tǒng)的零階優(yōu)化技術(shù)基準(zhǔn)研究。我們的研究不僅揭示了這些方法在精度和效率方面的表現(xiàn)，還深入探討了任務(wù)對齊和前向梯度的關(guān)鍵作用。利用這些實驗分析，我們提出了諸如分塊優(yōu)化、零階與一階混合訓(xùn)練、梯度稀疏化等技術(shù)，以進(jìn)一步增強基于零階優(yōu)化的大模型微調(diào)。這些改進(jìn)技術(shù)旨在在保持內(nèi)存效率的同時，提高微調(diào)的準(zhǔn)確性。
我們堅信，這些發(fā)現(xiàn)和技術(shù)的應(yīng)用可以大力降低大模型研究對硬件資源的要求，使得大模型微調(diào)在低端 GPU 也成為可能，從而進(jìn)一步推動學(xué)術(shù)研究并在工業(yè)界產(chǎn)生實際而有價值的影響。我們鼓勵廣大研究人員和技術(shù)開發(fā)者關(guān)注我們的研究成果，并探索更多利用 ZO 優(yōu)化的可能性。未來的研究將繼續(xù)探索這一領(lǐng)域的深層問題，以解鎖 LLM 微調(diào)中的更多潛力。
了解更多內(nèi)容請參考論文與 GitHub 倉庫，獲取更多信息和資源。
本文轉(zhuǎn)自 機器之心 ，作者：機器之心
