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

論文標(biāo)題：SparseLLM: Towards Global Pruning of Pre-trained Language Models

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

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

情景導(dǎo)入

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

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

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

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

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

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

全局剪枝

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

局部剪枝

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

方法缺陷

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

技術(shù)方法

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

動機(jī)

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

剪枝問題的重新表述

SparseLLM優(yōu)化過程

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

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

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

OPT 模型中的子問題求解

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

LLaMA 模型中的子問題求解

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

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

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

實(shí)驗(yàn)部分

實(shí)驗(yàn)設(shè)置

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

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

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

模型與剪枝策略

實(shí)驗(yàn)中，我們選擇了兩類主流的預(yù)訓(xùn)練語言模型進(jìn)行剪枝實(shí)驗(yàn)：

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

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

比較基線方法

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

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

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

結(jié)果分析

剪枝效果分析

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

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

模型大小與性能權(quán)衡

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

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

不同任務(wù)的通用性

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

實(shí)驗(yàn)小結(jié)

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

結(jié)論

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

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

未來研究方向

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

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

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