Efficient Sparse Attention needs Adaptive Token Release

近年来，大型语言模型（LLM）在各种以文本为中心的任务中展示了显著的能力。然而，它们的“大规模”特征引入了重大的计算和存储挑战，特别是在管理变压器的键值状态方面，这限制了它们的广泛适用性。因此，我们提出了一种自适应地从缓存中释放资源并重建必要的键值状态的方法。特别地，我们通过一个轻量级的控制器模块来近似理想的前K个稀疏注意力。该模块保留具有最高前K个注意权重的标记，并同时重建被丢弃但未来解码可能变得重要的标记。在自然语言生成和建模方面的全面实验表明，我们的方法不仅在性能方面与全注意力竞争，而且还实现了高达221.8%的吞吐量改进。可用于复制的代码在 https://github.com/WHUIR/ADORE 上。

论文提出了一种自适应释放缓存资源并重建必要的键值状态的方法，以解决大型语言模型的计算和存储挑战。

论文的关键思路是通过轻量级控制器模块来近似理想的前K个稀疏注意力，保留具有最高前K个注意力权重的标记，并同时重建被丢弃但未来可能成为解码必需的标记。

论文的实验结果表明，该方法不仅在性能上与全注意力竞争，而且还实现了高达221.8%的吞吐量改进。研究使用了自然语言生成和建模，代码已经在GitHub上开源。

最近在这个领域中，还有一些相关的研究，如《Sparse and Continuous Attention Mechanisms》、《Learning to Remember Rare Events》等。