本文來自于微信公眾號 機器之心(ID:almosthuman2014)，作者:機器之心。

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

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

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

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

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

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

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

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

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

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

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

　　方法介紹

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

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

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

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

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

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

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

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

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

　　動態(tài)批量加載

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

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

　　實驗及結(jié)果

　　模型配置:本文將 LLaMA2-7B 模型作為源模型，然后進行結(jié)構(gòu)化剪枝實驗，他們將 LLaMA2-7B 壓縮成兩個較小的目標(biāo)尺寸2.7B 和1.3B 參數(shù)，并將剪之后的模型與相同尺寸的模型進行了性能比較，包括 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)細節(jié)。

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

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

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

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

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

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

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

　　分析

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

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

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

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

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

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

　　與其他剪枝方法的對比

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

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

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

　　其他分析

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