提出kNN注意力的理论框架，让Transformer处理长序列更高效
$k$NN Attention Demystified: A Theoretical Exploration for Scalable Transformers
- 将自注意力重构成软最大分布期望，用kNN索引实现高效近似
- 设计新型亚二次算法，显著降低训练与推理时的计算开销
- 理论严谨且实证有效，适合做长序列建模的研究者参考
尽管功能强大，但传统Transformer在处理长序列时受限于自注意力的二次复杂度。为突破此瓶颈，研究者引入了k近邻（kNN）注意力机制，使每个标记仅关注其k个最近邻标记。尽管该方法在实践中表现良好，但其近似保证尚缺乏理论分析。本文建立了一个kNN注意力的理论框架，将自注意力重新表述为软最大分布的期望，并结合懒惰Gumbel采样与kNN索引实现高效近似。在此基础上，我们进一步提出新型亚二次算法，通过马尔可夫链等采样技术近似自注意力梯度。实验表明，这些算法在训练与推理阶段均具显著效率优势。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。