让视觉令牌剪枝更精准:直接评估信息损失,提升模型效率
DiffPrune: differentiable information throttling for token pruning in vision-language models

- 通过信息弱化机制直接优化令牌评分,避免间接代理损失
- 在10个基准上保留96.5%准确率,推理加速2.85倍,延迟仅0.69毫秒
- 适合追求高效视觉语言模型的开发者和研究者
视觉令牌剪枝通过移除冗余视觉令牌降低视觉语言模型(VLM)的计算开销。核心在于学习一个衡量令牌有用性的评分。现有方法通常使用Gumbel-Softmax近似训练中的离散选择,导致评分依赖于松弛剪枝算子的行为,而非实际的信息损失后果。本文提出DiffPrune,使令牌评分具有直接意义:训练时保留所有令牌,并根据评分削弱每个令牌的信息;若削弱某令牌影响任务性能,则评分被推动保护该令牌;反之则可降低其评分。由于损失通过实际的信息弱化路径反向传播,评分避免了松弛选择带来的不稳定性。DiffPrune采用信息弱化器,在视觉令牌中注入保持方差的噪声:高分令牌接近原始表示,低分令牌携带较少原始信息。推理时移除弱化器,使用学习到的评分进行硬性Top-K剪枝。在十个VLM基准上,DiffPrune在保留96.5%全模型精度的同时,将大语言模型预填充速度提升2.85倍,推理额外开销仅为0.69毫秒。代码将公开。
原文摘要 · Abstract (English)
Visual token pruning reduces the computational cost of Vision-Language Models (VLMs) by removing redundant visual tokens. The key is to learn a score that measures whether a token is useful. Existing methods typically rely on Gumbel-Softmax to approximate discrete selection during training. Such selectors make the score depend on the behavior of a relaxed pruning operator, not directly on the consequence of information loss. In this paper, we propose DiffPrune, which gives token scores a direct meaning. During training, DiffPrune keeps all tokens and weakens each token's information according to its score. If weakening a token hurts the task, the scorer is pushed to protect it; if not, the token can receive a lower score. Because the loss is differentiated through this actual information-throttling path, the scorer avoids the unstable surrogate path of relaxed token selection. DiffPrune implements this idea with an Information Throttler, which injects variance-preserving noise into visual tokens, where high-score tokens remain close to their original representations, while low-score tokens carry less original information. At inference, the throttler is removed, and hard top-K pruning is applied using the learned scores. Across ten VLM benchmarks, DiffPrune retains 96.5% of full-model accuracy while accelerating LLM prefill by 2.85x, with only 0.69 ms inference overhead. Code will be publicly available.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。