Loki：基于低秩键的高效稀疏注意力
机器学习
2024-11-11 v2
摘要
大语言模型（LLM）的推理在序列长度较大时会涉及高昂的计算和内存成本。特别是，LLM 推理中使用的自注意力机制在这些成本中占据重要比重，这催生了对自注意力计算进行近似以降低成本的兴趣。本文提出通过聚焦注意力块中计算的键向量维度来近似自注意力。我们的分析表明，键向量在多个数据集和模型中均存在显著低维空间。利用这一发现，我们提出了 Loki，这是一种新颖的稀疏注意力方法，基于注意力得分在低维空间中对 KV 缓存中的标记进行排序和选择。我们的评估显示，Loki 由于数据移动（加载/存储）和计算成本的减少，能够加快注意力计算，同时比其他流行的近似方法在保持模型效果方面更优。
引用
@article{arxiv.2406.02542,
title = {Loki: Low-rank Keys for Efficient Sparse Attention},
author = {Prajwal Singhania and Siddharth Singh and Shwai He and Soheil Feizi and Abhinav Bhatele},
journal= {arXiv preprint arXiv:2406.02542},
year = {2024}
}
备注
Proceedings of the Thirty-Eighth Annual Conference on Neural Information Processing Systems (Main Conference Track)