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

TL;DR

This paper provides a theoretical analysis of $k$NN attention in Transformers, introduces efficient approximation algorithms, and demonstrates their practical benefits in training and inference for scalable models.

Contribution

It establishes a theoretical framework for $k$NN attention, proposes novel sub-quadratic algorithms, and empirically validates their effectiveness.

Findings

Theoretical guarantees for $k$NN attention approximation

Development of sub-quadratic gradient approximation algorithms

Empirical improvements in training and inference efficiency

Abstract

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.