How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation

要約

古典的なトランスフォーマー アテンション スキームでは、3 つの $n \times d$ サイズ行列 $Q、K、V$ (クエリ、キー、および値のトークン) が与えられ、目標は新しい $n \times d を計算することです。
$ サイズ行列 $D^{-1} \exp(QK^\top) V$ ここで、$D = \mathrm{diag}( \exp(QK^\top) {\bf 1}_n )$。
この研究では、三重の相関関係を捉える注意の一般化を研究します。
この一般化により、変圧器では不可能であることが示されていた三重接続の検出に関する問題を解決できます。
この一般化の潜在的な欠点は、単純なアルゴリズムでは $n$ の 3 次時間が必要となるため、計算がさらに難しくなったように見えることです。
しかし、我々は、制限付きエントリ設定 (実際に発生し、理論と実践の両方でよく研究されています) では、実際にはほぼ線形の時間アルゴリズムが存在することを示します。
より正確には、一般化された計算を迅速に実行するには、制限されたエントリが必要かつ十分であることを示します。 $\bullet$ 良い点として、入力行列のすべてのエントリが $o(\sqrt[3]{\log によって上に制限されている場合)
n})$ 次に、$n^{1+o(1)}$ 時間で「テンソル型」アテンション行列を近似する方法を示します。
$\bullet$ マイナス面としては、入力行列のエントリが $\Omega(\sqrt[3]{\log n})$ ほど大きい場合、これより高速に実行されるアルゴリズムはないことがわかります。
$n^{3-o(1)}$ (きめの細かい複雑さ理論からの強力な指数時間仮説を仮定)。
また、構築、アルゴリズム、および下限が高次のテンソルと相関に自然に一般化されることも示します。
興味深いことに、テンソルの次数が高くなるほど、効率的なアルゴリズムを実現するには、エントリの境界を低くする必要があります。
したがって、私たちの結果は、エントリの境界性と、より表現力豊かで効率的なアテンション計算に使用できるテンソルの次数との間の自然なトレードオフをもたらします。

要約(オリジナル)

In the classical transformer attention scheme, we are given three $n \times d$ size matrices $Q, K, V$ (the query, key, and value tokens), and the goal is to compute a new $n \times d$ size matrix $D^{-1} \exp(QK^\top) V$ where $D = \mathrm{diag}( \exp(QK^\top) {\bf 1}_n )$. In this work, we study a generalization of attention which captures triple-wise correlations. This generalization is able to solve problems about detecting triple-wise connections that were shown to be impossible for transformers. The potential downside of this generalization is that it appears as though computations are even more difficult, since the straightforward algorithm requires cubic time in $n$. However, we show that in the bounded-entry setting (which arises in practice, and which is well-studied in both theory and practice), there is actually a near-linear time algorithm. More precisely, we show that bounded entries are both necessary and sufficient for quickly performing generalized computations: $\bullet$ On the positive side, if all entries of the input matrices are bounded above by $o(\sqrt[3]{\log n})$ then we show how to approximate the “tensor-type” attention matrix in $n^{1+o(1)}$ time. $\bullet$ On the negative side, we show that if the entries of the input matrices may be as large as $\Omega(\sqrt[3]{\log n})$, then there is no algorithm that runs faster than $n^{3-o(1)}$ (assuming the Strong Exponential Time Hypothesis from fine-grained complexity theory). We also show that our construction, algorithms, and lower bounds naturally generalize to higher-order tensors and correlations. Interestingly, the higher the order of the tensors, the lower the bound on the entries needs to be for an efficient algorithm. Our results thus yield a natural tradeoff between the boundedness of the entries, and order of the tensor one may use for more expressive, efficient attention computation.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google