本篇工作已被NeurIPS（2024 Conference on Neural Information Processing Systems）會議接收，文章第一作者為美國埃默里大學的博士生白光霽，通訊作者為趙亮教授，來自埃默里大學計算機系。共同作者包括美國阿貢國家實驗室的Yijiang Li和Kibaek Kim，以及埃默里大學的博士生凌辰。該研究主要探討了大語言模型的全局剪枝方法，旨在提高預訓練語言模型的效率。該成果的發(fā)表為大模型的剪枝與優(yōu)化研究提供了新的視角，并在相關領域具有重要的應用潛力。

論文標題：SparseLLM: Towards Global Pruning of Pre-trained Language Models

論文鏈接：https://arxiv.org/abs/2402.17946

• 代碼：https://github.com/BaiTheBest/SparseLLM

情景導入

隨著大型語言模型（LLMs）如 GPT 和 LLaMA 在自然語言處理領域的突破，現(xiàn)如今的模型能夠在各種復雜的語言任務中表現(xiàn)優(yōu)異。然而，這些模型往往包含數(shù)十億參數(shù)，導致計算資源的需求極為龐大。為了讓LLMs在更多的實際應用中變得可行，研究人員進行了大量的模型壓縮工作，其中包括剪枝、量化、知識蒸餾和低秩分解等方法。

剪枝作為一種重要的壓縮策略，通過引入稀疏性提升了內(nèi)存和計算效率。盡管剪枝的有效性已在視覺任務和較小規(guī)模的語言模型中得到驗證，但全局剪枝由于需要將整個模型加載到同一個GPU中，對于如今的大規(guī)模LLMs而言是不現(xiàn)實的。因此，近年來出現(xiàn)了局部剪枝方法比如SparseGPT和Wanda，盡管這些方法簡單粗暴地局部化每個layer的剪枝，從而在效率上有了提升，但局部剪枝帶來的效果往往是次優(yōu) (suboptimal) 的。

我們提出的 SparseLLM 框架，通過將全局剪枝問題分解為更易管理的子問題，從而在高稀疏度下也能實現(xiàn)高效的優(yōu)化和優(yōu)異的性能。SparseLLM 在實現(xiàn)了內(nèi)存和計算效率的同時，超越了當前最先進的剪枝方法，為LLMs的壓縮和應用樹立了新的標桿。

Figure 1: 此圖展示了全局剪枝、局部剪枝與我們提出的 SparseLLM 框架的對比。全局剪枝（左）：由于需要將整個模型加載到同一GPU中，內(nèi)存消耗過大，難以適用于大規(guī)模LLMs。局部剪枝（中）：只考慮每一層的局部誤差，盡管內(nèi)存需求較低，但會導致次優(yōu)的全局性能。SparseLLM（右）：通過引入輔助輸入和輸出，將全局剪枝問題分解為可管理的子問題，實現(xiàn)低內(nèi)存開銷下的全局剪枝，保持全局最優(yōu)性能。

現(xiàn)有方法及其缺陷

在模型剪枝中，剪枝方法主要分為兩大類：全局剪枝（Global Pruning）和局部剪枝（Local Pruning）。

全局剪枝

全局剪枝旨在對整個模型應用統(tǒng)一的稀疏化掩碼（sparsity mask），以最小化未壓縮模型與壓縮模型之間的全局損失。雖然理論上全局剪枝可以為模型提供最優(yōu)的性能，尤其是在高稀疏度情況下，但它的主要缺點在于其內(nèi)存開銷巨大。為了實現(xiàn)全局剪枝，整個模型必須被加載到同一個GPU中，這對于如今規(guī)模巨大的LLMs（如GPT和LLaMA）來說是不現(xiàn)實的。

局部剪枝

為了規(guī)避全局剪枝的內(nèi)存瓶頸，局部剪枝通過將模型壓縮分解為每一層的子問題來減少內(nèi)存消耗。局部剪枝的方法通常會對每一層的輸出進行獨立的稀疏化，并構造局部損失來衡量未壓縮和壓縮層之間的差異。雖然局部剪枝方法在資源使用方面更加高效，但由于它只關注每一層的局部最優(yōu)解，導致了全局上次優(yōu)的模型性能，尤其是在高稀疏度下表現(xiàn)尤為顯著。

方法缺陷

局部剪枝的局限性在于它無法很好地捕捉模型層之間的相互依賴關系，只針對各層的稀疏化進行優(yōu)化。這種過度約束每層激活值的行為，容易導致全局性能的下降。因此，盡管局部剪枝在某些情況下能帶來效率提升，但它往往無法達到全局剪枝所能帶來的最優(yōu)性能。

技術方法

在SparseLLM框架中，我們的目標是實現(xiàn)全局剪枝。SparseLLM通過將全局剪枝目標分解為多個子問題，每個子問題可以使用較少的資源解決，并且可以協(xié)同實現(xiàn)全局剪枝目標。SparseLLM的優(yōu)勢在于它能夠在內(nèi)存消耗較低的情況下實現(xiàn)全局剪枝。

動機

SparseLLM的開發(fā)基于以下觀察：LLMs可以被表述為一個復合函數(shù)，后一個模塊的輸出是下一個模塊的輸入。這使得可以通過輔助變量將全局剪枝目標重新表述為等價形式，從而將其分解為多個子問題。然后，開發(fā)了一種高效的算法，通過交替優(yōu)化每個子問題來實現(xiàn)全局最優(yōu)解。

剪枝問題的重新表述

SparseLLM優(yōu)化過程

原始的全局剪枝問題包含全局約束條件，為了避免直接求解帶約束的目標函數(shù)，我們引入了輔助變量和懲罰函數(shù)。通過引入  懲罰項，原始的全局目標函數(shù)被轉化為無約束問題，具體表達如下：

此時，每一層的剪枝問題被表述為線性投影模塊（上投影和下投影）的優(yōu)化問題，并且可以針對每一個子問題進行交替優(yōu)化。這一過程通過引入輔助變量和懲罰函數(shù)，使得我們能夠在不消耗過多內(nèi)存的前提下，求解全局最優(yōu)解。

Figure 2: 該圖展示了 SparseLLM 框架在不同LLM架構中的工作原理。左側描繪了在 OPT模型中的剪枝流程。SparseLLM通過引入輔助變量，將全局剪枝問題分解為可管理的子問題，利用線性層的上下投影模塊（Up Proj 和 Down Proj）以及ReLU激活函數(shù)來實現(xiàn)模型壓縮。右側描繪了在 LLaMA模型 中的剪枝流程。此處，SparseLLM除了使用線性層的上下投影外，還利用了SiLU激活函數(shù)與Gate Proj模塊來實現(xiàn)更復雜的剪枝操作。圖中標注了剪枝層（藍色）、輔助變量（綠色）和預訓練模型的固定輸入/輸出（粉色），以幫助理解各模塊在剪枝過程中的相互關系與作用。

OPT 模型中的子問題求解

在 OPT 模型中，我們的目標是對每一層的FFN模塊進行全局剪枝。根據(jù)公式5，我們首先考慮交替優(yōu)化子問題，對于每一個子問題的閉式解，我們分以下步驟進行：

LLaMA 模型中的子問題求解

在LLaMA模型中，我們的剪枝過程與OPT模型類似，唯一不同的是LLaMA模型中包含額外的門控投影層，并使用了SiLU激活函數(shù)。我們針對LLaMA的各層子問題進行如下求解：

由于SiLU函數(shù)的非凸特性，問題可以通過查找預先計算好的解決方案的查找表來快速求解。

通過上述過程，SparseLLM能夠在LLaMA和OPT兩種模型架構中實現(xiàn)高效的全局剪枝，最大化壓縮效果的同時，保持模型的全局性能。

實驗部分

實驗設置

為了驗證 SparseLLM 框架的有效性，我們在多個大規(guī)模語言模型上進行了廣泛的實驗。我們使用了 PyTorch 框架并結合 HuggingFace 的 Transformers 庫來處理模型和數(shù)據(jù)集。所有剪枝實驗均在 NVIDIA A100 GPU 上執(zhí)行，以確保實驗結果具備較好的可擴展性和適用性。

實驗的數(shù)據(jù)集選取了 WikiText2 (WT2)、PTB（Penn Treebank），以及 C4 的驗證集。這些數(shù)據(jù)集是語言模型壓縮任務中的標準基準集，確保了結果的廣泛適用性。為了保持公平性，我們在所有實驗中使用了相同的數(shù)據(jù)子集，并確保剪枝過程為零樣本設置，即不引入額外的訓練數(shù)據(jù)，僅依賴于模型的預訓練權重進行剪枝。

我們采用了模型的 困惑度（Perplexity） 作為評估指標，這是一種常用的衡量語言模型生成能力的指標。困惑度越低，模型的表現(xiàn)越好。所有剪枝后的模型均通過困惑度指標評估其在不同數(shù)據(jù)集上的性能，比較剪枝前后的性能變化以及不同稀疏度下的性能差異。

模型與剪枝策略

實驗中，我們選擇了兩類主流的預訓練語言模型進行剪枝實驗：

• OPT 模型：我們針對不同規(guī)模的 OPT 模型（如 OPT-1.3b、OPT-2.7b 等）進行全局剪枝，評估 SparseLLM 在不同大小的模型上的效果。
• LLaMA 模型：我們也對 LLaMA 系列模型進行了實驗，特別是 LLaMA-2 13b 模型，考察其在使用 SparseLLM 框架下的剪枝效果。

對于每一個模型，我們使用了不同的稀疏度設定，從 70% 到 90% 的非結構化稀疏度，以及 3:4 的半結構化稀疏度。在每個稀疏度水平下，我們測試了模型在各個基準數(shù)據(jù)集上的性能表現(xiàn)，以驗證 SparseLLM 在不同稀疏度條件下的剪枝效果。

比較基線方法

為了全面評估 SparseLLM 的效果，我們與現(xiàn)有的幾種剪枝方法進行了對比：

• 局部幅度剪枝：一種傳統(tǒng)的剪枝方法，逐層根據(jù)權重幅度進行獨立的局部剪枝，不考慮層間的依賴性。
• SparseGPT：一種最新的局部剪枝方法，專門針對預訓練語言模型進行高效的權重剪枝，能夠在高稀疏度條件下維持較好的性能。
• Wanda：另一種基于局部剪枝的方法，通過最小化剪枝后的模型輸出誤差，達到較好的壓縮效果。

我們分別在不同稀疏度條件下與上述方法進行了對比，詳細記錄了各模型在不同數(shù)據(jù)集上的困惑度結果。

結果分析

剪枝效果分析

在不同稀疏度條件下，SparseLLM 在大多數(shù)情況下能夠顯著超過基線方法，特別是在高稀疏度（≥ 70%）的情況下，SparseLLM 能夠維持較低的困惑度，甚至在一些任務中超越了未剪枝的模型。相比之下，局部幅度剪枝方法和 SparseGPT 的性能在稀疏度較高時下降明顯，證明了全局剪枝方法在模型壓縮中的優(yōu)勢。

此外，我們還分析了模型剪枝后的參數(shù)分布情況。SparseLLM 的剪枝策略能夠有效地保持模型重要參數(shù)的位置不變，減少了剪枝過程中重要信息的丟失，這也解釋了其在高稀疏度下性能優(yōu)越的原因。

模型大小與性能權衡

我們進一步分析了不同大小模型在使用 SparseLLM 剪枝時的表現(xiàn)。實驗表明，較大規(guī)模的模型在剪枝后的性能更為穩(wěn)定，尤其是 LLaMA 13b 等大規(guī)模模型，即使在90%稀疏度下，困惑度的增幅也非常有限。這一結果表明，隨著模型規(guī)模的增加，SparseLLM 能夠更有效地在剪枝中找到最優(yōu)的權重保留策略，降低剪枝對模型性能的負面影響。

Figure 3: Fast convergence of SparseLLM. Training loss per epoch for pruning layer 3 of OPT-125m at 80% sparsity (Left) and layer 6 of LlaMA-2 13b at 70% sparsity (Right).

收斂速度分析

除了剪枝性能，我們還對 SparseLLM 的收斂速度進行了分析。在實際訓練過程中，SparseLLM 在最初的幾個 epoch 內(nèi)即可快速達到較低的訓練損失，表明其剪枝過程非常高效。相比之下，其他局部剪枝方法在稀疏度較高時，需要更多的訓練步驟才能收斂到類似的性能水平。這進一步證明了全局剪枝策略在保持模型性能的同時，能夠有效減少訓練時間和資源消耗。

不同任務的通用性

為了驗證 SparseLLM 在零樣本設置下的通用性，我們在多個零樣本任務上測試了剪枝后的模型性能。實驗結果表明，SparseLLM 在大多數(shù)任務上保持了優(yōu)異的性能，特別是在涉及推理和生成任務時，SparseLLM 剪枝后的模型仍然能夠生成高質(zhì)量的文本輸出，而其他基線方法的性能在高稀疏度下顯著下降。

實驗小結

實驗結果表明，SparseLLM 框架能夠在不同規(guī)模的預訓練語言模型上實現(xiàn)高效的全局剪枝，同時保持良好的模型性能。無論是在較小的 OPT 模型上，還是在更大規(guī)模的 LLaMA 模型上，SparseLLM 均表現(xiàn)出色，特別是在高稀疏度的條件下表現(xiàn)尤為突出。此外，SparseLLM 的收斂速度和剪枝后的通用性也為其在實際應用中的高效性和適用性提供了強有力的支持。

結論

在本研究中，我們提出了 SparseLLM，一個針對大規(guī)模預訓練語言模型的全局剪枝框架。通過引入輔助變量和局部子問題，我們成功解決了全局剪枝在計算和內(nèi)存方面的瓶頸。實驗結果表明，SparseLLM 能夠在保持模型性能的同時，在高稀疏度下有效地壓縮模型。與現(xiàn)有的剪枝方法相比，SparseLLM 在困惑度、稀疏度和通用性方面展現(xiàn)出了顯著的優(yōu)勢，特別是在處理大型模型（如 LLaMA 和 OPT 系列）時，SparseLLM 的全局剪枝策略展現(xiàn)了更好的性能保持能力。

通過本研究，我們證明了全局剪枝在大規(guī)模語言模型中的潛力，不僅能夠顯著降低模型的計算和存儲成本，還能在高稀疏度下維持優(yōu)異的性能。這為實際應用中對資源敏感的任務提供了一個有效的解決方案。

未來研究方向

雖然 SparseLLM 取得了優(yōu)異的結果，但仍有若干值得深入研究的方向：

1. 動態(tài)剪枝策略的研究：當前的 SparseLLM 是基于靜態(tài)的全局剪枝策略，即在訓練前決定剪枝的層和稀疏度。未來可以探索動態(tài)剪枝策略，結合任務需求和模型訓練過程，實時調(diào)整剪枝比例和剪枝區(qū)域，以進一步提升模型壓縮效率。
2. 稀疏性與硬件加速的結合：盡管 SparseLLM 能夠大幅減少模型的計算和存儲開銷，但在實際部署中如何高效利用硬件加速器（如GPU、TPU）的稀疏性支持仍是一個開放問題。未來的工作可以研究如何更好地結合稀疏矩陣乘法和硬件加速技術，以實現(xiàn)更大的性能提升。
3. 剪枝后的模型微調(diào)：盡管我們的實驗已經(jīng)展示了 SparseLLM 在高稀疏度下的性能保持能力，但剪枝后的模型是否可以通過微調(diào)進一步提升性能仍然值得探索。未來的研究可以結合剪枝和微調(diào)技術，探討如何在稀疏化后進一步優(yōu)化模型性能。

通過探索這些方向，SparseLLM 有望在大規(guī)模語言模型的壓縮與高效推理中發(fā)揮更加廣泛的作用，為各類NLP應用提供更多的支持與創(chuàng)新。