AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
- 論文標題：SparseLLM: Towards Global Pruning of Pre-trained Language Models
- 論文連結：https://arxiv.org/abs/2402.17946
- 程式碼：https://github.com/BaiTheBest/SparseLLM
情景匯入
隨著大型語言模型（LLMs）如 GPT 和 LLaMA 在自然語言處理領域的突破，現如今的模型能夠在各種複雜的語言任務中表現優異。然而，這些模型往往包含數十億引數，導致計算資源的需求極為龐大。為了讓LLMs在更多的實際應用中變得可行，研究人員進行了大量的模型壓縮工作，其中包括剪枝、量化、知識蒸餾和低秩分解等方法。
現有方法及其缺陷
在模型剪枝中，剪枝方法主要分為兩大類：全域性剪枝（Global Pruning）和區域性剪枝（Local Pruning）。
全域性剪枝
全域性剪枝旨在對整個模型應用統一的稀疏化掩碼（sparsity mask），以最小化未壓縮模型與壓縮模型之間的全域性損失。雖然理論上全域性剪枝可以為模型提供最優的效能，尤其是在高稀疏度情況下，但它的主要缺點在於其記憶體開銷巨大。為了實現全域性剪枝，整個模型必須被載入到同一個GPU中，這對於如今規模巨大的LLMs（如GPT和LLaMA）來說是不現實的。
區域性剪枝
為了規避全域性剪枝的記憶體瓶頸，區域性剪枝透過將模型壓縮分解為每一層的子問題來減少記憶體消耗。區域性剪枝的方法通常會對每一層的輸出進行獨立的稀疏化，並構造區域性損失來衡量未壓縮和壓縮層之間的差異。雖然區域性剪枝方法在資源使用方面更加高效，但由於它只關注每一層的區域性最優解，導致了全域性上次優的模型效能，尤其是在高稀疏度下表現尤為顯著。
方法缺陷
區域性剪枝的侷限性在於它無法很好地捕捉模型層之間的相互依賴關係，只針對各層的稀疏化進行最佳化。這種過度約束每層啟用值的行為，容易導致全域性效能的下降。因此，儘管區域性剪枝在某些情況下能帶來效率提升，但它往往無法達到全域性剪枝所能帶來的最優效能。
技術方法
在SparseLLM框架中，我們的目標是實現全域性剪枝。SparseLLM透過將全域性剪枝目標分解為多個子問題，每個子問題可以使用較少的資源解決，並且可以協同實現全域性剪枝目標。SparseLLM的優勢在於它能夠在記憶體消耗較低的情況下實現全域性剪枝。
動機
SparseLLM的開發基於以下觀察：LLMs可以被表述為一個複合函式，後一個模組的輸出是下一個模組的輸入。這使得可以透過輔助變數將全域性剪枝目標重新表述為等價形式，從而將其分解為多個子問題。然後，開發了一種高效的演算法，透過交替最佳化每個子問題來實現全域性最優解。
剪枝問題的重新表述
表示模型的總層數，
表示需要剪枝的層集合，
和
分別代表原始預訓練模型的中間變數值。 是模型的最終輸出結果 。Remark：
時，解決該問題等同於全域性剪枝。而當
時，問題簡化為區域性剪枝，獨立地考慮每一層。透過調節
的大小，我們可以在全域性和區域性剪枝之間無縫切換。SparseLLM最佳化過程
懲罰項，原始的全域性目標函式被轉化為無約束問題，具體表達如下：
和
是超引數。這樣，我們將每一層的剪枝問題解耦為區域性子問題，同時保留了層之間的依賴性。
設定為
，其中
表示每個解碼器層FFN模組中線性層的索引集（具體數學細節詳見論文）。也就是說，同一個FFN模組中的線性層被全域性剪枝，而剩下的線性層仍按照區域性剪枝。對於每一層的FFN模組，我們的目標是最小化以下無約束的目標函式：
此時，每一層的剪枝問題被表述為線性投影模組（上投影和下投影）的最佳化問題，並且可以針對每一個子問題進行交替最佳化。這一過程透過引入輔助變數和懲罰函式，使得我們能夠在不消耗過多記憶體的前提下，求解全域性最優解。
OPT 模型中的子問題求解
剪枝權重的最佳化：首先，我們固定其他變數，最佳化權重 
和 
。具體來說，最佳化目標是最小化以下表示式：
為了求解該問題，我們首先對 
進行分解：
，其中 
（ 表示偽逆）。透過將分解後的 帶入原始損失函式中，可以得到權重剪枝問題的閉式解，類似於 SparseGPT 等剪枝方法。啟用值的更新：對於啟用值 
的最佳化，這是一個類似於權重剪枝的簡單最小二乘問題。啟用值 
的更新可以透過如下公式求解：
輸出 的更新：對於輸出 
，我們最小化以下非凸的損失函式：
由於ReLU函式的逐元素操作，問題可以透過簡單的if-then邏輯分為兩種情況： 
LLaMA 模型中的子問題求解
剪枝權重的最佳化：對於LLaMA模型中的FFN模組，權重剪枝的最佳化過程與OPT模型類似。我們透過最小化以下表示式來最佳化權重： 
並使用類似於OPT模型的分解方法求解。 啟用值 的更新：在LLaMA模型中，我們使用SiLU啟用函式，因此啟用值的更新公式如下： 
輸出 的更新：對於LLaMA模型中的輸出 
，我們最小化以下損失函式：
透過類似的數學操作，我們可以求得 
的解析解為：
門控投影輸出 
的更新：對於門控投影層的 ，我們透過最小化以下表示式來更新：
由於SiLU函式的非凸特性，問題可以透過查詢預先計算好的解決方案的查詢表來快速求解。
實驗部分
實驗設定
模型與剪枝策略
OPT 模型：我們針對不同規模的 OPT 模型（如 OPT-1.3b、OPT-2.7b 等）進行全域性剪枝，評估 SparseLLM 在不同大小的模型上的效果。 LLaMA 模型：我們也對 LLaMA 系列模型進行了實驗，特別是 LLaMA-2 13b 模型，考察其在使用 SparseLLM 框架下的剪枝效果。
比較基線方法
區域性幅度剪枝：一種傳統的剪枝方法，逐層根據權重幅度進行獨立的區域性剪枝，不考慮層間的依賴性。 SparseGPT：一種最新的區域性剪枝方法，專門針對預訓練語言模型進行高效的權重剪枝，能夠在高稀疏度條件下維持較好的效能。 Wanda：另一種基於區域性剪枝的方法，透過最小化剪枝後的模型輸出誤差，達到較好的壓縮效果。
結果分析
剪枝效果分析
模型大小與效能權衡
收斂速度分析
不同任務的通用性
為了驗證 SparseLLM 在零樣本設定下的通用性，我們在多個零樣本任務上測試了剪枝後的模型效能。實驗結果表明，SparseLLM 在大多數任務上保持了優異的效能，特別是在涉及推理和生成任務時，SparseLLM 剪枝後的模型仍然能夠生成高質量的文字輸出，而其他基線方法的效能在高稀疏度下顯著下降。
實驗小結
實驗結果表明，SparseLLM 框架能夠在不同規模的預訓練語言模型上實現高效的全域性剪枝，同時保持良好的模型效能。無論是在較小的 OPT 模型上，還是在更大規模的 LLaMA 模型上，SparseLLM 均表現出色，特別是在高稀疏度的條件下表現尤為突出。此外，SparseLLM 的收斂速度和剪枝後的通用性也為其在實際應用中的高效性和適用性提供了強有力的支援。
結論
未來研究方向
動態剪枝策略的研究：當前的 SparseLLM 是基於靜態的全域性剪枝策略，即在訓練前決定剪枝的層和稀疏度。未來可以探索動態剪枝策略，結合任務需求和模型訓練過程，實時調整剪枝比例和剪枝區域，以進一步提升模型壓縮效率。 稀疏性與硬體加速的結合：儘管 SparseLLM 能夠大幅減少模型的計算和儲存開銷，但在實際部署中如何高效利用硬體加速器（如GPU、TPU）的稀疏性支援仍是一個開放問題。未來的工作可以研究如何更好地結合稀疏矩陣乘法和硬體加速技術，以實現更大的效能提升。 剪枝後的模型微調：儘管我們的實驗已經展示了 SparseLLM 在高稀疏度下的效能保持能力，但剪枝後的模型是否可以透過微調進一步提升效能仍然值得探索。未來的研究可以結合剪枝和微調技術，探討如何在稀疏化後進一步最佳化模型效能。