yoheikikuta / paper-reading

Notes about papers I read (in Japanese)
156 stars 4 forks source link

[2019] ALBERT: A Lite BERT for Self-supervised Learning of Language Representations #38

Open yoheikikuta opened 4 years ago

yoheikikuta commented 4 years ago

論文リンク

https://arxiv.org/abs/1909.11942

公開日(yyyy/mm/dd)

2019/09/26

概要

parameter sharing と embedding の factorization によりパラメタ数を少なくした BERT モデルの提案。その上で hidden dim を増やし、パラメタ数はオリジナルの BERT より 1/10~1/5 程度に抑えつつより高い性能を発揮するモデルを構築した。 Next Sentence Prediction (NSP) が簡単なタスク過ぎてモデルの性能向上に理解していないのでは、という仮説のもと、連続する文章を正例はそのまま負例は順番を入れ替えて正しい順番かを当てる Senetnce Order Prediction (SOP) を導入。これはある意味で NSP を包含するより難しいタスクで、モデルの性能向上に寄与することを示した。 20191013 現在、GLUE 、SQuAd、RACE の leaderboard でトップを走るモデルとなっている。

yoheikikuta commented 4 years ago

BERT 以降のモデルのアップデートはそれなりに watch していて、この issue でも BERT, Transformer-XL, RoBERTa などのメモを残している。

性能向上だけでなく軽くて使いやすいモデルも研究が進んでいくといいなと思っていたところに軽いモデルであるという ALBERT の話を聞いたので読んでみた。 軽いというのは parameter sharing に依るところが大きいので、確かにモデルサイズや学習時間に関してはだいぶ改善しているが、inference time には特に改善がない(はず、論文では議論してない)。

実験がメインの論文だが、公式実装が公開されてないのがちといただけない。

yoheikikuta commented 4 years ago

モチベーションはモデルサイズと性能の関係にある。

NLP のモデルの性能はここ数年で大きく向上しているが、それに応じてモデルサイズも大きくなってきている。この辺は画像認識モデルと同じような歴史を辿っていると思う。

そしたら可能な限りモデルを大きくすりゃいいのではないかと思うが、(画像認識モデルでも同じように)物事はそう単純にはなっていない。 それを表す一つの例として、オリジナルの BERT-large モデルの hidden dim を単純に二倍にしたモデル (BERT-xlarge) での振る舞いを提示している。それが以下だ。

training loss は上下動が少し激しくなるくらいだが、Masked Language Model (MLM) の dev accuracy に関しては明白に悪くなっている。

このような背景を踏まえて、単純に学習可能なパラメタ数を増やすのではなく、効率的にしかし高い性能を誇るような BERT モデルを構築したいというのが目的となる。

yoheikikuta commented 4 years ago

この論文の貢献は以下の二点。

これらに加えて 概要 のところで書いた SOP も pre-training に採用している。 細かい議論はあるが、技術的にはもうこれで尽きている。

学習時にスケールアップをするためには微分情報をうまく取り扱うことでメモリを節約するものが提案されたりしているが、この論文のように parameter sharing してしまえば(メモリ節約の観点からは)簡単だ。 ちなみに Transformer でも parameter sharing をしているが、これはあくまで機械翻訳のモデルで encoder-decoder での parameter sharing である。

NSP の必要性に関しては登場以降懐疑的な見方をする研究がいくつか出ていたが、この論文では解くのが簡単過ぎるから効いてないのだとして、より難度が高い SOP を導入している。

yoheikikuta commented 4 years ago

もう少し具体的に書く。

embedding の factorization は、dim(E) = dim(H) のために $ O(V × H) $ となっていたところを、一旦 dim(E) を小さくしてそこから H につなぐことで $ O(V × E + E × H) $ とするというものだ。

実にシンプル。 こうしても良いだろうという理由として、モデリングの観点では context-independent な token embedding に対して hidden は context-dependent なものを学習する(つまり後者の方が含む情報が多い)ので、dim(E) << dim(H) がより効率的な設定だろう、という話。

parameter sharing は、一つの理由づけとして Deep Equilibrium Models を挙げていて、これはまだ詳しく読んでいないが parameter sharing をした設定で層を無限に深くする極限で embedding の in/out が固定点に収束する(そしてそれは性能の面でも良いモデルになっている)ということらしい。 特に良い理由づけになっている感じはないが、それを踏まえてこの論文ではえいやと全部 parameter sharing をしてしまうことにしている。

ちなみに embedding の in/out の振る舞いは、L2 distance や cosine similarity を見てみると以下のようになっていて、これは Deep Equilibrium Models とは違う振る舞いをしているとコメントしている(以下のグラフでは何かしらの平衡に到達している感じはなく、layer ID が大きいところでちょっと振動しているということっぽい?)。

NSP が簡単である、というのはこの問題が topic prediction と coherence prediction を合わせたようなものだからだ、と主張している。 余談だが、この coherence prediction ってのは見たことがなくて、広く使われて定着している言葉ではなくて単に文章の一貫性を見極める予測問題、という意味で使ってるだけっぽいかな。 なんにせよ、NSP は negative example を違う文章から持ってくるので、文章の連続性とかじゃなくてトピックが同じか否かで当てられる可能性がある(例えばスポーツの話題にあとに政治の話題が来たりするわけで、それを区別するのは簡単)。

ということで文章の一貫性により注目できるタスクとして Sentence Order Predition (SOP) を導入。 これは同じ文章から取った segment A, B に対して、正例はそのまま(なのでこれは NSP の正例と同じ)で、負例は B, A と並べ替えて正しい順番か否かを当てる問題である。 この場合、先に述べた topic prediction とはならないのでうまく機能することが期待でき、実装上もラクだ。

yoheikikuta commented 4 years ago

モデルのセットアップは以下。 embedding dim が 128 で、parameter-sharing してるので parameter 数が少ない点がポイント。

MLM は n-gram masking というのを使っている。これ知らなかったが、単純な mask ではなくて最大 3-gram になるように確率的に mask の作り方を変えるというものらしい。実装上細かいところは少し気になるが、アイデアとしては難しいところはない。

あとは lr = 0.00176 の LAMB optimizer で batch size = 4096 で Cloud TPU V3 で 125,000 steps 学習しましたとのこと。

LAMB optimizer ってのは知らなかったけど この論文 で導入されたものらしい。全然読んでないが batch size が大きい場合に layer-wise に adaptive に optimize するものっぽい。

yoheikikuta commented 4 years ago

あとは各種実験が並んでいる。

提案モデルの性能を示す主要な結果は以下。

BERT-large と ALBERT-xxlarge を比べると、パラメタ数の比較の意味で確かに性能が大きく向上している(average で 3.6 ポイント)。

parameter sharing はしているものの、ALBERT-xxlarge は hidden dim = 4096 なのでパラメタ数は決して少なくはない。 そしてあまり意味のない BERT-xlarge をベースにして training Speedup を比較しているのは流石にひどいと思う。ここは BERT-large と比べるべきだろう。BERT-large と比べると学習時間は 3 倍くらいだろうし、inference time はかなり遅いだろう。あとで学習時間を同じくらいとして比べたときにも良い、という実験もしているがそんなに impressive な結果には思えない。 ALBERT は軽くてしかも性能も良い!的な話を目にするが、実態は parameter-sharing によって保持すべき独立なパラメタ数が少なくなる、ということなので注意すべきと思う。practical な観点からも素晴らしい、というところまではまだ到達してない感じがする。

とはいえ parameter-sharing をしてこのように性能が出せるというのは興味深いし、研究が進んでいけばもっと良くなりそうな期待は持てる感じ。

yoheikikuta commented 4 years ago

あとはパラメタを変えて性能を比較したり ablation study をしたり、という構成。 自分的に気になったものだけ取り上げよう。

まずは parameter sharing のもの。 一つ目は parameter sharing の対象を変えることでどうなるかと、二つ目は parameter sharing をして layer 数を増やしていくとどうなるか。 前者はコスパを考えると全部 share してしまう方が良い、後者は深くするほど良いのではなくどこか適当なところで止めておくのが良い、という結果。

この段階だと単なる実験結果の羅列だけど、Deep Equilibrium Models との関連で言えば layer を増やしたときに性能が converge しているという結果になったりしないかは興味がある。 この論文で見た振る舞いは報告されているものとは異なると主張しているけど、最終層に近いところで変な振る舞いをしているので SOP の影響とかもありそうだし、まあまあ掘れるところはあるのではないか。

yoheikikuta commented 4 years ago

hidden size も変えて実験して 4096 を選んでいるが、これは増やしすぎは良くないことは想像できるのでそんなに面白い結果ではない(ので載せない)。

もう一つ面白いのは pre-training のデータを増やしたときの振る舞いの変化である。 MLM の性能は大きく上がっているが、down stream のタスクに関しては大して伸びていない。RoBERTa とかではデータ量が結構効いていたので、ここはどのように解釈できるのかは自分としてはよく分からない。 あまりはっきりと物が言えそうな領域とも思えないが、なかなか気になるところだ。

yoheikikuta commented 4 years ago

NSP と SOP の比較は以下。これはなかなか面白い結果だと思った。 まず Sentence Prediction (SP) 系のタスクを入れると、MLM の性能は落ちる。これは解こうとしているタスクが異なるので完全には両立はせずに少し MLM が犠牲になっているように見える。 次に、NSP で学習したものは SOP タスクを解けないが、SOP で学習したものは NSP を解ける。これは SOP がある程度 NSP を包含したより広いタスクになっていることを示唆していて、論文の狙い通りといったところ。

この結果は、MLM は pre-training として大成功を収めたと言って良いだろうが、downstream のタスクを考えればそれ以外の pre-training も重要であることを示唆している。

そのため、downstream タスクに合わせた pre-training 手法の考案(そのままやるとモデル学習のコストが高すぎるので、MLM だけの pre-trained model をベースに SOP のような pre-training はむしろ fine-tuning の方に含めてしまうとか)や、可能であれば互いに協調するような pre-training 手法の考案、などができれば面白いと思う。

yoheikikuta commented 4 years ago

ということで一通り読んだ。

テクニカルには、token embedding の次元を小さくしてから hidden で大きく戻すことと、layer 間の parameter sharing でパラメタ数を抑えることができて(SOPも導入して)性能は落とさない、というものだった。

特に発展が早い領域なので、どんどん良いものが出てくることを期待しよう。