LearnPrunerV2: Learnable Progressive Token Pruning for Vision-Language Models

1 Introduction

• Existing pruning criteria can miss task-specific regions through attention biases or query-agnostic selection, while static or heuristic schedules overlook redundancy's evolution across layers.Early pruning may be computationally efficient but blind to query relevance, whereas late pruning offers marginal speedup.Source: §1 · pp. 1–2
• LearnPrunerV2 places lightweight learnable pruning modules in every VLM layer, combining local visual saliency with query relevance from a shadow token to predict token importance.End-to-end optimization aligns pruning decisions with layer-specific visual significance and lets modules determine each token's exit depth.Source: §1 · p. 2
• Budget-Deterministic Pruning assigns fixed per-layer token budgets using offline-calibrated statistics and diversity-aware selection for hardware-friendly inference.The stated deployment aims include deterministic inference latency and maximized batching throughput.Source: §1 · pp. 2–3
• Across multiple VLMs and VQA benchmarks, LearnPrunerV2 outperforms prior state-of-the-art; pruning 94.4% of visual tokens in LLaVA-NeXT-7B preserves 96.2% of original performance with 9.7× fewer TFLOPs.The method is described as model-agnostic and applicable to architectures including LLaVA and Qwen-VL.Source: §1 · p. 2

2.1 Vision Language Models

• Visual tokens can number in the hundreds or thousands, increasing VLM computational costs; compression methods reduce computation but often require retraining and may degrade the performance upper bound.The paper therefore focuses on pruning techniques applicable directly to existing VLMs.Source: §2.1 · p. 3, §2.2 · p. 3

2.2 Vision Token Reduction for VLMs

• Pruning criteria divide broadly into importance-based attention selection and diversity-based token selection, with attention-map methods incompatible with Flash Attention.Importance-based examples include FastV and SparseVLM [49]; diversity-based examples include DivPrune [2], DART [38], and CDPruner [48].Source: §2.2 · p. 3
• Most pruning schedules are static or layer-specific, while multi-stage approaches such as PDrop [40], AutoPrune, and LearnPruner use heuristic schedules that may not match dynamic visual-information processing.The cited approaches include uniform pruning ratios and intervals, sample-specific policies, and two-stage pruning criteria.Source: §2.2 · p. 3

3.1 Preliminaries

• The LLM backbone is a stack of Transformer decoder layers that apply causal self-attention and feed-forward refinement before final hidden states support next-token prediction.The model input concatenates visual tokens, text tokens, and system prompts.Source: §3.1 · p. 3

3.2 Learnable Pruning Criterion

• Each decoder layer scores visual tokens by combining local visual features with query relevance derived from a shadow token that aggregates preceding context.The shadow token can attend to preceding tokens but remains invisible to other tokens, avoiding perturbation to the model’s original predictions.Source: Figure 2 · p. 4, §3.2 · p. 4
• Training uses straight-through thresholding and cumulative masks so pruned tokens remain excluded in later layers, while a pruning loss constrains the number preserved.The training mask simulates inference by cutting off interactions between pruned and retained tokens.Source: §3.2 · pp. 4–5

3.3 Budget-Deterministic Pruning

• Threshold-based pruning produces consistent, low-variance staircase retention trajectories across sampled inputs on LLaVA-v1.5 and Qwen2.5-VL.For LLaVA-v1.5, the trajectory drops sharply at the first layer, plateaus in early layers, declines again in middle layers, and approaches zero in deep layers.Source: §3.3 · p. 5, Figure 3 · p. 5
• BDP scales calibrated layer-wise retention ratios to a target average token budget, preserving their relative distribution while fixing each layer’s token count.The target average budget is converted into a total budget across layers, then used to set each layer’s budget.Source: §3.3 · pp. 5–6
• When more above-threshold candidates qualify than a layer’s budget allows, BDP combines importance ranking with diversity-aware selection to fill the fixed capacity.It selects high-importance pivot tokens, then fills remaining slots with candidates least similar to those pivots.Source: §3.3 · p. 6

4.1 Experimental Settings

• Evaluations cover nine image and three video benchmarks; base VLM weights remain frozen while pruning modules train on LLaVA-mix665k.Experiments include LLaVA-v1.5, LLaVA-NeXT, Video-LLaVA, and Qwen2.5-VL.Source: §4.1 · p. 6

4.2 Main Results

• LLaVA-NeXT retains 99.7% of original performance with 22.2% of visual tokens; Qwen2.5-VL exceeds 100% average relative accuracy at 66.7% and 77.8% pruning.Source: Table 2 · p. 7, Table 3 · p. 7, §4.2 · p. 7
• At 128 retained tokens, LearnPrunerV2 maintains 99.8% of vanilla performance and beats prior state of the art by 1.3%; the gap reaches 2.1% at 32 tokens.Source: §4.2 · p. 7

4.3 Ablation Study

• On a fixed 64-token budget, gating, shadow-token relevance, and hybrid parameter sharing each lead their respective ablations at 97.8% average performance.Input-only and input-plus-relevance score 96.6% and 97.2%; last-token attention and logits score 96.5% and 97.4%, while independent and fully shared schemes score 97.3% and 96.9%.Source: Table 4 · p. 8, §4.3 · p. 8
• At 32 tokens, diversity-aware selection beats Top-K by 0.9%; at 64 tokens, their average performances are 97.8% and 97.7%.Source: Table 5 · p. 8, §4.3 · p. 8

4.4 Efficiency Analysis

• Reducing LLaVA-NeXT visual tokens from 2,880 to 160 yields 9.7× fewer TFLOPs, 6.2× faster prefill, 8.7× smaller KV cache, and 3.7× faster total inference.Source: Table 6 · p. 8, §4.4 · p. 8

• LearnPrunerV2 is a learnable progressive token-pruning framework for accelerating VLMs.Source: §5 · p. 8
• It estimates token importance using intrinsic visual saliency and query relevance, with Budget-Deterministic Pruning assigning fixed per-layer budgets.Source: §5 · p. 8
• Experiments across multiple VLM architectures and benchmarks report strong accuracy-efficiency trade-offs and gains over prior methods.Source: §5 · p. 8

A.1 Results on Video-LLaVA

• LearnPrunerV2 consistently outperforms FastV on TGIF-QA, MSVD-QA, and MSRVTT-QA, especially under lower retention ratios.Source: §A.1 · p. 11
• Video-LLaVA uniformly samples 8 frames per video, producing 2,048 visual tokens; evaluation uses the first 1,000 samples per benchmark and GPT-4.1-assisted scoring.Source: §A.1 · p. 11
• The reported video results demonstrate LearnPrunerV2’s generalizability from image-based to video-based VLMs.Source: §A.1 · p. 11

A.2 Data Efficiency

• 98.5% relative performance is achieved on LLaVA-v1.5 with 10% of training data under the 128-token budget.Source: §A.2 · p. 11
• On LLaVA-v1.5 and LLaVA-NeXT, performance largely converges by 50% of the LLaVA-mix665k training data, matching full-dataset training.Source: §A.2 · p. 11
• Lower token budgets generally require more training data to reach convergence, which the authors attribute to the greater precision required by aggressive pruning.Source: §A.2 · p. 11

A.3 Visualization

• At shallow layers, LearnPrunerV2 removes many uninformative tokens, including padding regions and background areas.Source: §A.3 · p. 11, Figure 5 · p. 12, Figure 6 · p. 13
• Retained tokens converge toward query-relevant regions in middle layers, and nearly all visual tokens are pruned by deeper layers.Source: §A.3 · p. 11, Figure 5 · p. 12, Figure 6 · p. 13
• LearnPrunerV2 identifies query-relevant regions by layer 12, earlier than last-token attention, which typically requires deeper layers such as layer 16.Source: §A.3 · p. 11, Figure 7 · p. 14, Figure 8 · p. 15
• Compared with last-token attention, LearnPrunerV2’s importance scores show stronger cross-layer consistency, while last-token attention has positional bias and substantial fluctuations.Source: §A.3 · p. 11, Figure 7 · p. 14, Figure 8 · p. 15