Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction

要約

アテンション オペレータはおそらく、さまざまなタスクで最先端のパフォーマンスを実証してきた変圧器アーキテクチャの重要な特徴です。
ただし、トランスフォーマー アテンション演算子は、多くの場合、計算の複雑さがトークンの数に応じて二次関数的に増加するため、重大な計算負荷を課します。
この研究では、計算の複雑さがトークンの数に比例して増加する新しいトランスフォーマー アテンション オペレーターを提案します。
私たちは、トランスフォーマー スタイルのアーキテクチャが「ホワイト ボックス」アーキテクチャ設計によって自然に生成されることを示した以前の研究を拡張することによってネットワーク アーキテクチャを導き出しました。このアーキテクチャでは、ネットワークの各層は、最大の符号化レート削減目標の段階的な最適化ステップを実装するように設計されています (
MCR$^2$)。
具体的には、MCR$^2$ 目標の新しい変分形式を導出し、この変分目標の展開された勾配降下法から生じるアーキテクチャが、Token Statistics Self-Attendee (TSSA) と呼ばれる新しいアテンション モジュールにつながることを示します。
TSSA は線形の計算とメモリの複雑さを持ち、トークン間のペアごとの類似性を計算する一般的なアテンション アーキテクチャから根本的に逸脱しています。
視覚、言語、および長いシーケンスのタスクに関する実験では、TSSA をトークン統計トランスフォーマー (ToST) と呼ぶ標準的なセルフアテンションに交換するだけで、従来のトランスフォーマーと同等のパフォーマンスを実現しながら、計算効率と解釈性が大幅に向上することが示されています。
また、私たちの結果は、ペアごとの類似性スタイルの注意メカニズムが変圧器アーキテクチャの成功にとって重要であるという従来の通念に多少の疑問を投げかけています。
コードは https://github.com/RobinWu218/ToST で入手できます。

要約(オリジナル)

The attention operator is arguably the key distinguishing factor of transformer architectures, which have demonstrated state-of-the-art performance on a variety of tasks. However, transformer attention operators often impose a significant computational burden, with the computational complexity scaling quadratically with the number of tokens. In this work, we propose a novel transformer attention operator whose computational complexity scales linearly with the number of tokens. We derive our network architecture by extending prior work which has shown that a transformer style architecture naturally arises by ‘white-box’ architecture design, where each layer of the network is designed to implement an incremental optimization step of a maximal coding rate reduction objective (MCR$^2$). Specifically, we derive a novel variational form of the MCR$^2$ objective and show that the architecture that results from unrolled gradient descent of this variational objective leads to a new attention module called Token Statistics Self-Attention (TSSA). TSSA has linear computational and memory complexity and radically departs from the typical attention architecture that computes pairwise similarities between tokens. Experiments on vision, language, and long sequence tasks show that simply swapping TSSA for standard self-attention, which we refer to as the Token Statistics Transformer (ToST), achieves competitive performance with conventional transformers while being significantly more computationally efficient and interpretable. Our results also somewhat call into question the conventional wisdom that pairwise similarity style attention mechanisms are critical to the success of transformer architectures. Code will be available at https://github.com/RobinWu218/ToST.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google