通过不对称分块加速自注意力，显著降低长文本推理的内存与时间开销。
Inference-time sparse attention with asymmetric indexing
- 为查询和键分别设计不同分块，实现数据自适应的稀疏注意力
- 在50万词长文本下，内存查找量减少20倍，推理提速60%
- 仅需离线训练小型查询分类器，兼容预训练语言模型
Transformer模型中的自注意力机制是一种增量式关联记忆，将键向量映射到值向量。现有方法常采用k-means等标准划分方式结合GPU兼容的向量搜索算法来加速自注意力，但此类方法在此场景下效果不佳，主要因为（1）键和查询分布不同，（2）RoPE位置编码干扰桶分配。本文提出Saap（Self-Attention with Asymmetric Partitions），一种非对称索引技术，为键和查询分别使用独立分块，从而实现数据自适应的稀疏注意力。该方法适用于预训练语言模型，只需离线训练一个小型查询分类器。在长上下文Llama 3.1-8b模型上，序列长度达10万至50万词时，Saap通常使需查找的内存占比降低20倍，相较FlashAttention-v2可节省60%计算时间。
原文摘要 · Abstract (English)
Self-attention in transformer models is an incremental associative memory that maps key vectors to value vectors. One way to speed up self-attention is to employ GPU-compatible vector search algorithms based on standard partitioning methods such as k-means. However, such partitioning methods yield poor results in this context because (1) the keys and queries follow different distributions, and (2) the RoPE positional encoding hinders the bucket assignment. This paper introduces Saap (Self-Attention with Asymmetric Partitions), which overcomes these problems. It is an asymmetrical indexing technique that employs distinct partitions for keys and queries, thereby approximating self-attention with a data-adaptive sparsity pattern. It works on pretrained language models and only requires to train (offline) a small query classifier. On a long context Llama 3.1-8b model, with sequences ranging from 100k to 500k tokens, Saap typically reduces by a factor of 20 the fraction of memory that needs to be looked-up, which translates to a time saving of 60\% when compared to FlashAttention-v2.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。