Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers

在自回归Transformer中高效地容纳长序列，特别是在扩展上下文窗口内，由于自注意机制中固有的二次计算复杂度和大量KV内存需求，会带来重大挑战。在本文中，我们介绍了SPARSEK Attention，一种新颖的稀疏注意机制，旨在克服这些计算和内存障碍，同时保持性能。我们的方法集成了一个评分网络和一个可微分的top-k掩码操作器SPARSEK，为每个查询选择一个恒定数量的KV对，从而实现基于梯度的优化。因此，SPARSEK Attention在生成过程中提供了线性时间复杂度和恒定的内存占用。实验结果表明，SPARSEK Attention优于以前的稀疏注意方法，并在训练和推理过程中提供了显着的速度提升，特别是在语言建模和下游任务中。此外，我们的方法可以轻松地集成到预训练的大型语言模型（LLMs）中，只需进行最少的微调，为在各种应用中有效管理长程依赖提供了实用的解决方案。

本论文旨在解决在自回归Transformer中容纳长序列的计算和内存开销问题，提出了SPARSEK Attention，旨在通过选择每个查询的常数个KV对来克服这些计算和内存障碍，同时保持性能。

SPARSEK Attention是一种新颖的稀疏注意机制，通过集成得分网络和可微分的top-k掩码操作器，选择每个查询的常数个KV对，从而实现梯度优化，从而提供线性时间复杂度和生成期间的常数内存占用。

SPARSEK Attention在语言建模和下游任务中表现良好，比以前的稀疏注意方法表现更好，并提供了显着的速度改进。此外，该方法可以无缝地集成到预训练的大型语言模型（LLM）中，提供了一个有效地管理长距离依赖关系的实用解决方案。

近期在这个领域中的相关研究包括《Sparse Transformer》和《Compressive Transformers for Long-Range Sequence Modelling》。