Source-linked AI summary
TopV: Compatible Token Pruning with Inference Time Optimization for Fast and Low-Memory Multimodal Vision Language Model
Cheng Yang, Yang Sui, Jinqi Xiao, Lingyi Huang, Yu Gong, Chendi Li, Jinghua Yan, Yu Bai, Ponnuswamy Sadayappan, Xia Hu, Bo Yuan
TL;DR
VLM inference is burdened by large visual-token sequences and prior pruning methods’ reliance on attention scores and incompatible KV-cache behavior. TopV instead optimizes visual-token contribution with a visual-aware cost, prunes once during prefilling, and maintains compatibility with FlashAttention and KV cache. Across reported experiments, it preserves accuracy with reduced computation and memory, outperforming prior methods under comparable reductions.
Problem
VLM inference requires substantial computation and memory because visual tokens can comprise most input tokens, while prior pruning methods use attention-based criteria and can conflict with FlashAttention or KV caching.
Method
TopV formulates visual-token pruning as an optimization problem using feature similarity, spatial relationships, and central distance, then prunes once during prefilling without additional training.
Results
TopV outperforms previous token-pruning methods while reducing computational demand, memory requirements, and KV-cache storage.
Takeaways & Limitations
TopV provides a fast, low-memory VLM inference approach that remains compatible with FlashAttention and KV cache.
Takeaways & Limitations
The method’s memory behavior is discussed under the assumption that sufficient remaining GPU memory and non-reusable retained memory can affect CUDA allocation strategies.
Abstract
from arXiv · showhide
Vision-Language Models (VLMs) demand substantial computational resources during inference, largely due to the extensive visual input tokens for representing visual information. Previous studies have noted that visual tokens tend to receive less attention than text tokens, suggesting their lower importance during inference and potential for pruning. However, their methods encounter several challenges: reliance on greedy heuristic criteria for token importance and incompatibility with FlashAttention and KV cache. To address these issues, we introduce \textbf{TopV}, a compatible \textbf{TO}ken \textbf{P}runing with inference Time Optimization for fast and low-memory \textbf{V}LM, achieving efficient pruning without additional training or fine-tuning. Instead of relying on attention scores, we formulate token pruning as an optimization problem, accurately identifying important visual tokens while remaining compatible with FlashAttention. Additionally, since we only perform this pruning once during the prefilling stage, it effectively reduces KV cache size. Our optimization framework incorporates a visual-aware cost function considering factors such as Feature Similarity, Relative Spatial Distance, and Absolute Central Distance, to measure the importance of each source visual token, enabling effective pruning of low-importance tokens. Extensive experiments demonstrate that our method outperforms previous token pruning methods, validating the effectiveness and efficiency of our approach.
1. Introduction
VLM inference is costly because visual inputs add many tokens and substantial memory use. TopV addresses pruning limitations by optimizing visual-token importance without explicit attention scores, while preserving FlashAttention and KV-cache compatibility.
- Motivation: Visual tokens can constitute up to 87% of LLaVA-v1.5 OCR inputs and up to 95% of InternVL inputs, increasing inference computation and memory.LLaVA-v1.5 uses 576 visual tokens, while InternVL produces 256–1792 depending on image resolution.
- Limitations of prior work: Prior pruning methods use attention, merging, or layer-based withdrawal, but the introduction identifies greedy importance criteria and systems incompatibility as drawbacks.FastV uses attention scores, LLaVA-PruMerge merges less important tokens, and VTW removes vision tokens from selected layers.
- TopV: TopV formulates token selection as an optimization problem based on token contribution rather than attention scores, requiring no additional training.Its visual-aware cost incorporates feature similarity, spatial relationships, and central distance.
- TopV: TopV prunes once during prefilling, maintains the reduced token set during decoding, and remains compatible with FlashAttention and KV cache.Avoiding explicit attention-score computation preserves FlashAttention efficiency while reducing KV-cache size and memory usage.
2. Related Work
Related work shows that VLMs have expanded across image, video, audio, and language tasks, but increasing resolution and frame counts create severe token-growth challenges. This motivates more efficient multimodal computation.
- Vision-Language Models: VLMs such as CLIP, LLaVA, and BLIP-2 support image-language tasks including captioning and zero-shot image-to-text generation.Video-LLaMA and LLaVA-NeXT extend multimodal processing to video and audio.
- Vision-Language Models: Higher-resolution images and multiple video frames require substantially more tokens, creating computational bottlenecks for multimodal models.The passage notes that LLaVA can process high-resolution images into thousands of tokens, while video models allocate still more across frames.
3. Method
TopV estimates visual-token importance by optimally matching source tokens to target tokens with a visual-aware cost, then prunes low-contribution tokens and optionally recovers a uniform subset. The method uses Post-LN outputs as targets and combines feature, spatial, and central-distance factors.
- 3.1. Token Importance Formulation: TopV formulates visual-token pruning as an optimization problem that retains source tokens contributing most to subsequent-layer target tokens.The method is inspired by optimal transport and uses the resulting contribution scores to identify removable tokens.
- 3.1. Token Importance Formulation: The transport plan P produces a contribution matrix P* whose source-token scores determine which tokens are discarded.In experiments, importance is measured at Li = 2 and corresponding tokens are removed from subsequent layers.
- 3.5. Token Recovery: After pruning, TopV uniformly samples some discarded tokens for recovery, helping preserve balanced visual coverage and reduce visual collapse.The recovery strategy is intended especially for tasks requiring attention to diverse image regions, such as OCR.
- 3.2. Target Token Selection: TopV selects Post-LN outputs as target tokens because they capture attention processing while avoiding residual-induced similarity that can obscure source-target distinctions.The method empirically compares Pre-LN, Attention, Post-LN, and MLP positions.
- 3.3. Visual-Aware Cost Function: The visual-aware cost combines normalized feature similarity, relative spatial distance, and absolute central distance with weights α, β, and γ.A smaller γ limits excessive pruning of tokens near image edges.
- 3.3. Visual-Aware Cost Function: Relative spatial distance assigns lower transport costs to tokens representing the same or neighboring image regions than to distant tokens.This encodes the stronger expected correlation between nearby visual tokens.
4. Experiments
Experiments evaluate TopV across diverse VLM tasks and models using performance, memory, throughput, latency, and FLOPs-reduction measures. TopV maintains accuracy with substantial efficiency gains, supports KV-cache memory reuse, and selects foreground-focused tokens while token recovery helps generation-heavy tasks.
- Experimental Setup: Experiments span LLaVA and InternVL models across recognition, question answering, captioning, OCR, and video understanding tasks.The evaluation includes larger-scale LLaVA-13B and InternVL2-26B models in addition to the main LLaVA-v1.5-7B and InternVL2-2B settings.
- Main Results: For LLaVA, 35% vision-token FLOPs reduction causes negligible accuracy loss and raises average inference efficiency to 1.08×.At roughly 50% reduction, TopV exceeds FastV by 0.39% average accuracy while reducing dynamic memory by 49% and improving efficiency by 1.68×.
- Memory Analysis: KV-cache compatibility enables memory reuse across generated tokens, whereas without KV cache memory repeatedly spikes, especially for multi-token OCR generation.The profiling compares TopV, baseline, and FastV on AI2D and OCRBench for InternVL2-2B; static memory is separate from dynamic runtime memory.
- Ablations and Analysis: TopV concentrates selected tokens in foreground regions rather than background or padded areas, and token recovery benefits tasks requiring many generated tokens.The target-token-position and visual-aware-cost-function studies further analyze design choices, while Table 5 evaluates video question answering.
5. Conclusion
TopV optimizes VLM inference by pruning redundant visual tokens without additional training or fine-tuning. It remains compatible with FlashAttention and reduces KV cache storage by pruning once during prefilling.
- TopV formulates token pruning as an optimization problem to reduce computational demand and memory requirements.
A. Memory Analysis
GPU profiling shows that TopV substantially reduces dynamic memory usage compared with Baseline and FastV, while the model’s static memory remains unchanged by pruning.
- GPU profiling monitored peak memory usage after warmup on a subset of the dataset because profiling all tasks was computationally demanding.
- TopV reduces dynamic memory usage by 45.6% versus Baseline and 52.7% versus FastV on AI2D.Dynamic memory usage is 0.98GB for TopV, 1.8GB for Baseline, and 2.07GB for FastV.
- TopV reduces dynamic memory usage by 43.2% versus Baseline and 43.7% versus FastV on OCRBench.Dynamic memory usage is 0.71GB for TopV, 1.25GB for Baseline, and 1.26GB for FastV.
- InternVL2-2B consumes approximately 4.5GB of static GPU memory, unaffected by token pruning techniques.
B. Discussion
The discussion identifies combining TopV with model compression or KV cache compression as an important direction for further improving multimodal inference efficiency.
- Future work should examine how token pruning can be balanced with model compression and KV cache compression.These approaches are presented as orthogonal methods for improving inference speed and reducing memory usage.
- TopV’s reported AI2D GPU memory figures include 5.21 GiB maximum allocated and 5.53 GiB maximum reserved.
- TopV’s reported OCRBench GPU memory figures include 6.30 GiB maximum allocated and 8.19 GiB maximum reserved.
- The AI2D baseline figure provides GPU memory usage for InternVL2-2B without reporting maximum allocated or reserved values in the passage.
C. Evaluation Examples
Across dialogue and evaluation examples, TopV reduces vision-token computation while maintaining performance and often producing outputs closer to the ground truth than Baseline or FastV.
- TopV reduces vision-token FLOPs by 51% for LLaVA-v1.5-7B and 48% for InternVL2-2B while maintaining performance.
- In the OCR example, TopV recognizes “Jump,” while Baseline fails and FastV detects only the letter “J.”
- In the Writing task, TopV and Baseline produce similar, more detailed descriptions than FastV.
- In the Captioning task, TopV correctly identifies the rabbit, whereas Baseline and FastV classify it as a baby.
- The evaluation examples compare TopV, Baseline, and FastV in OCR, Writing, and Captioning tasks.