Search papers, labs, and topics across Lattice.
This paper introduces DiffPrune, a novel approach for visual token pruning in Vision-Language Models (VLMs) that directly associates token scores with their utility by employing an Information Throttler during training. Unlike traditional methods that rely on Gumbel-Softmax for discrete selection, DiffPrune weakens token information based on its score, allowing the model to learn the importance of each token more effectively. The results show that DiffPrune achieves 96.5% of full-model accuracy while accelerating LLM prefill by 2.85x, with minimal inference overhead of just 0.69 ms.
Token pruning can be done more effectively by directly linking token scores to their utility, achieving high accuracy with significantly reduced computational costs.
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.
