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