LearnPruner: Rethinking Attention-based Token Pruning in Vision Language Models
作者: Rinyoichi Takezoe, Yaqian Li, Zihao Bo, Anzhou Hou, Mo Guang, Kaiwen Long
分类: cs.CV
发布日期: 2026-04-27
备注: Accepted to ICLR 2026
💡 一句话要点
LearnPruner：重新思考视觉语言模型中基于注意力的Token剪枝
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 视觉语言模型 Token剪枝 注意力机制 模型加速 计算效率
📋 核心要点
- 现有视觉语言模型因视觉序列过长导致计算负担重，而基于注意力的token剪枝方法在确定token重要性方面存在不足。
- LearnPruner通过两阶段剪枝策略，首先移除冗余视觉token，然后保留LLM中间层中与任务相关的token。
- 实验表明，LearnPruner在显著减少计算量的同时，能够保持较高的模型性能，实现了更好的精度-效率平衡。
📝 摘要（中文）
视觉语言模型(VLM)最近在视觉理解和推理方面表现出了卓越的能力，但由于长的视觉序列输入，它们也带来了巨大的计算负担。目前的工作通过剪枝不重要的视觉token来解决这个问题，在保持模型性能的同时，实现了显著的计算量减少。Token剪枝的核心在于确定token的重要性，当前的方法主要依赖于视觉编码器或大型语言模型(LLM)的注意力分数。在本文中，我们分析了视觉编码器和LLM中注意力机制的有效性。我们发现视觉编码器存在注意力下沉的问题，导致对信息丰富的前景区域关注不足；而在LLM中，尽管之前的研究已经确定了注意力对token位置的偏见，但文本到视觉的注意力表现出对这种偏见的抵抗力，并在中间层实现了有效的剪枝指导。基于这些观察，我们提出了LearnPruner，这是一个两阶段的token剪枝框架，它首先通过视觉编码器后的可学习剪枝模块移除冗余的视觉token，然后在LLM的中间层仅保留与任务相关的token。实验结果表明，我们的LearnPruner可以在仅使用5.5%的视觉token的情况下，保留约95%的原始性能，并实现3.2倍的推理加速，从而实现了卓越的精度-效率权衡。
🔬 方法详解
问题定义：视觉语言模型（VLM）在处理长视觉序列时面临巨大的计算负担。现有的token剪枝方法依赖视觉编码器或LLM的注意力分数来判断token的重要性，但视觉编码器存在注意力下沉问题，无法有效关注前景信息，而LLM的注意力机制又存在对token位置的偏见，影响剪枝效果。
核心思路：LearnPruner的核心思路是结合视觉编码器和LLM的优势，设计一个两阶段的token剪枝框架。第一阶段利用可学习的剪枝模块去除视觉编码器输出中的冗余token，第二阶段利用LLM中间层文本到视觉的注意力来保留与任务相关的token。这种设计旨在克服视觉编码器的注意力下沉问题和LLM的位置偏见，从而更准确地确定token的重要性。
技术框架：LearnPruner框架包含两个主要阶段：1) 视觉token剪枝阶段：在视觉编码器之后，引入一个可学习的剪枝模块，该模块学习预测每个视觉token的重要性得分，并根据得分进行剪枝。2) LLM token选择阶段：利用LLM中间层的文本到视觉注意力，选择与文本描述最相关的视觉token。这两个阶段协同工作，首先去除冗余信息，然后保留关键信息。
关键创新：LearnPruner的关键创新在于其两阶段剪枝策略以及对LLM中间层文本到视觉注意力的有效利用。与现有方法相比，LearnPruner不仅考虑了视觉信息，还考虑了文本信息，从而能够更准确地确定token的重要性。此外，可学习的剪枝模块能够自适应地学习token的重要性，避免了手动设计或选择注意力机制的困难。
关键设计：在视觉token剪枝阶段，可学习剪枝模块的具体实现方式未知，但可以推测其可能采用类似MLP或卷积神经网络的结构，以预测每个token的重要性得分。在LLM token选择阶段，关键在于如何有效利用文本到视觉的注意力。论文提到LLM中间层对位置偏见具有抵抗力，因此可能选择中间层的注意力权重作为token选择的依据。具体的损失函数和参数设置未知。
🖼️ 关键图片
📊 实验亮点
LearnPruner在实验中表现出色，仅使用5.5%的视觉token，就能保留约95%的原始性能，并实现3.2倍的推理加速。这表明LearnPruner在精度和效率之间取得了显著的平衡，优于现有的token剪枝方法。这些结果验证了LearnPruner两阶段剪枝策略的有效性以及对LLM中间层注意力的合理利用。
🎯 应用场景
LearnPruner具有广泛的应用前景，可用于各种需要处理长视觉序列的视觉语言任务，如图像描述、视觉问答、视频理解等。通过降低计算成本，LearnPruner可以使VLM在资源受限的设备上运行，并加速模型的推理速度，从而提高用户体验。此外，该方法还可以应用于其他类型的序列数据，如文本和音频。
📄 摘要（原文）
Vision-Language Models (VLMs) have recently demonstrated remarkable capabilities in visual understanding and reasoning, but they also impose significant computational burdens due to long visual sequence inputs. Recent works address this issue by pruning unimportant visual tokens, achieving substantial computational reduction while maintaining model performance. The core of token pruning lies in determining token importance, with current approaches primarily relying on attention scores from vision encoders or Large Language Models (LLMs). In this paper, we analyze the effectiveness of attention mechanisms in both vision encoders and LLMs. We find that vision encoders suffer from attention sink, leading to poor focus on informative foreground regions, while in LLMs, although prior studies have identified attention bias toward token positions, text-to-vision attention demonstrates resistance to this bias and enables effective pruning guidance in middle layers. Based on these observations, we propose LearnPruner, a two-stage token pruning framework that first removes redundant vision tokens via a learnable pruning module after the vision encoder, then retains only task-relevant tokens in the LLM's middle layer. Experimental results show that our LearnPruner can preserve approximately 95% of the original performance while using only 5.5% of vision tokens, and achieve 3.2$\times$ inference acceleration, demonstrating a superior accuracy-efficiency trade-off.