手把手教你剪「羊駝」，陳丹琦團(tuán)隊(duì)提出LLM-Shearing大模型剪枝法

　　本文來(lái)自于微信公眾號(hào) 機(jī)器之心(ID:almosthuman2014)，作者:機(jī)器之心。

　　給 Llama2(羊駝)大模型剪一剪駝毛，會(huì)有怎樣的效果呢?今天普林斯頓大學(xué)陳丹琦團(tuán)隊(duì)提出了一種名為L(zhǎng)LM-Shearing的大模型剪枝法，可以用很小的計(jì)算量和成本實(shí)現(xiàn)優(yōu)于同等規(guī)模模型的性能。

　　自大型語(yǔ)言模型(LLM)出現(xiàn)以來(lái)，它們便在各種自然語(yǔ)言任務(wù)上取得了顯著的效果。不過(guò)，大型語(yǔ)言模型需要海量的計(jì)算資源來(lái)訓(xùn)練。因此，業(yè)界對(duì)構(gòu)建同樣強(qiáng)大的中型規(guī)模模型越來(lái)越感興趣，出現(xiàn)了 LLaMA、MPT 和 Falcon，實(shí)現(xiàn)了高效的推理和微調(diào)。

　　這些規(guī)模不等的 LLM 適用于不同的用例，但從頭開(kāi)始訓(xùn)練每個(gè)單獨(dú)的模型(即使是10億參數(shù)小模型)還是需要大量計(jì)算資源，這對(duì)于大多數(shù)科研機(jī)構(gòu)而言仍是很大的負(fù)擔(dān)。

　　因此在本文中，普林斯頓大學(xué)陳丹琦團(tuán)隊(duì)試圖解決以下問(wèn)題:能否利用現(xiàn)有預(yù)訓(xùn)練 LLM 來(lái)構(gòu)建一個(gè)規(guī)模更小、通用且在性能上有競(jìng)爭(zhēng)力的 LLM，同時(shí)比從頭開(kāi)始訓(xùn)練需要的計(jì)算量少得多?

　　研究者探索利用結(jié)構(gòu)化剪枝來(lái)實(shí)現(xiàn)目標(biāo)。這里的問(wèn)題是，對(duì)于通用 LLM，剪枝后的模型會(huì)出現(xiàn)性能下降，尤其是在剪枝后沒(méi)有大量計(jì)算投入的情況。他們使用的高效剪枝方法可以用來(lái)開(kāi)發(fā)規(guī)模更小但仍具有性能競(jìng)爭(zhēng)力的 LLM，并且與從頭開(kāi)始訓(xùn)練相比，訓(xùn)練需要的計(jì)算量也大大減少。

　　在對(duì) LLM 進(jìn)行剪枝之前，研究者確定了兩個(gè)關(guān)鍵技術(shù)挑戰(zhàn)，一是如何確定最終的性能強(qiáng)大、推理高效的剪枝結(jié)構(gòu)?LLM 目前的結(jié)構(gòu)化剪枝技術(shù)沒(méi)有指定的目標(biāo)結(jié)構(gòu)，導(dǎo)致剪枝后模型在性能和推理速度方面不理想;二是如何繼續(xù)預(yù)訓(xùn)練剪枝后的模型以達(dá)到預(yù)期性能?他們觀察到，與從頭開(kāi)始訓(xùn)練模型相比，使用原始預(yù)訓(xùn)練數(shù)據(jù)來(lái)訓(xùn)練會(huì)導(dǎo)致不同域出現(xiàn)不同的損失減少。

　　針對(duì)這兩個(gè)挑戰(zhàn)，研究者提出了「LLM - shearing」算法。這種新穎的剪枝算法被稱(chēng)為「定向結(jié)構(gòu)化剪枝」，它將源模型剪枝為指定的目標(biāo)架構(gòu)，該結(jié)構(gòu)通過(guò)現(xiàn)有預(yù)訓(xùn)練模型的配置來(lái)確定。他們表示，該剪枝方法在源模型中搜索子結(jié)構(gòu)，并在資源受限的情況下最大程度地保持性能。此外設(shè)計(jì)一種動(dòng)態(tài)批量加載算法，它能根據(jù)損失減少率按比例加載每個(gè)域的訓(xùn)練數(shù)據(jù)，從而高效利用數(shù)據(jù)并加速整體性能的提升。

　　最終，研究者將 LLaMA2-7B 模型剪枝成了兩個(gè)較小的 LLM，分別是 Sheared-LLaMA-1.3B 和 Sheared-LLaMA-2.7B，證實(shí)了其方法的有效性。

　　他們僅僅使用500億個(gè) token(即 OpenLLaMA 預(yù)訓(xùn)練預(yù)算的5%)進(jìn)行剪枝和繼續(xù)預(yù)訓(xùn)練，但對(duì)于11個(gè)代表性下游任務(wù)(如常識(shí)、閱讀理解和世界知識(shí))以及開(kāi)放式生成的指令調(diào)整，這兩個(gè)模型的性能仍然優(yōu)于其他同等規(guī)模的流行 LLM，包括 Pythia、INCITE 和 OpenLLaMA。

　　不過(guò)要提到一點(diǎn)，在這篇論文發(fā)布 Sheared-LLaMA-3B 的時(shí)候，最強(qiáng)3B 開(kāi)源模型的紀(jì)錄已經(jīng)被 StableLM-3B 打破了。

　　此外，下游任務(wù)性能軌跡表明，使用更多 token 來(lái)進(jìn)一步訓(xùn)練剪枝后的模型，將帶來(lái)更大的收益。研究者只對(duì)最多70億參數(shù)的模型進(jìn)行了實(shí)驗(yàn)，但 LLM-shearing 具有高度通用性，可以在未來(lái)的工作中擴(kuò)展到任何規(guī)模的大型語(yǔ)言模型。

　　方法介紹

　　給定一個(gè)現(xiàn)有的大模型 M_S(源模型)，本文目標(biāo)是研究如何有效地生成一個(gè)更小、更強(qiáng)的模型 M_T(目標(biāo)模型)。該研究認(rèn)為這需要兩個(gè)階段來(lái)完成:

　　第一階段將 M_S 剪枝為 M_T，雖然這樣減少了參數(shù)數(shù)量，但不可避免地導(dǎo)致性能下降;

　　第二階段持續(xù)預(yù)訓(xùn)練 M_T，使其性能更強(qiáng)。

　　結(jié)構(gòu)化剪枝

　　結(jié)構(gòu)化剪枝可以去除模型大量參數(shù)，從而達(dá)到壓縮模型并加速推理的效果。然而，現(xiàn)有的結(jié)構(gòu)化剪枝方法會(huì)導(dǎo)致模型偏離常規(guī)架構(gòu)的配置。例如 CoFiPruning 方法產(chǎn)生的模型具有不統(tǒng)一的層配置，與標(biāo)準(zhǔn)的統(tǒng)一層配置相比，這樣會(huì)產(chǎn)生額外的推理開(kāi)銷(xiāo)。

　　本文對(duì) CoFiPruning 進(jìn)行了擴(kuò)展，以允許將源模型剪枝為指定的任何目標(biāo)配置。例如，本文在生成2.7B 模型時(shí)使用 INCITE-Base-3B 架構(gòu)作為目標(biāo)結(jié)構(gòu)。

　　此外，本文還在不同粒度的模型參數(shù)上學(xué)習(xí)一組剪枝掩碼( pruning mask)，掩碼變量如下所示:

　　每個(gè)掩碼變量控制是否剪枝或保留相關(guān)的子結(jié)構(gòu)。例如，如果對(duì)應(yīng)的 z^layer=0，則需要?jiǎng)h除這個(gè)層。下圖2說(shuō)明了剪枝掩碼如何控制被剪枝的結(jié)構(gòu)。

　　剪枝之后，本文通過(guò)保留與每個(gè)子結(jié)構(gòu)中的掩碼變量相關(guān)的最高得分組件來(lái)最終確定剪枝后的架構(gòu)，并繼續(xù)使用語(yǔ)言建模目標(biāo)對(duì)剪枝后的模型進(jìn)行預(yù)訓(xùn)練。

　　動(dòng)態(tài)批量加載

　　該研究認(rèn)為對(duì)剪枝后的模型進(jìn)行大量預(yù)訓(xùn)練是很有必要的，這樣才能恢復(fù)模型性能。

　　受其他研究的啟發(fā)，本文提出了一種更有效的算法，即動(dòng)態(tài)批量加載，其可以根據(jù)模型性能簡(jiǎn)單地動(dòng)態(tài)調(diào)整域比例。算法如下:

　　實(shí)驗(yàn)及結(jié)果

　　模型配置:本文將 LLaMA2-7B 模型作為源模型，然后進(jìn)行結(jié)構(gòu)化剪枝實(shí)驗(yàn)，他們將 LLaMA2-7B 壓縮成兩個(gè)較小的目標(biāo)尺寸2.7B 和1.3B 參數(shù)，并將剪之后的模型與相同尺寸的模型進(jìn)行了性能比較，包括 OPT-1.3B、Pythia-1.4B、OPT-2.7B、 Pythia-2.8B、INCITE-Base-3B、OpenLLaMA-3B-v1、OpenLLaMA-3B-v2。表8總結(jié)了所有這些模型的模型體系結(jié)構(gòu)細(xì)節(jié)。

　　數(shù)據(jù):由于 LLaMA2的訓(xùn)練數(shù)據(jù)并不是公開(kāi)訪問(wèn)的，因此本文使用了 RedPajama 數(shù)據(jù)集 。表1提供了本文模型和基線模型使用的預(yù)訓(xùn)練數(shù)據(jù)。

　　訓(xùn)練:研究者在所有實(shí)驗(yàn)中最多使用了16個(gè) Nvidia A100GPU (80GB)。

　　SHEARED-LLAMA 優(yōu)于同等大小的 LM

　　本文表明，Sheared-LLaMA 明顯優(yōu)于現(xiàn)有的類(lèi)似規(guī)模的 LLM，同時(shí)只使用一小部分計(jì)算預(yù)算來(lái)從頭開(kāi)始訓(xùn)練這些模型。

　　下游任務(wù):表2展示了 Sheared-LLaMA 和類(lèi)似大小的現(xiàn)有預(yù)訓(xùn)練模型的零樣本和少樣本在下游任務(wù)上的性能。

　　指令調(diào)優(yōu):如圖3所示，與同等規(guī)模的所有其他預(yù)訓(xùn)練模型相比，指令調(diào)優(yōu)的 Sheared-LLaMA 實(shí)現(xiàn)了更高的獲勝率。

　　圖4顯示了 INCITEBase-3B 模型開(kāi)始時(shí)的精度要高得多，但其性能在持續(xù)的預(yù)訓(xùn)練過(guò)程中趨于穩(wěn)定。

　　分析

　　最后，研究者對(duì)本文方法的優(yōu)勢(shì)進(jìn)行了分析。

　　動(dòng)態(tài)批量加載的有效性

　　其中，研究者從以下三個(gè)方面的影響來(lái)分析動(dòng)態(tài)批量加載的有效性:(1) 跨域的最終 LM 損失，(2) 整個(gè)訓(xùn)練過(guò)程中每個(gè)域的數(shù)據(jù)使用情況，(3) 下游任務(wù)性能。結(jié)果均基于 Sheared-LaMA-1.3B 算法。

　　跨域損失差異。動(dòng)態(tài)批量加載的目的是平衡各域的損失降低率，使損失在大致相同的時(shí)間內(nèi)達(dá)到參考值。圖5中繪制了模型損耗(原始批量加載和動(dòng)態(tài)批量加載)與參考損耗之間的差異，相比之下，動(dòng)態(tài)批量加載能均勻地減少損失，各域的損失差異也非常相似，這表明數(shù)據(jù)使用效率更高。

　　數(shù)據(jù)使用情況。表3對(duì)比了 RedPajama 的原始數(shù)據(jù)比例和動(dòng)態(tài)加載的域數(shù)據(jù)使用情況(圖7展示了整個(gè)訓(xùn)練過(guò)程中域權(quán)重的變化)。與其他域相比，動(dòng)態(tài)批量加載增加了 Book 和 C4域的權(quán)重，這表明這些域更難恢復(fù)剪枝模型。

　　下游性能。如圖6所示，與在原始 RedPajama 分布上訓(xùn)練的模型相比，使用動(dòng)態(tài)批量加載訓(xùn)練的剪枝模型獲得了更好的下游性能。這表明，動(dòng)態(tài)批量加載所帶來(lái)的更均衡的損失減少可以提高下游性能。

　　與其他剪枝方法的對(duì)比

　　此外，研究者將 LLM-shearing 方法與其他剪枝方法進(jìn)行了比較，并報(bào)告了驗(yàn)證困惑度，它是衡量整體模型能力的一個(gè)有力指標(biāo)。

　　由于計(jì)算上的限制，下面的實(shí)驗(yàn)控制了所有比較方法的總計(jì)算預(yù)算，而不是將每種方法運(yùn)行到最后。

　　如表4所示，在相同稀疏度下，本文的目標(biāo)剪枝模型的推理吞吐量比非均勻剪枝 CoFiPruning 模型更高，但困惑度略高。

　　其他分析

　　表5顯示，在控制 token 總量的情況下，增加剪枝開(kāi)銷(xiāo)可以持續(xù)改善困惑度。然而，由于剪枝比持續(xù)的預(yù)訓(xùn)練更昂貴，研究者將0.4B 的 token 分配給剪枝。