Search papers, labs, and topics across Lattice.
This paper introduces DiffPrune, a fully differentiable token pruning method for VLMs that replaces discrete token selection with continuous information throttling using variance-preserving noise conditioned on learned importance scores. By directly operating on token representations, DiffPrune avoids surrogate gradients and enables more reliable learning of token importance. Experiments across ten VLM benchmarks demonstrate that DiffPrune achieves a 2.85x speedup in LLM prefill with minimal accuracy loss (96.5% retention).
Ditch Gumbel-Softmax: DiffPrune's fully differentiable information throttling prunes VLMs 2.85x faster with negligible accuracy loss.
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.
