Transformer大規模化へのブレークスルーとなるか!? 高効率化したReformer登場

3つの要点

✔️Attentionの計算量をO(n^2)からO(n log n)へと劇的に削減

✔️アクティベーションなどのメモリ使用量を大幅に削減

✔️速度・メモリ共に実装効率を大きく改善しながらも、Transformerの性能を維持

Reformer: The Efficient Transformer

written by Nikita KitaevŁukasz KaiserAnselm Levskaya

(Submitted on 13 Jan 2020, ICLR 2020)

subjects : Machine Learning (cs.LG); Computation and Language (cs.CL); Machine Learning (stat.ML)

大規模応用が支えるTransformerの記録的成果

Transformerを使った研究は次々にstate-of-the-artな結果を生み出していますが、その応用はどんどん大規模になっています。そしてこの大規模化の流れは、大きな研究機関以外では研究すること自体が難しくなることを意味します。現在それらが問題として指摘され、議論になっているようです。
ではどれくらいの規模になってきているのでしょうか?
  • Transformer層(以下Trm)ひとつのパラメーター数0.5B(=500,000,000)を保持するのに、2GBのメモリが必要(32ビット浮動小数点)。
  • トークン長64,000、Embeddingのサイズ1,024、バッチサイズ8のときのActivation(順伝播の計算結果)は、同じく 64K x 1K x 8 = 0.5B で2GB。
これらのサイズは、実際に計算するとき致命的ともいえる問題になります。
  • Trmが12層になったら、Activationは 2GB x 12 = 24GB。学習時には逆伝播するまで保持が必要。
  • Attentionの計算は、トークン長Lに対して計算量・メモリ量ともにO(L^2)、つまりたとえバッチサイズ1でもL=64KBのとき 64KB^2 x 4 (32ビット浮動小数点) = 16GB。
この規模になると数枚のGPU構成レベルでは全く取り扱えなくなります。
ではそんな用途の具体例を見てみましょう。

大規模なプロジェクトの例

論文 “Mesh-TensorFlow: Deep Learning for Supercomputers”1)[Mesh-TensorFlow 2018] Shazeer et al., 2018, “Mesh-TensorFlow: Deep Learning for Supercomputers

Googleの機械学習用チップTPUをメッシュで扱う研究論文では、機械翻訳を例題にしています。
Trm層のフィードフォワードを262,144次元、アテンションヘッドを256個にしたTransformerを計算させるのに、TPUを512個メッシュにしたそうです。
とても莫大で、Googleのような規模でないと不可能です。

論文 “Image Transformer”2)[Image Transformer 2018] Parmar et al., 2018, “Image transformer

画像にTransformerを応用した研究例で、中途半端な画像から残りの生成、超解像の応用があります。
超解像では入力8×8ピクセルから出力32×32ピクセルへ、この規模ですら計算量に苦慮し、「Local Self-Attention」という方法で計算利用を落としているほどです。
[Image Transformer 2018]より

論文 “Music transformer: Generating music with long-term structure”3)[Music transformer 2018] Huang et al., 2018, “Music transformer: Generating music with long-term structure

この音楽の生成では、1分の時間を本当は10msごとに区切って6,000個のデータで扱いたいところを、省データな表現で2,000次元まで落としています。
イメージとして文章では2,000単語で一つのデータとするところを、2,000音で一つの曲のデータとして扱うことになります。
それでも、2,000次元だと8ヘッドの場合8.5GBのメモリ消費になって扱えないため、「Relative Positional Self-Attention」でダイナミックにメモリ消費量を削減しています。
この方式を使うことで2,000次元という長い表現を取り扱えて、良い生成ができるようになっている論文ですが、これ以上のデータ量を扱いたいものの、そこがネックになっている現状も同時に表しています。
[Music transformer 2018]より
論文から引用したこのFigure 3からは、「Relative Positional Self-Attention」を使うことで(右側)、ケーデンス(終止、音楽用語)がフレーズの終わりに確認できるということです。つまり通常の楽曲に現れるような進行がうまく生成されるよう改善されたようです。
余談ですがこの図の見方を説明しましょう。これは「ピアノロール」と呼ばれる表現に近いものです。縦軸は鍵盤の位置、横軸は時間です。
音を鳴らしている間をバーで表現していて、下の図では2つの音が鳴らされている様子です。低い方は一小節ごとに音が進み、高い音でメロディを奏でているようです。

Attentionを近似計算で省リソース化〜LSH Attention

キーポイントの一つが、Attentionの処理を速く・小さくすることです。
これを実現する省リソース化の近似計算を行うために、「局所性鋭敏型ハッシュ (Locality-sensitive hashing)」を使いました。1999年から進化を続けたアルゴリズム(「LSH Algorithm and Implementation (E2LSH)」(MIT)より)で、類似した入力データが同じバケット(ハッシュ表の各要素)に入る性質を利用しています。
(ちなみにこのLSHは、機械学習アルゴリズムの一種k近傍法でも使われるものです)
似たデータをまとめて計算量を減らす仕組みのポイントは、Attentionに含まれるSoftMaxにあります。
[Transformer 2017] 4)[Transformer 2017] Vaswani et al., 2017, “Attention is all you needFigure 2 より “Scaled Dot-Product Attention”
SoftMaxに与えられる経路でMatMul(Q, K)のサイズが非常に大きくなるわけですが、SoftMax自体の演算結果は入力のうち大きな値だけがほとんど支配的です。ですからこの大きな値以外の演算を省略しよう、というアイデアなのです。
そしてその大きな値になるのは、QとKが近いときです。そこで、この近いデータを選ぶのに、LSHを使います。そうすることで、元の計算は、
下の計算式のようになります。
ここで i は Q を要素ごとに扱っている意味で、z(・, ・)は元々のSoftMaxの分母だと捉えてください。
j ∈ P_i がポイントで、P_iはq_iに近いk_jだけのjの集合です。つまり、q_iに近くないk_jの計算は省略されるのです。
※かなり省略して説明する方針で書かせて頂きました。正確には論文の計算式を御覧ください。
論文の図でイメージを確認しましょう。
[Reformer 2020] よりFigure 2
左側の図に注目しましょう。上から「LSH bucketing」をすることで、似たデータが集まります。
同じ色が同じバケットを表しているとして、それらの間だけで計算される様子が描かれています。
右側の図は、対角線上に似たデータだけが集まり、それらの間だけで計算されているイメージと捉えることができます。全ての要素間で計算するのではなく、計算する要素を省略しています。
これで、計算量もメモリ量もO(n^2)からO(n log n)になります。

Multi-round LSH attention

実際にはハッシュが同じものだけでなく近接n_r個(round)のバケットまで計算対象とすることで、ハッシュがうまくヒットしないケースもカバーできるようにしています。
roundの値を増やすことで、計算量・メモリ量は殆ど増えずに、近似計算の質を上げることが出来るようです。

アルゴリズムの改善はすごい、O(n^2)からO(n log n)へのインパクト

どのくらいの改善インパクトなのでしょうか。
例えばスピードに関しては、CPUに対してGPUで計算した場合、ざっくり10倍程度速く処理できます。逆に言えば、GPUにしたところでたかだか速度10倍にしかならない、ともいえます。
n O(n^2) O(n log n)
10 100 33
30 900 147
100 10,000 664
1,000 1,000,000 9,966
10,000 100,000,000 132,877
これに対して、O(n^2)とO(n log n)の主な関係を見てみると、n=100ですでにO(n log n)の方がO(n^2)に比べて10,000/664=約15倍です。
Transformerとして一文の単語数100規模は普通の使い方だと思われますが、すでにCPU対GPUに匹敵する、あるいはそれ以上の効果があります。
もしn=1,000単語扱いたいとしたら、O(n log n)の9,966に対してO(n^2)では1,000,000なので、約100倍の差があります。
これはすでにハードウェアでの改善規模を超えてきていて、まさにアルゴリズムでしかなし得ない領域の改善です。
O(n log n) は “linearithmic”、ほぼリニアO(n)と同じです。

リバーシブル化・チャンク化によるメモリ量削減

ここまで進めた省リソース化はAttention周辺でした。残るは、逆伝播に備えて保持しておくActivation(計算結果)と、フィードフォワードの計算のためのメモリ量削減です。

リバーシブル化で保持メモリ量の削減

計算結果を「再現」する先行研究に、ResNets5)[ResNets 2015] He et al., 2015, “Deep Residual Learning for Image Recognitionを基にしたRevNets6)[RevNets 2017] Gomez et al., 2017, “The reversible residual network: Backpropagation without storing activationsがあります。
通常のニューラルネットワークの学習では、一旦順伝播をネットワークの入力から順に出力に向けて計算すると、各所で計算結果を保持します。これらはその後の逆伝播の計算で使われます。
ところが、RevNetsでは各所での計算を保持する必要はなく、出力側から順次再計算して再現することができるのです。
Trm層に含まれるAttentionとFeedForwardのブロックはそれぞれResNets同様skip connectionしています。そこでRevNetsを応用して、Reversible Transformerを構成しています。
以上でリバーシブルになり、途中の計算結果を保持しておく必要がなくなりました。

チャンク化でフィードフォワードの計算時メモリ量削減

フィードフォワード(FeedForward)も大きなサイズで使われるため、この計算で使われるメモリの量も抑える必要があります。
この場合簡単で、各要素の計算に相関がないので小分けにして計算しよう(チャンク化)、というアイデアです。
以上でこの点についてもメモリ使用量の増加は無視できるようになります。

実験

実験の詳細は論文では探すことができなかったので記事「Google & UC Berkeley ‘Reformer’ Runs 64K Sequences on One GPU (medium)」を参照しました。
  • imagenet64 の画像生成タスクで、入力トークン長12K
  • enwik8 というテキストのタスクで、入力トークン長64K
どちらも非常に長い入力サイズです。これらのタスクで、同じ設定を使った通常のTransformerと比較されました。
Reformerを省リソースにする変更点一つ一つについてTransformerとの比較をして、性能に影響が出ないことが確認された上で、LSH Attentionでの近似の影響は次のようなグラフになりました。
通常の「full attention」対して、roundを8まで上げるとほぼ同じ性能になっています。つまり、8つのハッシュ値に合致するものまでを計算対象とすることで、性能を落とさず、計算量を抑えられるようです。roundを上げれば性能が落ちなくなるものの、この辺りは利用できる計算リソースとの兼ね合いで、という論文のコメントでした。
次に大きなReformerでの実験結果です。20層までのReformerでもGPU一つに乗り(!)、長い入力データで速い学習が実現できたということです。
ここでは右側のグラフにも注目です。「full attention」では入力シーケンス長が大きくなるに従い指数関数的に右肩上がり(※logスケール)で時間がかかるのに対して、LSH attemtionでは一定です。不気味な凄みがあります。
精度はどうだったのか、これは論文では前半で語られていて、ほぼTransformerと匹敵することが確認されています。
最後に計算量とメモリ量の比較の表を見ると、Reformerは l (入力データの長さ)、 nl (層の数)、 d_ff (フィードフォワードの数)に対してメモリ量が改善され、計算量は l が一つ減っていることがわかります。

まとめ

これまでご紹介したように、キーとなる改善手法自体はこの論文で発明されたものではありません。それらを道具として使ってTransformerの問題点を解決、応用の活路を開いていることに大きな意味があります。
機械学習の技術はITサービスと似て、あるタスクの解決法になるだけではなく、その解決法がスケールすることが重要です。
Reformerならば、TPUを512個メッシュにできるような大手でなくても画像など大きな入力を扱えるようになる可能性があり、AIの民主化には大きな意味があります。
アルゴリズム改善によるインパクト、もしかすると機械学習時代だからこそ、重要性を増しているかもしれません。
ICLR2020では殆ど完璧なスコア”8, 8, 6″で採択され、今後の研究に大きなインパクトがあるだろうと認められています(OpenReview)。要注目ですね。
今後の発展に、とても期待が高まります。

References   [ + ]

1. [Mesh-TensorFlow 2018] Shazeer et al., 2018, “Mesh-TensorFlow: Deep Learning for Supercomputers
2. [Image Transformer 2018] Parmar et al., 2018, “Image transformer
3. [Music transformer 2018] Huang et al., 2018, “Music transformer: Generating music with long-term structure
4. [Transformer 2017] Vaswani et al., 2017, “Attention is all you need
5. [ResNets 2015] He et al., 2015, “Deep Residual Learning for Image Recognition
6. [RevNets 2017] Gomez et al., 2017, “The reversible residual network: Backpropagation without storing activations

この記事をシェアする