arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2402.06262cs.CLcs.AI
收藏

关于键值受限生成式语言模型推理中驱逐策略的有效性

On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference

Shanghai Jiao Tong University(上海交通大学);University of Texas at Arlington(德克萨斯大学阿灵顿分校)

Siyu Ren, Kenny Q. Zhu

AI总结:

本文针对大型语言模型部署成本高及键值缓存内存开销大的问题,提出基于时间注意力分数和鲁棒性度量的缓存省略策略RoCo,并发布软件包EasyKV,实验验证其在预填充和自回归解码阶段的优越性。

摘要:

尽管大型语言模型(LLMs)近期取得了成功,但由于其过高的内存和计算需求,在资源受限的环境中部署它们的成本极高。除了模型参数外,键值缓存也存储在GPU内存中,并随批大小和序列长度线性增长。作为补救,近期的研究提出了各种驱逐策略,以将键值缓存的开销维持在给定预算内。本文探讨了现有驱逐策略在重要性分数计算和驱逐范围构建方面的有效性。我们指出了先前策略在这两个方面的不足,并提出了RoCo,一种基于时间注意力分数和鲁棒性度量的鲁棒缓存省略策略。涵盖预填充和自回归解码阶段的广泛实验验证了RoCo的优越性。最后,我们发布了EasyKV,一个专用于用户友好的键值受限生成式推理的通用软件包。代码可在 https://github.com/DRSY/EasyKV 获取。

英文摘要:

Despite the recent success associated with Large Language Models (LLMs), they are notably cost-prohibitive to deploy in resource-constrained environments due to their excessive memory and computational demands. In addition to model parameters, the key-value cache is also stored in GPU memory, growing linearly with batch size and sequence length. As a remedy, recent works have proposed various eviction policies for maintaining the overhead of key-value cache under a given budget. This paper embarks on the efficacy of existing eviction policies in terms of importance score calculation and eviction scope construction. We identify the deficiency of prior policies in these two aspects and introduce RoCo, a robust cache omission policy based on temporal attention scores and robustness measures. Extensive experimentation spanning prefilling and auto-regressive decoding stages validates the superiority of RoCo. Finally, we release EasyKV, a versatile software package dedicated to user-friendly key-value constrained generative inference. Code available at https://github.com/DRSY/EasyKV.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑