📝 LEARNPRUNER: RETHINKING ATTENTION-BASED TOKEN PRUNING IN VISION LANGUAGE MODELS - ICLR'26

Conference: ICLR 2026

1. Abstract (摘要)

【核心内容与总结】

• 研究背景： 视觉语言模型（VLMs）在视觉理解和推理上表现出色，但由于长视觉序列输入（如高分辨率图像或视频），带来了显著的计算负担。现有的“Token剪枝”技术通过删除不重要的视觉 Token 来缓解该问题，但其核心在于“如何决定 Token 的重要性”，现有方法通常直接依赖于视觉编码器（Vision Encoder）或大语言模型（LLM）的注意力分数（Attention Scores）。
• 核心发现：
  1. 视觉编码器端： 存在“注意力下沉（Attention Sink）”现象，导致其难以有效聚焦于包含丰富信息的前景区域。
  2. 大语言模型端（LLM）： 尽管先前的研究表明 LLM 存在对 Token 位置的注意力偏差（位置越后分数越高），但本文发现，文本到视觉（Text-to-Vision）的注意力机制能够抵抗这种偏差，并在 LLM 的**中间层（Middle Layers）**提供了有效的剪枝指导。
• 本文贡献： 提出了一种双阶段（Two-stage）Token 剪枝框架 LearnPruner。
  • 第一阶段： 在视觉编码器之后，通过一个可学习的剪枝模块 (Learnable Pruning Module, LPM) 剔除视觉冗余 Token。
  • 第二阶段： 在 LLM 的中间层，保留与任务（Query）相关的 Token。
• 实验效果： LearnPruner 在仅使用原始视觉 Token 数量 $ 5.5% $ 的情况下，保留了约 $ 95% $ 的原始性能，实现了 $ 3.2 \times $ 的推理加速，展示了良好的精度-效率权衡。

---

2. Introduction (引言)

【背景引入与动机分析】

• 重新审视注意力机制（结合 Figure 1）：
  • 如图 1 所示，Vision Encoder 使用的 [CLS] token 注意力虽然能部分覆盖前景目标，但通常会将过多注意力分配给信息量较低的背景区域。这与视觉 Transformer 倾向于在均匀背景区域生成高范数离群 Token（Artifacts）的现象一致。
  • 在 LLM 端，尽管存在“注意力偏移（Attention Shift，偏向图像下半部分）”现象，但这种偏差主要存在于视觉到视觉（V2V）或全 Token 注意力中。相反，**文本到视觉（Text-to-Vision）**的注意力能够有效地聚焦于与查询（Query）相关的区域。
• 现有剪枝方法的局限性：
  • (1) 在视觉编码器中，[CLS] token 无法充分关注显著前景，导致剪枝结果次优（尤其在 Token 预算有限时）。
  • (2) 在 LLM 中，文本到视觉的注意力具有抵抗注意力偏移的鲁棒性，能提供可靠指导，但如果将剪枝推迟到中间层进行，前期浅层网络仍需处理完整的视觉序列，导致加速收益有限。
• LearnPruner 解决方案：
  • 提出双阶段框架顺序移除冗余 Token。首先，采用轻量级可学习模块（取代传统的 [CLS] 分数）预测 Token 重要性，以消除固有的视觉冗余，并保留少量多样性 Token 提供补充信息。随后，剩余 Token 进入 LLM，在中间层基于文本查询进行第二次剪枝，进一步丢弃与指令无关的 Token。

---

3. Method (方法设计)

本节首先通过分析验证现有注意力机制在剪枝中的局限性，随后详细阐述 LearnPruner 的算法架构。

3.1 STUDY OF ATTENTION IN VLMS (VLM 注意力机制分析)

1. Vision Encoder 中的 [CLS] 注意力分析：

• 问题提出： 视觉编码器通常使用 [CLS] token 聚合全局信息，现有工作自然地将其注意力分数作为 Patch token 的重要性估计。但 [CLS] token 是否充分关注了显著的前景区域？
• 实验设计： 使用 LangSAM 等视觉 grounding 工具提取图像的“前景掩码 (Foreground Mask)”。在 LLaVA-1.5-7B 上对比了三种策略：基于 [CLS] 的全局选择（[CLS]_{all}）、受限于前景的 [CLS] 选择（[CLS]_{fg}），以及在前景中的随机选择（Random_{fg}）。
• 分析结论（参考 Figure 2a）： 限制在信息丰富区域的 [CLS]_{fg} 始终优于全局选择 [CLS]_{all}。值得注意的是，前景中的随机选择（Random_{fg}）也能达到与 [CLS]_{all} 相当的性能。这表明 [CLS] token 未能有效聚焦于显著的前景区域，其注意力得分与 Token 的真实重要性对齐较差。

2. LLM 中的注意力分析：

(Figure 3: 不同网络层中特定文本 Token 到图像的注意力热力图)

• 注意力偏移现象（参考 Figure 2b）： 将 LLM 接收到的注意力分解为视觉注意力和文本注意力。虽然两者都表现出注意力偏移（索引越大的 Token 得分越高），但文本注意力的上升趋势远比视觉注意力平缓。这是因为视觉注意力受位置编码衰减和因果掩码的影响更大，而文本注意力的模式在不同实例间差异较大，且受偏置影响较小。
• 不同剪枝指标的对比（参考 Figure 2c）： 对比了纯视觉注意力、纯文本注意力、双模态平均以及最后一个指令 Token 注意力的剪枝效果。结果显示，文本注意力 (Text Attn) 和最后一个 Token 注意力始终优于其他标准；引入视觉注意力反而会导致性能显著下降。
• 层级选择分析（参考 Figure 2c 与 Figure 3）： 注意力的可靠性从浅层到中间层逐渐增加，在中间层（如第 8 层以后）表现稳定，但在更深层则显著下降。如图 3 所示，在浅层和深层，无信息区域往往吸收了文本 Token 的大部分注意力；而在中间层，文本 Token 能够准确聚焦于与其语义相关的图像区域。因此，LLM 的中间层是执行剪枝的最佳位置。

---

3.2 LEARNPRUNER (LearnPruner 算法详解)

(Figure 4: LearnPruner 框架概览。包括视觉冗余移除和文本无关内容移除两个主要阶段)

基于上述分析，LearnPruner 采用了一种渐进式剪枝策略：在视觉编码器后使用可学习标准进行剪枝，并在 LLM 内部进行文本引导的剪枝。

阶段一：去除视觉冗余 (Remove Visual Redundancy)

目的： 图像本身存在视觉冗余，此阶段旨在视觉编码器输出后立即用更紧凑的 Token 表示图像信息。考虑到 [CLS] 注意力无法准确定位显著区域，该阶段引入了一个可学习剪枝模块 (LPM, Learnable Pruning Module)。

• LPM 模块设计与训练 (End-to-end Optimization):

  • 网络结构： 将视觉编码器输出的 Token 特征 $ X_v^{(0)} $ 输入到一个轻量级的多层感知机 (MLP) 中，进行二分类以决定每个 Token 是保留还是修剪。
  • 直通估计器 (STE)： 由于离散的二值决策是不可导的，模型采用 Straight-Through Estimator (STE) 技巧以支持端到端反向传播： $$ M_{soft} = \text{Softmax}(\text{MLP}(X_v^{(0)})) $$ $$ M_{hard} = \text{argmax}(M_{soft}) $$
  • 前向与反向传播： 前向传播时，利用二值掩码 $ M_{hard} $ 切断被剪枝 Token 与保留 Token 之间的注意力交互（通过修改 Attention 的 Softmax 掩码实现）；反向传播时，使用软掩码 $ M_{soft} $ 提供梯度。
  • 推理阶段： 在推理时，直接使用 $ M_{soft} $ 作为每个 Token 的重要性得分进行排序。

• 多样性 Token 选择 (Diversity-based Token Selection):

  • 动机： LPM 倾向于关注语义丰富的前景区域，可能会忽略在某些 VQA 任务中同样重要的背景信息。
  • 算法逻辑： 为了补充全面的视觉上下文，引入了基于多样性的选择。在 LPM 选出一组高信息量 Token 后，对于剩余的每一个 Token，计算其与已选集合中所有 Token 的余弦相似度 (Cosine Similarity) 并找出最大相似度值。然后，迭代地将具有最小的最大相似度值的 Token 加入保留集合（即选择与当前集合差异最大的 Token）。此过程持续进行，直到满足设定的第一阶段 Token 预算 $ R_1 $。

阶段二：去除文本无关内容 (Remove Text-Irrelevant Content)

目的： 虽然第一阶段保留了有价值的视觉信息，但并非所有视觉内容对特定的查询（Query）都是必要的。该阶段在 LLM 内部进一步移除与文本无关的 Token。

• 设计原则： 依据前文分析，文本注意力不易受注意力偏移影响，且对相关区域响应强烈。因此，直接利用文本注意力来指导 Token 剪枝。
• 计算公式： 文本相关性得分为所有文本 Query Token 对某一视觉 Token 注意力的平均值。对于第 $ k $ 层的注意力计算： $$ \tilde{A}^{(k)} = \frac{1}{N_q} \sum_{i=1}^{N_q} A^{(k)}(X_{q,i}^{(k)}, X_v^{(k)}) $$ 其中 $ X_{q}^{(k)} $ 和 $ X_{v}^{(k)} $ 分别表示第 $ k $ 层的 Query Token 和 Vision Token。$ N_q $ 是 Query Token 的数量。
• 剪枝执行： 根据计算出的平均注意力得分 $ \tilde{A}^{(k)} $，仅保留得分最高的 Top-k 个视觉 Token（预算为 $ R_2 $），它们将继续参与 LLM 较深层的跨模态交互，其余 Token 则被丢弃。

---

4. Experiments (实验评估)

本节在广泛使用的 VLM 基准上评估 LearnPruner 的性能，并进行了详尽的计算效率和消融分析。

4.1 实验设置

• 训练设置： 使用 10% 的 LLaVA-665K 数据集训练 LPM，期间冻结基础 VLM 的权重以保持原始性能。
• 推理设置： 第一阶段保留 $ R_1 $ 个 Token，其中 $ \lambda = 10% $ 的预算分配给多样性 Token。第二阶段在 LLM 的第 $ k = 12 $ 层进行，保留 $ R_2 $ 个 Token。为了与基线方法公平比较，实验中固定了不同阶段的保留比例（$ R_1:R_2 = 3 $）。

4.2 主要结果 (Main Results)

(Table 1: LLaVA-v1.5-7B 上的性能对比)

(Table 2: 高分辨率模型 LLaVA-NeXT-7B 上的性能对比)

• 在 LLaVA-v1.5-7B 上的结果 (Table 1):
  • 在 GQA, POPE, MME, TextVQA 等 8 个图像理解基准上评估。
  • 当视觉 Token 从 576 压缩到平均 128 个（保留 22.2%）时，LearnPruner 的平均准确率仅下降 1.5%，表现与需要额外训练 Decoder 层的 TwigVLM 相当，但本方法更轻量。
  • 在极端的压缩预算下（保留 64 或 32 个 Token），LearnPruner 维持了原始性能的 96.9% 和 94.8%，显著优于 FastV、SparseVLM 等现有方法，说明该方法能更有效地识别并保留最关键的视觉信息。
• 在 LLaVA-NeXT-7B 上的结果 (Table 2):
  • LLaVA-NeXT 将图像分割为多个子图，Token 序列长达 2880 个。在去除 88.9% 的 Token（仅保留 320 个）的情况下，LearnPruner 仍能保持 97.5% 的原始性能。

(Table 3: Video-LLaVA-7B 上的视频理解性能对比)

(Table 4: Qwen2.5-VL-7B 上的性能对比)

• 在 Video-LLaVA-7B 上的泛化能力 (Table 3):
  • 在 TGIF-QA, MSVD-QA 等视频基准上，LearnPruner 取得比 FastV 更高的准确率，证明了该方法在视频处理场景中的有效性。
• 在 Qwen2.5-VL-7B 上的表现 (Table 4):
  • 为了验证在不同架构上的通用性，在 Qwen 架构上进行测试。去除 88.9% 的视觉 Token 后，FastV 性能骤降至 71.4%，而 LearnPruner 仍保持 94.1% 的相对准确率。

4.3 效率分析 (Efficiency Analysis)

(Table 5: 计算效率对比，包含 TFLOPs、Prefill Time、总时间和 KV Cache 内存使用)

• 计算开销评估： 在 NVIDIA A100-80GB 上进行评估。
  • 在 LLaVA-v1.5-7B 中，将 Token 从 576 减少到 32 时，Prefill 时间和总推理时间分别实现了 $ 5.4 \times $ 和 $ 2.3 \times $ 的加速，同时 KV Cache 显存占用减少了 $ 6.8 \times $。
  • 对于长序列的 LLaVA-NeXT-7B，将 Token 从 2880 减少到 160 时，效率提升更为显著，Prefill 和总时间分别加速了 $ 6.0 \times $ 和 $ 3.2 \times $。由于 LPM 的轻量化设计，其带来的额外计算和内存开销可忽略不计。

4.4 消融实验 (Ablation Studies)

(Table 6: 不同重要性标准的有效性消融)

• 重要性标准的有效性：
  • 在第一阶段（去除视觉冗余），使用可学习模块（LPM）比使用 [CLS] 注意力性能提升了 1.7%。
  • 尝试在 LLM 中也插入 LPM（结合特征和注意力分布进行打分），结果显示没有进一步的性能提升（96.9% vs 96.9%）。这表明中间层的文本注意力信号已经足够可靠，直接利用注意力结果即可，无需训练多个 LPM。

(Table 7a: 剪枝位置 k 的消融; Table 7b: 第一、二阶段保留比例 $R_1:R_2$ 的消融)

• 剪枝位置与保留比例：
  • 剪枝层 $ k $ (Table 7a): 将第二阶段剪枝从第 8 层移至 12 层时性能逐渐稳定，证实了浅层注意力分布不可靠的结论。最终选择 $ k=12 $。
  • 比例 $ R_1:R_2 $ (Table 7b): 测试表明 $ R_1:R_2 = 3 $ 取得最佳效果，说明浅层需要保留较多 Token 以维持视觉完整性，而中深层仅需少量与文本高度相关的 Token 即可完成推理。

(Figure 5: 多样性 Token 比例 $\lambda$ 及训练数据规模的消融)

• 多样性 Token 比例 $ \lambda $ (Figure 5a)： $ \lambda $ 控制多样性 Token 的比例。过低会导致背景信息丢失，过高则会占用前景目标的预算。实验确定 $ \lambda = 0.1 $ 为最佳平衡点。
• 数据效率 (Figure 5b)： 使用 5% 到 100% 的 LLaVA-665K 数据集训练 LPM。结果显示，仅使用 5% 的数据，模型即可保持一致且优异的性能，证明了 LPM 设计的稳健性和数据效率。

4.5 可视化结果 (Qualitative Results)

(Figure 6: 剪枝结果可视化。黄色斑块表示被保留的 Token，平均保留 32 个)

(Figure 7: 剪枝结果可视化及失败案例分析)

• 对比分析： 与 VisionZip（基于 [CLS]，未能聚焦前景甚至选择了无效 Padding）和 FastV（偏向保留图像下半部分，对 Query 不敏感）相比，LearnPruner 能在第一阶段准确识别语义丰富的区域及背景上下文，并在第二阶段通过文本引导精确保留与问题最相关的 Token。
• 失败案例探讨 (Figure 7 底部)： 在极其复杂的视觉场景中（如寻找黑白色的食物、辨别右侧是否有消防栓），由于 Token 预算的严格限制，第一阶段可能未能充分保留关键的微小视觉特征；或者即使保留了相关 Token，基础 VLM 固有的推理能力限制也会导致预测错误。

---

5. Conclusion (结论)

本文对视觉编码器和 LLM 中的注意力机制进行了深入分析。基于上述观察，提出了一种两阶段剪枝框架 LearnPruner。该框架首先通过视觉编码器后部的可学习剪枝模块消除冗余的视觉 Token，随后在 LLM 的中间层利用文本引导进一步丢弃与指令无关的 Token。广泛的实验结果表明，LearnPruner 优于先前的 SOTA 剪枝方法，实现了更好的精度-效率权衡（Accuracy-efficiency trade-off）。