丸ごと理解できるAIとしてReformerモデルが発表され話題になっています。今回はこのReforerモデルが発表された論文の解説を行います。自然言語や音楽、動画などのSequentialデータを理解するには広範囲における文脈の依存関係を理解する必要があり困難なタスクです。need"の論文で紹介されたTransformerモデルは広くこれらの分野で用いられ、優秀な結果を出しています。機械翻訳などで有名なBERTはTransferモデルが基になっており、数千語にも及ぶコンテキストウィンドウが使われLSTMに比べて長い文脈を考慮することができます。このようにモデルの規模が大きくなってくるとリソースの問題が発生し、大きな研究機関以外はモデルの訓練が行えない状況です。ReformerはTransferを改良し軽量化することで、1つのアクセラレータ、16GBのメモリで100万ワードに及ぶ文章を扱えるようにしました。この記事ではTransformerモデルがどう改善されたかを解説していきますので、Transformerモデル自体の解説については過去の記事を参考にして下さい。大規模なTransformerモデルが抱える問題まずは従来のTransformerモデルを大規模化したときにどういった問題が起こるか説明します。attentionが使われています。$$トークンの長さを\(L\)とすると内積\(QK^T\)は計算量、メモリ量ともに\(O(L^2)\)となります。$$とメモリ消費が激しくなります。このように小説のような非常に長い文章をTransformerで扱いたい場合に、Attention層で必要な計算量、メモリ量が問題になります。Activationを保持しておくためのメモリ量の問題もう1つの問題が逆伝播に必要な各層のActivationを保持しておくためのメモリ量の問題です。保持するのに約2GB必要です。$$となるので、1つの層についてActivationを保持するためにさらに2GMのメモリが必要になります。典型的なTransformerモデルが12個以上の層を持つので、Activationを保持しておくだけで24GB以上必要になります。このようにモデルに与える文章が長くなるとすぐにメモリを使い果たしてしまいます。Attention)Attention層での処理で問題となるのは、内積\(QK^T\)の処理でした。(LSH)によって解決します。内積\(QK^T\)の処理が問題ですが、結局のところ知りたいのは\(softmax(QK^T)\)の結果です。softmaxの結果は内積の値が大きい要素に寄与するため、queryとkeyの内積が大きくなるペアの計算結果だけを用いて近似値を求めることができます。queryとkeyのすべてのペアについて内積を計算するのではなく、queryに対して類似したkeyだけを考慮すればよいので処理が効率化されます。このとき、queryと類似したkeyを選ぶのにLSHを使います。詳しい解説はここでは避けますが、LSHは高次元のデータを確立的処理によって次元圧縮する手法です。Reformerではハッシュ値の計算は次のようにして求めます。以下の画像で示すようにベクトルをランダムに回転させてどの領域に移るかによってどのバケットに入れるかを決定します。以外、xとyの移る先が異なるためそれぞれ違うバケットに入ります。一方で下の例では、xとyが3回の回転ですべて同じ領域に移るため同じバケットに入ります。次にLSHを利用してどのようにAttentionの計算を効率化するか具体的なステップを以下の図とともに説明します。図中の色はどのハッシュに割り当てられたかを示しており、似た単語は同じ色で表されます。LSHによってハッシュを割り当てて、ハッシュ値によって単語を並べ替えます。並べ替えた後、小さな塊に分割することで並列処理を可能にします。Attentionの処理を同じ塊の中および1つ前の塊に対して行うため大幅な計算不可の軽減ができます。Qの\(i\)番目の要素を表しており、\(\mathcal{P}_i\)はi番目のqueryに近いkeyの集合を表しています。\(\mathcal{z}\)はsoftmaxの分母の部分だと考えて下さい。n)\)まで削減することができます。これでAttentionの問題を解決することができました。Transformer)Attentionの問題は解決されましたが、学習時に逆伝播のためにActivationを保持しておかなければならずメモリ消費が大きくなる問題が残っています。Reformerではこの問題を解決するために順伝播時にActivationをメモリ内に保持しておくのではなく、逆伝播時に再計算する方法を取りました。(RevNet)を応用します。RevNetでは出力側から順次1つ先の層の結果を元にActivationを再計算します。通常のネットワークでは(a)のようにベクトルが通過するスタックに各層が追加されていきます。一方でRevNetでは各層において2つのActivaltionを持ちます。(b)のようにそのうち1つだけが通常のネットワークと同じように更新されます。1つのActivationは(c)のようにもう一方のActivationとの差分を捉えるために使われます。これにより順伝播時にActivationを保持する必要がなくなり、逆伝播時に出力側から順次再計算することでActivationを再現することができます。これでActivationを保持するためのメモリを削減することができました。最後にReformerのパフォーマンスを見てみましょう。実験には以下の入力のサイズが非常に大きいタスクに対して行われました。Transformerと通常のTransformerの性能を比べてみましょう。Transformerが通常のTransformerとほぼ等しい性能を見せました。Transformerによってメモリ消費を抑えても性能が犠牲になることはないとわかります。AttentionがTransformerの性能にどう影響するかみていきましょう。Attentionとの比較です。hashesはLSHでのハッシュの割り当て処理を何回行うか示しています。割り当てを複数回行った後、それぞれのラウンドでqueryと同じバケットに入ったすべてのkeyをAttentionの計算に利用します。このようなことをするかというと、LSHでは確率的にハッシュを割り当てるので、確率は低いですが類似した要素がちがうバケットに入る可能性があるからです。attentionと同じ性能になっています。最後に、以下のグラフは入力シーケンスの長さに対する学習にかかる時間を示しています。Attentionはほぼ一定になっていることが見てとれます。まとめReformerはTransformerモデルをリソース面で改善することで16GBのメモリ、単一のアクセラレータで最大で100万語の文章の処理を可能にしました。文章だけでなく画像や動画を扱うタスクへの応用も期待できます。大きな研究機関以外でも非常に長いシーケンスを扱えるようになる可能性も秘めており、AIの民主化という観点でも今後期待が高まりそうです。