基于位置持久稀疏注意力的快速准确 LLM 解码：TidalDecode
机器学习
2024-10-08 v1 人工智能
计算与语言
摘要
大型语言模型 (LLM) 已在 diverse NLP 任务上取得显著进展，尤其是处理长上下文的模型日益受到关注。然而，Transformer 架构所需的扩大 key-value (KV) 缓存大小在解码阶段加剧了内存约束，成为重要瓶颈。现有稀疏注意力机制 designed to address this bottleneck 有两个局限：(1)它们往往难以可靠地识别注意力中最相关的标记，(2)它们忽略了跨连续 Transformer 层的标记选择的空间一致性，这可能导致性能下降和显著的标记选择开销。本文引入 TidalDecode，一种简单而有效的算法和系统，通过位置持久稀疏注意力实现快速准确的 LLM 解码。TidalDecode 利用现有稀疏注意力方法所选标记的空间一致性，并引入少量标记选择层执行完整注意力以识别注意力得分最高的标记，而所有其他层执行稀疏注意力以使用预选标记。这种设计使 TidalDecode 能够在不牺牲生成结果质量的前提下，大大降低稀疏注意力的标记选择开销。评估在 diverse LLMs 和 tasks 上表明，TidalDecode 在保持与完整注意力方法相近的生成性能的同时，将 LLM 解码延迟降低最高可达 2.1 倍。
关键词
引用
@article{arxiv.2410.05076,
title = {TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention},
author = {Lijie Yang and Zhihao Zhang and Zhuofu Chen and Zikun Li and Zhihao Jia},
journal= {arXiv preprint arXiv:2410.05076},
year = {2024}
}