MM-ShiftKV: 解码感知的预填充阶段KV选择用于多模态大语言模型

# MM-ShiftKV: 面向多模态大语言模型的解码感知预填充阶段KV选择

来源：https://arxiv.org/html/2607.22586

Jinsong Shu¹, Chenyang Wu¹,¹, Zhongle Xie¹,✉, Baokun Wang², Lidan Shou³,⁴
¹浙江大学 ²蚂蚁集团 ³浙江大学区块链与数据安全全国重点实验室 ⁴杭州高新区（滨江）区块链与数据安全研究院

###### 摘要

键值（KV）缓存对于多模态大语言模型（MLLMs）的高效推理至关重要，但其内存占用随上下文长度线性增长，且由于大量视觉标记的存在，成为主要瓶颈。近年来的预填充阶段KV选择方法通过预填充统计量估计KV重要性，隐含假设预填充时的查询能够代表解码过程中遇到的查询。我们表明，这一假设在多模态推理中并不成立，因为解码时的查询方差显著大于预填充阶段的表示，导致在严格的缓存预算下KV重要性估计不稳定。因此，微小的排序错误可能会不成比例地丢弃语义关键的视觉标记，从而损害基础能力（grounding）和推理性能。我们提出MM-ShiftKV，一种无需训练、解码感知且严格预填充阶段的KV选择方法。MM-ShiftKV通过在预填充期间近似解码时的查询行为，构建方差扩展的*查询代理*，并根据其聚合的注意力质量来估计提示KV的重要性。在多模态基准测试上的实验表明，在严格的KV缓存预算下，MM-ShiftKV始终优于现有方法。我们的代码可在https://github.com/zjuDBxAI/MM-ShiftKV获取。

MM-ShiftKV: 面向多模态大语言模型的解码感知预填充阶段KV选择

参照标题(a)特征均值/标准差 (OCRBench)  参照标题(b)注意力质量覆盖率

图1: 预填充-解码统计量与解码时注意力覆盖率。(a) OCRBench上预填充和解码期间隐藏状态表示的逐层均值和方差。(b) 注意力质量覆盖率（保留的提示），测量为解码时分配给预填充后保留的提示KV标记的注意力概率质量比例。

## 1 引言

多模态大语言模型（MLLMs）扩展了纯文本语言模型，使其能够生成基于视觉输入的语言，从而实现光学字符识别（OCR）、文档理解和视觉问答等应用（Li等人，2024a；Bai等人，2025）。在推理过程中，这些模型处理简短的文本提示以及高分辨率视觉输入，这些输入在预填充阶段被编码为以视觉标记为主的长多模态序列（Arif等人，2025），随后通过自回归解码生成文本输出。高效解码依赖于*键值（KV）缓存*，其内存占用随编码序列长度线性增长，使得KV缓存大小和访问成本成为内存消耗和解码效率的主要瓶颈。

为缓解这一瓶颈，最近的工作提出了*预填充阶段KV缓存选择*，即在预填充后保留一部分KV状态，并在解码期间重复使用（Xiao等人，2023；Li等人，2024b；Devoto等人，2025；Park等人，2025）。与解码时缓存驱逐或自适应缓存压缩（Xiao等人，2024）相比，这些纯预填充方法具有吸引力，因为它们无需训练（Li等人，2024b），不引入解码时干预，并且与诸如FlashAttention（Dao，2023）等高级注意力内核保持兼容。大多数方法通过预填充期间观察到的统计量来估计KV重要性，隐含假设预填充阶段的表示和注意力行为能够代表解码阶段。由于多模态输入的异质性，这一隐含假设在多模态推理中变得脆弱。大量视觉冗余标记与少量语义关键标记共存（Tao等人，2025；Chen等人，2025），使得KV排序中的微小错误可能不成比例地移除关键表示。因此，现有的预填充阶段KV选择方法往往导致语言基础能力下降、推理不稳定，以及多模态任务上的显著性能下降（Devoto等人，2025；Li等人，2024b；Park等人，2025；Devoto等人，2024）。

在更基本的层面上，预填充和解码对应于多模态推理的不同功能阶段。预填充主要强调视觉感知和跨模态对齐，而解码则转向语言生成以及基于先前生成标记的推理。尽管两个阶段共享相同的模型参数，但它们可能导致隐藏状态和注意力查询的不同分布，我们在第2节中对此进行了系统分析。因此，仅从预填充阶段统计量得出的重要性估计可能与解码时的行为不一致，导致在严格的缓存预算下KV选择不可靠。

在这项工作中，我们提出MM-ShiftKV，一种无需训练且严格预填充阶段的KV选择框架，用于多模态推理。我们的核心思想是通过校准重要性估计，使其反映解码时查询的分布特性，而不是仅仅依赖预填充阶段统计量，从而使预填充阶段的KV选择明确地*感知解码*。这里的“解码感知”并不意味着在解码期间执行KV驱逐或重新排序，而是调整基于预填充的查询代理，以更好地近似解码时的行为。具体来说，MM-ShiftKV在预填充结束时通过从当前输入的统计量中采样方差扩展的查询代理，并根据聚合的注意力质量估计KV重要性，执行一次性KV选择。生成的紧凑提示KV缓存在解码期间保持不变，从而在严格预算下实现高效且鲁棒的多模态推理。

贡献。
- • 我们识别出多模态推理中持续存在的预填充-解码*尺度不匹配*，其中解码时的*隐藏状态表示*方差显著大于预填充阶段观察到的方差。
- • 我们表明，这种不匹配导致基于预填充的查询代理尺度不足，从而扭曲查询-键相关性估计，并在严格的KV缓存预算下导致KV选择不稳定。
- • 我们提出MM-ShiftKV，一种无需训练且严格*预填充阶段*的KV选择框架，构建方差扩展、解码感知的*查询代理*来估计提示KV重要性。
- • 我们在紧凑KV缓存约束下，在代表性的OCR、基础能力和长上下文VQA基准上展示了改进的准确率-内存-延迟权衡。

## 2 观察：多模态推理中的预填充-解码尺度不匹配

近期工作表明，大语言模型中的激活统计量具有可以利用的结构化属性，可以在无需训练的方式下使用（Liu等人，2024a）。受这些属性的启发，大多数*预填充阶段*KV选择方法假设预填充期间观察到的统计量在解码期间仍然代表KV使用情况。正如第1节所讨论的，这一假设对于预填充阶段KV选择至关重要，但在多模态推理中尚未被仔细检验，我们在附录A中提供了额外的理论分析。在本节中，我们表明该假设在经验上被违反，并描述了在OCRBench（Liu等人，2023）和Qwen2.5-VL-7B-Instruct上一致存在的*预填充-解码统计不匹配*。更多结果见附录C。

具体来说，我们观察到预填充和解码期间的隐藏状态表示在统计*尺度*上存在显著差异。图1(a)报告了Qwen2.5-VL-7B-Instruct在OCRBench上评估的隐藏特征的逐层均值和标准差。尽管预填充和解码共享相同的模型参数，但解码阶段的表示在跨层中始终表现出更大的方差，而均值偏移相对适中。这表明预填充阶段的统计量系统地低估了解码时表示的尺度，这一现象在先前的预填充阶段KV选择工作中缺乏专门研究。

这种统计不匹配对预填充阶段KV选择的稳定性有直接影响。现有方法通常依赖预填充阶段的信号，如局部注意力行为、序列级激活统计量和KV相似性，来估计标记的重要性。当这些信号相对于真实的解码时查询尺度不足时，产生的重要性估计就会失真。图1(b)使用*注意力质量覆盖率*量化了这种效应，该指标测量解码时分配给预填充后保留的提示KV标记的注意力概率质量比例。遵循标准的纯预填充协议，KV选择在预填充结束时执行一次，压缩后的提示KV缓存在解码期间保持固定。在不同方法中，注意力覆盖率持续降低，并且在严格的缓存预算下表现出高方差，表明从预填充统计量中选择的KV集未能可靠地捕捉解码时的KV使用情况。

综合来看，这些结果表明，仅基于预填充阶段统计量的预填充阶段KV选择在多模态设置中本质上是不稳定的。观察到的预填充-解码统计不匹配凸显了在严格保持纯预填充模式的同时，明确考虑解码时查询行为的必要性。这一见解构成了下一节中开发的解码感知预填充阶段KV选择方法的基础，该方法在图1(b)中实现了最高的注意力覆盖率。

参照标题
图2: MM-ShiftKV概述。该方法计算当前输入完整预填充序列的统计量，并在预填充期间构建方差扩展、解码感知的*查询代理*。这些查询代理通过注意力质量聚合和基于最后标记锚点的组别投票来估计提示KV重要性，从而在固定预算下生成紧凑的KV缓存。

## 3 方法

##### 核心思想。MM-ShiftKV通过明确近似未来解码查询对提示键的访问方式，来解决多模态推理中预填充阶段KV选择的不稳定性。由于解码时的查询在预填充期间不可用，核心思想是在预填充阶段构建一组合成*查询代理*，以近似解码查询的分布属性。这些查询代理持续关注的提示键更有可能在解码过程中重要，因此应在有限的KV缓存预算下保留。

### 3.1 问题定义

在描述我们的方法之前，我们形式化多模态推理中的预填充阶段KV选择问题。

##### 输入序列。给定多模态输入序列
\[
X = (x_1, x_2, \ldots, x_T),
\]
其中标记包括视觉和文本模态，\(T\)表示*预填充阶段*的序列长度。

##### 提示KV缓存。在预填充期间，模型在每个Transformer层\(\ell\)和KV头\(h\)计算所有提示标记的KV表示。得到的KV对集合构成*prompt KV cache*:
\[
\mathcal{C}^{(\ell,h)} = \{ (k_t^{(\ell,h)}, v_t^{(\ell,h)}) \}_{t \in \mathcal{T}},
\]
其中\(\mathcal{T} = \{1,\ldots,T\}\)索引提示标记。这些提示KV在解码期间被重复访问，并主导推理中的内存消耗和注意力计算。

##### 预填充阶段KV选择。对于每个层\(\ell\)和KV头\(h\)给定的缓存预算\(C_{\ell,h}\)，预填充阶段KV选择的目标是保留一个子集
\[
\mathcal{C}'^{(\ell,h)} \subseteq \mathcal{C}^{(\ell,h)}, \quad |\mathcal{C}'^{(\ell,h)}| \leq C_{\ell,h},
\]
该子集在整个解码过程中重复使用，以减少内存使用和注意力成本，同时保持生成质量。在实践中，KV选择独立于每个层和头，并在*预填充*结束时执行一次。请注意，选定的提示KV缓存在解码期间保持不变，而来自新生成标记的KV对以标准的自回归方式追加。

### 3.2 MM-ShiftKV

图2提供了MM-ShiftKV的概述。该方法包括三个步骤，均在预填充期间执行：(1) 构建解码感知查询代理，(2) 通过基于质量的投票估计提示KV重要性，(3) 基于聚合的重要性分数执行预算化KV选择。所有步骤严格限于预填充阶段，不引入解码时干预。更多实现细节见附录B。

#### 3.2.1 查询代理构建

为了近似解码时的查询行为，MM-ShiftKV使用预填充阶段的统计量构建一组合成*查询代理*。

##### 预填充隐藏状态统计量。设\(h^{(\ell)} \in \mathbb{R}^d\)表示预填充期间在层\(\ell\)产生的隐藏状态。我们使用逐元素统计量总结这些表示：
\[
\mu_{\text{pre}}^{(\ell)} = \mathrm{mean}(h^{(\ell)}), \quad \sigma_{\text{pre}}^{(\ell)} = \mathrm{std}(h^{(\ell)}),
\]
在提示标记上对每个特征维度计算。这些统计量分别针对每个输入样本在其完整预填充序列上计算，而不是跨不同样本共享。

##### 方差扩展采样。如第2节所示，解码时的查询方差显著大于预填充阶段的表示。为近似这一效应，我们引入方差扩展因子\(\gamma > 1\)，并从
\[
\tilde{H}_i^{(\ell)} \sim \mathcal{N}\!\left(\mu_{\text{pre}}^{(\ell)}, \mathrm{diag}\!\left((\gamma \sigma_{\text{pre}}^{(\ell)})^2\right)\right)
\]
中采样\(N\)个*代理隐藏状态*\(\{\tilde{H}_i^{(\ell)}\}_{i=1}^N\)。