Stepwise Token Selection for Efficient Multimodal Large Language Models

Abstract: In multimodal LLMs (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.

• The paper introduces StepPrune, a sequential pointer decoder that conditionally selects visual tokens and adapts token budgets based on input complexity.
• It employs a differentiable noise gating mechanism to blend selection weights with Gaussian noise, facilitating end-to-end training without fixed token counts.
• Empirical results show that StepPrune outperforms fixed-token pruning baselines, achieving speed-ups and high accuracy retention even with aggressive pruning.

Stepwise Token Selection for Efficient Multimodal LLMs

Motivation and Core Contribution

Efficient inference in multimodal LLMs (MLLMs) is bottlenecked by the length of the visual token prefix produced by vision towers such as CLIP-ViT, not the parameters of the LLM backbone itself. Existing visual token pruning approaches treat token retention as a rank-and-select operation—visual tokens are scored independently for relevance and a fixed number of top tokens are kept, assuming token independence and a uniform compression ratio across samples. This paper introduces StepPrune, a method reframing visual token selection as a sequential conditional decision process: a pointer-style decoder iteratively selects visual tokens conditioned on prior selections and dynamically determines when to stop via a learned termination action. This permits joint optimization of which tokens are retained and how many, adapting token budgets to input complexity. The pointer policy is made fully differentiable for end-to-end training via variance-preserving noise interpolation.

Sequential Visual Token Selection Strategy

Unlike one-shot scoring—where each token’s contribution is assessed independently—StepPrune’s autoregressive pointer mechanism selects tokens in sequence, each decision conditioned on the already chosen subset. The candidate set is augmented with a learned stop action, allowing the model to terminate token selection adaptively. By doing so, StepPrune abandons the fixed hyperparameter of “tokens to retain” ($K$) and instead produces an adaptive per-input $K$, determined by the interplay of visual content and user instruction.

Figure 1: One-shot scoring (top) versus stepwise sequential selection (bottom). Each row shows the retained tokens and the caption that the LLM produces from the resulting visual prefix.

The model’s architecture couples a cross-modal input encoder with a pointer decoder, illustrated in the following figure:

Figure 2: Architecture of StepPrune: module placement atop LLaVA-style pipeline and internal pointer mechanism, with tokens and text jointly encoded.

Differentiable Training via Noise-Gated Token Relaxation

End-to-end optimization through discrete selection is challenging, requiring either high-variance REINFORCE or surrogate relaxations. StepPrune utilizes a deterministic, variance-preserving noise gating mechanism: soft selection weights are derived from the sequential pointer distribution and a differentiable top-$k$ operator, then used to blend each token with isotropic Gaussian noise. A locally-masked denoiser ensures noise-gating remains effective, preventing tokens from recovering content through neighboring tokens. The resulting formulation aligns well with the downstream language modeling loss, allowing the pointer network to be trained alongside a light length penalty encouraging compression without prescribing token count a priori.

Empirical Results and Efficiency Gains

StepPrune consistently outperforms fixed-ratio pruning baselines (e.g., FastV, VisionZip, SparseVLM) across multiple compression levels on LLaVA-v1.5-7B and Qwen2.5-VL-7B. At aggressive pruning settings (88.9% token removal), StepPrune recovers 94.6% of full-model accuracy and achieves a 1.88$\times$ speed-up in prefill latency. Strongest gains emerge on dense, content-rich benchmarks such as MME, POPE, and TextVQA, validating that stepwise conditional selection is superior where visual cues are interdependent.

Figure 3: Adaptive budget analysis on MME (LLaVA-v1.5-7B): sample count and accuracy plotted against retained token count $K$.

StepPrune’s token count distribution is broad, adapting from sparse to dense inputs. Buckets of small $K$ do not suffer markedly worse accuracy than those with large $K$, indicating effective allocation of compute. Adaptive selection length is content-dependent, imaged in qualitative overlays revealing that the policy focuses on dominant evidence when present, and spreads coverage when the scene is more complex.

Efficiency analysis confirms that the FLOPs and latency costs of StepPrune are lower than fixed-ratio baselines at equivalent compression, and the pointer decoder overhead is negligible compared to the savings from a shortened prefix. Ablations verify the importance of bidirectional encoder context and robust score aggregation for performance retention.

Implications, Limitations, and Future Directions

StepPrune demonstrates that sequential, conditional, and adaptive token selection is not only feasible but empirically superior for efficient MLLMs. The pointer policy transfers across backbones and underlying vision towers, generalizing by construction. The adaptive policy enables per-input compute allocation without external oracle supervision or fixed budget tuning, paving the way for scalable, content-aware deployment of multimodal generative models under real-world applications (e.g., medical imaging, text-rich VQA).

Theoretically, StepPrune challenges the independence assumption of token pruning, highlighting the role of sequential decision policies in compositional model architectures. Practically, its differentiable relaxation mechanism bridges the gap to full end-to-end learning without surrogate gradients or auxiliary losses. Future directions for AI system design may expand on sequential selection mechanisms for other modalities, investigate cross-modal conditioning, or integrate fine-grained visual explainability. Exploration of hybrid selection strategies, runtime policy adaptation, or dynamic vision tower adjustments based on inferred complexity may further improve efficiency and robustness.

Conclusion

StepPrune replaces fixed, independent token scoring in multimodal LLMs with a sequential pointer-based selection mechanism that is conditioned and adaptive, coupled to a differentiable noise-based surrogate for end-to-end optimization. Empirical evidence across diverse benchmarks validates its superior accuracy-cost trade-off and cross-model generalization. The method advances the practical efficiency and adaptive reasoning capacity of MLLMs, suggesting a promising avenue for further research in content-adaptive computation and multimodal inference.