SLICEGPT: COMPRESS LARGE LANGUAGE MODELSBY DELETING ROWS AND COLUMNS

發表：ICLR24

機構：ETH Zurich

連接：https://arxiv.org/pdf/2401.15024

ABSTRACT

大型語言模型（Large Language Models, LLMs）已成為自然語言處理的基石，但其使用伴隨著在計算和內存資源方面的高昂代價。稀疏化（sparsification）提供了一種緩解這些資源限制的解決方案，近期的研究表明，訓練好的模型可以在訓練后進行稀疏化。然而，現有的稀疏化技術面臨一些挑戰：它們通常需要額外的數據結構，并且在現有硬件上只能帶來有限的加速效果。

在本文中，我們提出了一種新的訓練后稀疏化方法 SliceGPT，它用一個更小的（稠密）矩陣替換每一個權重矩陣，從而降低了網絡的嵌入維度。通過大量實驗證明，SliceGPT 可在不顯著損害性能的前提下，去除多達 25% 的模型參數（包括嵌入部分），適用于 LLAMA-2 70B、OPT 66B 和 Phi-2 模型，在零樣本任務中分別保留了 99%、99% 和 90% 的原始性能。

我們切片后的模型可在更少的 GPU 上運行，且無需任何額外的代碼優化即可運行更快：在 24GB 的消費級 GPU 上，我們將 LLAMA-2 70B 模型的推理總計算量降低為稠密模型的 64%；在 40GB 的 A100 GPU 上則降低至 66%。

我們還提出了一個新的洞察：Transformer 網絡中的計算不變性（computational invariance）。該性質使 SliceGPT 成為可能，并希望這一理念能激發并推動未來在減少預訓練模型內存與計算需求方向的研究。

代碼地址：
🔗 https://github.com/microsoft/TransformerCompression

1 INTRODUCTION

大型語言模型（Large Language Models, LLMs）是包含數十億參數的神經網絡，在數萬億 token 上進行訓練（Zhao et al., 2023）。LLM 的訓練成本極高，因此已出現將同一個預訓練模型用于多個任務的趨勢，即“基礎模型范式（foundation model paradigm）”。然而，LLM 的體積龐大，使得部署一個預訓練模型成為一項代價高昂的工作。許多模型需要多個 GPU 才能完成一次預測，而且由于模型是自回歸的，生成文本響應時需要多次執行神經網絡的前向傳播。因此，降低這些模型的計算需求成為廣泛關注的課題，通常通過訓練后技術來實現，這類技術統稱為模型壓縮（model compression）。

大多數模型壓縮方法可歸為以下四類之一：蒸餾（distillation）、張量分解（tensor decomposition，包含低秩分解）、剪枝（pruning）和量化（quantization）（Hoefler et al., 2021；Gholami et al., 2021；Zhu et al., 2023；Gupta & Agrawal, 2021）。在本研究中，我們專注于剪枝，但也希望我們的方法論能對其他方向的研究有所啟發。

雖然剪枝方法已提出多年，但許多方法在剪枝之后需要進行“恢復微調（recovery fine-tuning, RFT）”以維持模型性能，這使得整個流程成本高昂、難以擴展。而我們的 SliceGPT 方法能夠使用單張 GPU 在短短幾個小時內對大型模型進行壓縮，且即使不進行 RFT，也能在生成任務和下游任務上保持較強的性能。

剪枝方法的原理是將 LLM 中部分權重矩陣元素設為零，并（可選地）更新矩陣中其余元素以做出補償。這種操作產生稀疏模式，使得在神經網絡前向傳播中進行矩陣乘法時可以跳過某些浮點操作。計算加速的程度取決于稀疏程度及其模式：結構化稀疏通常帶來更大的加速。

與其他剪枝方法不同，SliceGPT 剪掉的是整個權重矩陣的行或列（即“切片”掉）。在切片之前，我們對網絡進行一次變換，這種變換不會改變預測結果，但能使切片的影響最小化。

最終的結果是，權重矩陣變得更小，block 之間傳遞的信號也變得更小：我們減小了神經網絡的嵌入維度。

圖 1 展示了我們的方法與現有稀疏化方法的比較。我們的貢獻如下：

我們引入了“計算不變性（computational invariance）”的概念：我們證明了可以對 Transformer 中的每個權重矩陣施加正交矩陣變換，而不會改變模型行為。
我們使用這種不變性來修改 Transformer 架構中的每個 block，使 block 之間傳遞的信號矩陣被投影到其主成分上。我們刪除變換后權重矩陣的部分列或行以減小模型規模。我們將這種變換和權重刪除的過程稱為 SliceGPT。
我們在 OPT（Zhang et al., 2022）和 LLAMA-2（Touvron et al., 2023）模型上進行了大量實驗證明， SliceGPT 可以將這些模型壓縮最多達 30%，在 WikiText-2 上的困惑度（perplexity）優于現有 2:4 稀疏化方案。在下游任務中，我們還在 Phi-2 模型上進行了實驗，顯示所有模型最多可被切片 30%，同時保留超過 90% 的稠密模型性能。

2 BACKGROUND

在本節中，我們首先描述一些關于 Transformer 架構的必要背景知識，這使我們能夠引入將在證明主要結果中使用的符號。然后我們將描述用于壓縮此類架構的稀疏化相關工作。

2.1 TRANSFORMER NETWORKS

2.2?RELATED WORK

在最簡單的設置中，可以使用基于幅值的稀疏化（magnitude-based sparsification），其方法是將模型中幅值最小的權重置為零（Han et al., 2016；Zhu & Gupta, 2017；Gale et al., 2019）。雖然這種方法具有良好的可擴展性，但應用于 LLM 時，會導致性能嚴重下降（Frantar & Alistarh, 2023）。

一種更為復雜的方法是 Optimal Brain Surgeon (OBS)（Hassibi et al., 1993；LeCun et al., 1989），它以系統化的方式移除對損失函數影響最小的權重。該方法通過使用 Hessian 矩陣的逆，對未被剪除的權重進行調整，從而補償權重移除帶來的誤差。不幸的是，對于擁有幾百萬參數的模型而言，OBS 并不實用，因為它需要計算并存儲 Hessian 矩陣的逆。

為了解決 OBS 所帶來的計算瓶頸，近期研究提出了兩個方向：

一是對 Hessian 矩陣的逆進行近似，例如 WoodFisher（Singh & Alistarh, 2020）；
二是對每一層分別應用該方法，例如 Optimal Brain Compression (OBC)（Frantar & Alistarh, 2022），也稱為逐層剪枝（layer-wise pruning）。

雖然這些技術在中等規模網絡上表現良好，但并不適用于 LLM，因為單層的權重矩陣通常包含超過 10810^8108 個參數。

GPTQ（Frantar et al., 2022）通過一種列按列（column-by-column）的策略，對 LLM 的權重矩陣進行量化（即以更低精度表示參數），并在下一列更新尚未量化的權重，從而解決了上述問題。

SparseGPT（Frantar & Alistarh, 2023）將這種思想應用于剪枝，對 LLM 進行非結構化（unstructured）和半結構化（semi-structured）稀疏化。Sun 等人（2023）則簡化了該思路，僅使用了 Hessian 的對角線部分。由于通過非結構化剪枝實現端到端速度提升是一項困難任務，他們還嘗試使用如 2:4、4:8（Mishra et al., 2021）等半結構化模式來誘導稀疏性。然而，采用這些結構會損害模型的準確率。

另一種壓縮策略是低秩近似（low-rank approximation），即將每個權重矩陣替換為兩個內部維度更小的矩陣之積，通常還需后續微調（Hu et al., 2021；Mahabadi et al., 2021；Noach & Goldberg, 2020；Tukan et al., 2020）。為了實現壓縮，內部維度必須小于原始維度的一半。相比之下，我們的方法將每個權重矩陣直接替換為一個更小的單個矩陣，通過減小嵌入維度達到目的，且無需微調。
我們提出的方法是刪除權重矩陣的行和列，這與卷積網絡文獻中**濾波器和通道的剪枝（filter/channel pruning）**相似。在該方向上，研究通常在 BatchNorm 參數（Liu et al., 2017）或網絡結構（Huang & Wang, 2018）上加入誘導稀疏性的正則化，然后訓練或微調，從而實現通道或結構的剪除。

與我們方法最相似的是 ThiNet（Luo et al., 2017；He et al., 2017），其通過在層之間引入線性操作（正如我們也將采用的），并結合更多的微調與正則化來實現壓縮。在這類研究中，模型通常比 LLM 小幾個數量級，例如 VGG16 網絡僅有 1.38 億個參數，差不多相當于我們考慮的最小 OPT 模型。
而 LLM 的巨大參數規模使得那些依賴大量微調的方案變得不現實，尤其是需要外部循環來選擇正則化參數的情況。

近期已有一些研究嘗試對 LLM 進行結構化剪枝（structured pruning），然后繼續訓練（或微調）以恢復性能。例如 LLM-pruner（Ma et al., 2023a）在進一步訓練之前剪除了 LLM 中的連接結構；與我們工作同時進行的 LLM Surgeon（van der Ouderaa et al., 2023）則將恢復微調與剪枝過程交替進行。

我們在本文中提供了 SliceGPT 作為一種單次剪枝（single-shot）方法，并同時展示了結合**剪枝后微調（post-slicing recovery fine-tuning）**的實驗結果。

3 SLICEGPT

我們的方法 SliceGPT 依賴于 transformer 架構中固有的一種計算不變性（computational invariance）。我們指的是：只要在某個組件的輸出中應用一個正交變換（orthogonal transformation），并在下一個組件中“撤銷”該變換，整個網絡的功能保持不變。我們關鍵的洞察是，在網絡的各個 block 之間執行的 RMSNorm 操作并不會影響這種變換：這兩個操作是可交換的（commute）。在本節中，我們首先描述這種不變性如何在 RMSNorm 連接的 transformer 網絡中成立，然后說明如何將使用 LayerNorm 的網絡轉換為 RMSNorm。接著，我們介紹如何使用主成分分析（PCA）計算每一層的變換矩陣，使得 block 之間傳遞的信號被投影到其主成分上。最后，我們描述如何刪除較弱的主成分，從而對應于在修改后的網絡中切除行或列。

3.1 COMPUTATIONAL INVARIANCE IN TRANSFORMER NETWORKS

3.2 LAYER NORM TRANSFORMERS CAN BE CONVERTED TO RMSNORM

3.3 A TRANSFORMATION PER BLOCK

3.4 SLICING

4 EXPERIMENTAL VALIDATION

Setup 我們使用 Hugging Face Transformers（Wolf et al., 2019）與 PyTorch（Paszke et al., 2019）實現了我們的代碼。正交變換矩陣 Q的計算是在一塊 80GB 顯存的 H100 GPU 上完成的，以 LLaMA-2 70B 模型為例，計算過程大約耗時 3.5 小時。

在進行 PCA 的過程中，我們使用**雙精度（double precision）**來計算協方差矩陣的特征向量。我們發現，若在 PyTorch 中使用單精度計算特征向量，最終模型的精度會出現偏差，詳細分析見附錄 A.2。

我們使用兩個不同的校準數據集進行了實驗：WikiText-2 訓練集（Merity et al., 2016）和 Alpaca 訓練集（Taori et al., 2023）。關于校準集大小與序列長度的消融實驗，詳見附錄 A.3。

Models, Tasks, and GPUs 我們在多個模型族上驗證了 SliceGPT 的效果，包括 OPT（Zhang et al., 2022）、LLaMA-2（Touvron et al., 2023），并在零樣本任務中額外評估了 Phi-2 模型。

我們排除了 OPT 175B 模型，因為它在性能上已經被更小的 LLaMA-2 模型超越。然而我們預計，這個更大的模型在 SliceGPT 上仍然能取得更好的壓縮效果，因為更大的模型通常有更多的冗余，壓縮潛力更大（詳見第 4.1 節）。

我們在以下兩類任務中評估我們的方案：

語言生成任務
熱門的零樣本（zero-shot）任務

為全面展示 SliceGPT 在推理速度上的提升效果，我們在不同類型的 GPU 上進行了實驗：

Quadro RTX6000（24GB 顯存）：代表消費級 GPU；
A100（40GB）與 H100（80GB）：代表數據中心級別 GPU。

Baseline Setup 我們最初計劃將 SliceGPT 與一種基線方案進行對比：該方案基于最小范數剪除列（或行）。然而我們發現，這種基線效果極差——在僅剪除少數列后，模型在 WikiText-2 上的困惑度（perplexity）飆升至上千。

因此，我們改為將 SliceGPT 與 SparseGPT（Frantar & Alistarh, 2023）進行比較，后者采用了 2:4 的稀疏化比率，這是目前唯一已知的能在實際硬件上帶來加速的稀疏化方案（Mishra et al., 2021）。

4.1?RESULTS

生成任務（Generation Task）

我們首先使用 WikiText-2 數據集展示我們的發現。在該實驗中，我們評估了 OPT 與 LLaMA-2 模型系列在不同模型規模下，在該數據集上進行剪枝的性能。表 1 展示了不同剪枝比例下模型的困惑度（perplexity）。

在應用于 OPT 模型時，SliceGPT 展現出優于 LLaMA-2 模型的性能，這與我們對這些模型譜分析的直覺一致（詳見附錄 A.4）。SliceGPT 的性能隨著模型規模的增大而提升。

將 SliceGPT 與 SparseGPT 相比，我們發現：在所有 LLaMA-2 模型中，SliceGPT 在剪除 25% 參數的情況下，其性能優于 SparseGPT 采用 2:4 稀疏模式的結果。對于 OPT 模型，除 2.7B 外，其余所有模型在 30% 剪枝比例下的表現也優于 SparseGPT。

零樣本任務（Zero-shot Tasks）

我們在五個知名的零樣本任務上評估 SliceGPT，包括：

PIQA（Bisk et al., 2020）
WinoGrande（Sakaguchi et al., 2021）
HellaSwag（Zellers et al., 2019）
ARC-e 和 ARC-c（Clark et al., 2018）

遵循相關工作（Frantar & Alistarh, 2023；Dettmers et al., 2022；Frantar et al., 2022；Dettmers et al., 2023），我們使用 LM Evaluation Harness（Gao et al., 2021）進行評估，采用其默認參數。

圖 5 展示了剪枝模型在這些任務中的平均得分：

上方圖為使用 WikiText-2 校準數據的得分；
下方圖為使用 Alpaca 校準數據的得分。

我們觀察到的模式與生成任務類似：

OPT 模型更容易壓縮，且準確率下降更小；
模型越大，壓縮后的性能下降越小。

在該實驗中我們也納入了 Phi-2 模型：我們發現，剪枝后的 Phi-2 模型在性能上與剪枝后的 LLaMA-2 7B 模型相當。OPT 與 LLaMA-2 的最大規模模型均可以被有效壓縮，即使剪除 30% 參數，也僅帶來幾個百分點的性能損失。

恢復微調實驗（Recovery Fine-Tuning）

我們在此還進行了恢復微調（Recovery Fine-Tuning, RFT）實驗：

我們使用 LoRA（Hu et al., 2021） 對剪枝后的 LLaMA-2 與 Phi-2 模型進行少量微調，方法參照 Ma et al. (2023a)。
對于使用 WikiText-2 剪枝的模型，我們使用約 1000 條序列進行微調；
對于使用 Alpaca 數據剪枝的模型，我們使用 5000 條序列進行微調。

所有 RFT 使用如下參數：

lora_r = 32，
lora_alpha = 10，
序列長度 = 1024，
其余超參數采用 Hugging Face PEFT 包（Mangrulkar et al., 2022）默認設置。

圖 6 展示了微調結果。我們觀察到，使用 Alpaca 數據集進行 RFT 的效果顯著優于使用 WikiText-2。這一差異可歸因于 Alpaca 與評估任務的相似性。

對于被剪除 30% 的最大 LLaMA-2 70B 模型，使用 Alpaca 做 RFT 后，其平均準確率為 74.3%，而原始稠密模型為 76.6%；
該剪枝模型的參數量約為 51.6B，推理吞吐量也有顯著提升（詳見后文）；
Phi-2 模型若僅用 WikiText-2 做 RFT 無法恢復準確率，但使用 Alpaca 可恢復若干百分點；
- 剪除 25% 后的 Phi-2 平均準確率為 65.2%，而稠密模型為 72.2%；
- 剪枝模型約含 2.2B 參數，保留了約 90.3% 的原始模型性能；

這表明，即使是小型語言模型，也能從訓練后剪枝中獲益。各任務的準確率表格見附錄 A.5。

吞吐量基準測試（Benchmarking Throughput）

與傳統稀疏方法只對 Win 與 Wout引入稀疏性不同，SliceGPT 還對中間特征 X引入了結構化稀疏性：直接剪除 X的整列，減小了嵌入維度，從而提高了壓縮模型的 FLOPs 效率和內存移動效率。

我們將剪枝比例為 25% 和 50% 的模型在 80GB H100 GPU 上與稠密模型進行了吞吐量對比。設定序列長度為 128，持續加倍 batch size，直到 GPU 顯存耗盡或吞吐量下降為止。

結果表明：

25% 剪枝模型的 token 吞吐量提升最多可達 1.55 倍；
50% 剪枝后，最大模型只需 1 張 GPU 就可運行（原需 2 張），吞吐量提升巨大：
- OPT-66B 提升為 3.13×
- LLaMA2-70B 提升為 1.87×
固定 GPU 數量下，相當于達到了：
- 6.26×（3.13× × 2）
- 3.75×（1.87× × 2） 的吞吐率增益。

雖然 SliceGPT 在 50% 剪枝下的 WikiText-2 perplexity 差于 SparseGPT（2:4），但在吞吐量上遠高于任何不剪中間特征 XXX 的稀疏方法。

對于 13B 模型，batch size 帶來的吞吐提升不明顯，因為模型本身占用顯存不多。但對于消費級 GPU（顯存小），小模型的吞吐量也會因剪枝而受益。詳見附錄 A.6。

推理時間（Inference Time）

我們接著評估了 SliceGPT 壓縮模型的端到端運行時間。

表 2 比較了 OPT-66B 與 LLaMA2-70B 模型在 Quadro RTX6000 與 A100 GPU 上生成單個 token 所需時間：

在 RTX6000 上，剪除 25% 后的模型推理時間縮短了 16~17%；
在 A100 上，推理時間減少了 11~13%；
同時也減少了所需 GPU 數量，降低了能耗與部署成本。

例如，LLaMA2-70B 在 RTX6000 上，推理一個 token 所需計算由 1764 GPU-ms 降至 1075 GPU-ms（64%）。

我們將這種提升歸因于：

替換為更小權重矩陣；
全程使用稠密計算內核（dense kernel），這是其他剪枝方法做不到的。

目前，SparseGPT（2:4）還無法實現端到端推理加速。因此，我們改為比較 transformer layer 內各操作的相對執行時間。我們發現：對于大模型，SliceGPT（25%）在推理速度與困惑度方面可與 SparseGPT（2:4）競爭甚至超越。詳見附錄 A.7。

計算成本（Compute cost）

所有 LLaMA-2、OPT 和 Phi-2 模型都可以在 單張 GPU 上在 1 到 3 小時內完成剪枝。

如果加上恢復微調，總時間也僅需 1 到 5 小時。詳見表 3。

5 CONCLUSION AND FUTURE WORK

我們提出了 SliceGPT，一種支持大型語言模型結構化剪枝的新方法。
該方法使我們無需任何額外代碼優化，即可將 LLaMA-2 70B 模型在 40GB A100 GPU 上的推理成本降低到稠密模型的 66%，所需 GPU 數量也從 4 塊減少到 3 塊，同時其保留集（held-out set）困惑度優于 SparseGPT 的 2:4 稀疏方案。
在 24GB RTX6000 GPU 上，推理成本被降至稠密模型的 64%，所需 GPU 數量從 7 塊減少到 5 塊。

在零樣本下游任務中，分別對 OPT-66B、LLaMA-2-70B 和 Phi-2 進行 25% 剪枝，其性能可分別保留稠密模型的 99%、96% 和 87%；
在進行恢復微調后（recovery fine-tuning），這兩個模型的表現進一步提高，LLaMA-2-70B 和 Phi-2 可分別恢復至 99% 和 90% 的稠密性能。

我們的方法仍有進一步發展的空間：

當前，在同等參數規模下，小型但稠密的語言模型往往優于被剪到同樣大小的大模型，但我們認為這種情況不會持續太久；
雖然 SliceGPT 剪枝后的模型參數量通常多于 SparseGPT，但我們的方法能使更大的 batch size 載入 GPU 顯存，且無需稀疏結構支持的額外開銷；
因此，一種可能的改進方向是將兩種方法結合，發揮各自優勢；
此外，還可以嘗試其他方式計算變換矩陣 QQQ，進一步提升剪枝效果；
若想繼續減少推理時間和 GPU 數量，還可以結合量化（quantization）（Xiao et al., 2023；Dettmers et al., 2022；Ashkboos et al., 2023；Frantar et al., 2022）以及其他結構化剪枝方法（如 Ma et al., 2023b）。