OTPrune: Distribution-Aligned Visual Token Pruning via Optimal Transport

Xiwen Chen; Wenhui Zhu; Gen Li; Xuanzhao Dong; Yujian Xiong; Hao Wang; Peijie Qiu; Qingquan Song

cs.CVFeb 22, 2026

OTPrune: Distribution-Aligned Visual Token Pruning via Optimal Transport

Authors: Xiwen Chen, Wenhui Zhu, Gen Li, Xuanzhao Dong, Yujian Xiong, Hao Wang, Peijie Qiu, Qingquan Song,

Organizations: Morgan Stanley · Clemson University · Arizona State University · Washington University in St. Louis ·

Abstract

Multi-modal large language models (MLLMs) achieve strong visual-language reasoning but suffer from high inference cost due to redundant visual tokens. Recent work explores visual token pruning to accelerate inference, while existing pruning methods overlook the underlying distributional structure of visual representations. We propose OTPrune, a training-free framework that formulates pruning as distribution alignment via optimal transport (OT). By minimizing the 2-Wasserstein distance between the full and pruned token distributions, OTPrune preserves both local diversity and global representativeness while reducing inference cost. Moreover, we derive a tractable submodular objective that enables efficient optimization, and theoretically prove its monotonicity and submodularity, providing a principled foundation for stable and efficient pruning. We further provide a comprehensive analysis that explains how distributional alignment contributes to stable and semantically faithful pruning. Comprehensive experiments on wider benchmarks demonstrate that OTPrune achieves superior performance-efficiency tradeoffs compared to state-of-the-art methods. The code is available at https://github.com/xiwenc1/OTPrune.

Figures & tables

Figure 1: Overview of the proposed OTPrune framework. Given visual and textual inputs, a vision encoder with a projector produces image tokens that are fed into a large language model together with text tokens and reasoning prompts. Existing methods, such as DivPrune , select tokens by maximizing diversity, which may overlook global representativeness. In contrast, OTPrune formulates pruning as distribution alignment by minimizing the 2-Wasserstein (optimal transport) distance between the full and pruned token distributions, thereby preserving both local diversity and global structure for efficient and semantically faithful multimodal reasoning.

Figure 2: Correlation between OT distance and downstream performance. We evaluate several manually designed selection strategies ( First-K , Last-K , Uniform , and Random ) along with the SOTA method DivPrune [ 3 ] . For each strategy, we compute the OT distance between the selected subset and the full token set and measure downstream performance across 11 multimodal benchmarks. Both OT distance (lower distance

\Rightarrow

lower rank) and task performance (higher score

\Rightarrow

lower rank) are ranked to compute Spearman’s correlation. A strong positive correlation confirms that smaller OT distance corresponds to higher performance.

	Method	TFLOP	COCO	Flickr	GQA	MMB	MME	MMMU	Nocaps	OKVQA	POPE	SQA	SEEDB	Avg.	One-shot
	Method	(ratio %)	CIDEr	CIDEr	EM	Acc	P-score	Acc	CIDEr	EM	F1	EM	Acc	Rank $\downarrow$	Rank $\downarrow$
LLaVA 1.5-7B	Original	3.228 (100.00)	1.10	0.75	61.96	64.09	1506	36.44	1.06	53.39	85.84	69.41	66.17	–	–
	Category: One-shot Pruning Methods
	VTW [ 30 ]	0.603 (18.46)	0.05	0.03	38.94	21.31	681	32.60	0.03	18.64	25.35	65.29	36.13	8.23	3.41
	FastV [ 8 ]	0.514 (15.69)	0.06	0.03	38.73	20.62	696	32.00	0.04	18.32	32.84	65.15	35.69	8.41	3.59
	DivPrune	0.512 (15.63)	0.96	0.62	56.85	59.19	1328	35.89	0.92	46.98	86.02	68.27	59.47	3.23	1.64

Table 1: Comparison of OTPrune with diverse baselines across 11 multimodal benchmarks. DivPrune ∗ aligns its token retention rate with PruMerge. For fair evaluation, all one-shot methods strictly retain 9.8% of visual tokens, using TFLOPs to indicate computational efficiency. The Avg. Rank

\downarrow

assesses all evaluated methods, while the One-shot Rank

\downarrow

strictly compares one-shot techniques. Our OTPrune consistently delivers the best performance. (

{\bullet}

: requires finetuning;

{\bigtriangleup}

: requires calibration).

Configuration				Performance
$m$	$d$	$k$	Method	Win Rate	Optimality Ratio
20	10	5	DivPrune	60.14	74.53
			OTPrune	94.56	87.81
30	20	10	DivPrune	84.62	76.48
			OTPrune	99.49	93.98
100	50	30	DivPrune	88.77	92.20

Table 2: Comparison between OTPrune and DivPrune on synthetic datasets with varying numbers of samples (

m

), feature dimensions (

d

), and selected subset sizes (

k

). The first two configurations are evaluated exhaustively over all

\binom{m}{k}

combinations, while the last three configurations use Monte Carlo estimation with 10 million different subsets to approximate the full combinatorial space.

Comparison	p -value
OTPrune vs DivPrune	0.028

Table 5

Figure 3: Comparison with diversity-based methods using LLaVA 1.5-13B. We evaluate OTPrune , DivPrune [ 3 ] , DPP [ 22 ] , and Random sampling across 11 multimodal benchmarks under different pruning ratios $\{0.05,\,0.098,\,0.15,\,0.2\}$ . We report both absolute performance and relative performance , defined as the ratio between pruned and original model performance ( pruned/original ). We also compute the OT distance between each subset and the full token set and visualize the relative distance normalized by OTPrune, i.e., $f(\mathcal{C}_{\text{method}})/f(\mathcal{C}_{\text{OTPrune}})$ .

Figure 4: Relative OT distance comparison across pruning methods. We evaluate OTPrune , DivPrune [ 3 ] , DPP [ 22 ] , and Random sampling using LLaVA 1.5-13B across 11 multimodal benchmarks under different token ratios $\{0.05,\,0.098,\,0.15,\,0.2\}$ . Each bar shows the relative OT distance normalized by OTPrune, i.e., $f(\mathcal{C}_{\text{method}})/f(\mathcal{C}_{\text{OTPrune}})$ , where lower values indicate closer alignment with the original token distribution. OTPrune consistently achieves the smallest OT distance across all ratios, demonstrating superior distributional fidelity. The last panel reports the overall average across all 11 datasets.

Figure 5: Sensitivity analysis of the balancing coefficient $\gamma$ .

Visual Token Pruning

Visual token pruning reduces the image and video patches processed by transformers to lower inference cost while preserving task-relevant information.

See all 269 papers

S $^2$ Prune: Spatially Structured Visual Token Pruning for Multimodal Large Language Models
Sep 1, 2026Yuanyuan Jia, Shunpu Tang, Qianqian YangEfficient VLM Inference Multimodal Large Language Models
College of Information Science and Electronic Engineering, Zhejiang University Hangzhou, China
SinkPruner: Sink-Free Visual Token Pruning for Multimodal Large Language Models
Sep 1, 2026Shiyu Li, Zi-Yuan Hu, Shijia Huang +3Efficient VLM Inference Multimodal Large Language Models
1The Chinese University of Hong Kong · 2Weitu AI · 3Peking University
ToDRE: Effective Visual Token Pruning via Token Diversity and Task Relevance
May 24, 2025Duo Li, Zuhao Yang, Xiaoqin Zhang +2Efficient VLM Inference Large Vision-Language Models
CCDS, NTU, Singapore · CCST, ZJUT, China · Terminus AI Lab, UCAS, China

Sep 1, 2026J/K move · Enter open · S save

S $^2$ Prune: Spatially Structured Visual Token Pruning for Multimodal Large Language Models

Yuanyuan Jia, Shunpu Tang, Qianqian Yang

College of Information Science and Electronic Engineering, Zhejiang University Hangzhou, China

Efficient VLM Inference Multimodal Large Language Models Visual Token Pruning LVLM Pruning

Open paper AskPDF arXiv

Visual token pruning reduces the inference overhead of multimodal large language models (MLLMs) by retaining only a subset of visual tokens. Existing methods usually select tokens based on importance or redundancy. However, we observe that these criteria produce stable spatial biases across inputs and do not always outperform simple Uniform Grid sampling, highlighting the value of broad spatial coverage. Motivated by this, we propose S

^2

Prune, a training-free pruning method that preserves spatial coverage while adapting token density to local image structure. We first divide the image into regions and assign at least one token to each region to preserve coverage. The remaining token budget is then distributed according to Laplacian variation, giving more tokens to regions with richer structure. We then use Early Representation Change (ERC), computed from the first decoder block, to select representative tokens within each region. We evaluate S

^2

Prune across diverse settings and two MLLM architectures. On Qwen2.5-VL-7B-Instruct, it achieves the highest average accuracy among the evaluated training-free pruning methods. With only 32 of the original 576 visual tokens, it still retains 79.3% of the full-model performance. Code is available at https://github.com/yuanyuanjia71-spec/S2Prune.

Abstract

Visual Token Pruning

Explore similar work

S2^22Prune: Spatially Structured Visual Token Pruning for Multimodal Large Language Models

SinkPruner: Sink-Free Visual Token Pruning for Multimodal Large Language Models

ToDRE: Effective Visual Token Pruning via Token Diversity and Task Relevance

S $^2$ Prune: Spatially Structured Visual Token Pruning for Multimodal Large Language Models