Balanced Token Pruning (BTP) Overview

• Balanced Token Pruning (BTP) is a set of algorithmic strategies that reduce the number of visual tokens in transformers, ensuring balanced retention of class context, prompt alignment, and diversity.
• BTP employs multi-objective formulations—such as per-class quotas and local-global joint pruning—to maintain both local output consistency and global information coverage across layers.
• Empirical results demonstrate that BTP achieves substantial token reductions (up to 88.9%) with minimal accuracy loss (<1%), significantly cutting computational cost in segmentation and vision-language tasks.

Balanced Token Pruning (BTP) refers to a set of algorithmic methodologies for accelerating vision transformers and large vision-LLMs (LVLMs) via structured reduction of visual tokens, while explicitly balancing multiple competing objectives such as class/context coverage, prompt alignment, and overall representation diversity. Unlike naïve or static pruning based solely on singular criteria (e.g., attention magnitude), BTP strategically manages the preservation and removal of tokens to maintain both local output consistency at each layer and global information content across depths and modalities. Modern BTP techniques appear in both semantic segmentation and multi-modal transformer inference settings, where excessive token counts lead to prohibitive computational cost.

1. Theoretical Foundations and Objectives

Balanced Token Pruning is characterized by its multi-objective formulation. Theoretical advances formalize the trade-off between preserving task-relevant information and achieving computational efficiency. In vision-LLMs, this is captured by minimizing the loss in model outputs measured under the Hausdorff distance, with the output perturbation bounded by the maximal covering radii of pruned tokens relative to prompt and visual sets (Li et al., 15 May 2025):

$$\|\mathcal{F}(X)-\mathcal{F}(X_s)\|\;\le\;C_\ell\, \max \Bigl\{ \min\{d_H(S, V),\,d_H(V,P)\}, \min\{d_H(S,V),\,d_H(S,P)\}\Bigr\}$$

where $S$ is the pruned token set, $V$ is the original set of visual tokens, $P$ is the set of prompt tokens, $d_H$ is the Hausdorff distance, and $C_\ell$ is the Lipschitz constant for the model block. This bound highlights the necessity for BTP to maintain balance between "prompt alignment" and "visual preservation" objectives, since exclusive pursuit of either introduces substantial output error.

The BTP constraint, by its various implementations, enforces a minimum level of coverage for critical semantic or spatial information -- either by per-class quotas or multi-objective budget allocation.

2. Algorithmic Methodologies in BTP

There are two primary algorithmic frameworks in Balanced Token Pruning:

• Per-Class Quota Pruning (Semantic Segmentation): Within Dynamic Token Pruning (DToP) (Tang et al., 2023), the algorithm divides the ViT into $M$ stages, each terminated by an auxiliary head $\mathcal{O}_m$ that estimates per-token class confidence. At every stage $m$, all tokens with predicted confidence $p_{m,i}\geq p_0$ are finalized except for the top-$k$ tokens per semantic class, guaranteeing minimal contextual representation for each class. Tokens are selected into

$$E_m^{(c)} = \{ i \in E_m : \arg\max P_m[i, \cdot] = c \}$$

the top-$k$ with largest $p_{m,i}$ per class $c$ are retained, ensuring context is not erased for "easy" classes. The set of survivors $S_m$ passed to the next stage aggregates residual "hard" tokens and these quotas.

• Multi-Stage Local-Global Joint Pruning (Vision-LLMs): In LVLMs, BTP operates via a multi-stage pruning schedule, partitioning layers into semantic-shift stages identified on a calibration set. At each stage, tokens are selected based on a composite criterion:

$$\text{score}_j = \lambda_i\, S_{\text{img}}^{(\ell_i)}(j) + (1 - \lambda_i)\, \text{diversityContribution}(j)$$

where $S_{\text{img}}^{(\ell_i)}$ measures attention-based importance and $\text{diversityContribution}(j)$ reflects marginal coverage for the set diversity function $F_{\text{dis}}$. $\lambda_i$ is monotonically increased at deeper stages, prioritizing diversity (global coverage) early and local score (attention) later (Li et al., 28 May 2025).

• Bi-Objective Covering via Greedy Radius Trading: MoB (Multi-Objective Balanced Covering) formulates pruning as a joint covering problem:

$$(\mathcal{S}_p^\star, \mathcal{S}_v^\star) = \arg\min_{|\mathcal{S}_p| + |\mathcal{S}_v| = K} \max\{ d_H(\mathcal{S}_p, P), d_H(\mathcal{S}_v, V) \}$$

where the total token budget $K$ is allocated via two-phase greedy covering: prompt alignment by $k$-NN, visual diversity via farthest point sampling. This delivers provable error bounds and $O(N(L+K)d)$ time complexity (Li et al., 15 May 2025).

3. Architectural Integration and Inference Workflow

Semantic Segmentation (DToP-based): The ViT encoder is partitioned into $M$ consecutive stages, each concluded by an auxiliary segmentation head. At every stage, the head yields per-token confidence, and tokens classified as "easy" are exited and finalized unless they are within the top-$k$ class quota. Only the survivor tokens are propagated to higher layers, resulting in dynamic reduction of the input sequence length and computation at each subsequent stage. The final dense prediction map is constructed by combining early-exited tokens at their respective exit stages with tokens predicting at the deepest head (Tang et al., 2023).

Vision-LLMs (LVLMs): Pruning layers are selected using calibration set statistics (cosine similarity, attention shift). Pruning proceeds across $S$ stages, each with a retention schedule. Tokens surviving at each stage are selected by maximal composite local-global score. Inference applies the same layer/stage configuration learned on the calibration set, yielding a compressed and more efficient sequence (Li et al., 28 May 2025).

4. Empirical Results and Efficiency Gains

Model	Dataset	Token Reduction (%)	Accuracy Retained (%)	Key Benefits
ViT-Base+ATM (DToP)	ADE20K	21 (GFLOPs)	+0.1 (mIoU)	No drop in mIoU, 21% FLOPs reduction
ViT-Base+FCN (DToP)	ADE20K	25 (GFLOPs)	0.0 (mIoU)	No accuracy drop at 25% compute saving
SegViT-L (DToP)	ADE20K	33 (GFLOPs)	-0.5 (mIoU)	33% FLOPs reduction, mIoU drop < 0.5
LLaVA-v1.5-7B (BTP)	Multi-Benchmark	78 (tokens)	98	3x TFLOPs cut, 67–75% KV-cache cut
LLaVA-Next-7B (MoB)	Multi-Benchmark	88.9 (tokens)	97.9	1.3–1.5x speed-up, <0.5% accuracy loss
Qwen2-VL-7B (MoB)	Multi-Benchmark	66.7 (tokens)	98.4	Outperforms single-objective and prior methods

BTP consistently preserves >96% of the original model's performance with only 11–22% of tokens retained, depending on the model and task, and delivers end-to-end latency and memory reductions of 3x or better. Purely attention-only or diversity-only schedules yield lower retained accuracy or do not yield practical end-to-end gains due to $O(n^2)$ overheads in GPU computation. BTP formulations circumvent these limitations with spatial initialization, calibration-guided stage selection, and balanced score functions (Tang et al., 2023, Li et al., 28 May 2025, Li et al., 15 May 2025).

5. Multi-Objective Trade-offs and Hyperparameter Tuning

Empirical ablations demonstrate that naïve pruning, even with high confidence thresholds, can induce accuracy drops by as much as 1–2% absolute, particularly via context starvation for overrepresented "easy" classes. The introduction of per-class quotas (e.g., top-$k$ tokens per class) or budget allocation (prompt vs. visual coverage) restores and often increases mIoU by 0.8–1.0% at equivalent FLOPs budgets (Tang et al., 2023).

In practice, $k=5$ (per-class quota) or carefully tuned $\lambda_i$ schedules (local/global weighting) are transferable across architectures and tasks. Lower $k$ values marginally increase computational savings but begin to erode mIoU, while higher $k$ values provide diminishing returns for accuracy. Calibration-based selection of pruning layers and adaptive weighting strategies are essential for optimal compression-accuracy trade-offs (Li et al., 28 May 2025, Li et al., 15 May 2025).

6. Theoretical and Practical Advances

The introduction of closed-form error bounds—such as those derived from covering theory and the Hausdorff distance—establishes the mathematical optimality and robustness of BTP and related MoB algorithms (Li et al., 15 May 2025). The two-phase greedy radius dynamic allocation ensures linear scalability with respect to the number of tokens and comprehensive adaptability to different pruning objectives. These theoretical results are borne out in extensive empirical investigations across vision-language and segmentation benchmarks.

Moreover, BTP functions as a drop-in module for existing architectures and can operate in both task-agnostic and task-aware modes, depending on the structure of the calibration set and the pruning schedule. Its flexibility and provable performance guarantees under multi-objective criteria position BTP as a foundational paradigm for practical transformer acceleration in both unimodal and multi-modal domains.