GoogleがTransformerの計算を軽量化したReformerを提案:AIで小説『罪と罰』を全て理解?

課題

言語、音楽、映像などを理解することは、広範囲にわたる文脈に依存している場合、困難なタスクになります。たとえば、映像内で人や物体が表示されなくなったかなり後に再び表示される場合、多くのモデルはその外観を忘れてしまいます。機械翻訳などで主流となっているTransformerモデルはBERTの基になったモデルであり、一度に使用できる文章表現(コンテキストウィンドウ)は数百から数千語に及びます。また、同じ仕組みを用いて、ピクセルや音符などテキスト以外のアプリケーションにも使用でき、それに基づいて音楽や画像の生成も可能です。

しかし、Transformerを数千語よりさらに大きな範囲(コンテキストウィンドウ)に拡張すると、問題が生じます。多数のレイヤーを使用したTransformerモデルは、非常に大きな計算コスト(データ量Nに対してN^2の計算量)が必要なため、段落数が少ないテキストの要約や短い音楽を生成する場合しか使用できませんでした。

解決方法

今回Googleが発表したReformerは、最大100万ワードのコンテキストウィンドウを処理するように設計されたTransformerモデルで、さらに使用メモリは16GBのみです。
Transformerモデルを非常に大きなテキスト・シーケンスに適用する場合の最初の課題は、アテンション層の処理方法です。Transformerモデルでも使用されているアテンションは、Scaled Dot-Product Attention(正規化内積型注意)とも呼ばれます。これは、ある「単語」と周辺単語から計算されたKeyベクトルと特定の処理のために計算されたQueryベクトルの類似度を内積として計算し、その内積の値が高い(つまり、KeyとQueryの値が近い)要素に関して、値が高くなる(似ていない関係では0)ようなマップを作り、これを注意マップとして、この注意マップにより、単語の特徴量の重み付き平均を計算し、この値をその後の処理に使用するといったもです。

Reformerはこの内積型注意におけるKeyとQueryの類似度の計算を局所性鋭敏型ハッシュ(LSH: Locality Sensitive Hashing)に置き換えたものです。元々、LSHは、高次元ベクトルを確率的処理によって次元圧縮を行います。このハッシュ計算により値が確率的に「バケット」に振り分けられることにより、効率的に類似性を判断することができます。

ReformerではLSHを用いて、KeyとQueryの可能なすべてのペアを計算するのではなく、類似の特徴ベクトル(ワード)に対して一致するようなハッシュ関数を計算します。それぞれのワードにハッシュ値が割り当てられると、値が再配置されて同じハッシュを持つ要素がまとめられ、セグメントに分割されて並列処理が可能になります。これにより、どのワードが重要かを示すアテンションが計算され、計算負荷が大幅に軽減されます。

次の図では、異なる色が異なるハッシュを表し、類似する単語は同じ色を持ちます。

https://techable.jp/archives/115600

さらに、Reformerでは、各レイヤの入力をメモリに格納するのではなく、学習時の逆伝搬中(バックプロパゲーション)において、Reversible Layerと呼ばれる仕組みを用いて、オンデマンドで再計算します。これは、各レイヤで適用されたアクティベーションを単純に減算するもので、データをメモリにバックアップすることなく、メモリ使用を効率化できます。

Reformerにおけるこれらふたつのアプローチは、単一のGPUアクセラレータ上で100万語までの長さのテキストシーケンスを処理することを可能にします。

どうなったか

Reformerは画像に関しては、断片的な画像を復元することができるようになりました。このように画像処理や映像処理への応用は大きな可能性を示していますが、テキストへの応用はさらにエキサイティングだとGoogleは述べています。Reformerは小説全体を一度に単一のデバイス上で処理することができ、トレーニング例では『罪と罰』の全文を処理する方法がGoogle Colabのノートブックで示されました。周囲の単語から情報を集約し、文脈における言語の特定の意味を決定することができます。

https://ai.googleblog.com/2020/01/reformer-efficient-transformer.html?m=1

まとめ

Reformerは、未来のTransformerモデルの基礎を提供するだろうとGoogleは述べています。すでに、さらに長いシーケンスに適用する方法と位置エンコーディングの処理を改善する方法を探し始めているようです。Reformer解析などを大規模にできることが特徴であり、大規模化が新しい応用に繋がるでしょう。

参考資料

(蒲生由紀子・森裕紀)