BERTでseq2seqしたいときのモデルの理解

BERT generation周り学習時の備忘用.
huggingfaceの翻訳とほしい部分の要約.
以下の順:

1. Encoderの定式化
2. Decoderの定式化
3. EDモデルの定式化
4. pre-trained BERTでEDモデルするにはどうするか

まとめ

• BERTはencoder onlyのため, せっかくがっつり事前学習してweightを保存していてもseq2seqには向かない (事前にsequence lengthが欲しい)
• decoder onlyのGPT2はseq2seqには向くものの, $X_{i+1}$の情報を$X_i$に使えないのがもったいない.
• Rothe et al. (2020)ではBERTやGPT2といった事前学習モデルをwarm-startingなencoder-decoder (ED)モデルとして活用する際のTIPSをまとめた.

---

seq2seqの定式化

$X_{1:n}$を長さnの入力の単語列(センテンス, 以降sequence), $Y_{1:m}$を長さmの出力sequenceとして, seq2seqモデルは以下の条件付き確率分布として定式化される:

$$p_{\theta_{model}}(Y_{1:m}|X_{1:n})$$

このとき各単語ベクトルを用いて$X_{1:n}=[\bm{x_{1}},...,\bm{x_{n}}]$などと書けることに注意.

BERTの定式化

Encoder側の代表.
BERTは$X_{1:n}$を長さnの入力seqから, コンテキストを反映したseq(要はRNNでいうhidden state) $\overline{X_{1:n}}$へと射影する.
その後, 適宜pooling layerとclassification layerを追加して, $\overline{X_{1:n}}$をクラス$c$へと対応付けるなり様々なタスクを解く.

$$f_{\theta_{BERT}}:X_{1:n} \to \overline{X_{1:n}}$$

$$f_{\theta_{pooling,c}}:\overline{X_{1:n}} \to c$$

重要な点は, BERTのようなencoder-onlyモデルは入力seqを長さが決まっている出力seqへと射影するモデルであり, 出力seqの長さは入力seqに依存しない こと.
そのためBERTは素直に使うとseq2seqには向かない.

GPT2の定式化

Decoder側の代表.
GPT2はunidirectional self-attentionを用いて$Y_{1:m - 1}$の入力seqから, 次の単語についてのlogitベクトル列$L_{1:m}=[\bm{l_{1},...\bm{l_{m}}}]$への射影として定式化される:

$$f_{\theta_{GPT2}}:Y_{1:m - 1} \to \overline{L_{1:m}}$$

logitベクトルをsoftmaxに通せば単語列$Y_{1:m}$の確率分布が得られる.
この確率分布は$m - 1$の条件付き「次の単語」分布に分解できるので以下のように見るとわかりやすい:

$$p_{\theta_{GPT2}}(Y_{1:m})=\prod_{i=0}^m p_{\theta_{GPT2}}(\bm{y_{i}}|Y_{0:i - 1})$$

,where

$$p_{\theta_{GPT2}}(\bm{y_{i}}|Y_{0:i - 1}) = \bm{Softmax}(\bm{l_i}) = \bm{Softmax}(p_{\theta_{GPT2}}(Y_{0:i - 1})$$

GPT2は文章の生成に向くが, unidirectionalな性質のため, $\bm{x_i}$が$\bm{x_{i + 1}}$に依存できず, EDモデルに比べて劣るとされている(参考).

EDモデルの定式化

EDモデル一般の話.
まずEncoderにより, 長さnの入力seq $X_{1:n}$から, コンテキストを反映したseq $\overline{X_{1:n}}$へと射影する.
その後, Decoderにより, $\overline{X_{1:n}}$と一つ前の出力seq $Y_{0:m - 1}$からlogitベクトル列$L_{1:m}$へと射影する.

$$p_{\theta_{enc,dec}}(Y_{1:m}|X_{1:n})=p_{\theta_{dec}}(Y_{1:m}|\overline{X_{1:n}})=\prod_{i=1}^m p_{\theta_{dec}}(\bm{y_i}|Y_{0:i-1},\overline{X_{1:n}})$$

, where $\overline{X_{1:n}}=f_{\theta_{enc}}(X_{1:n})$

$\overline{X_{1:n}}$と一つ前の出力seq $Y_{0:m - 1}$で条件づけた次の単語分布は, GPT2と同様にlogitベクトル列$L_{1:m}$にsoftmaxをかければ得られる.

$$p_{\theta_{dec}}(\bm{y_i}|Y_{0:i-1},\overline{X_{1:n}})=\bm{Softmax}(\bm{l_i})$$

BERTによるwarm-startingなseq2seqのアプローチ

いよいよpre-trained BERTによるseq2seq.
pre-trained BERTのweightをseq2seqに利活用するにあたり, encoderとdecoderに分けて考える.

encoder側

encoder側はBERTと同じ構造なので, そのままpre-trainedのweightをセットすればいい.

decoder側

decoderは3箇所BERTと異なる.
唯一$\overline{X_{1:n}}$を反映するためのcorss-attentionのみ乱数での初期化だが, 他にはBERTのweightを利活用できる:

1. cross-attentionの追加
   • ここが 最も異なる
   • decoderはコンテキストを反映したseq $\overline{X_{1:n}}$ によって条件付けされる必要がある
   • 各BERTブロック内self-attentionとfeed-forwardの間に乱数で初期化した, $\overline{X_{1:n}}$とのcross-attentionを入れる
2. bi-directional self-attentionをuni-directionalに変更
   • BERTブロックのself-attentionはbi-directionalだが, decoderにするためにはここがuni-directionalである必要がある
   • ただquery, key, valueの関係性は同一なので, weight自体はBERTのものをセットする
   • (雑感) 少し不思議な気もするが, 感覚的にweightのmeanやvarがそれっぽい値になるだけでもfine-tuningする上では捗る印象なので, 奏効するのか
3. LM-headの追加
   • decoder出力$L_{1:m}$はロジットのベクトル列なので, 欲しい条件付き確率$p_{\theta_{dec}}(Y_{1:n} | \overline{X})$に変換するLM headがdecoderブロックの最後に来る
     • LM headはlanguage model headの略, GPT2由来
   • LM headのweightはword embedding $W_{emb}$に対応するので, BERTモデルの$W_{emb}$をdecoderのLM headにセットする

(オマケ) weight-sharing

Raffel et al.(2020)は, seq2seqではencoder/decoderのweightをshareしても(同じにしても)それほど性能が下がらないことを指摘している.
この辺りとwarm-startingは相性が良さそう.

テスト

テストは本家に丁寧に載っているので割愛, ページのちょうど真ん中ら辺からpracticeを載せてくれている.
ここにコードだけをまとめたバージョンもある.
transformersのEncoderDecoderModelにて, encoderとdecoderにそれぞれpre-trainedのBERTを指定すれば後は普段通りという感じ.
スクラッチの際には渡せる形のpre-trainedを作ることが少し手間になるが, 構造を変えない限りは入れ替えるだけなので簡単.