arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2310.01801cs.CL

模型告诉你该丢弃什么:面向大语言模型的自适应KV缓存压缩

Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

University of Illinois Urbana-Champaign(伊利诺伊大学厄巴纳-香槟分校);Microsoft(微软公司)

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, Jianfeng Gao

AI总结:

本研究提出即插即用的大语言模型自适应KV缓存压缩方法FastGen,通过剖析注意力结构差异化构建缓存,无需重训练即可大幅降低GPU内存占用且几乎不损失生成质量。

摘要:

本研究提出自适应KV缓存压缩方法,这是一种即插即用的方法,可降低大语言模型(LLMs)生成式推理的内存占用。与保留所有上下文 token 的键、值向量的传统KV缓存不同,我们通过针对性剖析来识别注意力模块的内在结构。基于识别出的结构,我们以自适应方式构建KV缓存:在侧重局部上下文的注意力头上驱逐长距离上下文,在以特殊 token 为中心的注意力头上丢弃非特殊 token,仅对广泛关注所有 token 的注意力头使用标准KV缓存。此外,由于采用轻量级注意力剖析来指导自适应KV缓存的构建,FastGen 无需耗费大量资源的微调或重新训练即可部署。在我们针对各类任务的实验中,FastGen 大幅降低了GPU内存消耗,且生成质量损失可忽略不计。我们将发布代码及兼容的CUDA内核以保证可复现性。

英文摘要:

In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.

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

↑