Googleは、入力シーケンス長に対して線形スケールするTransformerディープラーニングアーキテクチャのPerformerをオープンソースとして公開した。この特徴によってPerformerは、画素推定(pixel-prediction)やタンパク質配列のモデリングといった、長いシーケンスを必要とするタスクに使用することができる。
Google Research内のチームがarXivで公開された論文の中で、このモデルといくつかの実験について説明している。PerformerはFAVOR+(Fast Attention Via positive Orthogonal Random)と呼ばれる汎化アテンション機構(attention mechanism)を使用して、Transformerモデルで一般的に使用される標準的なソフトマックスアテンション(softmax attention)の高精度な近似を行うことで、空間的および時間的複雑性を2次から線形に低減する。複雑性を低減することにより、Performerは、通常のTransformerでサポートされているよりも長いシーケンス長を必要とするアプリケーション内での使用が可能になっている。さらに、FAVOR+アテンション機構は、既存のTransformerモデルと完全な後方互換性を持つことにより、スパースアテンション(sparse attention)など他の高効率アテンションスキームに対して優位な位置にある。チームメンバのKrzysztof Choromanski、Lucy Colwell両氏が、Googleのブログに次のように記している。
私たちの研究は、アテンションやTransformerアーキテクチャ、さらにはカーネルメソッドについて考える新たな方法を開拓したと、私たちは考えています。
Transformerニューラルネットワークアーキテクチャはシーケンス学習( sequence learning)、特に自然言語処理(NLP)の分野においては一般的な選択肢である。その理由になっているのが、再帰型ニューラルネットワークアーキテクチャ(RNN)などに対するアドバンテージの存在だ。中でも、シーケンス内の前方の項目を"記憶"できるセルフアテンション機構は、シーケンス全体の並列動作を可能にすることで、トレーニングや推論の高速化に寄与している。一方でセルフアテンションは、シーケンス内の各項目が他のすべての項目とリンク可能であるため、その演算上およびメモリの複雑性は、処理される可能性のある最大シーケンス長をNとした場合、\(O(N^2)\)に達する。これが1,024アイテム程度のシーケンス長という、GPUのメモリ制約による実用上の限界を設けることになっているのだ。
Transformerのオリジナルのアテンション機構は、サイズNxNの行列と、それに続くソフトマックス演算によって実装されている。行と列がそれぞれクエリとキーに対応し、アテンション行列に入力シーケンスが乗算されることで、一連の類似性(similarity)値が出力される。PerformerのFAVOR+アルゴリズムはこの行列を"ランダムフィーチャ(random features)" — クエリとキーのランダムな非線形関数を含む2つの行列に分解(decompose)する。研究チームは、この分解によってオリジナルのアテンションを任意の精度で近似すると同時に、計算ストレージの複雑度を\(O(N)\)まで削減できることを示した。さらにこのアルゴリズムは、ソフトマックス以外の他の類似度演算(similarity operation)も可能であるため、さらに一般化されたアテンション定義を生成することができる。
より長いシーケンスでのトレーニングの有用性を実証するために、チームはPerformerを使用してタンパク質シーケンス"言語モデル"を開発した。このモデルでは、たんぱく質の"単語"がアミノ酸の"文字"の線形シーケンスで表現される。これらのシーケンスでトレーニングされたモデルは、結果として得られるタンパク質分子に関する幾何学的情報の予測に使用することができる。Performerが長いシーケンスをサポートすることにより、いくつかのシーケンスを結合して、タンパク質間の相互作用を予測することが可能になった。最大8,192のアミノ酸からなるこれらの長いシーケンスは、一般的には大規模なTransformerのメモリでもオーバーロードになるものだ。サイズの小さなTransformerでもこのデータでトレーニングすることは可能だが、達成可能な精度は、Performerの24パーセントに対して19パーセント程度に留まる。
最近では、アテンションの複雑性を軽減するためのスキームが、他にもいくつか開発されている。例えば、OpenAIは昨年、アテンション行列のスパース因数分解(sparse factorization)を開発し、ネットワークの複雑性を\(O(N^2)\)から\(O(N\sqrt(N))\)に低減している。Google Researchは先日、局所性鋭敏型ハッシュ(locarity-sensitive hashing、LSH)によってメモリ要件を\(O(N\log(N))\)に低減した近似アテンション演算を使用するReformerを発表した。さらにGoogleは、3つのより小さなアテンション機構を組み合わせて使用するBigBirdも開発している。BigBirdはPerformerと同じく線形複雑度を持っているが、Performerの開発者たちによると、BigBirdの"スタック型"アテンションサブレイヤは既存のトレーニング済モデルで使用することが難しいため、トレーニングの再実施と"非常に多くのエネルギ消費"を必要とする、ということである。加えて、スパース手法は特殊なスパース行列の乗法演算を必要とすることが多く、すべてのハードウェアプラットフォーム上で実行可能ではない場合があるのだ。
Perfomerの高速アテンションモジュールとタンパク質言語モデルのコードはGitHub上で公開されている。