Each Complexity Deserves a Pruning Policy

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

AutoPrune is a training-free, plug-and-play visual token pruning framework for large vision-language models (LVLMs) that assigns each input a personalized pruning schedule instead of a fixed one. It measures cross-modal mutual information in early decoder layers to estimate sample/task complexity, then maps that scalar to a budget-constrained logistic retention curve that controls how aggressively tokens are dropped at each layer. On LLaVA-1.5-7B it prunes 89% of visual tokens and cuts inference FLOPs by 76.8% while retaining 96.7% of full-model accuracy.

Problem

Existing training-free visual token pruning methods apply a fixed pruning schedule uniformly across all inputs and tasks—either a predetermined layer-wise ratio or a hand-crafted heuristic. This rigidity is mismatched to the empirical observation that simple queries cause cross-modal attention to collapse early while complex, reasoning-intensive queries maintain diffuse attention across many layers. Fixed schedules therefore either over-prune complex inputs (hurting accuracy) or under-prune simple ones (wasting compute), and most existing layer-wise heuristics also lack a mechanism to guarantee adherence to a global FLOPs or token budget.

Method

Complexity estimation via mutual information. At early decoder layers, AutoPrune computes the mutual information I(V,T) between visual tokens V and textual tokens T by treating the softmax-normalized cross-attention weights as a probability distribution. High I(V,T) signals strong text-vision alignment (simple task); low I(V,T) signals weak alignment (complex task).

Budget-constrained logistic retention curve. For each input q, a logistic function fq(x) = N_init / (1 + exp(k_q(x − x⁰_q))) defines how many visual tokens survive at each decoder layer x. The slope k_q and inflection point x⁰_q are set linearly as functions of I(V,T): high MI → steeper slope and earlier inflection (aggressive front-loaded pruning); low MI → shallower slope and later inflection (conservative, late-onset pruning). The curve is analytically integrated and rescaled so that its area equals the user-specified token or FLOPs budget c_max, guaranteeing budget adherence without manual tuning. Token counts per layer are obtained by evaluating the rescaled curve at integer layer indices via binary search.

Token selection. Within each layer, tokens are ranked by attention saliency and the bottom-ranked tokens (up to the layer's budget) are dropped; discarded tokens can be revived via a token-revival policy ρ. The overall complexity is O(N_h N_t N_v), dominated by the MI estimation term and negligible relative to full model inference.

Key Contributions

• Neuroscience-grounded analysis linking sample/task complexity to cross-modal attention decay and inter-layer saliency fluctuation patterns in LVLMs.
• AutoPrune: a training-free, plug-and-play complexity-adaptive pruning framework using mutual information as a complexity proxy and logistic retention curves as per-sample pruning policies.
• Closed-form budget normalization that guarantees any user-specified token or FLOPs budget is met exactly, without heuristic approximation.
• Generalization to Vision-Language-Action (VLA) models for autonomous driving (Senna on nuScenes) with zero hyper-parameter re-tuning.

Results

• LLaVA-1.5-7B, 89% pruning (64 tokens retained): AutoPrune retains 96.7% of full-model accuracy vs. 87.6% for PDrop (CVPR'25), a 9.1% improvement; FLOPs reduced by 76.8%.
• LLaVA-1.5-7B, 78% pruning (128 tokens): 98.1% accuracy vs. 95.6% (PDrop), >93% for others.
• LLaVA-1.5-7B, 66% pruning (192 tokens): 99.0% accuracy (near-lossless), FLOPs reduced >57%.
• LLaVA-NeXT-7B, 160 tokens retained: 94.9% accuracy vs. 86.7% (FasterVLM, CVPR'25), over 8 pp gap.
• LLaVA-NeXT-7B, 320 tokens: 98.2% accuracy, best among all compared methods.
• Senna VLA (nuScenes), 25% token retention: 111.23% planning accuracy relative to unpruned baseline (pruning removes noisy tokens); PyramidDrop achieves only 98.89% at the same level.
• Ablation (Table 4): Mutual information outperforms cosine similarity and average-attention as complexity indicators; logistic curve outperforms linear, tanh, and exponential schedules on GQA and TextVQA.

Limitations

• The authors note that token importance does not strictly decrease monotonically with decoder depth: deeper layers occasionally retain more critical tokens than shallower ones, which the current logistic schedule cannot fully capture.
• The pruning curve shape is parameterized by a small number of fixed hyper-parameters (k₀, γ, β, x₀) whose robustness across radically different architectures is not fully characterized.
• The pruned model occasionally outperforms the baseline on VLA tasks, suggesting the presence of noisy tokens—a phenomenon left as future work rather than actively exploited.
• Evaluation is limited to image-based LVLMs and one VLA system; video-based models with extremely long token sequences are not evaluated.

Relevance to Vision-Language Models

AutoPrune directly addresses the inference efficiency bottleneck in LVLMs caused by long visual token sequences, a key obstacle to real-time deployment. By demonstrating that a simple mutual-information signal computed from cross-attention is sufficient to differentiate easy from hard inputs and drive nearly lossless pruning at 89% sparsity, the paper adds a practically important, architecture-agnostic tool to the VLM efficiency toolkit. Its extension to VLA models shows the principle generalizes beyond standard vision-language tasks, which is relevant for embodied-AI applications of VLMs. The neuroscience-grounded analysis of how cross-modal attention evolves across decoder depth also offers interpretability insights applicable to understanding VLM reasoning processes more broadly.