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

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

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

1 背景介紹

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

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

2 團隊工作介紹：TRIME、CofiPruning

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

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

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

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

上述兩個工作實際上只是在測試階段采用了緩存分布和k近鄰分布，在訓(xùn)練過程中只是延續(xù)了傳統(tǒng)的語言模型，在推斷階段語言模型并沒有優(yōu)化緩存和數(shù)據(jù)存儲區(qū)的結(jié)合。

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

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

因此整個TRIME的訓(xùn)練目標(biāo)函數(shù)包括兩個部分：

（1）基于輸出詞嵌入矩陣的預(yù)測任務(wù)。

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

算法希望最終訓(xùn)練的網(wǎng)絡(luò)能夠?qū)崿F(xiàn)最終預(yù)測的單詞盡可能準(zhǔn)確，同時同一訓(xùn)練批內(nèi)共享同一個待預(yù)測單詞的文本盡可能相似，以使得正在訓(xùn)練過程中讓所有的文本記憶表示通過反向傳播實現(xiàn)端到端的神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)。算法的實現(xiàn)思想在很大程度上受到2020年提出的稠密檢索（dense retrieval）所啟發(fā)，稠密檢索在訓(xùn)練階段對齊詢問和正相關(guān)文檔并且利用同一批內(nèi)的文檔作為負(fù)樣本，在推斷階段從大型數(shù)據(jù)存儲區(qū)中檢索相關(guān)文檔。

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

為了能夠盡可能減小訓(xùn)練和測試階段的不一致性，需要采取一定的數(shù)據(jù)處理策略來更好地構(gòu)建訓(xùn)練記憶。局部記憶指的是在同一個數(shù)據(jù)片段中的先前標(biāo)記（tokens），使用代價極其低廉?？梢圆捎秒S機取樣的批處理方式就能直接在訓(xùn)練階段和測試階段同時利用局部記憶，這就得到了基于局部記憶的基礎(chǔ)版TrimeLM模型。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

（1）針對整層剪枝可以獲得速度上的提升，相關(guān)工作指出大概50%左右的神經(jīng)網(wǎng)絡(luò)層是可以被剪枝的，但是粗粒度的剪枝對準(zhǔn)確率的影響比較大。

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

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

在實驗階段，在針對句子分類任務(wù)的GLUE數(shù)據(jù)集和針對問答任務(wù)的SQuAD1.1數(shù)據(jù)集上，可以發(fā)現(xiàn)CofiPruning在相同的速度和模型規(guī)?；A(chǔ)之上比所有的蒸餾和剪枝基線方法表現(xiàn)更好。

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

綜上所述，CofiPruning是一種非常簡單有效的模型壓縮算法，通過對粗粒度單元和細(xì)粒度單元聯(lián)合剪枝，結(jié)合逐層蒸餾的目標(biāo)函數(shù)可以聯(lián)通結(jié)構(gòu)剪枝和知識蒸餾這兩種算法的優(yōu)勢，從而實現(xiàn)更快的處理速度和更小的模型結(jié)構(gòu)。針對模型壓縮的未來趨勢，陳丹琦還重點提及了是否能夠?qū)鏕PT-3這樣的大型語言模型進(jìn)行剪枝，同時是否能夠?qū)ι嫌稳蝿?wù)進(jìn)行剪枝，這些都是未來可以重點關(guān)注的研究思路。

3 總結(jié)與展望

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