The Model Knows Which Tokens Matter: Automatic Token Selection via Noise Gating

TL;DR

AutoSelect is a novel method that dynamically selects important visual tokens in vision-language models, reducing inference cost while maintaining high accuracy by using a lightweight importance scorer and noise gating during training.

Contribution

It introduces AutoSelect, a training framework that learns to identify and select critical tokens without auxiliary objectives, enabling efficient inference across various vision-language models.

Findings

Retains 96.5% of full model accuracy on ten benchmarks.

Speeds up inference by 2.85x with minimal latency overhead.

Transfers effectively across different model architectures.

Abstract

Visual tokens dominate inference cost in vision-language models (VLMs), yet many carry redundant information. Existing pruning methods alleviate this but typically rely on attention magnitude or similarity scores. We reformulate visual token pruning as capacity constrained communication: given a fixed budget K, the model must allocate limited bandwidth to maximally preserve visual information. We propose AutoSelect, which attaches a lightweight Scorer and Denoiser to a frozen VLM and trains with only the standard next token prediction loss, without auxiliary objectives or extra annotations. During training, a variance preserving noise gate modulates each token's information flow according to its predicted importance so that gradients propagate through all tokens; a diagonal attention Denoiser then recovers the perturbed representations. At inference, only the Scorer and a hard top-K selection remain, adding negligible latency. On ten VLM benchmarks, AutoSelect retains 96.5% of full model accuracy while accelerating LLM prefill by 2.85x with only 0.69 ms overhead, and transfers to different VLM backbones without architecture-specific tuning. Code is available at https://github.com/MedHK23/AutoSelect.