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

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

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

1 背景介紹

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

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

2 團(tuán)隊(duì)工作介紹：TRIME、CofiPruning

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3 總結(jié)與展望

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