用神經架構搜尋給 LLM 瘦身，同等準確度可讓 LLaMA2-7B 模型大小降 2 倍。
大型語言模型（LLM）的一個主要特點是「大」，也因此其訓練和部署成本都相當高，如何在保證 LLM 準確度的同時讓其變小就成了非常重要且有價值的研究課題。
對此，研究社群已經提出了多種多樣的方法，比如剪枝、稀疏化、量化等，它們的效果也各不一樣。
近日，Intel Labs 釋出了一項研究成果，宣稱可使用神經架構搜尋（NAS）高效地為 LLM 「瘦身」。他們基於 LLaMA2-7B 模型的實驗表明，該技術不僅能降低模型大小，有時甚至還能讓模型的準確度獲得提升！
論文標題：LLaMA-NAS: Efficient Neural Architecture Search for Large Language Models
論文地址：https://arxiv.org/pdf/2405.18377
神經架構搜尋（NAS）是一種讓演算法根據樣本集自動設計神經網路架構的技術。之前曾被用來探索各式的新型神經網路架構，甚至可能發現人類從未構建過的網路架構，比如谷歌就曾透過 NAS 發現了一個新的特徵金字塔架構，參閱《谷歌大腦提出 NAS-FPN：這是一種學會自動架構搜尋的特徵金字塔網路》
Intel Labs 的這個研究團隊猜想：也許可以透過神經架構搜尋（NAS）來降低模型的大小和複雜性。單次 NAS 方法往往需要一次訓練一個超級網路（super-network）並允許使用權重共享原理進行子網路取樣。但是，從頭開始訓練這些超級網路成本高昂，因此 LLM 領域少有探索。InstaTune 等技術則是利用現成的預訓練權重，在微調階段創造一個超級網路。
為了緩解與 LLM 相關的大小和複雜度問題，該團隊做了一番研究探索：
該團隊首次高效地使用單次 NAS 來降低 LLM 的大小和計算複雜度。他們的研究表明：對於某些標準的基準評測任務而言，LLaMA2-7B 龐大且複雜得沒有必要。
他們發現，一種使用單次 NAS 的 LLM 壓縮方法不僅效能優於剪枝和稀疏化技術，而且還無需剪枝和稀疏化通常需要的額外恢復微調步驟。
他們分析了網路引數，以便更好地理解新方法發現的網路的架構特徵。他們證明，沒有一套單一的架構啟發式方法可以應用於多個標準基準測試任務。
他們提出了一種框架，無需專門的軟體和或硬體就能直接用於產生壓縮版 LLM。他們還發現，使用不做任何修改的標準版量化技術就能進一步壓縮這些網路。
方法
搜尋方法
為了讓大型語言模型（LLM）適應多樣化的硬體平臺和效能要求，該團隊的選擇是使用 InstaTune 方法，這是一種全新的單次神經架構搜尋（NAS）正規化，參閱論文《Instatune: Instantaneous neural architecture search during fine-tuning》。
InstaTune 對傳統的微調過程進行了擴充套件，其做法是讓模型的架構具有彈性，使其可以探索固定結構之外更寬廣的設計空間。這種靈活性讓模型可以適應不同的任務、資料集和計算資源。InstaTune 增強 NAS 的方式是利用已有的模型權重，從而無需再從頭開始訓練超級網路。透過將 NAS 過程嵌入到微調階段，InstaTune 能在節省計算資源的同時確保派生出的子網路針對目標任務進行了專門最佳化，而不是針對預訓練目標進行了最佳化。
在這項研究中，研究團隊使用 InstaTune 中提出的技術和 Alpaca 資料集對預訓練後的 LLaMA2-7B 模型進行了微調，得到了一個超級網路和一個搜尋空間 —— 它們可用於搜尋過程。
但是，不同於 InstaTune，該團隊並未在 LLaMA2-7B 上執行基於強大教師模型或超級網路的知識蒸餾，原因則主要是計算和記憶體的限制。
此外，雖然 InstaTune 會在同一次迭代中為超級網路和隨機取樣的子網路同時計算損失，但該團隊為了降低記憶體使用量，選擇的做法是交替地為超級網路和隨機取樣的子網路計算損失。在這個具體案例中，超級網路就是基礎 LLaMA2-7B 模型，沒有任何額外的層，MLP 模組的中間層大小也沒有擴充套件。
遵循 InstaTune 的方法，該團隊也採用了一種進化搜尋框架來在多目標設定中對架構進行最佳化，從而在給定下游任務上對模型大小和準確度進行最佳化。
具體來說，他們採用的演算法是輕量迭代式神經架構搜尋（LINAS）演算法。LINAS 組合了 NSGA-II 搜尋和網路效能預測器，可以高效地找到帕累托最優的網路配置。它能在真實資料上迭代式地評估子網路，從而為效能預測器提供資訊，然後它們就能預測大量子網路的效能。之後，將最有潛力的子網路選出來，再在下一輪迭代中進行評估。這個過程持續進行，直到在真實資料上完成一定的評估次數。
搜尋空間
如前所述，執行最優子網路搜尋的演算法是 LINAS。使用 InstaTune 在 Alpaca 資料集上對預訓練 LLaMA2-7B 模型進行微調之後，該團隊定義了一組允許引數值，以便在搜尋最優子網路時使用。這個搜尋空間的建立方式是改變整個網路的層數以及每個 MLP 模組的中間大小。
表 1 給出了這些引數值，最終得到的搜尋空間大小約為 1.3 × 10^10。
評估實驗和結果
任務
該團隊選用了以下任務來評估新方法的表現：
AI2 推理挑戰（ARC）
大規模多工語言理解（MMLU）
TruthfulQA
WinoGrande
搜尋分析
如前所述，該團隊使用了 LINAS 演算法來搜尋 LLaMA2-7B 的最優子網路，使用的搜尋空間見表 1。每個任務的搜尋過程都執行 250 次評估。
AI2 推理挑戰
針對 AI2 推理挑戰（ARC）任務，是在模型大小 / ARC-c 準確度和模型大小 / ARC-e 準確度目標空間中執行的搜尋。圖 1 展示了在這兩個目標空間中的帕累託邊界。
執行完搜尋後，該團隊找到了一些能提供更高準確度和 / 或更小規模的子網路架構。比如某些子網路的 ARC-c 準確度比預訓練的 LLaMA2-7B 更高，如圖中紅點所示。舉例來說：一個子網路的準確度與預訓練 LLaMA2-7B 網路一樣，都規模卻小 1.1 倍；另一個子網路的大小一樣，但準確度高 1.9%。
大規模多工語言理解
該團隊也使用表 1 中的引數搜尋了在大規模多工語言理解（MMLU）任務上的帕累托最優子網路。該搜尋的目標空間是模型大小 / MMLU 準確度。此外，為了瞭解在推理速度方面的增益，該團隊還評估了搜尋得到的帕累托最優子網路的吞吐量。
圖 2 展示了在模型大小 / MMLU 準確度以及吞吐量 / MMLU 準確度這兩個目標空間中的帕累託邊界。
可以看到，在這些帕累托最優子網路中，某些架構在模型大小和吞吐量這兩方面都有優勢。比如某些子網路的 MMLU 準確度比預訓練的 LLaMA2-7B 更高，如圖中紅點所示。在一個案例中，子網路的準確度比預訓練 LLaMA2-7B 高 1.1%，同時又比它小 1.5 倍，快 1.3 倍。
TruthfulQA
對於 TruthfulQA 任務，該團隊發現對於 TruthfulQA MC1，預訓練 LLaMA2-7B 明顯是過度引數化了。如圖 3 所示，新方法發現的許多網路架構都明顯優於預訓練 LLaMA2-7B，同時大小也小得多。
WinoGrande
在 WinoGrande 上得到的結論與在 ARC-c 上得到的結果很相近。圖 4 給出了在模型大小 / WinoGrande 搜尋空間中的帕累託邊界。
搜尋找到了一些有更高準確度和 / 或更小 guim 的子網路架構。舉例來說：一個子網路的準確度與預訓練 LLaMA2-7B 網路一樣，都規模卻小 1.1 倍；另一個子網路的大小一樣，但準確度高 1.4%。
標準基準效能總結
表 2 給出了各任務的一些比較突出的子網路，並總結了它們的標準基準效能。
剪枝和稀疏化效能比較
在透過新方法找到帕累托最優子網路之後，該團隊選出了那些規模同等或更小的子網路（它們是透過剪枝或稀疏化得到的），並評估了它們在不同任務上的效能。然後，該團隊比較了它們與 LLM-Pruner 和 SliceGPT 的標準基準效能，結果見表 3 和 4。
在表 3 中，ARC-c 和 MMLU 的某些子網路的準確度高於新方法找到的更大子網路。
原因是網路架構中不同中間大小的位置可能不會影響模型大小，但會影響準確度。因此，較小的子網路也可能取得更高的準確度，因為它選擇了更好的中間大小。對於新方法找到的大小相近的子網路，在每個任務上的準確度均高於 LLM-Pruner 和 SliceGPT。
量化
為了進一步提升新方法找到的帕累托最優子網路的效能，該團隊使用了定點（INT8）量化，使用的工具是 bitsandbytes。具體來說，首先量化搜尋找到的帕累托最優子網路，然後在模型大小 / 準確度目標空間中再次評估它們。解碼器層中的所有線性運算都被量化到了 INT8，而剩餘層（比如嵌入層）依然是 FP16。
圖 5 展示了量化前後在模型大小 / 準確度目標空間中的帕累託邊界。
在為帕累托最優子網路使用了定點（INT8）量化之後，在目標空間中的帕累託邊界「向左移了」。很顯然，量化對準確度的影響很小，但卻能顯著降低模型大小。事實上，取得了高 1.1% MMLU 準確度且模型大小降低了 1.5 倍的非量化子網路在量化後不僅維持著原來的準確度，而且大小還降低了 2.5 倍。
表 5 比較了表 2 中非量化版子網路與其定點（INT8）量化後的版本。
可以看到，對於 MMLU，相比於非量化版的預訓練 LLaMA2-7B 網路，量化版子網路的準確度幾乎沒有損失，而模型大小卻分別降低了 2.0 和 2.5 倍。對於 ARC-e，量化版網路的準確度還比非量化版更高一點。而對於 WinoGrande，量化版子網路的大小不僅分別降低了 1.8 和 2.0 倍，同時還取得了 0.2% 的少量準確度增益。
搜尋空間分析
該團隊最後分析了層數和中間大小方面的性質。在層數方面，圖 6 展示了在不同的任務上，一個被選擇的網路對於不同的準確度百分位數具有層數 l ∈ {24, 28, 32} 的機率。
可以看到，層數越多，準確度往往就越高。但是，對於 ARC-c、MMLU 和 TruthfulQA 等某些任務而言，情況卻並非如此，這說明 LLaMA2-7B 在這些任務有過度引數化的現象。
在中間大小方面，圖 7 給出了在 ARC-c 任務上，選擇出的所有 32 層網路架構，一個特定層的中間大小為 s ∈ {5504, 11008} 的機率。
可以看到，取得高 ARC-c 準確度的網路只在某些層有較高機率選擇 5504 的中間大小。
具體來說，在 3、4、18、25、26、28、29 和 32 層，中間大小 s = 5504 的機率明顯高於 s = 11008。評估 3、4、18、25、26、28、29 和 32 層的中間大小為 s = 5504 且其它層的中間大小為 s = 11008 的網路架構，得到的 ARC-c 準確度為 45.9%，與預訓練 LLaMA2-7B 的準確度相當，從而驗證了圖 7 的結果。
另外，該團隊還分析了在 MMLU 任務上評估的 24 層網路架構選擇以及在 WinoGrande 任務上評估的 32 層網路架構選擇，結果見圖 8 和 9。
總體而言，可以得出結論：對於某些任務來說，沒有必要使用完整的網路大小。
而這些結果的差異又表明，並不存在一種適用於所有任務的單一架構啟發式方法。