譯者 | 朱先忠

審校 | 重樓

簡(jiǎn)介

近年來(lái)，大型語(yǔ)言模型的發(fā)展突飛猛進(jìn)。BERT成為最受歡迎和最有效的模型之一，可以高精度地解決各種自然語(yǔ)言處理（NLP）任務(wù)。繼BERT模型之后，一組其他的模型也先后出現(xiàn)并各自展示出優(yōu)秀的性能。

不難看到一個(gè)明顯趨勢(shì)是，隨著時(shí)間的推移，大型語(yǔ)言模型（LLM）往往會(huì)因其訓(xùn)練的參數(shù)和數(shù)據(jù)數(shù)量呈指數(shù)級(jí)增加而變得更加復(fù)雜。深度學(xué)習(xí)研究表明，這種技術(shù)通常會(huì)帶來(lái)更好的運(yùn)行結(jié)果。然而，遺憾的是，盡管機(jī)器學(xué)習(xí)世界已經(jīng)克服了不少關(guān)于大型語(yǔ)言模型相關(guān)的問(wèn)題；但是，可擴(kuò)展性的問(wèn)題已經(jīng)成為有效訓(xùn)練、存儲(chǔ)和使用大型語(yǔ)言模型的主要障礙。

考慮到上述問(wèn)題，人們已經(jīng)開(kāi)發(fā)出不少的壓縮大型語(yǔ)言模型的特殊方法。在這篇文章中，我們將重點(diǎn)討論轉(zhuǎn)換器蒸餾方法，這種方法誕生了名為T(mén)inyBERT的一個(gè)迷你版本的BERT模型。此外，我們還將介紹TinyBERT模型的學(xué)習(xí)過(guò)程，以及使TinyBERT模型變得如此強(qiáng)大的幾個(gè)微妙原因。本文基于TinyBERT的官方論文整理而成。

主要思想

在最近的文章中，我們已經(jīng)討論了DistilBERT模型中蒸餾技術(shù)的工作原理。簡(jiǎn)而言之，蒸餾技術(shù)的主要思想是：修改損失函數(shù)目標(biāo)，以便使學(xué)生模型和教師模型的預(yù)測(cè)結(jié)果相似。在DistilBERT模型中，損失函數(shù)比較學(xué)生模型和教師模型的輸出分布，并兼顧兩個(gè)模型的輸出嵌入（針對(duì)相似性損失）。

有關(guān)DistilBERT模型的更多的細(xì)節(jié)，請(qǐng)參考文章《Large Language Models: DistilBERT — Smaller, Faster, Cheaper and Lighter》，地址是：
“https://towardsdatascience.com/distilbert-11c8810d29fc?source=post_page-----1a928ba3082b--------------------------------”。此文的主要內(nèi)容介紹了BERT模型壓縮的秘密，目標(biāo)是實(shí)現(xiàn)師生模型框架效率的最大化。

從表面上看，TinyBERT模型中的蒸餾框架與DistilBERT模型沒(méi)有太大變化：再次修改了損失函數(shù)，目標(biāo)是使學(xué)生模型模仿教師模型。然而，在TinyBERT模型的情況下，它更進(jìn)了一步：損失函數(shù)不僅考慮了師生兩個(gè)模型產(chǎn)生的結(jié)果，還考慮了如何獲得預(yù)測(cè)結(jié)果的問(wèn)題。根據(jù)TinyBERT模型論文作者介紹，TinyBERT損失函數(shù)由三個(gè)部分組成，它們涵蓋了師生兩個(gè)模型的不同方面：

1. 嵌入層的輸出
2. 從轉(zhuǎn)換器層導(dǎo)出的隱藏狀態(tài)和注意力矩陣
3. 預(yù)測(cè)層輸出的logits值

轉(zhuǎn)換器蒸餾損失函數(shù)示意圖

那么，比較師生兩種模型的隱藏狀態(tài)有什么意義呢？通過(guò)包括隱藏狀態(tài)和注意力的輸出結(jié)果，注意力矩陣使得學(xué)生模型有可能學(xué)習(xí)教師模型的隱藏層內(nèi)容，從而構(gòu)建與教師模型相似的層。這樣，提取的模型不僅可以模仿原始模型的輸出，而且模仿其內(nèi)部行為。

那么，為什么復(fù)制教師模型的行為很重要呢？研究人員聲稱(chēng)，通過(guò)BERT模型學(xué)習(xí)到的注意力權(quán)重有利于捕捉語(yǔ)言結(jié)構(gòu)。因此，它們對(duì)另一種模式的蒸餾也給了學(xué)生模型更多獲得語(yǔ)言知識(shí)的機(jī)會(huì)。

層映射

TinyBERT模型僅代表一種較小的BERT版本，具有較少的編碼器層。現(xiàn)在，不妨讓我們將BERT模型層數(shù)定義為N，將TinyBERT模型層數(shù)定義為M。鑒于層數(shù)不同，如何計(jì)算蒸餾損失值的問(wèn)題尚不明確。

為此，引入了一個(gè)特殊函數(shù)n=g（m）來(lái)定義哪個(gè)BERT模型層n用于將其知識(shí)提取到TinyBERT模型中的相應(yīng)層m。然后，所選擇的BERT層用于訓(xùn)練期間的損失值計(jì)算。

引入的函數(shù)n=g（m）具有兩個(gè)推理約束：

• g（0）=0。這意味著，BERT模型中的嵌入層被直接映射到TinyBERT模型中的嵌入圖層，這是有意義的。
• g（M+1）=N+1。該等式指示，BERT模型中的預(yù)測(cè)層被映射到TinyBERT模型中的預(yù)測(cè)層。對(duì)于所有其他TinyBERT模型中滿足條件1≤m≤m的那些層，需要映射n=g（m）的相應(yīng)函數(shù)值。現(xiàn)在，假設(shè)已經(jīng)定義好了這樣的函數(shù)。
  有關(guān)TinyBERT模型設(shè)置的問(wèn)題，將在本文稍后進(jìn)行研究。

轉(zhuǎn)換器蒸餾

1.嵌入層蒸餾

原始輸入在被傳遞到模型之前，首先被標(biāo)記化，然后被映射到學(xué)習(xí)的嵌入層。然后，這些嵌入層被用作模型的第一層。所有可能的嵌入層都可以用矩陣的形式表示。為了比較學(xué)生模型和教師模型的嵌入層有多大的不同，可以在他們各自的嵌入矩陣E上使用標(biāo)準(zhǔn)回歸度量。例如，轉(zhuǎn)換器蒸餾使用均方誤差（MSE）作為回歸度量。

由于學(xué)生模型和教師模型的嵌入矩陣具有不同的大小，因此不可能通過(guò)使用均方誤差來(lái)明智地比較它們的元素。這就解釋了為什么學(xué)生模型嵌入矩陣乘以可學(xué)習(xí)的權(quán)重矩陣W，從而導(dǎo)致結(jié)果矩陣與教師模型嵌入矩陣具有相同的形狀。

嵌入層蒸餾損失函數(shù)。

由于學(xué)生模型和教師模型的嵌入空間是不同的，矩陣W在將學(xué)生模型的嵌入空間線性轉(zhuǎn)換為教師模型嵌入空間方面也起著重要作用。

2.轉(zhuǎn)換器層蒸餾

轉(zhuǎn)換器層蒸餾損失函數(shù)可視化展示

2A. 注意力層蒸餾

轉(zhuǎn)換器中的多頭注意力機(jī)制的核心是生成包含豐富語(yǔ)言知識(shí)的多個(gè)注意力矩陣。通過(guò)轉(zhuǎn)移教師模型的注意力權(quán)重，學(xué)生模型也可以理解重要的語(yǔ)言概念。為了實(shí)現(xiàn)這一思想，使用損失函數(shù)來(lái)計(jì)算學(xué)生模型和教師模型注意力權(quán)重之間的差異。

在TinyBERT模型中，考慮了所有的注意力層，并且每一層的最終損失值等于所有頭部的相應(yīng)學(xué)生模型和教師模型注意力矩陣之間的均方誤差值之和。

注意層蒸餾損失函數(shù)計(jì)算公式

值得注意的是，用于注意力層提取的注意力矩陣A是未歸一化的，而不是它們的softmax輸出softmax（A）。根據(jù)研究人員的說(shuō)法，這種微妙之處有助于更快地收斂并提高性能。

2B. 隱藏層蒸餾

為了實(shí)現(xiàn)獲取豐富語(yǔ)言知識(shí)的想法，蒸餾操作也被應(yīng)用到轉(zhuǎn)換器層的輸出上。

隱藏層蒸餾損失函數(shù)計(jì)算公式。

這里，權(quán)重矩陣W起到與上述用于嵌入層蒸餾的權(quán)重矩陣相同的作用。

3.預(yù)測(cè)層蒸餾

最后，為了使學(xué)生模型再現(xiàn)教師模型的輸出結(jié)果，使用了預(yù)測(cè)層損失函數(shù)。它包括計(jì)算兩個(gè)模型預(yù)測(cè)的logit向量之間的交叉熵。

預(yù)測(cè)層蒸餾損失函數(shù)計(jì)算公式

值得注意的是，有些情況下，logits要除以控制輸出分布的平滑度的溫度參數(shù)T。在TinyBERT模型中，溫度參數(shù)T設(shè)置為1。

損失方程

在TinyBERT模型中，根據(jù)其類(lèi)型特征，每一層都有自己的損失函數(shù)。考慮到某些層或多或少的重要性作用，將相應(yīng)的損失值乘以常數(shù)a。最終的損失函數(shù)等于所有TinyBERT模型層的損失值的加權(quán)和。

TinyBERT模型中的損失函數(shù)計(jì)算公式

大量實(shí)驗(yàn)表明，在三種損失分量中，轉(zhuǎn)換器層蒸餾損失對(duì)模型性能的影響最大。

模型訓(xùn)練

需要注意的是，大多數(shù)自然語(yǔ)言處理模型（包括BERT）開(kāi)發(fā)過(guò)程可大致劃分為兩個(gè)階段：

1. 在一個(gè)大型數(shù)據(jù)語(yǔ)料庫(kù)上對(duì)模型進(jìn)行預(yù)訓(xùn)練，以獲得語(yǔ)言結(jié)構(gòu)的一般知識(shí)。
2. 在另一個(gè)數(shù)據(jù)集上對(duì)模型進(jìn)行微調(diào)，以解決特定的下游任務(wù)。

遵循與此同樣的思想，研究人員研發(fā)了一個(gè)新的框架TinyBERT，它的學(xué)習(xí)過(guò)程也是由類(lèi)似上面的兩個(gè)階段組成的。在這兩個(gè)訓(xùn)練階段中，使用轉(zhuǎn)換器蒸餾算法將BERT模型知識(shí)轉(zhuǎn)換成TinyBERT模型。

階段一：普通蒸餾。TinyBERT作為一個(gè)教師模型，通過(guò)預(yù)先訓(xùn)練（無(wú)需微調(diào)）的BERT模型獲得了豐富的語(yǔ)言結(jié)構(gòu)常識(shí)。通過(guò)使用更少的層和參數(shù)，在這個(gè)階段之后，TinyBERT模型的性能通常比BERT模型差一些。

階段二：特定任務(wù)的蒸餾。這一次，微調(diào)版的BERT模型扮演了教師模型的角色。為了進(jìn)一步提高性能，正如研究人員所提出的，在訓(xùn)練數(shù)據(jù)集上應(yīng)用了數(shù)據(jù)增強(qiáng)方法。實(shí)驗(yàn)結(jié)果表明，經(jīng)過(guò)特定任務(wù)的蒸餾操作后，TinyBERT模型在BERT模型方面取得了相當(dāng)?shù)男阅堋?/span>

TinyBERT模型訓(xùn)練流程示意圖

數(shù)據(jù)增強(qiáng)

針對(duì)特定任務(wù)的蒸餾，引入了一種特殊的數(shù)據(jù)增強(qiáng)技術(shù)。在這種數(shù)據(jù)增強(qiáng)技術(shù)中，首先從給定的數(shù)據(jù)集中提取序列，然后以下列兩種方式之一替換一定百分比的單詞：

• 如果某單詞被標(biāo)記為同一個(gè)單詞，則該單詞由BERT模型預(yù)測(cè)，并且用預(yù)測(cè)后的結(jié)果單詞替換序列中的原始單詞。
• 如果單詞被標(biāo)記為幾個(gè)子單詞，那么這些子單詞將被最相似的GloVe嵌入（全局向量的詞嵌入：Global Vectors for Word Representation）所取代。

盡管模型大小顯著減小，但是所描述的數(shù)據(jù)增強(qiáng)機(jī)制通過(guò)允許TinyBERT學(xué)習(xí)更多不同的示例，對(duì)其性能產(chǎn)生了很大影響。

數(shù)據(jù)增強(qiáng)技術(shù)示意圖

模型設(shè)置

由于只有14.5M個(gè)參數(shù)，TinyBERT模型比基礎(chǔ)型BERT模型小約7.5倍。它們的詳細(xì)比較如下圖所示：

基礎(chǔ)BERT模型與TinyBERT模型比較

對(duì)于層映射，論文作者提出了一種統(tǒng)一的策略。根據(jù)該策略，層映射函數(shù)將每個(gè)TinyBERT層映射到按序排序的每三個(gè)為一組的BERT層中的第一個(gè)：g（m）=3*m。論文作者還研究了其他的策略（如采用所有底部或頂部BERT層），但僅有統(tǒng)一策略顯示出最佳實(shí)驗(yàn)結(jié)果。這個(gè)結(jié)論似乎是比較合乎邏輯的，因?yàn)樗试S從不同的抽象層轉(zhuǎn)移知識(shí)，使轉(zhuǎn)移的信息更加多樣化。

基于不同的層映射策略情況：圖中展示了基于GLUE數(shù)據(jù)集的性能比較結(jié)果

就訓(xùn)練實(shí)現(xiàn)過(guò)程方面，TinyBERT模型是在英語(yǔ)維基百科（2500M個(gè)單詞）上訓(xùn)練的，其大多數(shù)超參數(shù)與BERT模型庫(kù)中使用的相同。

結(jié)論

轉(zhuǎn)換器蒸餾是自然語(yǔ)言處理中的一項(xiàng)重要措施?？紤]到基于轉(zhuǎn)換器的模型是目前機(jī)器學(xué)習(xí)中最強(qiáng)大的模型之一，我們可以通過(guò)應(yīng)用轉(zhuǎn)換器蒸餾來(lái)有效地壓縮它們來(lái)進(jìn)一步開(kāi)發(fā)利用它們。這方面最偉大的例子之一是TinyBERT模型，它在BERT模型基礎(chǔ)上壓縮了7.5倍。

盡管參數(shù)大幅減少，但實(shí)驗(yàn)表明，TinyBERT模型的性能與BERT基礎(chǔ)模型基本相當(dāng)：在GLUE基準(zhǔn)數(shù)據(jù)集上的測(cè)試結(jié)果表明，TinyBERT模型獲得77.0%的得分，與得分為79.5%的BERT模型相距并不遠(yuǎn)。顯然，這是一個(gè)驚人的成就！最后，其他的一些流行的壓縮技術(shù)，如量化或修剪等，都可以應(yīng)用于TinyBERT模型壓縮算法，從而使此模型體積變得更小。

除非另有說(shuō)明，否則本文中所有圖片均由作者本人提供。

參考資料

• TinyBERT: Distilling BERT for Natural Language Understanding：https://arxiv.org/pdf/1909.10351.pdf。

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專(zhuān)家博客、講師，濰坊一所高校計(jì)算機(jī)教師，自由編程界老兵一枚。

原文標(biāo)題：Large Language Models: TinyBERT — Distilling BERT for NLP，作者：Vyacheslav Efimov