Skip to content

LearnPruner: Rethinking Attention-based Token Pruning in Vision Language Models

🕒 Published (v1): 2026-04-27 01:56 UTC · Source: Arxiv · Venue: ICLR 2026 · link

Why this paper was selected

Attention-based token pruning rethought; attention scores don't indicate importance

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LearnPruner is a two-stage visual token pruning framework for VLMs that replaces unreliable attention-based importance criteria with a lightweight learnable module after the vision encoder, then applies text-guided pruning in the LLM's middle layers. It retains 94.8% of baseline performance using only 5.5% of vision tokens (32 out of 576), achieving up to 3.2× total inference speedup on LLaVA-NeXT-7B.

Problem

Current attention-based token pruning methods have two failure modes: (1) vision encoder [CLS] attention suffers from "attention sink" artifacts that assign high scores to uninformative background regions, making it a poor importance estimator; (2) LLM attention-based pruning is contaminated by "attention shift" bias (positional encoding + causal masking causes higher-indexed tokens to receive disproportionate attention), and pruning in shallow LLM layers incurs substantial redundant computation before the signal becomes reliable.

Method

LearnPruner operates in two sequential stages:

Stage 1 — Remove Visual Redundancy (post-vision-encoder): A lightweight learnable pruning module (LPM), a small MLP (~0.53M parameters) performing binary classification on vision token features, predicts per-token importance scores. Training uses the Straight-Through Estimator (STE) to backpropagate through the discrete binary pruning mask: \(M_\text{soft} = \text{Softmax}(\text{MLP}(X_v^{(0)}))\), \(M_\text{hard} = \text{argmax}(M_\text{soft})\). A diversity-based complement selects a small fraction (\(\lambda = 10\%\) of the retained budget) of tokens maximally dissimilar (minimum cosine similarity) to LPM-selected tokens, capturing background information the foreground-focused LPM may miss.

Stage 2 — Remove Text-Irrelevant Content (LLM middle layer \(k=12\)): Text-to-vision attention averaged over query tokens and attention heads, \(\tilde{A}^{(k)} = \frac{1}{N_q}\sum_{i=1}^{N_q} A^{(k)}(X_{q,i}^{(k)}, X_v^{(k)})\), guides pruning of query-irrelevant tokens. The ratio \(R_1 : R_2 = 3\) determines the split between the two pruning stages. The LPM is trained on 10% of LLaVA-665K with the base VLM frozen.

Key Contributions

  • Empirical analysis demonstrating [CLS] attention in vision encoders fails as an importance criterion: foreground-constrained random selection outperforms global [CLS]-based selection.
  • Empirical demonstration that text-to-vision attention in LLM middle layers is resistant to attention-shift bias and provides reliable pruning signal, whereas vision-to-vision attention degrades performance.
  • A lightweight (~0.53M param) learnable pruning module with STE training that directly predicts token importance without relying on attention maps.
  • A diversity-based complement mechanism to recover background information missed by the foreground-biased LPM.
  • Generalization across model families (LLaVA-1.5, LLaVA-NeXT, Video-LLaVA, Qwen2.5-VL) and input types (image, high-resolution, video).

Results

  • LLaVA-v1.5-7B, 32 tokens retained (94.4% reduction): 94.8% relative accuracy vs. 576-token baseline; nearest competitor VisionZip achieves 87.8%, DivPrune 90.5%.
  • LLaVA-v1.5-7B, 64 tokens retained: 96.9% relative accuracy; outperforms all training-based (TwigVLM: 96.0%) and training-free competitors.
  • LLaVA-v1.5-7B, 128 tokens retained: 98.5% relative accuracy; comparable to TwigVLM (99.0%) at ~1,170× fewer learnable parameters (0.53M vs. 610M).
  • LLaVA-NeXT-7B, 160 tokens retained (94.4% reduction): 94.0% relative accuracy vs. TwigVLM 93.3% and VisionZip‡ 93.3%.
  • Qwen2.5-VL-7B, 142 tokens retained (88.9% reduction): 94.1% average vs. FastV 71.4%.
  • Efficiency (LLaVA-NeXT-7B, 160 tokens): 6.1× FLOPs reduction, 3.2× total inference speedup, 8.9× KV cache reduction vs. vanilla; prefill time reduced 6.0×.
  • Video (Video-LLaVA-7B, 228 tokens): 97.0% relative accuracy vs. FastV 86.7% (implied from scores: LearnPruner avg 2.45 vs. FastV 2.21 score at 228-token budget).

Limitations

  • Training required (though lightweight): the LPM must be fine-tuned per base VLM on task data, unlike fully training-free methods.
  • The diversity-based token selection is greedy and iterative at inference time, adding overhead not fully characterized.
  • Generalization to architectures beyond the LLaMA family (only Qwen2.5-VL tested outside LLaMA) is not thoroughly evaluated.
  • Middle-layer pruning (\(k=12\)) is a fixed hyperparameter; its optimal value may vary across model architectures and was only ablated on LLaVA-1.5.
  • Performance on instruction-following or generation-heavy tasks (e.g., image captioning, complex reasoning chains) is not reported.

Relevance to Vision-Language Models

LearnPruner directly addresses the practical deployment bottleneck of VLMs: visual token length scales quadratically with image resolution and linearly with video frames, making aggressive yet accurate pruning critical. By systematically diagnosing why both vision-encoder [CLS] attention and LLM shallow-layer attention fail as importance signals, this work establishes a principled empirical foundation for redesigning pruning criteria in the VLM pipeline. The two-stage cascade — learned foreground selection then query-conditioned filtering — generalizes cleanly to high-resolution (LLaVA-NeXT tiling) and video inputs, which are the dominant frontier for VLM scale. For researchers tracking efficiency methods in VLMs, this paper provides both diagnostic tools (foreground mask ablations, attention decomposition by modality and layer depth) and a lightweight training recipe that is practical even under constrained compute budgets.