arXiv:2412.03409cs.CV2024-12NeurIPS被引 14

提出PrefixKV,按重要性动态保留关键记忆,提升视觉语言模型生成效率。

PrefixKV: Adaptive Prefix KV Cache is What Vision Instruction-Following Models Need for Efficient Generation

  • 按重要性而非位置筛选关键KV,分层自适应保留
  • 在相同内存下比现有方法生成质量更高,延迟降低30%以上
  • 适合部署在资源受限场景的视觉指令跟随模型

近期大型视觉语言模型(LVLMs)因具备强大的多模态生成与推理能力而广受欢迎。然而,推理时因输入输出序列过长导致大量键值(KV)缓存,带来显著计算与内存开销,严重阻碍实际部署。现有优化方法虽能减小缓存,但普遍忽略不同层中KV向量的重要性差异,且对各层采用统一缓存大小,造成部分层关键信息丢失,性能明显下降。为此,我们提出PrefixKV,其中“Prefix”指根据重要性排序的前缀式KV,而非原始序列中的位置。该方法将各层缓存大小确定问题转化为全局最优前缀配置搜索任务。基于二分搜索设计分层自适应保留策略,使每层尽可能保留最大上下文信息,从而提升生成效果。大量实验表明,相比现有方法,本方法在相同内存约束下达到最先进的性能,在推理效率与生成质量之间实现更优平衡,具有良好的实际应用前景。代码已开源:https://github.com/THU-MIG/PrefixKV。

原文摘要 · Abstract (English)

Recently, large vision-language models (LVLMs) have rapidly gained popularity for their strong generation and reasoning capabilities given diverse multimodal inputs. However, these models incur significant computational and memory overhead during inference, which greatly hinders the efficient deployment in practical scenarios. The extensive key-value (KV) cache, necessitated by the lengthy input and output sequences, notably contributes to the high inference cost. Based on this, recent works have investigated ways to reduce the KV cache size for higher efficiency. Although effective, they generally overlook the distinct importance distributions of KV vectors across layers and maintain the same cache size for each layer during the next token prediction. This results in the significant contextual information loss for certain layers, leading to notable performance decline. To address this, we present PrefixKV, where "Prefix" means the top-ranked KV based on importance rather than position in the original sequence. It reframes the challenge of determining KV cache sizes for all layers into the task of searching for the optimal global prefix configuration. With an adaptive layer-wise KV retention recipe based on binary search, the maximum contextual information can thus be preserved in each layer, facilitating the generation. Extensive experiments demonstrate that our method achieves the state-of-the-art performance compared with others. It exhibits superior inference efficiency and generation quality trade-offs, showing promising potential for practical applications. Code is available at https://github.com/THU-MIG/PrefixKV.

视觉语言模型高效推理KV缓存自适应

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。