Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attention

大型语言模型（LLMs）现在支持极长的上下文窗口，但是基础的注意力机制的二次复杂度导致了显著的Time-to-First-Token（TTFT）延迟。现有的解决这个复杂度的方法需要额外的预训练或微调，而且经常会牺牲模型的准确性。在本文中，我们首先提供了接近无损稀疏注意力的理论和实证基础。我们发现，运行时动态捕捉特定头部的稀疏模式对于低开销至关重要。为了解决这个问题，我们提出了SampleAttention，一种自适应的结构化和接近无损的稀疏注意力。利用观察到的显著稀疏模式，SampleAttention关注于相邻令牌的固定百分比，以捕捉本地窗口模式，并采用两阶段查询引导的键值过滤方法，自适应地选择一组最小的键值，以低开销地捕捉列条纹模式。全面的评估表明，SampleAttention可以无缝地替代现成的LLMs中的基础注意力，几乎不会损失准确性，并将TTFT与FlashAttention相比减少了高达$2.42\times$。

解决问题：论文提出了一种解决大语言模型中注意力机制复杂度问题的方案，旨在降低模型的TTFT延迟。

关键思路：论文提出了SampleAttention，一种自适应的结构化、接近无损的稀疏注意力机制，通过捕捉头部特定的稀疏模式来降低注意力机制的复杂度。

其他亮点：论文通过理论和实验给出了SampleAttention的优越性，可以在不牺牲模型准确度的情况下，将其应用于现有的大型语言模型中，从而大大减少TTFT延迟。实验使用了多个数据集，并且表明SampleAttention可以取代vanilla attention和FlashAttention。论文还提供了开源代码。

相关研究：最近的相关研究包括《Generating Long Sequences with Sparse Transformers》和《Reformer: The Efficient Transformer》。