ZETA: Leveraging Z-order Curves for Efficient Top-k Attention

要約

近年、Transformer はシーケンス モデリング アーキテクチャの基本的な構成要素となっています。
しかし、その中心となるのは自己注意の使用であり、そのメモリと計算コストは​​シーケンス長 $N$ に応じて二次関数的に増大し、長いシーケンスでは法外に高価になります。
有望なアプローチは、トップ $k$ アテンションです。これは、$k$ 最も関連性の高いトークンのみを選択し、スペースと計算需要を大幅に削減しながら、バニラのセルフ アテンションと同等のパフォーマンスを達成します。
ただし、因果マスクでは、現在のクエリ トークンが過去のトークンのみに対応する必要があるため、既存の上位 $k$ アテンション手法では最も関連性の高いトークンを並行して効率的に検索することができず、トレーニング効率が制限されます。
この研究では、\textbf{E}効率的な \textbf{T}op-$k$ \textbf{A}ttention の \textbf{Z}-Order Curves を活用して、全体の過去のトークンの並列クエリを可能にする ZETA を提案します。
シーケンス。
% 空間計算量と時間計算量の両方で $\mathcal{O}(N \log N)$ になります。
まず、キーとクエリの次元の選択には、次元の呪いと射影後の相対距離の保持との間のトレードオフが含まれることを理論的に示します。
この洞察を踏まえて、値とは対照的にキーとクエリの次元を削減し、さらに $Z$ オーダー曲線を活用して低次元のキーとクエリを \emph{1} 次元空間にマッピングすることを提案します。これにより、並列ソートが可能になります。
により、上位 $k$ のトークン選択の効率が大幅に向上します。
実験結果は、ZETA が \textsc{Multi-Query Associative Recall} 合成タスクでは標準的な注意のパフォーマンスに匹敵し、\textsc{Long Range Arena} および \textsc{WikiText-103} 言語モデリングでは注意とそのバリアントよりも優れていることを示しています。

要約(オリジナル)

Over recent years, the Transformer has become a fundamental building block for sequence modeling architectures. Yet at its core is the use of self-attention, whose memory and computational cost grow quadratically with the sequence length $N$, rendering it prohibitively expensive for long sequences. A promising approach is top-$k$ attention, which selects only the $k$ most relevant tokens and achieves performance comparable to vanilla self-attention while significantly reducing space and computational demands. However, causal masks require the current query token to only attend to past tokens, preventing the existing top-$k$ attention method from efficiently searching for the most relevant tokens in parallel, thereby limiting training efficiency. In this work, we propose ZETA, leveraging \textbf{Z}-Order Curves for \textbf{E}fficient \textbf{T}op-$k$ \textbf{A}ttention, to enable parallel querying of past tokens for entire sequences. % in both space and time complexity of $\mathcal{O}(N \log N)$. We first theoretically show that the choice of key and query dimensions involves a trade-off between the curse of dimensionality and the preservation of relative distances after projection. In light of this insight, we propose reducing the dimensionality of keys and queries in contrast to values and further leverage $Z$-order curves to map low-dimensional keys and queries into \emph{one}-dimensional space, which permits parallel sorting, thereby largely improving the efficiency for top-$k$ token selection. Experimental results demonstrate that ZETA matches the performance of standard attention on the synthetic \textsc{Multi-Query Associative Recall} task and outperforms attention and its variants on \textsc{Long Range Arena} and \textsc{WikiText-103} language modeling.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google