BERTを超えたXLNetの紹介
この記事では、「BERTを超えた！」と話題のXLNetとその基幹となっている技術について説明します。
概要
https://arxiv.org/abs/1906.08237
XLNetは2019/6/19に、”XLNet: Generalized Autoregressive Pretraining for Language Understanding”と題してArxivに投稿された論文です。一言(?)でいうと
Transformer-XLを単語の順序を入れ替えた（元の順序情報は保持）もので学習させることで、自己回帰モデルで双方向の意味依存関係を取得できるようにしたと主張。20を超えるタスクでBERT超え。
といったところでしょうか。この記事では、背景となる関連技術も含めてXLNetの着想と技術について順々に説明していきます。
自然言語処理のタスク学習の流れと事前学習の種類
深層学習を用いた自然言語処理は翻訳、QAタスク、文書分類など多岐にわたります。
深層学習でそれらのタスクを解く際は、タスクに跨って有用な表現を教師なしで取得する「事前学習」と、事前学習の結果をもとにタスク用に再学習させる「微調整(fine-tuing)」という２段階にわけるという手法が、深層学習が発展して以降メジャーなアプローチ方法となっています。
事前学習はBERTのような自己符号化(autoencoding)をもとにした手法と、LSTMを用いたSeq2Seq、GPTのような自己回帰を用いた言語モデルの２つが頻繁に使われるようです。
自己符号化をもとにしたBERTは、順方向・逆方向の情報をうまく扱えますが、予測対象の単語同士の依存関係を学習しにくいという特徴があります。
一方、LSTM等のRNN系ネットワークを用いたSeq2Seqのような自己回帰言語モデルは、順々に単語を読み込ませるため予測対象の単語同士の依存関係を学習できますが、順方向・逆方向の情報を同時に扱えないという問題がありました。
今回紹介するXLNetは、予測対象の単語同士の依存関係を学習できる自己回帰言語モデルでありながら、自己回帰言語モデルの弱点でありBERTの良いところである「順方向・逆方向の情報を同時に扱える」性質を持っています。つまり、両者のいいとこ取りをした手法になっていると著者たちは主張しています。
XLNetと関連研究
Transformer
https://arxiv.org/pdf/1706.03762.pdf
“Attention Is All You Need”というLSTM・CNNを愛用していた人たちに対する挑発的なタイトルでも話題になった論文で提案されたモデルです。
CNNでもLSTMでもないdot product Attentionという機構を提案し、それを積み重ねたモデル(Transformr)で既存手法を大きく上回る成果を上げました。この機構はXLNetだけでなく最近話題になったGPT、BERTにも使われています。
Transformerで使われる(dot-product) Attentionでは、Query, Key, Valueの３つの変数を使います。端的にいえば、Query単語とKey単語を関連性(Attention Weight)を計算し、それぞれのKeyに紐づくValueをかけるという仕組みです。
それをAttention Headを複数用いた（全結合でいうと”隠れ層の数”を増やした）Multi-Head Attentionは以下のように定義されます。上図の(Sigle Head) AttentionはQ,Kをそのまま使っていましたがMulti-Head Attentionでは、各Headに専用の射影行列W_i^Q, W_i^K, W_i^V がついており、それで射影した特徴量を用いてAttentionをかけます。
このdot product Attentionで使うQ,K,Vを同じデータからもってくるとSelf-Attentionと呼ばれます。TransformerのEncoder部分や、Decoder部分の最初のAttentionがそれにあたります。Decoder部分の上の方はQueryをEncoderから、K,VをDecoderから持ってきているのでSelf-Attentionではありません。
実際に適用したイメージを図で描くと以下の図にようになります。この図は、makingという単語をqueryとしてそれぞれのkey単語に対するAttention Weightを算出したものを可視化したものになっています。TransformerではMulti-Head Attentionを用いて後ろの層に伝播させており、それぞれのヘッドは異なった依存関係を学習しています。下図のkeyの単語に複数の着色がされていうますが、それぞれのヘッドのAttention Weightを表したものになっています。
Transformerの詳細な内容については、このブログが分かりやすいのでオススメです。
Transformer-XL
https://arxiv.org/pdf/1901.02860.pdf
しかし、Transformerはdot-product Attentionの性質上、固定長しか扱えません。理想的には長大な固定長を持つモデルを構築すれば対応できますが、Transformerは重量級モデルなので、あまり現実的な選択肢ではないでしょう。この問題を解決しようとしたのがTransformer-XLです。対象文書の前部分(Transformerでは参照できない部分)でgradientをとらないようにすることによって、計算グラフを固定したまま文書全体を参照しながら計算値を出せるようにしています。
その戦略を可能にするために、”Segment-Level Recurrence with State Reuse”と”Relative Positional Encodings”という２つの手法を提案しています。
Segment-Level Recurrence with State Reuse
雑にまとめると、「過去のシーケンスの情報の勾配をとらないことで計算グラフは固定長のままにするが、値はメモリに残しておいて参照する」手法です。この手法をとることで、通常のTransformerが固定長以上のシーケンス情報を扱えないのに対し、固定長以上のシーケンス情報を扱えるようになります。
Relative Positional Encodings
雑にまとめると、位置に対する学習パラメータを導入することで、位置関係の影響を”学習”できるようになったことが改善点のようです。
通常のTransformerとTransformer-XLのMulti-Head Attentionを定式化すると以下のように書けます。違いとしては、Key側の位置エンコーディング項U_jが相対値なっている点、、Query側の位置エンコーディング項U_iが学習パラメータu, vに置き換わっている点、c項・d項のW_qとW_kがの項W_(k,E),W_(k,R)に置き換わっている点、です。
c項で文字情報項の大きさを、d項で相対位置情報の大きさを学習できる仕組みになっています。
ELMo
https://arxiv.org/pdf/1802.05365.pdf
双方向の情報を情報を使って、文脈を考慮した分散表現を得る手法を提案した手法です。具体的には順方向・逆方向の単方向LSTMを別々に学習させたネットワークの中間層を重み付けし、結合したものを分散表現として用います。（BiLM:Bi-directional Language Modelという機構。双方向ですが,別々に学習させるためBi-directional LSTMではないことに注意）
順方向・逆方向両方の情報を用いているので、文脈を考慮した分散表現が取得できるとELMoの著者たちは主張してます。
BERT
https://arxiv.org/pdf/1810.04805.pdf
言わずと知れた2018年の有名モデルです。ELMoでは双方向の情報を使っているが、実際は「予測先が見えてはいけない」という制約がります。
具体的には、”I play tennis every day”という文の”tennis”部分を予測する場合は、順方向でtennisの次に予測対象になる”day”が逆方向では見えてしまいます。この状態では上手く予測モデルを作れないため、ELMoでは順方向と逆方向を別々に学習させた後に活用するという手法をとっていました。
BERTの著者達は、予測したい単語を[mask]にすることで順方向・逆方向の両方の情報を同時に使う手法を提案しました。（実際は[mask]だけで代替するわけではないですが、本筋ではないため詳細は原論文を参照してください）
入力文は下記のようになり、順方向から読み込んでも逆方向から読み込んでも、予測対象の情報が漏洩しないという特色があります。
さらに先行研究として存在していたOpenAIのGPTモデル のように双方向LSTMの代わりに双方向Transformer Encoderを使ってうことで、モデルの表現力を大幅に向上させています。
BERTは様々なタスクにおいてstate-of-the-artを更新しただけではありません。BERTには「事前学習した後にモデルはそのままで多様なタスクにそのまま応用できる」という優れた点があります。本筋ではないので詳細は省きますが、最終層の出力の仕方を変えるだけで、色々なタスクにそのまま応用できる工夫がなされています。この点もBERTが話題になった理由の１つでしょう。
XLNet
さて、やっと本題のXLNetに入ります。
多くの素晴らしい功績を残したBERTですが、XLNetの著者達は２つの問題点を指摘しています。
- [mask]という特殊な文字を使用していので、それらが出現しないfine-tuning時に問題がでてくる
- 後述するようにBERTのような自己符号化モデルでは、予測対象の単語が複数あった場合、自己回帰モデルのように予測対象の単語間の依存関係をとることができない
著者たちは、[mask]を必要としない自己回帰言語モデルを使いながらも双方向の情報を同時に扱える手法と、それにTransformer-XLの技術を加えたXLNetというモデルを提案しています。
次の段落から、いよいよXLNetの大きな特徴となっている【単語の予測順序の入れ替え】と【Two Stream Self-Attention】を説明していきます。
単語の予測順序の入れ替え
自己回帰言語モデルの弱点である「双方向の情報を同時に扱えない」という点を克服するために、「単語の予測順序を入れ替える」手法をXLNetでは使っています。しかし、単語を入れ替えた文書をそのまま入れるのではなく、Transformer-XLで提案されたRelative Positional Encodingsを使って元の文書における位置情報を残しておくことがポイントです。
例として”I play tennis today”という文を予測する場合を考えます。この文書を予測するとき、XLNetでは例えば[tennis play today I]というように並び替えた状態で予測を実施します。
この場合、例えば3番目に位置する”today”を予測するためには、入れ替えた後に１,2番目に位置（元の文書で３,2番目に位置）する”tennis”,”play”との情報を扱って予測をします。
次に4番目に位置する”I”を予測するためには、入れ替えた後に1,2,3番目に位置（元の文書で3,2,4番目に位置）する”tennis”,”play”,”today”の情報を扱って予測をします。
Two Stream Self-Attention
基本的にはTransformer-XLと同様にTransformer Encoder(Self-Attention)をモデルで用います。しかし入れ替えにより単語順序がランダムになっているため、このままでモデルを構築すると、次の単語の出現確率はどの単語でも同じになってしまい、上手く学習できません（入れ替わりがあるため、”次”という位置に意味がない）。
そのため著者達は通常のSelf-Attention(Content Stream Attention)の他にQuery Stream Attentionというものを導入しています。
下図がそれを模式的に表したもので、潜在変数hとgという２つを用いてContent Stream Attention、Query Stream Attentionを構築しています。
Content Stream AttentionとQuery Stream Attentionの違いは、前者は参照位置の単語情報を直接参照できるのに対し、後者はそれができないことです。
例として、”I play tennis today” の並びを入れ替えた[tennis, play, today, I]という文を考えます。(この場合の位置zのシーケンスは[3,2,4,1]となり、その情報をTransformer-XLで提案されたrelative positional encodingを使って付与されています）
(a)のようにContents Stream Attentionの場合は、読み込み済みの[tennis, play, today]（位置:[3,2,4]）に対応する潜在変数h_3, h_2, h_4(単語情報に相当)だけでなく、参照位置の単語である”I”（位置:[1])に対応する潜在変数h_1(単語情報に相当)を参照することができます。参照位置の単語である”I”を含む全ての情報を使ってKey, Valueを構成し、Queryは参照位置の単語である”I”に対応する潜在変数h_1を使ってAttentionをかけています。
一方、Query Stream Attentionの場合は参照位置以外の[tennis, play, today]（位置:[3,2,4]）に対応する潜在変数h_3, h_2, h_4(単語情報に相当)のみでKey, Valueを構成し、Queryを参照位置（位置:[1]）に対応する潜在変数g_1（位置のみ、単語情報なし）からとってきてAttentionをかけています。
そして単語を予測する際は、予測位置のQuery Stream Attentionの情報を使って予測を実施します。
実例
では、XLNetで”I play tennis today”という文を入れ替えた[tennis, play, today, I]という語順で予測を実施する場合を考えてみます。
XLNetでは、この文の単語順を入れ替えた[tennis, play, today, I]という順序で予測を実施しますが、元の文における位置情報[3,2,4,1]をTransformer-XLで提案されたRelative positional encodingを使って付与します。
まず、入れ替えた順序でいうと2番目に位置する”play”を予測する場合を考えてみます。入れ替えた順に予測するので、入れ替えた後に１番目に位置（元の文書で３番目に位置）する”tennis”の情報とメモリに保存した情報を扱います。メモリに保存されているのはTransformer-XLで”Segment-Level Recurrence with State Reuse”(勾配はとらないが情報は使う)の部分です。
上図では、簡単のため”play”を予測するために必要なg_2^(2)に関わる情報のみに焦点を当てて図示し、この時点では解禁されていない層g,hを薄く表示しています（薄い矢印は"play”予測のためには必要ないですが、情報が伝播していることを示すために表示しました）。
ここで注目してほしいのは、予測のために必要なg_2^(2)は、『”play”の位置情報[g_2^(1)]』と『”play”より前の時系列の単語情報[h_3^(1), mem]』の２種類をみていることです。Transformer、Transformer-XLでは後者の『”play”より前の時系列の単語情報』のみを見て”play”を予測しおり、前者の情報は使 っていません。ここが大きな違いになっています。
次に[tennis, play]から”today”を予測する場合を考えてみます(下図)。少し煩雑になりましたが、ここでも同様に、”today”を予測するために必要なg_4^(2)は、『”today”の位置情報[ g_4^(1)]』と『”today”より前の時系列の単語情報[h_3^(1), h_2^(1), mem]』の２種類をみています。
次に[tennis, play, today]から”I”を予測する場合は下図のようになります。かなり煩雑になりましたが、ここでも同様に”I”を予測するために必要なg_1^(2)は、『”I”の位置情報[ g_1^(1)]』と『”I”より前の時系列の単語情報[h_3^(1), h_2^(1), h_4^(1), mem]』の２種類をみています。
この文はここで終わりですが、文に続きがある場合はTransformer-XLのように、最初の部分を勾配を取らないパラメータとして固定（メモリに入れる）し、ずらして文の続きを学習させます。もちろん文の続きとは、元の文の順番の続きではなく「並び替えられた文の続き」です。
上の図では少し煩雑になりすぎたので、”I”を予測するという状態でg_1^(2)とh_2^(2)が何を見ているかを可視化してみます。
まずg_1^(2)です。１層目のhの依存関係を省略しました。予測対象時刻の位置情報gをqueryとして、それより前の時刻の単語情報[play tennis today]を見ていることがわかります。
次にh_2^(2)です。gはqueryとして使っている時刻の情報はKとVでは使わず、KとVはそれより前の時系列の単語情報を使っていました。しかし、hではqueryで使っている時刻の情報もKとVで使うことができます。そして、Q・K・V全て扱っているものは単語情報です。（通常のSelf-Attentionと同じ)
つまり、Two Stream Attentionでは、「予測したい時刻の位置情報と、それより前の時刻の単語情報の関係性」を扱うQuery Stremと、「予測したい時刻より前の単語情報同士の関係性」を扱うContent Streamで、バラバラにされた文章をうまく再構成している機構になっています。
BERTとの比較
まず、順方向・逆方向両方の情報を扱えるBERTとXLNetを比較を、具体例として”New York is a city”という文を推論する場合で考えてみます。
BERTでは、予測対象の単語にマスクをかける([mask]という特殊文字で置き換える）ので、入力文は以下のようなります。
[mask] [mask] is a city
一方XLNetでは、文の入れ替えを行うため以下のようになります。（ここでは例として”is a city New York”という入れ替え方になったとします）
New予測時の入力： is a city
York予測時の入力：is a city New
よって、BERT, XLNetの目的関数は以下のように書くことができます。ここで注目してほしいのは、XLNetの目的関数の第２項で予測対象の単語である”New”と"York”の依存関係が現れているところです。著者達は予測対象の単語同士の依存関係が学習できる点を、BERTと比較した際のXLNetの利点として挙げています。
自己回帰言語モデルとの比較
次に自己回帰言語モデルとXLNetを比較しています。
通常の自己回帰モデルは順方向、もしくは逆方向に文書を読み込ませるため、片方向の依存関係しか学習することができません。
例えば、”New York is a city”を順方向に読み込ませる場合だと、p(York | New) は学習できますが、p(New | York)は学習できません。ELMoのように順方向・逆方向両方の情報を扱う手法もありますが、前述のように事実上別々に学習させているため不十分そうです。
一方、XLNetでは学習時に単語の予測順序の入れ替えを行うことによって、順方向・逆方向両方の依存関係を学習することができます。
結果
多いので、一部を抜粋して掲載します。
まず、RACEというデータセットでは、BERT-Large(モデルが大きな高性能Ver.のBERT)を超える結果となっています。
SQuADではBERTのアンサンブルを単一モデルで超えています。
他にも色々なデータセットでstate-of-the-art更新しており、結果20タスクでBERTを超える結果を示しています。
最後に
GithubのREADMEに記載されていますが、XLNetではTransformer-XLのように過去のシーケンスを保持するためか大量の高性能GPUが必要のようです（著者達はTPUで実験）。BERTやBigGANsもそうですが、最近のSOTAモデルは「金(大量のGPU/TPU)で殴る」ものが多い気がします。家の低性能GPUでも学習して遊べるSOTAモデルでないかなぁと切実に思う今日この頃です。
解釈が間違っている点等あれば、指摘していただけると嬉しいです！
謝辞
Twitterを通して@stomohide様より誤っている部分のご指摘と、その箇所に関する議論をしていただきました。この場を借りて厚くお礼申し上げます。
