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 .