[論文レビュー] Sparse Sinkhorn Attention

Yi Tay, Dara Bahri+3 more|arXiv (Cornell University)|Feb 26, 2020

Natural Language Processing Techniques参考文献 24被引用数 77

ひとこと要約

Sparse Sinkhorn Attention は、メモリ効率の良いスパースアテンションを作成する微分可能なソーティング機構を導入し、ローカルウィンドウを介した準グローバルな文脈を可能にし、ベーシックな Transformers と競合する性能を実現します。

ABSTRACT

We propose Sparse Sinkhorn Attention, a new efficient and sparse method for learning to attend. Our method is based on differentiable sorting of internal representations. Concretely, we introduce a meta sorting network that learns to generate latent permutations over sequences. Given sorted sequences, we are then able to compute quasi-global attention with only local windows, improving the memory efficiency of the attention module. To this end, we propose new algorithmic innovations such as Causal Sinkhorn Balancing and SortCut, a dynamic sequence truncation method for tailoring Sinkhorn Attention for encoding and/or decoding purposes. Via extensive experiments on algorithmic seq2seq sorting, language modeling, pixel-wise image generation, document classification and natural language inference, we demonstrate that our memory efficient Sinkhorn Attention method is competitive with vanilla attention and consistently outperforms recently proposed efficient Transformer models such as Sparse Transformers.

研究の動機と目的

• ドット積アテンションの二乗的メモリコストを削減しつつ、 Sparse アテンション出力を学習する動機づけ。
• 入力ブロックを再配置する微分可能なソーティングベースのアテンション機構を提案し、ローカル計算で準グローバルな文脈を実現。
• オートレグレッシブデコーディングのための因果的バリアント（Causal Sinkhorn Balancing）と、シーケンスを切り詰める動的エンコーディング（SortCut）を導入。
• Sinkhorn アテンションと標準アテンションを組み合わせた混合モデルを検討し、性能を向上。
• 言語モデリング、画像生成、文書分類、自然言語推論において有効性を示す。

提案手法

• ブロック毎の置換行列を出力するパラメータ化されたメタソーティングネットワーク（SortNet）を導入。
• ソーティング行列を二重確率分布（ビルコホ Polytope）へ射影する微分可能な Sinkhorn バランシングを適用。
• 局所的な近傍性を維持しつつ準グローバルなアテンションを可能にするため、個々のトークンではなくブロックをソート。
• ソートされたブロック内でアテンションを計算し、キーとバリューに共通のソーティング演算子を適用。
• 離散的な置換を近似するためにガンベルノイズと温度を組み込む（Gumbel-Sinkhorn）。
• オプションとして Sinkhorn アテンションを標準のバニラアテンションと混合（Mixture model）。
• 因果設定へ拡張（i）累積和に基づく因果ソーティングネットワーク、（ii）将来マスキングを伴う因果 Sinkhorn バランシング。
• SortCut を提供し、ソート後にシーケンスを予算（N_k）で切り詰めて計算量をさらに削減。

実験結果

リサーチクエスチョン

• RQ1 Sparse Sinkhorn Attention は、メモリ複雑性を削減しつつ、ベースのアテンションと競合的な性能を達成するか？
• RQ2 提案されたソーティングベースのスパースアテンション変種（Causal Sinkhorn Balancing、SortCut）は、エンコーディング/デコoding の効率と精度にどのような影響を与えるか？
• RQ3 Sinkhorn とバニラアテンションのミックスは、各アプローチ単独よりも多様なタスクで上回るか？
• RQ4 Sparse Sinkhorn Attention のメモリ複雑性は、実践的には密結合アテンションおよび Sparse Transformers と比較してどうか？
• RQ5 ソーティングのハイパーパラメータ（温度、Sinkhorn の反復回数）が性能に与える影響は？

主な発見

• Sparse Sinkhorn Attention は自己注意のメモリを O(ell^2) から O(B^2 + N_B^2) に削減し、SortCut を用いると特定の設定で線形時間 O(ell N_k) にもなる。
• さまざまなタスク（アルゴリズム的ソーティング、言語モデリング、ピクセル単位生成、文書分類、NLI）において、Sinkhorn Transformers は vanilla Transformers に匹敵するか、上回ることが多く、Sparse Transformers や Local Attention を上回ることが多い。
• LM1B のサブワードモデリングでは、Sinkhorn Transformers はいくつかのブロックサイズで局所アテンションおよび Sparse Transformer よりも優れており、Sinkhorn Mixture が最良の結果を達成。
• CIFAR-10 の画像生成では、Sinkhorn Transformer が試験された効率モデルの中で最も良い Bpd（bytes-per-dimension）を達成。
• SortCut エンコーダは、 vanilla Transformers に匹敵する性能を大幅なメモリ削減で実現可能。
• アブレーションによると、ソーティングネットワークは単純な線形形式で最も良好に機能し、Sinkhorn 正規化（N_k > 0）の適用が性能にとって決定的である。