论文

每个密钥具有更多价值:非对称稀疏注意力可实现更快的LLM解码

More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding

模型推理推理加速

摘要

大语言模型 (LLM) 中的自动回归生成受到记忆和注意机制的计算需求的限制。稀疏注意力方法通过仅选择注意力矩阵的高概率条目来减轻这种成本。我们观察到,在许多此类方法中,这使得概率值乘法可以忽略不计,将瓶颈转移到查询键步骤。因此,可以减少关键头以加速推理,同时保留更多价值头可以在有限的额外解码成本下保留容量。我们引入了稀疏非对称组查询注意力(SAGA),它将键和值的头数解耦以利用这一原理,并将其与近似 top-N (Atop-N)注意力配对,这是一种简单的稀疏注意力方法,旨在研究稀疏性和头数不对称之间的相互作用。我们从理论上形式化了这种不对称性的好处,并通过对高达 1.5B 参数的模型进行延迟测量和质量评估来实证验证它们。 SAGA 和 Atop-N 在长上下文中共同实现了端到端解码速度,超过了我们的全注意力 GQA 基线上的 $2\times$。在评估基准上,使用 SAGA 从头开始训练的模型几乎与可比较的 GQA 变体的质量相匹配。为了促进采用,我们引入了一种高效的微调方法,该方法将预训练模型转换为 SAGA 架构,使从业者能够从我们的方法中受益,而无需进行昂贵的重新训练。