StepPrune: Adaptive Sequential Visual Token Selection across Multimodal Large Language Models
Authors: Hansen Zhang, Landi He, Mingde Yao, Lijian Xu
Organizations: Shenzhen University of Advanced Technology, Shenzhen, China · School of Computer Science, Nanjing University, Nanjing, China · CUHK MMLab, CPII under InnoHK, Hong Kong, China
Visual prefixes account for a major portion of the per-layer computation in multimodal large language models (MLLMs), making visual-token pruning a direct approach to accelerating inference. Existing top-K methods typically evaluate tokens independently and apply a uniform budget to all inputs, overlooking both selection-dependent interactions and variations in visual complexity across samples. In contrast, we propose StepPrune, which formulates visual-token pruning as an adaptive sequential decision process. Conditioned on previously selected tokens and textual context, StepPrune progressively constructs the retained subset and automatically determines its size through a learned STOP action. During training, a variance-preserving noise gate provides a differentiable surrogate for the discrete selection process, whereas during inference, unselected tokens are physically removed before language-model prefill. A grouped selection mechanism further extends StepPrune to high-resolution inputs. Experiments across LLaVA-1.5, LLaVA-NeXT, Qwen2.5-VL, and InternVL3 show that StepPrune achieves the best average normalized performance retention across all evaluated pruning rates on LLaVA-1.5, Qwen2.5-VL, and InternVL3, while remaining competitive on the substantially longer AnyRes prefixes of LLaVA-NeXT. On LLaVA-1.5, StepPrune retains 94.6% of the full-prefix normalized performance while pruning 88.9% of the visual tokens. At a mean retained count of 64, StepPrune reduces prefill latency from 59.95 ms to 40.05 ms, corresponding to a 1.50x prefill speed-up.
Figures & tables
Figure 1: Comparison between one-shot and sequential visual-token selection. One-shot selection retains redundant sign-level tokens but misses the street-name evidence required by the prompt. Sequential selection conditions each decision on previously selected tokens and progressively recovers the complementary visual evidence needed to produce the correct answer
Figure 2: Overview of StepPrune. Visual and text tokens are jointly encoded to produce cross-modal representations. The conditional pointer policy sequentially selects visual tokens according to the previously selected history and terminates with a learned STOP action. During training, the aggregated selection probabilities control a variance-preserving noise gate, while the locally masked denoiser prevents information leakage across visual tokens. The vision encoder and language model remain frozen.
Figure 3: Adaptive token allocation across multimodal benchmarks. (a) Distribution of the adaptively retained token count K on each benchmark, where the value above each distribution denotes the mean retained count. (b) Percentage of samples assigned to different retention intervals. (c) Difference in normalized performance between adaptive selection and fixed-budget selection under matched mean retained counts; positive values indicate an advantage for adaptive selection
Figure 4: Sequential selection trajectories on LLaVA-1.5-7B and LLaVA-NeXT-7B. Columns show the input, cumulative selections at early, middle, and final stages, the final retained subset, and the STOP trajectory. The upper example is predicted correctly, while the lower example is a failure case.
Method
Token (#)
FLOPs (TFLOPs)
E2E Latency (ms)
Prefill (ms)
Decode / Token (ms)
KV Cache (GB)
GPU Memory (GB)
Avg. (%)
Baseline: Full Model
LLaVA-1.5-7B
576
8.772
138.37
59.95
19.59
0.3350
15.296
100.0
Lightweight Methods (Retain 64 Tokens)
PDrop
64
4.690
135.84
49.37
21.46
0.1705
14.764
84.7
VisionZip
64
1.984
140.98
50.13
22.65
0.0666
14.588
92.0
SparseVLM
64
2.253
124.29
41.29
20.67
0.0762
18.919
85.6
Table 5: Efficiency and normalized performance retention on LLaVA-1.5-7B and LLaVA-NeXT-7B under different retained-token budgets. Within each retained-token setting, the best and second-best pruned results for each metric are shown in bold and underlined, respectively. Lower is better for FLOPs, latency, KV cache, and GPU memory, whereas higher is better for average performance retention.
Backbone
Input tokens
320
640
960
1280
1600
1920
2240
LLaVA-1.5-7B
Retained tokens
32
64
96
128
160
192
224
Full prefill
29.34
53.24
70.86
116.89
142.93
189.83
236.55
StepPrune prefill
30.00
37.14
37.75
40.91
46.13
52.28
57.73
Speed-up
0.98×
1.43×
1.88×
2.86×
3.10×
3.63×
4.10×
Selector latency
9.65
16.15
20.45
24.25
28.85
34.41
39.14
Selector share
32.16%
43.49%
54.17%
59.28%
62.54%
65.81%
67.80%
Table 6: Scaling with input token count at 10% retention on a single NVIDIA RTX PRO 6000 Blackwell GPU. Latencies are in milliseconds
Variant
Avg. Ret. (%)
(a) Encoder self-attention direction
Causal (unidirectional)
92.4
Bidirectional †
94.6
(b) Scorer aggregation (see Eq. ( 6 ))
Last-step pointer only
90.3
Uniform sum ( At≡1,β=1 )
83.9
Table 7: Component ablations on LLaVA-1.5-7B. Each row reports average performance retention (%) over the eight benchmarks in Table 1 . Rows marked with † indicate the default configuration
Variant
K≈192
K≈128
K≈64
Hard pruning
97.3
95.8
93.7
Noise-gated inference
97.0
94.9
92.5
Gap
0.3
0.9
1.2
Table 8: Train–inference diagnostic on LLaVA-1.5-7B. Values are average performance retention (%) over GQA, MMB, POPE, and MME
Figure 5: Hyperparameter sensitivity on LLaVA-1.5-7B. (a) Adaptive mean retained count K and matched- K retention under different length-penalty weights λ . (b) Matched- K retention under different geometric decay factors β . Retention is averaged over GQA, MMB, POPE, and MME at matched K≈64 .
In multimodal large language models (MLLMs), inference cost is largely dominated by the visual token prefix rather than the language backbone, making token reduction a key factor for improving efficiency. Existing approaches typically assign independent importance scores to visual tokens and retain a fixed number of top-ranked tokens, implicitly assuming token independence and a uniform compression ratio across inputs. In this work, we reformulate visual token pruning as a sequential decision-making process. Specifically, we introduce a pointer-style selection mechanism that iteratively chooses informative tokens, conditioning each decision on previously selected ones, and dynamically determines when to stop via a learned termination action. This enables joint optimization of both the selected subset and its size. To enable end-to-end training under standard language modeling objectives, we design a differentiable relaxation based on a variance-preserving noise interpolation scheme, allowing gradients to propagate through the discrete selection process. Extensive experiments on LLaVA-v1.5-7B and Qwen2.5-VL-7B demonstrate that our approach consistently outperforms fixed-ratio baselines across different compression levels. Under aggressive pruning that removes 88.9% of visual tokens, our method preserves 94.6% of the original accuracy while achieving a 1.88x speed-up in prefill latency.
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 S2Prune, 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 S2Prune 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.
Yuanyuan Jia, Shunpu Tang, Qianqian Yang
College of Information Science and Electronic Engineering, Zhejiang University Hangzhou, China
Visual-token pruning can substantially reduce the inference cost of multimodal large language models (MLLMs), yet existing methods largely rely on fixed, handcrafted heuristics and costly expert trial and error. As pruning objectives, budgets, and model architectures diversify, manually navigating the expanding design space becomes increasingly difficult. This paper aims to build an AI4AI framework for visual-token pruning by addressing a natural question: Can large language models automatically design effective visual-token reduction algorithms? Although LLMs possess broad algorithmic knowledge and strong reasoning capabilities, translating such general knowledge into effective solutions for a specialized task remains nontrivial. We argue that the key lies in designing an appropriate search-state representation that connects the internal knowledge of LLMs with the structural requirements and constraints of visual-token pruning. Based on this insight, we propose AutoPrune, a training-free framework for LLM-driven visual-token pruning policy design. At its core, AutoPrune introduces a Token Pruning Domain-Specific Language (TPDSL) comprising 131 reusable atoms for budget control, token scoring, selection constraints, and token reassembly. A key property of TPDSL is that it represents each search state as a residual modification of a strong base policy. This residual formulation narrows the search space and directs the LLM's attention toward the policy components that are most consequential for performance. Experiments on 14 multimodal benchmarks and three MLLM backbones demonstrate the effectiveness, efficiency, and transferability of AutoPrune. Even when removing 94.4% of visual tokens, AutoPrune preserves more than 99% of full-token performance while reducing FLOPs by 9.9x and prefill latency by 6.4x.
Zhen Liu, Wenli Huang, Wei Song +3
Xi’an Jiaotong University · Ningbo University of Technology · 3North China University of Technology +3