文獻閱讀筆記-ALBERT ： A lite BERT for self-supervised learning of language representations

0. 背景

機構：谷歌
作者：
發佈地方：ICLR 2020
面向任務：自然語言理解
論文地址：https://openreview.net/pdf?id=H1eA7AEtvS
論文代碼：暫未

0.1 摘要

預訓練自然語言表徵時，增加模型大小一般是可以提升模型在下游任務中的性能。但是這種純粹依賴模型尺寸進而期望大力出奇跡的想法在未來會越發困難。進一步增加模型大小將帶來以下困難：(1)GPU/TPU內存不足(2)訓練時間會更長(3)模型退化。
所以，爲了解決上述這些問題，本文提出通過兩種參數精簡技術來降低內存消耗，並加快BERT的訓練速度。此外，本文還引入一個自監督損失(self-supervised loss)，用於對句子連貫性(inter-sentence coherence)建模，並證明該損失函數能夠提升多句子作爲輸入的下游任務的性能。本文所提出的模型ALBERT在 GLUE、RACE 和 SQuAD 這3個基準上都取得了新的SOTA結果，且參數量還少於 BERT-large。

1. 介紹

過往的研究者們在諸多NLP任務上的實驗已經表明，模型規模在取得SOTA結果上至關重要。在應用場景中通常是預訓練一個大規模的模型，再對其進行蒸餾萃取出一個更小的模型。考慮模型大小的重要性，我們不禁要問：“擁有更好的NLP模型是否能夠和擁有更大的模型一樣容易？”

上述問題首要解決的便是
(1)內存受限。當下的各種SOTA模型動輒數億甚至數十億個參數，倘若要擴大模型規模，這個內存問題是無法迴避的。
(2)訓練速度上的限制。由於通信開銷與模型中參數的數量成正比，在分佈式訓練中訓練速度將成爲一大瓶頸。簡單地增加隱含層單元數，只會適得其反，從而降低性能。Table 1和Fig 1中可以看出將BERT-large 的隱藏層單元數增加一倍， 該模型（BERT-xlarge）在 RACE 基準測試上的準確率顯著降低。

Table 1：增加 BERT-large 的隱藏層單元數量，模型在 RACE 上的表現變差。

Fig 1：BERT-large 和 BERT-xlarge 的訓練損失（左）和 dev mask 的 LM 準確率（右）。模型增大之後，其mask LM準確率降低了， 同時沒有出現明顯的過擬合跡象。

爲解決上述問題，已有先賢們做了一些研究，比如模型並行化、智能內存管理等。這些解決方案只解決了內存受限問題，而沒有考慮模型通信開銷過大和模型退化問題。爲此，本文提出A Lite BERT(ALBERT)模型以解決上述三個問題。 該模型的參數量遠遠少於傳統的 BERT 架構。

ALBERT 引入2種參數精簡技術，克服了擴展預訓練模型面臨的主要障礙。

第一種技術是對嵌入參數進行因式分解(factorized embedding parameterization)。將一個大的詞彙嵌入矩陣分解爲兩個小矩陣，從而將隱藏層的大小與詞彙嵌入的大小分離開來。這種分離便於後續隱藏層單元數量的增加，怎麼說呢？就是增加隱藏層單元數量，並不顯著增加詞彙嵌入的參數量。

第二種技術是跨層參數共享。這一技術可以避免參數量隨着網絡深度的增加而增加。

這兩種技術都顯著降低了 BERT 的參數量，同時不顯著損害其性能， 從而提升了參數效率。ALBERT 的配置類似於 BERT-large，但參數量僅爲後者的 1/18，訓練速度卻是後者的 1.7 倍。 這些參數精簡技術還可以充當某種形式的正則化，可以使訓練更加穩定，且有利於泛化。

爲了進一步提升 ALBERT 的性能， 本文還引入了一個自監督損失函數，用於句子順序預測（SOP，sentence-order prediction）。SOP 主要聚焦於句間連貫，用於解決原版 BERT中下一句預測（NSP）損失的低效問題。因爲確實已有研究(Yang et al., 2019; Liu et al., 2019)表明NSP是可以去掉的。

基於上述的這3個設計，ALBERT 能夠擴展爲更大的版本，在參數量仍然小於 BERT-large的同時，性能可以顯著提升。本文在GLUE、SQuAD 和 RACE 這3個自然語言理解基準測試上都刷新了記錄：在 RACE 上的準確率爲 89.4%，在 GLUE 上的得分爲 89.4，在 SQuAD 2.0 上的 F1 得分爲 92.2。

GLUE上最新榜單，2019年9月27號結果

SQuAD 2.0上最新榜單，2019年9月27號結果

2. 相關工作

2.1 擴大模型表徵

自然語言的表徵學習其重要程度不言而喻，目前的趨勢已從前兩年預訓練詞向量如標準的Word2Vec抑或是基於上下文語境的ELMo轉變爲整個網絡的預訓練+下游任務微。調。一般上來說，模型更大能夠提升性能，比如Devlin et al. (2019)的工作就表明採用更多的隱含層單元數、更多的隱含層和更多的attention head能夠提升性能。但是，凡事都有其適用範圍。Devlin在隱含層單元數設置爲1024後，就沒有繼續增加了。本文繼續進一步，發現當隱含層單元數增加爲2048個後，模型性能是退化的。所以，擴大自然語言的表徵學習模型，並不能簡單地增加模型尺寸了事。

另外，面對有限的計算資源，特別是GPU/TPU的內存受限，大規模模型的訓練會顯得很艱難。之前研究者們的方案有以下幾種：(1)Chen et al. (2016) 提出的gradient checkpointing，以降低內存消耗，從而使得內存的佔用是次線性的，其代價是需要一個額外的forward pass。(2)Gomez et al. (2017)提出從下一層重構每一層的激活，如此便無需存儲中間激活。這兩種方法是以犧牲速度爲代價換取內存佔用的降低。反觀本文的方法，通過參數精簡技術降低內存佔用的同時還加速了訓練速度，着實令人驚喜，有木有。

2.2 跨層參數共享

跨層參數共享在之前的Transformer架構中就有，但是Transformer只注重訓練用於標準encoder-decoder的任務，而非預訓練+下游任務微調的設置。與本文的觀察所不同的是，Dehghani et al. (2018)發現跨層參數共享(Universal Transformer,UT)在語言建模和主謂一致方面能夠取得比標準Transformer更好的結果。近來， Bai et al. (2019)提出的Deep Equilibrium Model (DQE)表明DQE能夠在某一層取得輸入嵌入與輸出嵌入保持一致的平衡點。

2.3 句子次序目標

爲了學習句子之間的連貫性，ALBERT引入了一個預測兩個連續句子次序的目標函數。針對篇章中句子的連貫和銜接，此前就有學者在預訓練中採用類似操作，比如基於Skip thought (Kiros et al., 2015)和FastSent (Hill et al., 2016)是通過對句子進行encoding進而預測句子周邊詞彙，從而學習句子嵌入。Gan et al., 2017雖然也是學習句子嵌入，但是使用的是預測後續的句子，而非僅僅其周邊單詞。[Jernite et al., 2017; Nie et al., 2019]所採用的目標則是預測顯性話語標記語(explicit discourse markers，即話語末尾標記語)。本文所提出的損失函數是面向文本段落segments，而非句子。BERT中採用的一個損失函數是預測兩個segments是否來自相同文檔，更確切地說是，其正樣本是下一個句子，而負樣本是隨機的其他文檔句子。相比於上述的這些方法，本文的句子次序預測更具挑戰性，於下游任務更有益。

3. ALBERT的組成

3.1 模型架構的選擇

ALBERT 架構的主幹網絡與 BERT 相似，即使用 Transformer 編碼器和 GELU 非線性激活函數。爲便於論述，本文做如下約定：詞嵌入大小爲 $E$、編碼器層數爲$L$、隱藏層大小爲 $H$。與 Devlin 等人的研究一樣，本文將 前饋網絡/濾波器 大小設置爲$4H$，將注意力頭的數量設置爲$H/64$。

本文ALBERT框架與BERT相比有以下3點不同：

嵌入向量參數化的因式分解
在BERT 及後續的 XLNet 和 RoBERTa 中，WordPiece 詞嵌入大小$E$和隱藏層大小$H$是相等的，即$E \equiv H$。這在建模和實際使用中，看起來可能並不是最優的。

(1)從建模的角度來說，WordPiece 詞嵌入的目標是學習上下文無關的表示，而隱藏層嵌入的目標是學習上下文相關的表示。[Liu et al.,2019]通過上下文長度相關的實驗表明，BERT的表徵能力很大一部分來自於使用上下文以在學習過程提供上下文相關的表徵信息。因此，將 WordPiece 詞嵌入大小$E$ 從隱藏層大小$H$ 分離出來，可以更高效地利用總體的模型參數， 其中$H$遠遠大於$E$。
(2)從實踐的角度，自然語言處理使用的詞典大小$V$一般非常龐大，如果$E$恆等於$H$，即$E \equiv H$，那麼增加$H$ 將直接增大詞典嵌入矩陣的大小($V \times E$)。這會導致模型參數劇增，而模型訓練過程中大部分參數的更新是很稀疏的。

因此，本文所提出的ALBERT對詞嵌入參數進行了因式分解，將其分解爲兩個小矩陣。本文不再將 one-hot 向量直接映射到大小爲$H$ 的隱藏空間，而是先將它們映射到一個低維詞嵌入空間$E$，然後再映射到隱藏空間。通過這種分解，可以將詞嵌入參數從 $O(V \times H)$ 降低到 $O(V \times E+E \times H)$。這在$H$ 遠遠大於$E$的時候，參數量減少得非常明顯。

跨層參數共享

此外，ALBERT中還使用了一種跨層參數共享機制來進一步提升參數效率。其實目前有很多方式來共享參數，比如只共享前饋網絡不同層之間的參數，或者只共享注意力機制的參數，而 ALBERT 採用的是所有層共享所有參數。

在之前也有學者提出過類似機制，比如[Dehghani et al. (2018)]的(Universal Transformer, UT)和Bai et al. (2019)和Deep Equilibrium Models(DQE)。雖然上述方案都有一定提升效果，但在度量過程發現詞嵌入的 L2 距離和餘弦相似性是震盪而不是收斂。如下圖 2 展示了每一層輸入與輸出嵌入矩陣間的 L2 距離與餘弦相似性。

Figure 2展示了BERT-Large 與 ALBERT-Large(設置參見Table 2) 每一層輸入嵌入與輸出嵌入間的 L2 距離與餘弦相似性。從Figure 2可以發現 ALBERT 從一層到另一層的轉換要比 BERT 平滑得多。這說明權重共享有效地提升了神經網絡參數的魯棒性。儘管相比於 BERT，ALBERT在兩個評價指標都有所下降，但在 24 層以後，它們也不會收斂到 0。這也表明ALBERT參數的解空間與DQE是大不相同的。

句間連貫性損失
除了語言建模損失(masked language modeling，MLM)外，BERT 還使用了額外的下一句預測損失(NSP)。NSP損失本是爲了提升下游任務的性能，但是後來很多研究 (Yang et al., 2019; Liu et al.,
2019)發現這種機制並不是很高效，因此決定去除它。

據此我們猜測，NSP低效的原因，主要是它的難度太小。因爲下一句預測將主題預測和連貫性預測結合到單個任務中，然而主題預測比連貫性預測簡單得多，同時NSP與MLM損失函數學到的內容是有重合的。

由於句間建模在語言理解中非常重要，因此本文提出一種基於語言連貫性的損失函數。在ALBERT中本文使用了一個句子次序預測（SOP）損失函數，它會避免預測主題，而只關注建模句子之間的連貫性。SOP正樣本獲取方法與BERT相同，而負樣本僅僅是將正樣本的兩個segments次序對換。實驗結果表明SOP能夠在合理範圍內解決NSP任務。在使用了該損失函數後，ALBERT能顯著提升下游多句子編碼任務的性能。

3.2 模型設置

本文用到的BERT和ALBERT對應的超參數如Table 2所示。可以看出，ALBERT的模型參數量遠遠小於對應的BERT模型。例如，ALBERT-large的參數量僅爲BERT-large的1/18，即18M vs 334M。對於ALBERT-xxlarge，由於24層網絡結構與12層網絡結構在結果上接近，但是計算量更大，所以本文主要採用12層的網絡結構。

4. 實驗結果

4.1 實驗設置

爲了進行更公平的對比，本文一方面使用與原始 BERT相同的配置訓練試驗模型，另一方面採用 BOOKCORPUS 和 English Wikipedia 共計 16GB 的純文本作爲預訓練任務的數據。與BERT一樣，使用的詞典大小是30，000；此外還借鑑了XLNet中使用的SentencePiece。在MLM目標函數上使用$n-gram$的masking，隨機選用$n-gram$的mask遮蔽輸入。預測生成的n-gram的概率：
$p(n)=\frac{1 / n}{\sum_{k=1}^{N} 1 / k}$
本文設置的n-gram最大長度爲3，即MLM目標最多由3個全詞組成。

4.2 評估基準

略

4.3 BERT VS ALBERT

從Table 3看出，ALBERT-xxlarge 的參數量只有 BERT-Large 70% ，但性能卻能夠顯著超越BERT-large。具體表現在SQuAD v1.1上提升1.7%，在SQuAD v2.0上提升4.2%，MNLI上提升2.2%，在SST-2上提升3.0%，在RACE上提升8.5%。此外，還觀察到BERT-xlarge在全部的指標上全面潰敗於BERT-base。這說明形如BERT-xlarge的大參數模型相較於更小參數量的模型是更難訓練的。另一個有趣的發現是，相同訓練配置下(相同TPUs數量)訓練一次，數據的吞吐速度差異。由於ALBERT模型通信更少，計算量更低，所以相比於BERT有更高的數據吞吐量。這裏選用最慢的BERT-xlarge作爲速度參照基準。

接下來介紹ALBERT中3個主要部分對模型提升所帶來的影響。

4.4 嵌入向量的因式分解

Table 4 展示了在ALBERT-base(具體設置參見上述Table 2)上修改詞嵌入大小 E 帶來的影響，它們的參數量及下游任務效果也都展示在內。

Table 4：ALBERT-base 隨詞嵌入大小的改變，其性能與參數量的變化

對於non-shared下(BERT-style)，更大的嵌入尺寸能夠取得更好的結果，但是提升的幅度其實不大。對於all-shared(ALBERT-style)，嵌入大小128是最好的。基於上述這些結果，本文在後續的實驗中的嵌入大小統一選用$E=128$。

4.5 跨層參數共享

Table 5 展示了不同跨層參數共享的效果，同樣使用 ALBERT-base 作爲示例模型，此外還增加了嵌入大小爲768的結果。對比了所有all-shared策略(ALBERT-style)、not-shared 策略(BERT-style)及其介於二者之間的中間策略(僅注意力參數共享，FNN不共享；僅FNN參數共享，注意力參數不共享)。

從上述結果可以看出，all-shared型策略在E=768和E=128上都會一定程度上降低性能。但是，需要說明的是，下降幅度較小，對於E=128，平均下降1.5；對於E=768，平均下降2.5。再細看，共享FFN層的參數，應該是罪魁禍首；而注意力機制的參數共享帶來的影響不能一概而論，對於E=128反而在平均性能上提升了0.1，對於E=768平均性能下降0.7。

4.6 句子次序預測(SOP)

Table 6 展示了SOP與下一句預測損失（NSP）的對比效果。

本文這裏對比了3種策略：沒有句子間損失(比如XLNet和RoBERTa)、NSP(比如BERT)、SOP(ALBERT)。這裏採用的ALBERT也是ALBERT-base。對比過程，一方面對比自身任務中的準確率，另一方面是下游任務的性能表現。在自身任務這一維度，可以看出NSP損失對於SOP幾乎是沒有任何益處，NSP訓練後，在SOP上的表現只有52%，這跟瞎猜差不了多少。據此，可以得出結論：NSP建模止步於主題識別。反觀SOP損失，確實一定程度上能夠解決NSP任務，其準確率爲78.9%，而自身的準確率爲86.5%。更爲重要的是，在下游任務上SOP損失統統起到促進作用，具體表現在SQuAD1.1提升1%，SQuAD 2.0提升2%，RACE提升1.7%。

4.7 相同訓練時長下的對比

從Table 3中的提速結果看出BERT-large的數據吞吐量大概是ALBERT-xxlarge的3.17倍。一般而言，訓練越長性能越高，鑑於此，我們進一步對比相同訓練時長下不同模型的表現。

在訓練了差不多相同的時間之後，ALBERT-xxlarge 明顯優於 BERT-large。

4.8 引入額外訓練集和Dropout的影響

上述實驗都是在 Wikipedia 和 BOOKCORPUS 數據集上進行的，那麼，如果增加額外的數據會對結果產生怎樣的影響？這裏採用的額外數據與XLNet和RoBERTa中的相同。

Figure 3a 表明，添加額外數據後，模型的開發集 MLM 準確率顯著提升。
添加額外數據後模型在下游任務中的性能情況，如Table 8 所示：

我們還注意到，即使在訓練了 100 萬步之後，最大的模型仍然沒有過擬合。因此，嘗試刪除dropout，以進一步提高模型能力。如Figure 3b 所示，去掉 dropout 可以顯著提高 MLM 準確度。去掉dropout後在下游任務上的表現如Table 9所示：

4.9 當下SOTA模型在NLU任務上的對比

除了上述實驗之外，ALBERT 在 GLUE、SQuAD 和 RACE 基準測試中都取得了 SOTA 結果，如Figure 10、11 所示：

5. 總結

本文的貢獻主要在3個方面：
(1)對嵌入矩陣分解，解除詞嵌入和隱含層大小的關係，便於隱含層大小的擴展而不劇增模型參數
(2)跨層參數共享，雖然一定程度上微微降低性能，但是可以大大地降低模型參數，收益很高
(3)放棄NSP，引入SOP，更爲有力地學習句子間的連貫性