ATP-LLaVA: Adaptive Token Pruning for Large Vision Language Models

TL;DR

The paper tackles the high cost of processing many visual tokens in LVLMs by showing pruning effects vary across LLM layers and instances. It introduces ATP-LLaVA with an Adaptive Token Pruning module inserted between decoder layers and a Spatial Augmented Pruning policy, using learnable thresholds and differentiable masks. Budget-constrained training further guides pruning to balance computation with vision-language understanding, achieving 75% average token reduction with only 1.9% performance loss across seven benchmarks. This approach delivers substantial efficiency gains while maintaining near-original performance, enabling scalable LVLM deployment on limited hardware.

Abstract

Large Vision Language Models (LVLMs) have achieved significant success across multi-modal tasks. However, the computational cost of processing long visual tokens can be prohibitively expensive on resource-limited devices. Previous methods have identified redundancy in visual tokens within the Large Language Model (LLM) decoder layers and have mitigated this by pruning tokens using a pre-defined or fixed ratio, thereby reducing computational overhead. Nonetheless, we observe that the impact of pruning ratio varies across different LLM layers and instances (image-prompt pairs). Therefore, it is essential to develop a layer-wise and instance-wise vision token pruning strategy to balance computational cost and model performance effectively. We propose ATP-LLaVA, a novel approach that adaptively determines instance-specific token pruning ratios for each LLM layer. Specifically, we introduce an Adaptive Token Pruning (ATP) module, which computes the importance score and pruning threshold based on input instance adaptively. The ATP module can be seamlessly integrated between any two LLM layers with negligible computational overhead. Additionally, we develop a Spatial Augmented Pruning (SAP) strategy that prunes visual tokens with both token redundancy and spatial modeling perspectives. Our approach reduces the average token count by 75% while maintaining performance, with only a minimal 1.9% degradation across seven widely used benchmarks. The project page can be accessed via https://yxxxb.github.io/ATP-LLaVA-page/.

Figures (4)

• Figure 1: (a) Previous methods employ a fixed, pre-defined token pruning ratio. (b) Illustration of ATP-LLaVA, which dynamically selects the adaptive pruning ratio for each layer of the LLM decoder based on the instance-specific characteristics.
• Figure 2: Comparison of vision token pruning at different LLM decoder layers and pruning ratios across fine-grained and coarse-grained tasks of the SEED-Image li2023seed. Fine-grained tasks include instance counting, spatial relation, etc. Coarse-grained tasks include scene understanding, etc.
• Figure 3: Illustration of the Adaptive Token Pruning (ATP) module. The ATP module can be flexibly inserted between any two LLaMA decoder layers. It adaptively predicts pruning thresholds for current layer and instance. Redundant or text-irrelevant visual tokens are pruned at this stage, and they will be ignored by other tokens in subsequent LLaMA decoder layers.
• Figure 4: Visualized vision token pruning results of ATP-LLaVA. White tokens represents the pruned tokens. The uniform sampled tokens are pruned by spatial pruning threshold, while the sparse tokens are pruned by redundant pruning threshold. Zoom in to have a better view.