“Making large models smaller”這是很多語言模型研究人員的學術追求，針對大模型昂貴的環境和訓練成本，陳丹琦在智源大會青源學術年會上做了題為“Making large models smaller”的特邀報告。報告中重點提及了基于記憶增強的TRIME算法和基于粗細粒度聯合剪枝和逐層蒸餾的CofiPruning算法。前者能夠在不改變模型結構的基礎上兼顧語言模型困惑度和檢索速度方面的優勢；而后者可以在保證下游任務準確度的同時實現更快的處理速度，具有更小的模型結構。

陳丹琦  普林斯頓大學計算機科學系助理教授

陳丹琦于2012年畢業于清華大學姚班，2018年獲得斯坦福大學計算機科學博士學位，師從斯坦福大學語言學和計算機科學教授 Christopher Manning。

1 背景介紹

近年來，自然語言處理領域正在迅速被大語言模型所主導。自從GPT 3問世以來，語言模型的規模呈現指數級增長。大型科技公司不斷發布越來越大的語言模型。近期，Meta AI發布了OPT語言模型（一個蘊含了1750億參數的大型語言模型），并向公眾開放了源代碼和模型參數。

研究學者們之所以如此推崇大語言模型，是因為它們出色的學習能力和性能表現，但是人們對于大語言模型的黑盒性質仍了解甚少。向語言模型輸入一個問題，通過語言模型一步一步地推理，能夠解決非常復雜的推理問題，比如推導出計算題的答案。但與此同時，大型語言模型也存在著風險，特別是它們的環境和經濟成本，例如：GPT-3 等大規模語言模型的能源消耗和碳排放規模驚人。面對大語言模型訓練成本昂貴、參數量龐大等問題，陳丹團隊希望通過學術研究縮減預訓練模型的計算量并且讓語言模型更有效率地適用于下層應用。為此重點介紹了團隊的兩個工作，一個是一種語言模型的新型訓練方法稱之為TRIME，另一個是一種適用于下游任務的有效模型剪枝方法稱之為CofiPruning。

2 團隊工作介紹：TRIME、CofiPruning

論文地址：https://arxiv.org/abs/2205.12674

傳統語言模型的訓練流程如下：給定一段文檔，將其輸入到Transformer編碼器中得到隱向量，進而將這些隱向量輸送到softmax層，該層輸出為由V個詞嵌入向量組成的矩陣，其中V代表詞匯量的規模，最后可以用這些輸出向量對原先的文本進行預測，并與給定文檔的標準答案進行比較計算梯度，實現梯度的反向傳播。然而這樣的訓練范式會帶來以下問題：（1）龐大的Transformer編碼器會帶來高昂的訓練代價；（2）語言模型輸入長度固定，Transformer的計算量規模會隨著序列長度的變化呈平方級增長，因此Transformer很難處理長文本；（3）如今的訓練范式是將文本投影到一個固定長度的向量空間內來預測接下來的單詞，這種訓練范式實際上是語言模型的一個瓶頸。

為此，陳丹琦團隊提出了一種新的訓練范式——TRIME，主要利用批記憶進行訓練，并在此基礎之上提出了三個共享相同訓練目標函數的語言模型，分別是TrimeLM，TrimeLMlong以及TrimeLMext。TrimeLM可以看作是標準語言模型的一種替代方案；TrimeLMlong 針對長范圍文本設計，類似于Transformer-XL；TrimeLMext結合了一個大型的數據存儲區，類似于kNN-LM。在前文所述的訓練范式下，TRIME首先將輸入文本定義為，然后將輸入傳送到Transformer編碼器中，得到隱向量，經過softmax層之后得到需要預測的下一個單詞，在整個訓練范式中可訓練的參數為和E。陳丹琦團隊的工作受到了以下兩個工作的啟發：（1）2017年Grave等人提出的連續緩存（Continuous cache）算法。該算法在訓練過程中訓練一個普通的語言模型；在推斷過程中，給定輸入的文本，首先列舉給定文本先前出現的所有單詞和其中所有等于下一個需要預測單詞的標記位置，然后利用隱變量之間的相似度和溫度參數計算緩存分布。在測試階段，對語言模型分布和緩存分布進行線性插值可以得到更好的實驗效果。

（2）2020年Khandelwal等人提出的k近鄰語言模型（kNN-LM），該方法與連續緩存算法類似，二者之間最大的不同在于kNN-LM為所有的訓練樣例構建了一個數據存儲區，在測試階段將對數據存儲區的數據進行k近鄰搜索，從而選擇最佳的top-k數據。

上述兩個工作實際上只是在測試階段采用了緩存分布和k近鄰分布，在訓練過程中只是延續了傳統的語言模型，在推斷階段語言模型并沒有優化緩存和數據存儲區的結合。

除此之外，還有一些針對超長文本的語言模型工作值得關注，例如在2019年提出的結合注意力循環（Attention recurrence）機制的Transformer-XL和在2020年提出的基于記憶壓縮（Memory compression）的Compressive Transformers等。

在之前介紹的幾項工作基礎之上，陳丹琦團隊構建了一個基于批記憶的語言模型訓練方法，主要思想是針對相同的訓練批（training batch）構建一個工作記憶（working memory）。針對給定文本預測下一個單詞的任務，TRIME的思想與對比學習十分類似，不僅僅考慮利用softmax詞嵌入矩陣預測下一個單詞出現概率的任務，還新增了一個模塊，在這個模塊中考慮所有出現在訓練記憶（training memory）中且與給定文本需要預測的單詞相同的所有其他文本。

因此整個TRIME的訓練目標函數包括兩個部分：

（1）基于輸出詞嵌入矩陣的預測任務。

（2）在訓練記憶（training memory）中共享同一個待預測單詞文本的相似度，其中需要衡量相似度的向量表示是在通過最終前饋層的輸入，采用縮放點積衡量向量相似度。

算法希望最終訓練的網絡能夠實現最終預測的單詞盡可能準確，同時同一訓練批內共享同一個待預測單詞的文本盡可能相似，以使得正在訓練過程中讓所有的文本記憶表示通過反向傳播實現端到端的神經網絡學習。算法的實現思想在很大程度上受到2020年提出的稠密檢索（dense retrieval）所啟發，稠密檢索在訓練階段對齊詢問和正相關文檔并且利用同一批內的文檔作為負樣本，在推斷階段從大型數據存儲區中檢索相關文檔。

TRIME的推理階段幾乎與訓練過程相同，唯一的區別在于可能會采用不同的測試記憶，包括局部記憶（Local memory），長期記憶（Long-term memory）和外部記憶（External memory）。局部記憶指的是出現在當前片段中的所有單詞，并且已經被注意力機制向量化；長期記憶指的是由于輸入長度限制導致無法直接獲取但與待處理文本來源于相同文檔的文本表示，外部記憶指的是存儲所有訓練樣本或者額外語料庫的大型數據存儲區。

為了能夠盡可能減小訓練和測試階段的不一致性，需要采取一定的數據處理策略來更好地構建訓練記憶。局部記憶指的是在同一個數據片段中的先前標記（tokens），使用代價極其低廉。可以采用隨機取樣的批處理方式就能直接在訓練階段和測試階段同時利用局部記憶，這就得到了基于局部記憶的基礎版TrimeLM模型。

長期記憶指的是在同一個文檔先前片段中的標記，需要依賴于同一個文檔的先前片段。為此將同一個文檔中的連續片段（consecutive segments）放入同一個訓練批中，這就得到了集合長期記憶的TrimeLMlong模型。

外部記憶需要結合大型數據存儲區進行檢索。為此可以利用BM25將訓練數據中的相似片段放入同一個訓練批中，這就得到了結合外部記憶的TrimeLMext模型。

綜上所述，傳統的語言模型在訓練階段和測試階段都沒有利用記憶；連續緩存方法只在測試階段采用了局部記憶或者長期記憶；k近鄰語言模型在測試階段采用了外部記憶；而針對TRIME算法的三種語言模型，在訓練階段和測試階段都采用了記憶增強的方式，其中TrimeLM在訓練階段和測試階段都采用了局部記憶，TrimeLMlong在訓練階段針對相同文檔的連續片段放入同一批訓練，在測試階段結合了局部記憶和長期記憶，TrimeLMext在訓練階段針對相似文檔放入同一批訓練，在測試階段結合了局部記憶、長期記憶和外部記憶。

在實驗階段，在WikiText-103數據集上進行模型參數247M，切片長度3072的測試時，基于TRIME算法的三種版本的語言模型都能取得比傳統Transformer更好的困惑度（perplexity）效果，其中基于實際距離的TrimeLMext模型可以取得最好實驗效果。同時TrimeLM和TrimeLMlong也能保持和傳統Transformer 接近的檢索速度，同時兼具了困惑度和檢索速度的優勢。

在WikiText-103數據集上進行模型參數150M，切片長度150的測試時，可以看到由于TrimeLMlong在訓練階段針對相同文檔的連續片段放入同一批訓練，在測試階段結合了局部記憶和長期記憶，因此盡管切片長度只有150，但是在測試階段實際可利用的數據可以達到15000，實驗效果遠遠好于其他基線模型。

針對字符級別的語言模型構建，基于TRIME算法的語言模型在enwik8數據集上同樣取得了最好的實驗效果，同時在機器翻譯的應用任務中，TrimeMT_ext也取得了超過基線模型的實驗效果。

綜上所述，基于TRIME算法的語言模型采用了三種記憶構建的方式，充分利用同一批內的相關數據實現記憶增強，在引入記憶的同時卻沒有引入大量的計算代價，并且沒有改變模型的整體結構，相比于其他基線模型取得了較好的實驗效果。

陳丹琦還著重提到了基于檢索的語言模型，實際上TrimeLMext可以看作是k近鄰語言模型的一個更好的版本，但是在推斷過程中這兩種算法相較于其他的基線模型速度要慢接近10到60倍，這顯然是難以接受的。陳丹琦指出了基于檢索的語言模型未來可能的發展方向之一：是否可以利用一個更小的檢索編碼器和一個更大的數據存儲區，從而實現計算代價縮減到最近鄰搜索。

相比于傳統的語言模型，基于檢索的語言模型有顯著的優勢，例如：基于檢索的語言模型可以更好的實現更新和維護，而傳統的語言模型由于利用先前知識進行訓練無法實現知識的動態更新；同時基于檢索的語言模型還可以更好的利用到隱私敏感的領域中去。至于如何更好的利用基于檢索的語言模型，陳丹琦老師認為或許可以采用fine-tuning、prompting或者in-context learning的方式來輔助解決。

論文地址：https://arxiv.org/abs/2204.00408

模型壓縮技術被廣泛應用于大語言模型，讓更小的模型能夠更快地適用于下游應用，其中傳統的主流模型壓縮方法為蒸餾（Distillation）和剪枝（pruning）。對于蒸餾而言，往往需要預先定義一個固定的學生模型，這個學生模型通常是隨機初始化的，然后將知識從教師模型傳送到學生模型中去，實現知識蒸餾。

例如，從原始版本的BERT出發，經過通用蒸餾，即在大量無標注的語料庫上進行訓練之后，可以得到基礎版本的TinyBERT4，針對基礎版本的TinyBERT4，還可以通過任務驅動的蒸餾方法得到微調過的TinyBERT4，最終得到的模型在犧牲輕微的準確度基礎之上可以比原先的BERT模型更小且處理速度更快。然而這種基于蒸餾的方法也存在著一定的缺陷，例如針對不同的下游任務，模型的架構往往是固定不變的；與此同時需要利用無標注數據從零開始訓練，計算代價太大。

對于剪枝而言，往往需要從一個教師模型出發，然后不斷地從原始模型中移除不相關的部分。在2019年提出的非結構化剪枝可以得到更小的模型但是在運行速度方面提升很小，而結構化剪枝通過移除例如前饋層等參數組實現實際應用的速度提升，例如2021年提出的塊剪枝可以實現2-3倍的速度提升。

針對傳統蒸餾和剪枝方法存在的局限性，陳丹琦團隊提出了一種名為CofiPruning的算法，同時針對粗粒度單元和細粒度單元進行剪枝，并設計了一個逐層蒸餾的目標函數將知識從未剪枝模型傳送到剪枝后的模型中去，最終能夠在保持超過90%準確率的基礎之上實現超過10倍的速度提升，比傳統的蒸餾方法計算代價更小。

CofiPruning的提出建立在兩個重要的基礎工作之上：

（1）針對整層剪枝可以獲得速度上的提升，相關工作指出大概50%左右的神經網絡層是可以被剪枝的，但是粗粒度的剪枝對準確率的影響比較大。

（2）是對更小單元例如頭部進行剪枝可以獲得更好的靈活性，但是這種方法在實現上會是一個更有難度的優化問題，且不會有太大的速度提升。

為此陳丹琦團隊希望能夠在粗粒度單元和細粒度單元同時剪枝，從而兼具兩種粒度的優勢。除此之外，為了解決從原始模型到剪枝模型的數據傳送，CofiPruning在剪枝過程中采用逐層對齊的方式實現知識的傳送，最終的目標函數包括了蒸餾損失和基于稀疏度的拉格朗日損失。

在實驗階段，在針對句子分類任務的GLUE數據集和針對問答任務的SQuAD1.1數據集上，可以發現CofiPruning在相同的速度和模型規模基礎之上比所有的蒸餾和剪枝基線方法表現更好。

針對TinyBERT，如果沒有經過通用蒸餾，實驗效果會大打折扣；但是如果利用通用蒸餾，盡管實驗效果可以有所提升但是訓練時間代價會非常昂貴。而CofiPruning算法不但能夠獲得與基線模型近乎持平的效果，在運行時間和計算代價方面還都有很大提升，可以用更少的計算代價獲得更快的處理速度。實驗表明，針對粗粒度單元，第一層和最后一層前饋層在最大程度上會被保留而中間層更有可能會被剪枝；針對細粒度單元，上層神經網絡的頭部和中間維度更有可能會被剪枝。

綜上所述，CofiPruning是一種非常簡單有效的模型壓縮算法，通過對粗粒度單元和細粒度單元聯合剪枝，結合逐層蒸餾的目標函數可以聯通結構剪枝和知識蒸餾這兩種算法的優勢，從而實現更快的處理速度和更小的模型結構。針對模型壓縮的未來趨勢，陳丹琦還重點提及了是否能夠對例如GPT-3這樣的大型語言模型進行剪枝，同時是否能夠對上游任務進行剪枝，這些都是未來可以重點關注的研究思路。

3 總結與展望

大型語言模型如今取得了非常喜人的實際應用價值，但是由于昂貴的環境和經濟成本，隱私和公平性方面的困擾以及難以實時更新的問題，導致大型語言模型仍有很多待改進之處。陳丹琦認為，未來的語言模型或許可以用作大型的知識庫，同時在未來語言模型的規模需要大幅度削減，或許可以利用基于檢索的語言模型或者稀疏語言模型來代替稠密檢索，模型壓縮的工作也需要研究者們重點關注。