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

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

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

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

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

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

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

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

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

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

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

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

　　方法介紹

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

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

　　第二階段持續預訓練 M_T，使其性能更強。

　　結構化剪枝

　　結構化剪枝可以去除模型大量參數，從而達到壓縮模型并加速推理的效果。然而，現有的結構化剪枝方法會導致模型偏離常規架構的配置。例如 CoFiPruning 方法產生的模型具有不統一的層配置，與標準的統一層配置相比，這樣會產生額外的推理開銷。

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

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

　　每個掩碼變量控制是否剪枝或保留相關的子結構。例如，如果對應的 z^layer=0，則需要刪除這個層。下圖2說明了剪枝掩碼如何控制被剪枝的結構。

　　剪枝之后，本文通過保留與每個子結構中的掩碼變量相關的最高得分組件來最終確定剪枝后的架構，并繼續使用語言建模目標對剪枝后的模型進行預訓練。

　　動態批量加載

　　該研究認為對剪枝后的模型進行大量預訓練是很有必要的，這樣才能恢復模型性能。

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

　　實驗及結果

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

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

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

　　SHEARED-LLAMA 優于同等大小的 LM

　　本文表明，Sheared-LLaMA 明顯優于現有的類似規模的 LLM，同時只使用一小部分計算預算來從頭開始訓練這些模型。

　　下游任務:表2展示了 Sheared-LLaMA 和類似大小的現有預訓練模型的零樣本和少樣本在下游任務上的性能。

　　指令調優:如圖3所示，與同等規模的所有其他預訓練模型相比，指令調優的 Sheared-LLaMA 實現了更高的獲勝率。

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

　　分析

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

　　動態批量加載的有效性

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

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

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

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

　　與其他剪枝方法的對比

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

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

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

　　其他分析

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