[論文レビュー] Reformer: The Efficient Transformer

Nikita Kitaev, Łukasz Kaiser+1 more|arXiv (Cornell University)|Jan 13, 2020

Advanced Image and Video Retrieval Techniques参考文献 18被引用数 323

ひとこと要約

tldr: Reformerは可逆層、チャンク化されたフィードフォワード、 locality-sensitive hashing (LSH) アテンションを導入してメモリと計算を削減し、長いシーケンスでのTransformerみたいな性能をはるかに高い効率で達成します。

ABSTRACT

Large Transformer models routinely achieve state-of-the-art results on a number of tasks but training these models can be prohibitively costly, especially on long sequences. We introduce two techniques to improve the efficiency of Transformers. For one, we replace dot-product attention by one that uses locality-sensitive hashing, changing its complexity from O($L^2$) to O($L\log L$), where $L$ is the length of the sequence. Furthermore, we use reversible residual layers instead of the standard residuals, which allows storing activations only once in the training process instead of $N$ times, where $N$ is the number of layers. The resulting model, the Reformer, performs on par with Transformer models while being much more memory-efficient and much faster on long sequences.

研究の動機と目的

• 長いシーケンス上での大規模Transformerモデルに伴う高いメモリコストと計算コストを動機付ける。
• 性能を維持しつつ、メモリと計算量を削減するアーキテクチャと手法を提案する。
• 長いシーケンスタスクや標準ベンチマークで経験的に検証し、効率向上を示す。
• 共有QKアテンション、可逆層、LSHアテンションがトレーニングダイナミクスと精度に与える影響を評価する。

提案手法

• ドット積アテンションを locality-sensitive hashing (LSH) アテンションに置換し、アテンションの計算量を O(L^2) から O(L log L) に削減する。
• 各層の活性を保存することを回避するために reversible residual layers を使用し、ネットワーク深さに伴う N倍のメモリ増加を排除する。
• フィードフォワード層をチャンク化して、活性を小さなブロックで処理することでメモリ使用量を削減する。
• プロジェクション経路を簡素化するために shared-QK attention を採用し、性能への影響を分析する。
• LSHアテンションの精度を向上させ、バケット衝突を緩和するために複数ラウンドのハッシングを採用する。
• 理論的なメモリ・時間計算量の比較と、タスク間の経験的アブレーションを提供する。

実験結果

リサーチクエスチョン

• RQ1 可逆層はメモリ使用量を削減しつつ、Transformerの性能を維持できるか？
• RQ2 LSHベースのアテンションは長いシーケンスに対して完全なアテンションを十分に近似できるか、そしてハッシュパラメータは精度にどう影響するか？
• RQ3 shared-QK attention とチャンク化されたフィードフォワードはトレーニングダイナミクスと効率にどう影響するか？
• RQ4 非常に長いシーケンスに対するReformerの実用的なメモリと速度の利点は、標準のTransformerと比べてどの程度か？

主な発見

• 可逆Transformerは標準のTransformerの性能に匹敵しつつ、メモリ使用量を実質的に削減し、より深いモデルを可能にする。
• Shared-QK attentionは性能を損なわず、enwik8でのトレーニングを加速することさえある。
• LSHアテンションはハッシュ回数が増えるにつれて完全なアテンションに近づく（例えば、回数が多いほど完全なアテンションに近い結果となる）。
• チャンク化されたフィードフォワード層は、同じパラメータでのTransformerと数値的同値性を保ちながらメモリを削減する。
• Reformerは長いシーケンスでより大きく、メモリ効率の良いモデルを可能にする（例：12–20層構成まで）、enwik8やimagenet-64などの長文コンテキストタスクでより速いトレーニングとメモリ節約を示す。
• WMT14英語-ドイツ語翻訳では、可逆Transformerが異なる構成で競争力のあるBLEUスコアを達成している。例：ベース可逆モデル27.6 BLEU（sacreBLEU 27.4）、ビッグ可逆モデル29.1 BLEU（sacreBLEU 28.4）。