LLaVA-PruMerge: Adaptive Token Reduction for Efficient Large Multimodal Models

Yuzhang Shang; Mu Cai; Bingxin Xu; Yong Jae Lee; Yan Yan

Source-linked AI summary

LLaVA-PruMerge: Adaptive Token Reduction for Efficient Large Multimodal Models

Yuzhang Shang, Mu Cai, Bingxin Xu, Yong Jae Lee, Yan Yan

arXiv:2403.15388v6Mar 22, 2024cs.CVcs.AIcs.CL

TL;DR

LMMs face rising inference costs because visual-token prefixes lengthen Transformer contexts, whose attention complexity is quadratic. PruMerge adaptively selects important tokens using sparse class-token attention and merges them by similarity, reducing token counts while preserving comparable performance. Applied to LLaVA-1.5, it uses a small fraction of visual tokens across diverse tasks and also supports video inference without additional training.

Problem
LMMs use many visual tokens, and increasing input length raises Transformer attention costs quadratically, especially for high-resolution images and videos.
Method
PruMerge uses sparse attention between the class token and spatial tokens to prune unimportant tokens, then clusters and merges selected tokens by key similarity.
Results
6.9% of visual tokens on average maintained comparable performance across diverse visual question-answering and reasoning tasks when applied to LLaVA-1.5.
Takeaways & Limitations
Visual-token reduction offers substantial computational savings while preserving LMM reasoning capabilities and can complement other acceleration techniques.
Takeaways & Limitations
PruMerge introduces a marginal performance discrepancy relative to the original LLaVA model, motivating the refined PruMerge+ variant.

Abstract

from arXiv · show

Large Multimodal Models (LMMs) have shown significant visual reasoning capabilities by connecting a visual encoder and a large language model. LMMs typically take in a fixed and large amount of visual tokens, such as the penultimate layer features in the CLIP visual encoder, as the prefix content. Recent LMMs incorporate more complex visual inputs, such as high-resolution images and videos, which further increases the number of visual tokens significantly. However, due to the inherent design of the Transformer architecture, the computational costs of these models tend to increase quadratically with the number of input tokens. To tackle this problem, we explore a token reduction mechanism that identifies significant spatial redundancy among visual tokens. In response, we propose PruMerge, a novel adaptive visual token reduction strategy that significantly reduces the number of visual tokens without compromising the performance of LMMs. Specifically, to metric the importance of each token, we exploit the sparsity observed in the visual encoder, characterized by the sparse distribution of attention scores between the class token and visual tokens. This sparsity enables us to dynamically select the most crucial visual tokens to retain. Subsequently, we cluster the selected (unpruned) tokens based on their key similarity and merge them with the unpruned tokens, effectively supplementing and enhancing their informational content. Empirically, when applied to LLaVA-1.5, our approach can compress the visual tokens by 14 times on average, and achieve comparable performance across diverse visual question-answering and reasoning tasks. Code and checkpoints are at https://llava-prumerge.github.io/.

1. Introduction

LMM inference is costly because large visual-token prefixes increase Transformer attention costs, especially for high-resolution images and videos. PruMerge adaptively prunes and merges redundant visual tokens, preserving comparable performance while substantially reducing token count.

Method: PruMerge adaptively retains more tokens in information-dense images and fewer tokens in simpler images while leaving other LMM procedures unchanged.The selected tokens concentrate on regions with dense information.
Motivation: LMM inference cost is driven by both large language-model size and long visual-token contexts, whose attention cost grows quadratically with input length.LLaVA-1.5 appends 576 visual tokens, while Video-LLaVA uses 2048.
Method: PruMerge exploits sparse class-token attention to identify important visual tokens, then merges selected tokens using similarity-based clustering and weighted averaging.PruMerge+ additionally samples tokens spatial-uniformly to complement the selected tokens.
Results: 5.5% of visual tokens—about 32 tokens on average—maintained comparable performance to all 576 tokens across diverse LLaVA-1.5 benchmarks.The reported average is across the evaluated tasks and reflects adaptive token counts.
Results: PruMerge also accelerated Video-LLaVA inference without additional training and improved performance across multiple video benchmarks.The passage reports both faster processing and improved benchmark performance after integrating PruMerge during inference.

2. Related Work

Prior efficient-LMM approaches primarily reduce backbone size or use compression, while token-reduction work often targets internal vision-transformer computation. PruMerge instead reduces prefix visual tokens for overall LMM efficiency using sparsity-based selection and linear-time merging.

Efficient Large Multimodal Models (LMMs): Efficient LMM work commonly reduces backbone size or applies quantization, while compact architectures target mobile deployment.These approaches include MobileVLM, MobileVLM-v2, and quantized LLaVA-1.5.
Token Reduction: Existing token-reduction methods mainly reduce tokens inside vision-transformer blocks to lower internal ViT computation.They progressively reduce token numbers across transformer blocks.
Token Reduction: PruMerge targets numerous prefix visual tokens in LMM context rather than the internal transformer structure, achieving comparable performance with fewer than one-fourth of the original tokens.Its core selection mechanism identifies anchor tokens from sparse modality-encoder attention scores.
Token Reduction: PruMerge’s k-nearest-neighbor clustering has O(n) time complexity, compared with CrossGet’s O(n^2) graph-matching complexity.The comparison concerns the token-merging module.

3. Method: Token Pru-Merging

PruMerge reduces LMM visual-token computation by adaptively selecting informative tokens from sparse class-token attention and supplementing them through key-similarity clustering. PruMerge+ extends selection to improve the balance between token reduction and performance.

Adaptive Important Token Selection: PruMerge first selects important visual tokens using outlier detection on class-token attention, whose sparse distribution indicates that only a small subset carries key visual information.The class attention value is taken from the ViT penultimate layer.
Adaptive Important Token Selection: Selected tokens adapt to image content: complex images receive more attentive tokens, which are typically spatially aligned with important regions.This behavior is illustrated for images ranging from simple scenes to screens containing dense text.
Token Supplement via Similar Key Clustering: PruMerge clusters visual tokens by key similarity and merges pruned tokens into unpruned cluster centers through class-attention-weighted averaging.Final-layer key vectors serve as token representations, and k-nearest neighbors define the clusters.
PruMerge+: PruMerge+ adds spatially uniform sampling of additional tokens from less-critical regions, preserving substantial reduction while reducing the original method’s performance discrepancy.Its average visual-token reduction is fourfold with minimal performance degradation.

4. Experiments

The experiments evaluate PruMerge across performance, efficiency, generalization, and component effectiveness, including comparisons on six benchmarks.

Experimental Setup: The experiments separately examine LLaVA-1.5 performance, LMM efficiency, cross-modality generalization, and component effectiveness.
Experimental Setup: PruMerge and PruMerge+ are evaluated on six benchmarks while adaptively reducing visual tokens.They use 5.5% and 25.0% of visual tokens on average, respectively, while achieving competitive performance with LLaVA-1.5.

4.1. Main Results

On diverse visual question-answering and reasoning benchmarks, PruMerge achieves comparable performance to LLaVA-1.5 using a small fraction of its visual tokens.

Evaluation Setup: LLaVA-1.5 is evaluated both directly and after one epoch of LoRA fine-tuning with reduced visual tokens.
Benchmark Results: PruMerge achieves comparable performance with LLaVA-1.5 across six visual question-answering and reasoning benchmarks using an average of 32 visual tokens.The evaluated tasks are VQAv2, ScienceQA, TextVQA, POPE, MME, and MMBench.
Benchmark Results: PruMerge outperforms previous methods such as BLIP2 and InstructBLIP on the evaluated benchmarks.
Benchmark Results: PruMerge performs better than LLaVA-1.5 on POPE and ScienceQA.

4.2. Efficiency Analysis

PruMerge reduces visual-token-driven computational costs in LMMs by lowering LLM prefill operations and memory demands, while remaining compatible with other acceleration techniques.

Efficiency Analysis: 14.4 times fewer visual tokens reduces the LLaVA prefill input from 576 visual tokens to approximately 40, matching a 40-token text prompt.The analysis considers a 336 × 336 image processed by CLIP-ViT and a 40-token text prompt.
Efficiency Analysis: PruMerge reduces required FLOPs and computational memory demands while accelerating the LLM prefill process.
Compatibility: PruMerge can complement quantization and factorization because its token reduction strategy is described as orthogonal to those techniques.

4.3. Generalization on Video-LLM

The study extends PruMerge to Video-LLaVA, applying it during inference without additional training to reduce the visual-token burden of video inputs.

Video-LLaVA Setup: Video-LLaVA samples 8 frames and extracts 2048 visual tokens for video perception.
Video-LLaVA Setup: PruMerge and PruMerge+ are integrated into Video-LLaVA during inference without training a new model.
Video Evaluation: Table 3 compares Video-LLaVA variants with other LVMs on video reasoning benchmarks.

4.4. Ablation Study

The sampling analysis compares adaptive, sequential, and spatial token selection, showing that PruMerge performs better across downstream tasks with training-free token reduction. Ablations further show gains from Adaptive Important Token Selection and Token Supplement, while fine-tuning enhances performance.

Token Sampling Strategy Analysis: PruMerge consistently outperforms sequential and spatial sampling across all downstream tasks in training-free experiments.Its stronger TextVQA performance indicates that selected tokens preserve detailed information needed for OCR.
Ablation Studies: With 6.9% of visual tokens held constant, progressively adding AITS and TS further enhances downstream performance.The ablation isolates the contributions of Adaptive Important Token Selection and Token Supplement.
Training Analysis: Fine-tuning further improves LLaVA-PruMerge performance across ScienceQA, TextVQA, POPE, and MME.The reported results distinguish fine-tuned models from training-free models.
Comparison to Token Reduction Methods: PruMerge+ significantly outperforms unimodal token-reduction methods and CrossGet under the same reduction ratio on multimodal tasks.The comparison supports the method’s effectiveness for multimodal contexts.

5. Conclusion

The paper improves LMM efficiency by reducing visual-token quantity through spatial-redundancy-aware pruning and merging. On LLaVA-1.5, retaining 6.9% of visual tokens maintains comparable performance across diverse visual question-answering and reasoning tasks.

Conclusion: The proposed plug-and-play module reduces visual tokens using class-to-spatial-token similarity for pruning and merging.The method targets spatial redundancy while preserving visual information relevant to LMM reasoning.
Conclusion: 6.9% of visual tokens maintains comparable performance across a wide range of visual question-answering and reasoning tasks.The conclusion frames this reduction as enabling substantial computational savings without sacrificing reasoning capabilities.

7. Appendix

PruMerge selects informative visual tokens using sparse class attention, then merges similar tokens into selected centers. Visualizations indicate representative token coverage that supports substantial reduction while minimizing performance loss.

Algorithm: Class-attention sparsity guides selection of the most informative visual tokens at the end of the visual encoder.This is the first stage of both PruMerge and PruMerge+.
Algorithm: Selected tokens serve as cluster centers for grouping and merging similar tokens.The second stage supplements the retained tokens with information from grouped tokens.
Token Positions: PruMerge’s selected-token positions provide comprehensive and representative visual coverage.The visualization connects token placement with minimized performance loss during substantial reduction.