Skip to content

Balanced Token Pruning: Accelerating Vision Language Models Beyond Local Optimization

🕒 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

Existing visual token pruning methods for LVLMs optimize either the current layer's output (attention-based) or global token diversity (diversity-based), but not both simultaneously. Balanced Token Pruning (BTP) introduces a multi-stage pruning framework that shifts from diversity-dominant objectives in early layers to attention-dominant objectives in deeper layers, achieving 78% token compression while retaining 96.7% of average model performance.

Problem

Token pruning in LVLMs has two competing failure modes: attention-based methods (e.g., FastV, PyramidDrop) minimize local output error at each pruned layer but are globally suboptimal because the set of important tokens at layer l does not subsume important tokens at layer l+k; diversity-based methods (e.g., DivPrune) select semantically spread tokens that better cover future-layer needs but fail to preserve local output consistency, accumulating error in deeper layers. Neither approach jointly considers local and global objectives, and pruning layer selection has relied on manual tuning or dataset-specific validation.

Method

BTP formulates a local-global objective (Eq. 7) that combines a local attention-based term (minimize output divergence at the current pruned layer) with a diversity term (maximize semantic spread to anticipate token importance at downstream layers), weighted by a per-stage scalar λᵢ that increases monotonically across stages:

\[\mathcal{L}_{local\text{-}global} = -\sum_i \sum_{j \in P_i} \left(\lambda_i \cdot \text{Attn}^{(i)}(X_I, X_T) + (1-\lambda_i) \cdot F_{dis}(P_i)\right)\]
  • Early stages (large λ → more diversity weight): retain spatially diverse tokens to preserve representational quality for subsequent layers.
  • Deep stages (small λ → more attention weight): prioritize task-relevant tokens for local output fidelity.

Attention rebalancing: rather than taking raw top-k attention scores (biased by positional encoding toward late tokens), BTP over-selects top-k' > k tokens and fills the retained set with earlier-position tokens first, then later-position tokens.

Spatial diversity initialization: instead of solving the O(n²) Max-Min Diversity Problem (MMDP) in feature space, BTP initializes the diversity-selected set by solving MMDP over 2D Manhattan/Euclidean distances on the token grid, reducing overhead to O(n) levels and enabling GPU-friendly execution.

Calibration-based layer selection: using 64 unlabeled samples from LLaVA-655k, BTP computes layer-wise cosine similarity of image token hidden states and selects pruning layers just before/after layers where token semantics shift significantly (high cosine dissimilarity), avoiding trial-and-error layer tuning.

Key Contributions

  • Formal analysis identifying that attention-based pruning optimizes local optima and diversity-based pruning neglects local consistency, unifying them under a single local-global objective.
  • BTP: a plug-and-play, multi-stage pruning method that transitions smoothly from diversity- to attention-dominant objectives across depth.
  • Attention rebalancing operation to correct positional-encoding bias in importance score computation.
  • Spatial-position-based MMDP initialization that achieves O(n) diversity selection and avoids the latency blowup of feature-space MMDP.
  • Calibration-based, dataset-agnostic pruning layer selection using only 64 shared samples.

Results

  • LLaVA-1.5-7B at 128/576 tokens (78% compression): BTP 98% avg. performance vs. VTW 89%, PDrop 94%, FastV 96%, DivPrune 96%.
  • LLaVA-1.5-13B at 128/576: BTP 98% vs. VTW 89%, PDrop 96%, FastV 95%, DivPrune 96%.
  • LLaVA-1.6-7B at 22% token retention: BTP 94% vs. VTW 85%, PDrop 92%, FastV 94%, DivPrune 92%.
  • Qwen2.5-VL-7B at 25% token retention: BTP 97% vs. VTW 65%, PDrop 82%, FastV 96%, DivPrune 96%.
  • Latency (LLaVA-1.5-7B, 128 tokens): BTP 0.134s (7% reduction), DivPrune 0.224s (54% increase), FastV 0.122s (15% reduction) — BTP is practical where DivPrune is not.
  • KV cache (LLaVA-1.6-7B): BTP reduces from 1.11 GB to 0.28 GB (74.7% ↓), matching DivPrune.
  • On LLaVA-COCO captioning, BTP (LLaVA-1.6) scores 98.9 vs. original 106.6, vs. DivPrune 99.1 and FastV 92.6.

Limitations

  • Requires a calibration set (64 samples) at setup time; though lightweight, this is an extra preprocessing step not needed by purely training-free methods like FastV.
  • λ schedule and stage boundaries are determined empirically; the paper does not provide a closed-form derivation for optimal λᵢ values.
  • Spatial diversity initialization is an approximation of the true feature-space MMDP; the paper notes Euclidean outperforms Manhattan but does not evaluate against full feature-space MMDP quantitatively.
  • Evaluated primarily on static image QA benchmarks; performance on video or streaming modalities with even larger token counts is unexplored.
  • DivPrune's latency issue motivates BTP's spatial initialization, but BTP does not match FastV's raw latency at the same compression ratio.

Relevance to Vision-Language Models

BTP directly addresses one of the primary deployment bottlenecks in modern LVLMs — the quadratic KV-cache and attention cost from thousands of visual tokens, especially under dynamic high-resolution encoding (LLaVA-NeXT: up to 2,880 tokens). By unifying the two dominant pruning paradigms (attention-based and diversity-based) under a single stage-aware objective, this work advances understanding of why existing methods fail rather than just empirically outperforming them. The calibration-based, dataset-agnostic layer selection is particularly practical for plug-and-play deployment across different LVLM architectures (LLaVA, Qwen2.5-VL). This is relevant to anyone working on efficient LVLM inference, edge deployment of multimodal models, or KV-cache compression.