针对多模态大模型推理中的关键值缓存瓶颈,提出解码感知的预填充阶段选择方法。
MM-ShiftKV: Decode-Aware Prefill-Stage KV Selection for Multimodal Large Language Models

- 通过构建方差扩展查询代理,模拟解码阶段的查询行为。
- 在严格缓存预算下,显著提升视觉关键信息保留率和推理准确率。
- 无需训练,适用于资源受限的多模态模型部署场景。
键值(KV)缓存对多模态大语言模型(MLLMs)高效推理至关重要,但其内存开销随上下文长度线性增长,且因大量视觉标记而成为主要瓶颈。现有预填充阶段的KV选择方法基于预填充统计估计重要性,隐含假设预填充阶段查询具有代表性。我们发现,在多模态推理中,解码阶段查询的方差远大于预填充阶段表示,导致在紧缩缓存预算下,KV重要性估计极不稳定。微小排序误差可能错误丢弃语义关键视觉标记,严重损害定位与推理性能。为此,我们提出MM-ShiftKV,一种无需训练、解码感知且严格仅在预填充阶段运行的KV选择方法。该方法通过构造方差扩展的查询代理,在预填充阶段近似解码阶段的查询行为,并基于其聚合注意力质量评估提示中KV的重要性。在多模态基准测试中,MM-ShiftKV在严格缓存预算下持续优于现有方法。代码已开源:https://github.com/zjuDBxAI/MM-ShiftKV。
原文摘要 · Abstract (English)
Key-Value (KV) caching is essential for efficient inference in multimodal large language models (MLLMs), yet its memory footprint grows linearly with context length and becomes a major bottleneck due to the large number of visual tokens. Recent prefill-stage KV selection methods estimate KV importance from prefilling statistics, implicitly assuming that prefilling-time queries are representative of those encountered during decoding. We show that this assumption breaks down in multimodal inference, where decoding-time queries exhibit substantially larger variance than prefilling-stage representations, leading to unstable KV importance estimation under tight cache budgets. As a result, small ranking errors can disproportionately discard semantically critical visual tokens and degrade grounding and reasoning performance. We propose MM-ShiftKV, a training-free, decode-aware and strictly prefill-only KV selection method. MM-ShiftKV approximates decoding-time query behavior during prefilling by constructing variance-expanded query proxies and estimates prompt KV importance based on their aggregated attention mass. Experiments on multimodal benchmarks demonstrate that MM-ShiftKV consistently outperforms existing methods under strict KV-cache budgets. Our code is available at https://github.com/zjuDBxAI/MM-ShiftKV.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。