[論文レビュー] Combiner: Full Attention Transformer with Sparse Computation Cost

Hongyu Ren, Hanjun Dai+5 more|arXiv (Cornell University)|Jul 12, 2021

Topic Modeling参考文献 44被引用数 28

ひとこと要約

Combinerは注意を構造化された条件付き分布としてモデル化することでサブ二次コストで全注意を提供し、長いシーケンスのテキストおよび画像タスクで強力な結果を達成し、複数のベンチマークでベースラインに匹敵するか凌ぐ。

ABSTRACT

Transformers provide a class of expressive architectures that are extremely effective for sequence modeling. However, the key limitation of transformers is their quadratic memory and time complexity $\mathcal{O}(L^2)$ with respect to the sequence length in attention layers, which restricts application in extremely long sequences. Most existing approaches leverage sparsity or low-rank assumptions in the attention matrix to reduce cost, but sacrifice expressiveness. Instead, we propose Combiner, which provides full attention capability in each attention head while maintaining low computation and memory complexity. The key idea is to treat the self-attention mechanism as a conditional expectation over embeddings at each location, and approximate the conditional distribution with a structured factorization. Each location can attend to all other locations, either via direct attention, or through indirect attention to abstractions, which are again conditional expectations of embeddings from corresponding local regions. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention, resulting in the same sub-quadratic cost ($\mathcal{O}(L\log(L))$ or $\mathcal{O}(L\sqrt{L})$). Combiner is a drop-in replacement for attention layers in existing transformers and can be easily implemented in common frameworks. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach, yielding state-of-the-art results on several image and text modeling tasks.

研究の動機と目的

• 非常に長いシーケンスに対する全注意を用いたスケーラブルなシーケンスモデリングを動機づける。
• Combiner を、サブ二次コスト内で全注意を保持する構造化された条件付き期待値アプローチとして導入する。
• 既存のスパース注意パターンから Combiner を実装して表現力を維持する方法を示す。
• 画像の密度推定と長距離シーケンス課題で最先端の結果を実証しつつ、実装可能性を確保する。

提案手法

• 注意を、クエリ位置に与えられた値埋め込みの条件付き期待値として再解釈する。
• 条件付き分布を、直接の注意項と複数の局所的（抽象化ベースの）注意に分解する。
• 既存のスパースパターンに触発されたサブ二次の局所因数分解スキームを用いて、コスト O(L log L) または O(L sqrt L) で全注意を達成する。
• 局所領域の確率を、スケールドドット積と単純な抽象化（最大プーリングや DeepSets）でパラメータ化する。
• 既存のスパースパターンを全注意バリアントに変換する、Fixed・Logsparse・Axial・Learnable などのいくつかの Combiner の実装を提供する。
• 局所因数分解の表現力を高めるための任意のミクスチャー成分。

実験結果

リサーチクエスチョン

• RQ1 注意分布の構造化因数分解を用いて、サブ二次の計算量とメモリで全注意を実現できるか？
• RQ2 既存のスパース注意パターンを再利用して、効率を犠牲にせずに全注意を得るにはどうすればよいか？
• RQ3 テキストと画像領域における自己回帰および双方向タスクで、Combiner の実証的パフォーマンス向上は何か？

主な発見

• Combiner は各ヘッド内で全注意をサブ二次コストで達成し、スパースバリアントと同等の漸近的複雑さ（例: O(L log L) または O(L sqrt L)）に一致する。
• CIFAR-10 では、Combiner-Axial が 2.77 bits/dim を達成し、いくつかのベースラインを上回る。
• ImageNet-64 では、Combiner-Axial が 3.42 bits/dim を達成し、従来のスパース/トランスフォーマーのベースラインと競合するまたはそれを上回る。
• Wiki-40B の言語モデリングでは、Combiner-Fixed-2k および Combiner-Axial-2k がいくつかのベースラインより低い perplexity を達成し、8k シーケンスを用いるとさらなる向上がある。
• Long-Range Arena ベンチマークでは、Combiner バリアント（Fixed, Axial）は複数タスクで標準の Transformer の性能に匹敵するかそれを上回る。
• Combiner は強力な密度推定および自己回帰/双方向モデリング性能を示しつつ、標準的なフレームワークで実装可能である。