Loki: Low-Rank Keys for Efficient Sparse Attention

这篇文章讨论了大型语言模型的推理过程中计算和内存成本的问题，特别是当使用长序列长度时，其中的自注意机制会大大增加这些成本。为了解决这个问题，近期有一些论文提出了针对推理的稀疏注意力近似方法。本文提出了一种新的方法，通过关注注意力块中计算的关键向量的维度来近似自注意力计算。我们的分析表明，关键向量位于一个显著较低维度的空间中，在多个数据集和模型中都一致存在。利用这个观察结果，我们提出了一种名为Loki的新型稀疏注意力方法，它根据在低维空间中计算的注意力分数对KV缓存中的标记进行排名和选择。我们的评估结果表明，相比其他流行的近似方法，Loki能够更好地维持模型的有效性，同时由于减少了数据移动（加载/存储）和计算成本，加速了注意力计算。

本论文试图解决自注意力机制在大型语言模型推理中计算和内存成本高昂的问题，提出了一种稀疏注意力近似方法。

论文提出了一种基于注意力块中键向量维度的自注意力计算近似方法。研究发现，键向量在一个明显较低维度的空间中，这一观察结果被用来设计稀疏注意力方法Loki。

实验结果表明，Loki能够在减少数据移动和计算成本的同时，比其他流行的近似方法更好地维持模型的有效性。该论文使用了多个数据集进行评估，没有开源代码。值得深入研究的是，稀疏注意力的设计可以应用于其他自注意力模型中。

最近的相关研究包括使用稀疏注意力的其他模型，如Longformer和Sparse Transformer。