用神經架構搜索給LLM瘦身，模型變小，準確度有時反而更高
用神經架構搜索給 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。
總體而言，可以得出結論：對于某些任務來說，沒有必要使用完整的網絡大小。
而這些結果的差異又表明，并不存在一種適用于所有任務的單一架構啟發式方法。
