DiffPrune: differentiable information throttling for token pruning in vision-language models

2026-08-03

AI Summary

This work addresses the high computational cost of current vision-language models caused by redundant visual tokens and the instability of conventional pruning methods that rely on unreliable relaxed selection mechanisms to assess token importance. The authors propose a differentiable information throttling mechanism that dynamically attenuates token information through variance-preserving noise injection, enabling token scores to directly reflect their impact on task performance. This approach eliminates the need for relaxation approximations such as Gumbel-Softmax, ensuring stable training. At inference time, hard top-K pruning is applied, achieving a 2.85× speedup in LLM prefilling with only a 3.5% accuracy drop across ten vision-language benchmarks and an additional overhead of merely 0.69 milliseconds.

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.

Problem

• token pruning
• vision-language models
• information loss
• differentiable selection
• visual tokens

Innovation

• differentiable pruning
• information throttling
• vision-language models
• token scoring
• variance-preserving noise