DIFFPRUNE: DIFFERENTIABLE INFORMATION THROTTLING FOR TOKEN PRUNING IN VISION-LANGUAGE MODELS

Modern Vision-Language Models (VLMs) process high-resolution images that generate a large number of visual tokens, leading to significant computational overhead during the prefill stage. While token pruning is a common strategy to reduce this cost, existing methods often rely on Gumbel-Softmax or similar surrogate-gradient estimators to handle the discrete nature of token selection. These approaches suffer from a forward-backward mismatch, where the gradient updates do not accurately reflect the actual pruning operation, leading to unstable training and suboptimal token selection.

DiffPrune introduces a fully differentiable framework for visual token pruning. Instead of forcing a hard keep-or-drop decision during training, it employs an Information Throttler that continuously modulates the information content of each token. This throttler uses variance-preserving (VP) noise, where the intensity of the noise is controlled by a learned importance score. High-score tokens remain largely intact, while low-score tokens are suppressed by noise. A train-only Diagonal-Attention Block is used to ensure these noised representations remain compatible with the frozen downstream model. Because the loss is differentiated through this continuous throttling path, the Scorer receives consistent, stable gradients. At inference, the throttler is removed, and the model performs standard hard top-K pruning based on the learned scores.

DiffPrune demonstrates significantly higher gradient-direction consistency compared to traditional Gumbel-Softmax approaches, leading to more stable optimization. Across ten VLM benchmarks, including LLaVA-1.5, LLaVA-NEXT, and Qwen2.5-VL, DiffPrune consistently outperforms existing pruning methods. It retains 96.5% of full-model accuracy while accelerating LLM prefill by up to 2.85x, with minimal inference overhead (0.69 ms). The results show that the learned rankings transfer effectively to hard deployment, even at aggressive pruning ratios where only a small fraction of original tokens are retained.