微軟UNILM 2.0：優雅的統一預訓練模型

本文首發於微信公衆號：NewBeeNLP，歡迎關注獲取更多幹貨資源。

刷arxiv看到了之前比較喜歡的一個工作UNILM的續集，這不得提高優先級先來品品（雖然還有一大堆TODO）。關於UNILM 1.0 在之前的文章中已經有介紹了，見站在BERT肩膀上的NLP新秀們（PART II），這種做到NLU和NLG簡潔統一的框架真的是非常贊！

目前NLP預訓練模型主要採取兩大類預訓練目標來進行語言模型訓練：Auto-encoding（如BERT等）和Auto-regressive（如GPT/XLNEet等）。如下圖，簡單來說Auto-encoding就是類似BERT那樣的通過[MASK]標記來獲取前向和後向的信息；而Auto-regressive呢就是自左向右或者自右向左地對文本進行語言建模，它的好處是對生成任務有天然的優勢，缺點呢就是隻能得到單向的信息不能同時利用上下文的信息。

而本文使用的Partially Auto-Regressive則可以認爲是對Auto-regressive的擴展，將token級別擴展到了span級別，並且不限制方向性。

模型預訓練

UNILM 2.0的模型框架沒有什麼好說的，骨架網絡就是Transformer，跟主流的都差不多，論文的亮點在於預訓練目標。在一個統一的BERT式的網絡基礎上，設計了兩種預訓練目標，自編碼和部分自迴歸，分別對應雙向語言模型和seq2seq語言模型， 然後進行聯合訓練：
$\mathcal{L}=\mathcal{L}_{\mathrm{AE}}+\mathcal{L}_{\mathrm{PAR}}$

自編碼

上圖中左半部分，自編碼（Auto Encoding，AE）就像BERT一樣，利用已知的上下文信息來單獨地計算被MASK掉的token的概率。
$\mathcal{L}_{\mathrm{AE}}=-\sum_{x \in \mathcal{D}} \log \prod_{m \in M} p\left(x_{m} | x_{\backslash M}\right)$
其中，$M=\left\{m_{1}, \cdots, m_{|M|}\right\}$爲被mask詞集合，$x_{\backslash M}$表示出去mask詞集的其他詞，$x_{m}$則表示被mask詞，$D$表示訓練語料集合。

部分自迴歸

上圖中右半部分，部分自迴歸（Partially AutoRegressive，PAR ）在每一步因式分解的時候可以選擇預測單個token或多個tokens（爲連續的span），當所有步均選擇預測一個token時，則會變爲自迴歸模式。
$\begin{aligned} p\left(x_{M} | x_{\backslash M}\right) &=\prod_{i=1}^{|M|} p\left(x_{M_{i}} | x_{\backslash M_{\geq i}}\right) \\ &=\prod_{i=1}^{|M|} \prod_{m \in M_{i}} p\left(x_{m} | x_{\backslash M \geq i)}\right. \end{aligned}$
其中，$x_{M_{i}}=\left\{x_{m}\right\}_{m \in M_{i}}$表示第$i$步需要預測的token，$M_{\geq i}=\bigcup_{j \geq i} M_{j}$表示還沒有預測出的詞（每一步都不一樣）。所以部分自迴歸預訓練損失可以表示爲：
$\mathcal{L}_{\mathrm{PAR}}=-\sum_{x \in \mathcal{D}} \mathbb{E}_{M} \log p\left(x_{M} | x_{\backslash M}\right)$
其中，$\mathbb{E}_{M}$是分解的期望，但是在預訓練中只隨機採用一種分解。

關於Mask的策略

對於輸入，隨機採樣15%作爲mask集合，這其中40%的概率mask掉長度爲2到6的n-gram連續span，60%的概率mask單獨一個token。

僞掩碼語言模型

正如前面提到的，如果要用部分自迴歸的話，每一步預測時可用的tokens是不一樣的，因此如果直接使原始BERT中的MLM，必須爲每個分解步驟構建一個新的Cloze實例，這將導致部分自迴歸的預訓練不可行。

爲此提出了一種Pseudo-Masked LM(PMLM) 的訓練方式，簡單來說就是對masked token額外添加一個帶對應位置編碼信息的標記[P]，用於在部分自迴歸預訓練時預測的placeholder。

舉個栗子，對於第一幅圖中的最後一個例子，即4,5——>2，下圖爲PMLM進行部分自迴歸預測的流程。

• Step 1：通過計算$p\left(x_{4}, x_{5} | x_{\backslash\{2,4,5\}}\right)$來同時 預測$x_{4}$和$x_{5}$，爲了避免信息泄露，這一步attend的元素爲$x_{1}, x_{3}, x_{6}$以及$[P]_{4}$和$[P]_{5}$。
• Step 2：通過計算$p\left(x_{2} | x_{\backslash\{2\}}\right)$來預測$x_{2}$，這一步可以獲取上一步得出的結果，即這一步attend的元素爲$x_{1}, x_{3}, x_{4}, x_{5}, x_{6}$以及$[P]_{2}$
  

具體實現的話，基本思想還是跟UNILM 1.0一樣的，即Self-Attention Masks，來控制計算某一個token時其他元素可以attend的權限。以上圖爲例，其Self-Attention Masks矩陣爲，

可以看到，[M]和已知的$\left(x_{1}, x_{3}, x_{6}\right)$可以被全局attend。[M] 對AE 部分來說，因爲是bidirectional，所以互相都允許attend；對PAR部分，可以提供”完整“的context（比如對P4/P5來說，就可以看到M2），這樣避免了之前autoregressive不能看future context，導致一些位置被完全masked。

模型輸入

• 在預訓練階段，模型輸入形式爲：[SOS] S1 [EOS] S2 [EOS]，其中S1和S2是連續的文本，[SOS] 和[EOS] 是表示文本起始以及結束的特殊標記，輸入token的表示是word embedding, absolute position embedding以及 segment embedding的加和。
• 在微調階段，不同任務的輸入形式不同：
  • 對於NLU任務，輸入爲“[SOS] TEXT [EOS]，然後用[SOS]作爲text的表示用於下游任務；
  • 對於NLG任務，輸入爲[SOS] SRC [EOS] TGT [EOS] ，一旦解碼出[EOS]，則停止解碼，解碼階段使用beam search；

實驗部分

論文中也展示了很多具體實驗，QA、GLUE、Abstractive Summarization、Question Generation，就不展開，粗瞄一眼效果都很不錯的樣子，感興趣的自行閱讀~

還是消融性實驗比較有趣，作者在多個任務上對UNILN2.0的組件進行拆解，各模型表現如下表。

隨便嘮嘮

UNILM1.0剛出來的時候同期還有一個MASS，也是微軟的，MASS的做法是顯式地設計一個Encoder和一個Decoder來做seq2seq生成。而UNILM1.0是使用了一個更爲優雅的設計，即通過attention mask矩陣來達到在一個單獨的BERT中實現seq2seq任務，不顯式區分encoder和decoder。

• UNILM2.0是對1.0中的seq2seq預訓練任務進行了優化，引入了部分自迴歸訓練目標，可能是受了XLNet的啓發，PMLM中的特殊標記[P]的作用有點類似於XLNet中的雙流自注意力機制。

來自論文作者董神：有點兒 三流注意力（真實、mask、pseudo mask）的意思，motivation就是最大限度重用計算的hidden states，來實現 unified language model pretraining 裏的 bidirectional AE (bert) 和 seq2seq PAR (bart)

• 另外在預訓練過程中，輸入文本的上下文表示會被AE和PAR重複利用，避免了無關計算。