$k$NN Attention Demystified: A Theoretical Exploration for Scalable Transformers

要約

トランスフォーマーは、そのパワーにもかかわらず、自己注意の二次的な複雑さのため、長いシーケンスという課題に直面しています。
この制限に対処するために、$k$-Nearest-Neighbor ($k$NN) 注目のような手法が導入され [Roy, ​​Saffar, Vaswani, Grangier, 2021]、各トークンがその $k$ 最も近いトークンのみに注目できるようになりました。
$k$NN の注目は、トランスフォーマーの効率を高めることに経験的に成功していることを示していますが、その正確な近似の保証は理論的に分析されていません。
この研究では、$k$NN アテンションの理論的枠組みを確立し、自己注意をソフトマックス分布に対する期待として再定式化し、効率的な近似のために $k$NN インデックスを使用した遅延ガンベル サンプリング [Mussmann、Levy、Ermon、2017] を活用します。
このフレームワークに基づいて、マルコフ連鎖ベースの推定などの効率的なサンプリング手法を活用して自己注意勾配を近似する新しい二次二次アルゴリズムも提案します。
最後に、経験的な実験を通じてこれらのアルゴリズムの実際的な有効性を実証し、トレーニングと推論の両方における利点を示します。

要約(オリジナル)

Despite their power, Transformers face challenges with long sequences due to the quadratic complexity of self-attention. To address this limitation, methods like $k$-Nearest-Neighbor ($k$NN) attention have been introduced [Roy, Saffar, Vaswani, Grangier, 2021] enabling each token to attend to only its $k$ closest tokens. While $k$NN attention has shown empirical success in making Transformers more efficient, its exact approximation guarantees have not been theoretically analyzed. In this work, we establish a theoretical framework for $k$NN attention, reformulating self-attention as expectations over softmax distributions and leveraging lazy Gumbel sampling [Mussmann, Levy, Ermon, 2017] with $k$NN indices for efficient approximation. Building on this framework, we also propose novel sub-quadratic algorithms that approximate self-attention gradients by leveraging efficient sampling techniques, such as Markov Chain-based estimation. Finally, we demonstrate the practical effectiveness of these algorithms through empirical experiments, showcasing their benefits in both training and inference.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google