Balanced Token Pruning: Accelerating Vision Language Models Beyond Local Optimization

TL;DR

Balanced Token Pruning (BTP) tackles the token-burden problem in LVLMs by jointly optimizing pruning across current layers and downstream layers. It uses a calibration-set to stage pruning, applying diversity-focused objectives early to preserve downstream representations and attention-focused objectives later to maintain local output fidelity, guided by a local-global objective. Empirical results across multiple LVLMs show substantial token reduction with minimal performance loss and notable improvements in end-to-end latency and memory usage, with ablations validating the importance of the balance strategy, spatial initialization, and attention rebalancing. This approach enables efficient, scalable deployment of LVLMs on edge devices while preserving multimodal capabilities.

Abstract

Large Vision-Language Models (LVLMs) have shown impressive performance across multi-modal tasks by encoding images into thousands of tokens. However, the large number of image tokens results in significant computational overhead, and the use of dynamic high-resolution inputs further increases this burden. Previous approaches have attempted to reduce the number of image tokens through token pruning, typically by selecting tokens based on attention scores or image token diversity. Through empirical studies, we observe that existing methods often overlook the joint impact of pruning on both the current layer's output (local) and the outputs of subsequent layers (global), leading to suboptimal pruning decisions. To address this challenge, we propose Balanced Token Pruning (BTP), a plug-and-play method for pruning vision tokens. Specifically, our method utilizes a small calibration set to divide the pruning process into multiple stages. In the early stages, our method emphasizes the impact of pruning on subsequent layers, whereas in the deeper stages, the focus shifts toward preserving the consistency of local outputs. Extensive experiments across various LVLMs demonstrate the broad effectiveness of our approach on multiple benchmarks. Our method achieves a 78% compression rate while preserving 96.7% of the original models' performance on average. Our code is available at https://github.com/EmbodiedCity/NeurIPS2025-Balanced-Token-Pruning.

Figures (12)

• Figure 1: Layer-wise visualization of attention in LVLMs.
• Figure 2: Impact of different pruning strategies on layer-wise representations.
• Figure 3: Overview of BTP: We first use a calibration set to determine the pruning layers. In the early layers, we emphasize diversity-based pruning to preserve the output of subsequent layers. In the deeper layers, attention-based pruning is prioritized to maintain the output of the pruning layers. Due to the pruning strategy, we achieve an overall optimal pruning balance.
• Figure 4: Layer-wise image token hidden state dynamics and attention allocation in LVLMs.
• Figure 5: Effect of various pruned methods on the output of decoder layers.