DiffPrune: differentiable information throttling for token pruning in vision-language models
2026-08-03 • Computer Vision and Pattern Recognition
Computer Vision and Pattern Recognition
AI summaryⓘ
The authors introduce DiffPrune, a method to speed up Vision-Language Models by smartly deciding which visual parts (tokens) are important. Instead of guessing token importance with approximations like before, DiffPrune gently reduces the information from less important tokens during training and learns directly how much each token matters to the task. Tokens that don’t affect performance can be more ignored, while important ones are protected. This results in faster processing with almost no drop in accuracy, as shown across multiple benchmarks.
Vision-Language Modelstoken pruningGumbel-Softmaxinformation throttlingvariance-preserving noisetop-K pruningvisual tokensmodel accelerationinference speedtoken scoring
Authors
Landi He, Mingde Yao, Shawn Young, Lijian Xu
Abstract
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.