提出可微分的视觉令牌剪枝方法，让模型自动学习哪些视觉信息该保留。
Beyond Surrogate Gradients: Fully Differentiable Token Pruning for Vision-Language Models
- 用连续控制替代离散选择，通过信息抑制机制实现端到端可微优化
- 在10个基准上保持96.5%准确率，推理前缀加速2.85倍，延迟仅0.69毫秒
- 适合需要高效推理的视觉语言模型部署，尤其关注计算效率的场景
视觉令牌剪枝通过移除冗余视觉令牌降低视觉语言模型（VLM）的计算开销。现有方法通常依赖Gumbel-Softmax在训练中近似离散选择，但其优化依赖代理梯度而非真实选择过程，导致令牌重要性学习不可靠。本文提出DiffPrune，将剪枝重新定义为对令牌信息的连续控制而非离散选择学习。具体而言，引入信息调节器，利用基于重要性得分的方差保持噪声调制每个令牌，重要性越高则信息抑制越少。该设计直接作用于令牌表示，自然提供学习令牌重要性的完全可微优化路径。推理时通过硬阈值对学习到的得分进行剪枝。在十个VLM基准上，DiffPrune在仅0.69毫秒推理开销下，保持96.5%的全模型准确率，并使LLM预填充阶段提速2.85倍。
原文摘要 · Abstract (English)
Visual token pruning reduces the computational cost of Vision-Language Models (VLMs) by removing redundant visual tokens. Existing methods typically rely on Gumbel-Softmax to approximate discrete selection during training. However, the optimization is driven by surrogate gradients rather than the true selection process, leading to unreliable learning of token importance. In this paper, we propose DiffPrune, which reformulates pruning as continuous control of token information instead of discrete selection learning. Specifically, we introduce an Information Throttler that modulates each token using variance-preserving noise conditioned on importance scores, where higher scores induce less information suppression during training. This design directly operates on token representations, naturally providing a fully differentiable optimization path for learning token importance. At inference, tokens are removed via hard thresholding on 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 of inference overhead.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。