🧠

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

2022/09/27に公開

bert

seq2seq

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

Encoderの定式化
Decoderの定式化
EDモデルの定式化
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を利活用できる:

cross-attentionの追加
- ここが 最も異なる
- decoderはコンテキストを反映したseq $\overline{X_{1:n}}$ によって条件付けされる必要がある
- 各BERTブロック内self-attentionとfeed-forwardの間に乱数で初期化した, $\overline{X_{1:n}}$ とのcross-attentionを入れる
bi-directional self-attentionをuni-directionalに変更
- BERTブロックのself-attentionはbi-directionalだが, decoderにするためにはここがuni-directionalである必要がある
- ただquery, key, valueの関係性は同一なので, weight自体はBERTのものをセットする
- (雑感) 少し不思議な気もするが, 感覚的にweightのmeanやvarがそれっぽい値になるだけでもfine-tuningする上では捗る印象なので, 奏効するのか
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にセットする

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

テスト

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