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 高分接收，程式碼已開源；詳情如下。
- 論文地址：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 數學公式如下所示。
在此之前，零階最佳化已經大量應用於機器學習問題中，比如對抗樣本生成和防禦，黑箱模型解釋，強化學習和自動機器學習；詳細的演算法和應用介紹請參看 [1]。在大模型領域，MeZO [2] 率先提出將零階隨機梯度下降法（ZO-SGD）用作大模型微調，並展示了零階最佳化器的無限潛力。於此同時，ZO-SGD 是最簡單、基本的 BP-free 最佳化器，它的許多更高階的變種 [3] 能否在大模型微調領域帶給我們更多驚喜，是一個亟待研究的課題。本文系統地評測了以下無需反向傳播（BP-free）的最佳化演算法在大模型微調任務上的效能、效率和相容性，目的是向社群展示零階最佳化器在多種大模型任務上的廣泛潛力：- ZO-SGD-Sign：基於符號的（sign-based）零階隨機梯度下降 [5]
- ZO-SGD-MMT：帶有動量（momentum）的零階隨機梯度下降 [6]
- ZO-SGD-Cons：保守（conservative）梯度更新的零階隨機梯度下降 [7]
本研究還包括了 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 來說，使用模板 “<CLS>SENTENCE. It was [terrible|great].<SEP>” 能夠在 ZO-SGD 上帶來約 10% 的效能提升。然而對於一階最佳化器（如 FO-SGD），無論是否使用模板效能差異都不大，襯托出零階最佳化器的獨特性。SST2 作為一個較為基礎的任務，其實驗結果可以支援以下結論：- ZO-Adam 似乎是最有效的零階最佳化器：在 8 個微調設定中的 4 箇中表現最佳。
- Forward-grad 是一種競爭力強但以前被忽視的方法，特別是在全模型微調 (full fine-tuning) 中。
- ZO-SGD-Cons 和 ZO-SGD-MMT 也展示了強大的效能，而 ZO-SGD-Sign作為最簡單的零階最佳化器，往往是最弱的方法。
進一步，該研究使用更大的模型 OPT-13B，在更復雜、更困難的任務（COPA 和 WinoGrande）上進行實驗，得出以下結論：- 在更復雜的任務中，不同最佳化器的效能差異被進一步放大。
- ZO-Adam 和 ZO-SGD-MMT 在各種實驗下展示了非常好的穩定性，這可能歸因於減少了方差的設計。
- LoRA 微調對於零階演算法始終表現出強大的魯棒性，在各種實驗環境中穩定且可靠。
以 OPT-13B 模型在 MultiRC 資料集上微調為例，作者還進一步對比分析了不同零階和一階最佳化器的記憶體與時間成本。如下表所示：首先，從記憶體效率的角度看，ZO-SGD、ZO-SGD-Cons 和 ZO-SGD-Sign 顯示出了類似的高記憶體效率，只需要一個 A100 GPU 來進行大型語言模型的微調。這並不令人驚訝，因為這些零階最佳化器採用相對簡單的最佳化步驟，主要依賴於零階梯度估計器 RGE 的利用。其次，Forward-Grad 似乎是零階最佳化方法在記憶體效率方面超過一階方法的臨界點（例如與 ZO-Adam 相比）。最後，與一階方法相比，零階最佳化每次迭代的執行時間成本降低了約 41.9%（以 ZO-SGD 與 FO-SGD 為例）。
作者進一步深入比較了 ZO-SGD 與 FO-SGD 在不同序列長度下的記憶體效率。可以看到，ZO-SGD 的記憶體消耗保持一致，因為其峰值記憶體消耗僅由模型引數大小決定，相比之下，隨著序列長度的增加，FO-SGD 的峰值記憶體消耗先保持不變，然後開始增加。因此，在長上下文長度的設定中，ZO-SGD 將展示出更好的記憶體效率優勢。具體的記憶體理論值和實驗值可參見原論文。
零階最佳化器在應用於 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。
零階和一階混合訓練（Hybrid ZO and FO fine-tuning）反向傳播（BP）從深層至淺層神經網路依次計算權重梯度。由於零階最佳化器在記憶體佔用上有遠超傳統一階最佳化器的優勢，但一階最佳化器的效能往往更好。因此，採用零階和一階最佳化器的結合將達到一種記憶體使用和效能之間的平衡（trade-off）。具體而言，對於較深層次網路，可以利用一階最佳化器透過反向傳播精確計算梯度；對於淺層網路，則可以透過零階最佳化器進行梯度估算。實驗結果表明，在淺層部分（例如 OPT-1.3B 的前 8/24 層）採用零階最佳化器，而在剩餘的深層部分使用一階最佳化器，可以在節約大約三分之一的視訊記憶體的同時，達到與完全使用一階最佳化器相同的效能水平。
使用稀疏梯度的零階最佳化器（ZO with gradient pruning）在一階最佳化器中，梯度剪裁通常用於加速訓練過程；而在零階最佳化器中，透過梯度剪裁引入的稀疏梯度可以進一步降低梯度估計的方差，從而提高效能。本文研究了在零階最佳化器中應用基於幅值的剪裁策略來獲取每一層的稀疏率，然後根據這些稀疏率生成隨機的稀疏梯度掩碼（mask），並將其應用於隨機梯度估計的擾動上。實驗結果顯示，適度的梯度稀疏性（約 20% 左右）能給零階最佳化器帶來一定程度的效能提升。
在本文中，我們展示了零階最佳化器在大型語言模型微調中的有效應用。透過利用損失差分來近似梯度，零階最佳化方法避免了反向傳播和啟用儲存的需求，極大地節省了記憶體資源。我們透過擴大已有的研究範圍，將不同的零階最佳化方法、任務型別及評估指標容納到了本次評測中，進行了首次系統的零階最佳化技術基準研究。我們的研究不僅揭示了這些方法在精度和效率方面的表現，還深入探討了任務對齊和前向梯度的關鍵作用。利用這些實驗分析，我們提出了諸如分塊最佳化、零階與一階混合訓練、梯度稀疏化等技術，以進一步增強基於零階最佳化的大模型微調。這些改進技術旨在在保持記憶體效率的同時，提高微調的準確性。我們堅信，這些發現和技術的應用可以大力降低大模型研究對硬體資源的要求，使得大模型微調在低端 GPU 也成為可能，從而進一步推動學術研究並在工業界產生實際而有價值的影響。我們鼓勵廣大研究人員和技術開發者關注我們的研究成果，並探索更多利用 ZO 最佳化的可能性。未來的研究將繼續探索這一領域的深層問題，以解鎖 LLM 微調中的更多潛力。瞭解更多內容請參考論文與 GitHub 倉庫，獲取更多資訊和資源。[1] Liu, et al,. "A primer on zeroth-order optimization in signal processing and machine learning." IEEE Signal Processing Magazine 37, no. 5 (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 Zeroth-order Methods for Nonconvex Stochastic Programming.” [5] Liu, et al., “signSGD via Zeroth-Order Oracle. ” ICLR’ 2019.[6] Huang, et al., “Accelerated Zeroth-Order and First-Order Momentum Methods from Mini to Minimax Optimization.” JMLR’ 2022. [7] Kim, et al., “Curvature-Aware Derivative-Free Optimization.”[8] Chen, et al., “ZO-AdaMM: Zeroth-Order Adaptive Momentum Method for Black-Box Optimization.”[9] Baydin, et al., “Gradients without Backpropagation.”