Google AI研究者チームは先頃、Transformerディープラーニングモデルの効率改善バージョンであるReformerをオープンソースとして公開した。アテンション(attention)の計算にハッシュ技術を使用し、Reversible Residual Layer(可逆的残差レイヤ)を採用することにより、Reformerは、100万ワードのテキストシーケンスを16GBのメモリと単一のGPUアクセラレータのみで処理することが可能である。
ICLR 2020で発表予定の論文の中でチームは、同モデルの詳細といくつかの実験について報告している。標準的なTransformerモデルで大量のメモリを必要とする最も大きなソースは、アテンション計算とネットワークレイヤのアクティベーションである。チームはこの２つに、メモリ空間を削減するためのテクニックを適用した。アテンションに必要なメモリを削減するために、Reformerでは、LSH(locality-sensitive hashing)によるアテンション演算の最適化を行って、入力シーケンス長Nに対するメモリ使用量を\(O(N^2)\)から\(O(N\log{N})\)に削減した。また、Reversible Layerを使用することで、すべてのネットワークレイヤのアクティベーション保存を不要とし、最新のレイヤのアクティベーションのみを保存するようにした。これら2つのテクニックを組み合わせることで、大規模なGPUクラスタがなくても、これまでより長い入力シーケンスに対して、より深いモデルのトレーニングが可能になった。研究者によると、
このような大きなコンテキストウィンドウによって、テキストだけでなく、画素や楽譜を対象とするアプリケーションにTransformerを使用して、音楽やイメージを生成することが可能になりました。
Transformerはディープラーニングアーキテクチャとして、BERTやRoBERTA、GPT-2といった最新の自然言語処理(NLP)システムに採用されている。Transformerへの入力は、文章のワードのようなデータアイテムのシーケンスである。出力も同じくシーケンスで、例えば機械翻訳では入力文の翻訳結果である。Transformerのオペレーションで鍵となるのが、入力ワードのペアの重要性を計算するアテンションだ。アテンションの計算は、入力シーケンスの長さの２乗に比例して複雑になる。ネットワークは一般的に何十という数のレイヤを持ち、それぞれがアテンションのコンポーネントを含むため、トレーニングに必要なメモリが単一のGPUの容量を超過して、GPUクラスタによるモデルの並列トレーニングが必要になることが少なくない。
Reformerの最初の改善は、アテンションに必要なメモリ容量を削減したことだ。アテンションの操作は、ベクトルの正規化されたドット積(dot-product)とソフトマックス(softmax)によって行われる。正規化されたドット積は２つのベクタの角距離、つまり、同じ方向を示している程度を効率的に計算し、ソフトマックスがそれらの値の最大値に重み付けを行う。角度空間で互いに最も近いベクトルに対してのみ計算を行うことで、アテンションの計算を簡略化することが可能になる。研究者チームがここで発案したのは、局所性鋭敏型キャッシュ(locality-sensitive hashing)(LSH)関数を使用することだ。この関数は、近接したベクトルに対して同じハッシュコードを出力する。LSHを使用して近接するベクタをグループ化し、近傍間のアテンションのみを計算することによって、精度の低下を数パーセントに抑えながら、アテンションの複雑性を\(O(N^2)\)から\(O(N\log{N}\)に低減することが可能になった。
第2の改善は、表示対象外のレイヤに関するアテンションを保存する必要性を低減したことだ。ニューラルネットワークのトレーニングで使用される標準的なバックプロパゲーション技術では、これらの値を格納するストレージが必要になる。深いネットワーク — すなわち、より多くのレイヤを持つネットワークは、浅いものよりパフォーマンスに優れる傾向のあることが、研究によって明らかになっている。一方で、アクティベーションを格納するメモリコストは、レイヤの数に対して直線的に増加する。そのため、個々のレイヤが必要とするストレージは数GBであっても、数十のレイヤを持つ大規模なネットワークでは、単一GPUのオンボード容量を簡単に越えてしまう。Reformerはこれを、可逆的なレイヤを実装することで解決した。バックプロパゲーション中、可逆的レイヤは自身のアクティベーションを回復することができる。前のレイヤのアクティベーションに対して、逆の計算を実行すればよいのだ。この方法によってReformerでは、精度に影響を与えることなく、最終レイヤのアクティベーションのみを格納しておけばよくなった。
BERTやGPT-2といったTransformerのモデルは、最先端のNLPを実現しているが、これらのモデルのトレーニングコストが極めて高いことから、精度を維持したままその複雑性を低減するための研究が進められている。OpenAIが昨年導入したSparce Transformerでは、アテンションのメカニズムを分解することによって、複雑性を\(O(N^2)\)から\(O(N\sqrt{N})\)に低減している。Googleが先日リリースしたALBERTでは、BERTに比較してパラメータが89パーセント少なくなったが、精度もやや低くなっている。Hacker NewでのReformerに関する議論の中で、あるユーザが次のように話している。
これはすごいことだと思います。大規模なアテンションウィンドウ生成時のリソース消費の急増を漸近的に抑制できたことで、これまで以上に複雑なモデルの構築が可能になるはずです。
Reformerのソースコードは、Traxパッケージの一部としてGitHubで公開されている。