跳转至

TwinPrune: Density-Aware Two-Phase Token Pruning for Vision-Language Models

会议: NeurIPS 2026
arXiv: 2610.17530
代码: https://github.com/Kai-Liu001/TwinPrune
领域: 多模态效率
关键词: 视觉语言模型, token剪枝, KV缓存驱逐, 空间密度, 两阶段剪枝

一句话总结

针对视觉语言模型在长序列与高分辨率下的预填计算与解码显存瓶颈,提出基于局部空间余弦相异度的统一密度信号,在前向预填阶段执行密度倾斜的最远点二部合并,在自回归解码阶段复用相同密度进行KV缓存一步驱逐,实现计算、显存与精度的三维帕累托前沿提升。

研究背景与动机

视觉语言模型(VLM)在文档理解、具身智能体、多步视觉推理以及GUI控制等应用中取得了显著进展,但这些场景通常需要将成百上千个视觉token输入数十层Transformer中进行处理。随着输入图像分辨率提升与切片(tile)数量剧增,系统面临着两个并发收紧的效率瓶颈:一是由大语言模型主干网络处理大量视觉token导致的预填计算延迟;二是由长序列自回归生成阶段每层必须持久驻留的键值(KV)缓存所引发的显存膨胀。现有的token剪枝工作要么专注于预填阶段通过注意力或最远点采样丢弃token(如FastV、DivPrune、DART),要么在解码阶段维护重击者(heavy-hitter)或注意力汇流缓存(如H2O、SnapKV、StreamingLLM),且大多声称可以在压缩数十倍的同时几乎无损保留精度。

然而,这一繁荣的共识背后潜藏着严重的评估机制缺陷。主流文献广泛采用单token输出型基准(如MMBench、POPE、ScienceQA等预填探针),其评分仅依赖首个解码token的argmax预测(即 \(y_1=\arg\max_v [W_o h_T]_v\))。该预测在预填阶段结束时便已被完全决定,解码阶段的任何KV缓存驱逐操作在此刻尚未发生。因此,大量声称在解码阶段实现无损剪枝的算法,本质上测量的仅是预填输出质量而非解码生成质量;当面对TextVQA、MM-Vet和DocVQA等依赖完整自回归序列生成的全流程探针基准时,直接丢弃全部视觉KV缓存会导致预测剧烈崩塌,而预填探针分数却几乎纹丝不动。

因此,亟需建立真正将预填计算与解码显存解耦并在全流程基准上有效验证的统一压缩框架。核心 idea:提出一种基于局部空间余弦相异度的无训练空间密度信号,在预填阶段与最远点采样融合实现高覆盖的二部特征合并,并在解码阶段零额外开销复用该信号执行一步式KV缓存Top-\(k\)驱逐,以单一特征统计量同时主导预填计算加速与解码显存压减。

方法详解

整体框架

TwinPrune 将视觉token压缩视为连续贯穿预填与解码的双阶段协同过程。输入图像经视觉编码器网格化为 \(H \times W\) 的视觉token序列,并与文本输入拼接送入LLM主干网络。在设定的预填早期插入层 \(L\),模型通过计算各视觉token与其局部空间邻域的余弦相异度提取多尺度空间密度信号 \(\rho\);随后在预填阶段执行密度倾斜的最远点二部采样与特征合并,将压缩后的token序列继续前向传播以削减FLOPs;在预填结束进入自回归解码时,系统直接复用已计算的密度分数对残留视觉token的KV缓存执行Top-\(k\)驱逐,保留关键RoPE位置的视觉KV条目,从而在不增加额外注意力重计算的前提下大幅压缩自回归生成期的显存占用。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像与文本指令<br/>网格化视觉token"] --> B["LLM浅层前向计算<br/>传播至指定插入层 L"]
    B --> C["空间局部密度评估<br/>计算多尺度邻域余弦相异度"]
    C --> D["阶段一:预填密度倾斜最远点合并<br/>k-center多样性与密度打标二部合并"]
    D --> E["LLM深层预填完成<br/>获得首个解码token输出 y1"]
    E --> F["阶段二:解码单步KV缓存驱逐<br/>复用密度分数Top-k截断KV条目"]
    F --> G["自回归长序列生成<br/>全流程探针输出最终答案"]

关键设计

1. 空间局部密度评估:无训练的多尺度几何非冗余度度量

针对传统L2范数受图像亮度和对比度主导、注意力分数受文本提示词扰动与注意力汇流偏移的问题,本文提出基于空间邻域余弦相异度(Cosine Dissimilarity)的固有特征统计量。对于排列在 \(H \times W\) 规则网格上的视觉token,位置 \((i,j)\) 的局部密度定义为其与局部 \(w \times w\) 邻域内其余token的平均余弦距离: $\(\rho_{ij} = 1 - \frac{1}{|\mathcal{N}_w(i,j)|} \sum_{(i',j') \in \mathcal{N}_w(i,j)} \frac{\langle x_{ij}, x_{i'j'} \rangle}{\|x_{ij}\| \cdot \|x_{i'j'}\|}\)$ 为了同时捕获图像中的精细边缘高频纹理与较大尺度的语义区域边界,算法在多个尺度 \(w \in \{3, 5, 7\}\) 下分别计算相异度并取最大值 \(\rho = \max_{w} \rho^{(w)}\)。高 \(\rho_{ij}\) 代表该token与周围上下文存在显著差异、承载独特视觉信息,而低 \(\rho_{ij}\) 则反映其处于平滑冗余区域,这一度量完全无需额外训练与参数调节。

2. 预填阶段密度倾斜最远点合并:平衡多样性覆盖与信息丰富度

针对纯最远点采样(FPS)在极端剪枝比例下易被孤立噪声点带偏、而纯贪心密度选择又容易造成空间聚集导致几何覆盖塌缩的矛盾,设计了密度倾斜的最远点贪心循环(Density-Tilted Farthest-Point Selection)。算法将密度归一化为 \(\tilde{\rho} \in [0, 1]\),预计算未选token与已选集合的最小余弦距离 \(\text{mind}_j\),并以密度最高点作为初始种子。后续每一步迭代按照综合得分进行贪心选取: $\(s_t = \arg\max_{j \notin S} \big( \text{mind}_j + \beta \tilde{\rho}_j \big), \quad \text{随即更新} \quad \text{mind}_j \leftarrow \min(\text{mind}_j, D_{s_t j})\)$ 系统设定倾斜标量 \(\beta = 0.3\)。在维持覆盖半径(k-center cover radius)理论近似保证的同时,利用密度打碎距离饱和时的平局;选出的保留token充当聚类锚点,被舍弃的源token特征通过二部匹配加权合并注入最近锚点中,既消除冗余又避免了直接剪枝导致的上下文信息截断。工程上,该循环通过两路Triton核函数(argmax规约与 \(64\times 32\) 瓦片更新)融合并在FP16精度下捕获为模块级CUDA Graph缓存,消除内核发射开销。

3. 解码阶段单步KV缓存驱逐:零重算开销的跨阶段信号复用

针对自回归生成阶段视觉KV缓存线性吞噬显存且现有动态窗口驱逐方法需要频繁维护注意力矩阵的痛点,TwinPrune 在预填-解码交接边界执行单步KV驱逐。由于预填阶段提取的空间密度具有核心集(Coreset)代表性,它在期望上逼近所有后续注意力聚合操作的加权空间,因此无需引入解码期新的注意力热度统计,直接按预填期计算好的 \(\rho\) 降序排列,截取前 \(k_{\text{dec}} = \lceil N_{\text{post}} \cdot r_{\text{decode}} \rceil\) 个视觉KV条目保留,其余条目完全释放。幸存的KV缓存维持原始RoPE旋转位置编码,首个生成token的Query位置设置在原预填序列长度加一处,确保相对几何映射完全自洽。

实验关键数据

主实验

在InternVL-3.5-8B主干模型与全流程长序列基准上,TwinPrune分别在预填阶段与解码阶段展示出显著的精度与速度/显存帕累托优势。

表1:InternVL-3.5-8B预填剪枝阶段主流方法对比(不同保留率与插入层)

方法 会议 Token保留比 (TPL) MMBench POPE ScienceQA TextVQA Avg-4 Prefill加速比
Baseline (无压缩) - 100.0% 0.8707 0.8903 0.9539 0.6493 0.8411 \(1.00\times\)
Block 1: \(L=15\), 剪枝 50%
Random - 65.2% 0.8862 0.9155 0.9782 0.6202 0.8500 \(1.28\times\)
ToMe ICLR 2023 65.2% 0.8584 0.8913 0.9346 0.5325 0.8042 \(1.16\times\)
FastV ECCV 2024 65.2% 0.8584 0.8868 0.9326 0.5150 0.7982 \(1.28\times\)
PruMerge ICCV 2025 65.2% 0.8584 0.8943 0.9410 0.5937 0.8219 \(1.21\times\)
DivPrune CVPR 2025 65.2% 0.8630 0.8923 0.9430 0.6017 0.8250 \(1.17\times\)
DART EMNLP 2025 65.2% 0.8467 0.8877 0.9316 0.6026 0.8171 \(1.28\times\)
TwinPrune (Ours) NeurIPS 2026 65.2% 0.8909 0.9142 0.9792 0.7195 0.8759 \(1.16\times\)
Block 2: 极端压力 \(L=3\), 剪枝 85%
Random - 24.1% 0.7980 0.8398 0.8622 0.3285 0.7071 \(3.12\times\)
ToMe ICLR 2023 24.1% 0.8320 0.8798 0.8860 0.4524 0.7626 \(2.07\times\)
FastV ECCV 2024 24.1% 0.7423 0.7665 0.7833 0.1939 0.6215 \(3.11\times\)
PruMerge ICCV 2025 24.1% 0.8313 0.8578 0.8825 0.3866 0.7396 \(2.47\times\)
DivPrune CVPR 2025 24.1% 0.0000 0.0000 0.0000 0.0041 0.0010 \(2.98\times\)
DART EMNLP 2025 24.1% 0.0000 0.0000 0.0000 0.0041 0.0010 \(3.09\times\)
TwinPrune (Ours) NeurIPS 2026 24.1% 0.8282 0.8670 0.8929 0.5696 0.7894 \(2.67\times\)

表2:解码阶段视觉KV缓存驱逐对比(InternVL-3.5-8B跨三类保留率)

方法 TextVQA (50%) TextVQA (25%) TextVQA (12.5%) MM-Vet (50%) MM-Vet (25%) MM-Vet (12.5%) DocVQA (50%) DocVQA (25%) DocVQA (12.5%)
Baseline (无驱逐) 0.733 0.733 0.733 0.756 0.756 0.756 0.642 0.642 0.642
StreamingLLM 0.637 0.573 0.507 0.667 0.561 0.514 0.468 0.384 0.324
H2O 0.697 0.629 0.552 0.662 0.571 0.497 0.555 0.438 0.347
SnapKV 0.707 0.643 0.573 0.654 0.585 0.513 0.577 0.469 0.377
Density (Ours) 0.726 0.697 0.648 0.708 0.634 0.551 0.575 0.466 0.382

消融实验

在InternVL-3.5-8B固定配置(\(L=15\), 剪枝 50%)下评估不同重要性评分信号对预填与解码表现的影响:

评分信号配置 TextVQA准确率 MMBench准确率 说明
空间密度 \(\rho\) (完整多尺度) 0.616 0.860 兼顾局部边缘与区域异质性
单尺度密度 \(\rho\) (\(w=3\)) 0.613 0.859 仅差0.03%,多尺度略优但单尺度已极强
L2 范数 \(\|x\|_2\) 0.594 0.861 下降2.2%,受像素亮度与对比度高频噪点误导
注意力 V-norm 0.584 0.862 下降3.2%,受提示词Query偏置扰动
复合得分 \(\rho \cdot \|v\|\) 0.598 0.860 混合相关性弱的信号反而引入扰动噪声

关键发现

  • 评估基准结构性盲区:在纯首字argmax打分的预填探针(MMBench、POPE)上,即使在解码阶段丢弃100%的视觉KV缓存,预测翻转率仅有 \(\le 0.54\%\);而在依赖长序列自回归生成的TextVQA上,全丢KV会导致 \(45.5\%\) 的样本预测发生剧烈翻转。
  • 极端剪枝下的结构鲁棒性:在极浅层(\(L=3\))剪枝85%的高压工况下,纯选择型算法(DivPrune、DART)因为破坏了序列连续性且未合并上下文导致模型输出彻底崩溃(准确率跌至 0.001),而 TwinPrune 通过特征二部合并依然维持了 0.7894 的四项平均准确率。
  • 大模型规模缩放增益:从 2B 扩展至 38B(InternVL-3.5 系列),随着模型参数容量增大,同等压缩比例下的准确率损失不增反减,展现出优异的跨尺度泛化能力。

亮点与洞察

  • 揭示评测机制盲区:明确指出单token打分基准无法评测解码期KV压缩效果的结构性硬伤,为社区确立了必须使用全流程探针基准进行评估的诊断标准。
  • 一石二鸟的统一表征:仅需在浅层计算一次局部空间余弦相异度,即可无缝复用于预填期二部特征合并与解码期单步KV驱逐,无需二次计算注意力矩阵。
  • 高度优化的工程落地:利用 Triton 算子融合与模块级 CUDA Graph 缓存,将密集合并算法的调度开销降低近两个数量级,在 A6000 上实现端到端 \(3.24\times\) 的预填推理加速。

局限与展望

  • 单阶段预填压缩局限:当前预填剪枝仅在单一预设层(如 \(L=3\) 或 \(L=15\))进行一次性处理,尚未探索跨层递进式渐进剪枝策略。
  • 非规则网格拓扑限制:空间密度计算依赖于输入视觉token的二维网格邻接先验,对于采用Q-Former或抽象感知器(Perceiver)等非网格投影架构的VLM需要专门适配邻域拓扑。
  • 未来方向:探索动态自适应学习邻域聚合核函数,并将任务特定的文本提示偏好动态融入密度加权中。

相关工作与启发

  • vs DivPrune: DivPrune 使用纯最远点采样追求多样性,但在高压缩比或浅层插入时易陷入孤立噪声点;本文通过密度倾斜机制结合二部合并,消除了极端工况下的崩塌风险。
  • vs FastV: FastV 依赖深层自注意力权重直接截断token;本文在浅层提取固有空间几何密度且执行特征合并,保留了更完整的上下文语义并在 TextVQA 上大幅领先。
  • vs SnapKV / H2O: 传统解码KV压缩依赖复杂的注意力累积统计且需要维护滑动窗口;本文直接复用预填阶段的空间密度信号执行一次性静态驱逐,以零重算开销在极限保留率下全面胜出。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示预填探针盲区并提出空间密度贯穿预填与解码的统一剪枝范式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖3大VLM家族7种尺度,提供严格的预填/解码/全流程三维帕累托前沿分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 论证严密,数学与核心集理论支撑扎实,评测实验直指要害。
  • 价值: ⭐⭐⭐⭐⭐ 为大模型长视觉上下文部署提供了极具实用性与理论深度的加速方案。