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

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; rethinks attention-based token pruning in VLMs

Ask a follow-up

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

TL;DR¶

LearnPruner is a two-stage token pruning framework for VLMs that replaces unreliable attention-based importance scores with a lightweight learnable MLP after the vision encoder, followed by text-guided pruning in the LLM's middle layers. It retains only 5.5% of vision tokens while preserving ~95% of original model performance and achieving up to 3.2× inference acceleration.

Problem¶

Existing attention-based token pruning methods rely either on vision encoder [CLS] token attention or LLM attention scores, both of which are flawed: the [CLS] token exhibits attention sink, over-attending to uninformative background artifacts rather than salient foreground regions; meanwhile, LLM attention suffers from positional attention shift (bias toward higher-index tokens due to causal masking and positional encoding decay), particularly in vision-to-vision attention patterns. These limitations lead to suboptimal token selection, especially under aggressive pruning budgets.

Method¶

LearnPruner operates in two sequential stages:

Stage 1 — Remove Visual Redundancy (after vision encoder): A lightweight learnable pruning module (LPM), a small MLP (~0.53M parameters), takes vision encoder output features $X_v^{(0)}$ and performs binary classification per token via softmax over two classes. Since argmax is non-differentiable, a Straight-Through Estimator (STE) enables end-to-end training: the hard binary mask $M_\text{hard} = \text{argmax}(\text{Softmax}(\text{MLP}(X_v^{(0)})))$ is used in the forward pass while gradients flow through the soft mask $M_\text{soft}$. Only the LPM is trained (base VLM frozen) on 10% of LLaVA-665K. Additionally, $\lambda = 10\%$ of the retained budget is filled by diversity tokens selected greedily via minimum cosine similarity to already-selected tokens, compensating for background information the LPM may miss.

Stage 2 — Remove Text-Irrelevant Content (LLM middle layer $k=12$): Text-to-vision attention is computed as the average attention from all query tokens to each vision token across all heads at layer $k$: $$\tilde{A}^{(k)} = \frac{1}{N_q} \sum_{i=1}^{N_q} A^{(k)}(X_{q,i}^{(k)}, X_v^{(k)})$$ The top-$R_2$ vision tokens by this score are retained. The ratio $R_1 : R_2 = 3:1$ is fixed given a target budget.

The key insight motivating Stage 2 is that text attention resists positional bias (gradual increase with index vs. the sharp linear bias in visual attention) and is most semantically precise in middle LLM layers (layers 8–12 for LLaVA-1.5), degrading in both shallow and deep layers.

Key Contributions¶

Systematic empirical analysis demonstrating that [CLS] attention in vision encoders is corrupted by attention sink, and that vision-modality LLM attention is distorted by positional bias while text-to-vision attention is substantially more robust.
Learnable pruning module (LPM) with STE training to replace [CLS]-based importance scoring after the vision encoder, trained with only 0.53M parameters on 10% of instruction-tuning data.
Diversity-token supplement to recover potentially pruned background context not captured by the learned importance scores.
Two-stage progressive pruning (post-encoder + LLM middle layer) achieving superior accuracy-efficiency Pareto frontier compared to prior training-free and training-based methods.

Results¶

LLaVA-v1.5-7B (576 → 32 tokens, 94.4% reduction): 94.8% relative accuracy, outperforming all baselines (next best: DivPrune at 90.5%, VisPruner at 89.7%); 2.3× prefill speedup, 1.5× total speedup.
LLaVA-v1.5-7B (576 → 64 tokens): 96.9% relative accuracy vs. TwigVLM 96.0% (610M params) and DivPrune 93.9%.
LLaVA-v1.5-7B (576 → 128 tokens): 98.5% relative accuracy, comparable to TwigVLM (99.0%) but with far fewer added parameters (0.53M vs. 610M).
LLaVA-NeXT-7B (2880 → 160 tokens): 94.0% relative accuracy, 6.0× prefill speedup, 3.2× total speedup, 8.9× KV cache reduction.
Qwen2.5-VL-7B (1280 → 142 tokens, 88.9% reduction): 94.1% average accuracy vs. FastV's 71.4%.
Video-LLaVA-7B (2048 → 228 tokens): Average accuracy 43.9 vs. FastV 39.3; score 2.45 vs. FastV 2.21.

Limitations¶

Training is required (LPM fine-tuning), adding overhead compared to fully training-free methods, though the scale is small (0.53M params, 10% of one dataset).
The diversity token selection (greedy cosine similarity) is a heuristic that may not optimally capture all relevant background tokens in every scenario.
The fixed pruning layer ($k=12$) and ratio $R_1:R_2 = 3$ were set empirically for LLaVA-1.5 and may not transfer optimally to all VLM architectures.
Total inference speedup (1.5×–3.2×) is substantially lower than prefill speedup because decoding stage latency is less affected by vision token count.
Evaluation uses GPT-4.1 as a judge for video QA, introducing potential evaluation noise.

Relevance to Vision-Language Models¶

LearnPruner directly addresses the compute bottleneck from long visual token sequences in VLMs, a central scalability challenge as high-resolution and video inputs push token counts toward thousands. The paper's analytical contribution—characterizing the failure modes of both encoder-side [CLS] attention (artifact-driven attention sink) and LLM-side attention (positional shift in visual attention, reliability profile across layers)—provides principled grounding for why prior pruning methods degrade under aggressive budgets. The result that text-to-vision attention in middle LLM layers is most reliable for query-aware selection is a practically useful design rule for future efficient VLM architectures. For researchers tracking VLMs, this work offers a clean recipe for combining learned encoder-level redundancy removal with LLM-level query-aware selection, with strong empirical coverage across LLaVA-1.5, LLaVA-NeXT, Video-LLaVA, and Qwen2.5-VL.