DTP: A Simple yet Effective Distracting Token Pruning Framework for Vision-Language Action Models

TL;DR

Vision-Language-Action models often over-attend to task-irrelevant image tokens, hurting manipulation performance. The authors propose Distracting Token Pruning (DTP), an inference-time, plug-and-play framework that constructs a task-relevant region and an attention-based pruning mask to remove distracting tokens, refining action-token generation without changing model weights. They formalize an upper-bound analysis using conditional uncertainty $E(\alpha) = H(A^\ast \mid Z_{\alpha})$ and a normalized score $P(\alpha) = 1 - \frac{E(\alpha)}{H(A^\ast)}$, showing that sweeping the tolerance $\tau$ yields near-optimal attention patterns on multiple VLAs. Across SpatialVLA, Nora, and UniVLA, DTP achieves consistent task-success improvements and reveals a negative correlation between unimportant attention and success, motivating attention-aware robustness in future embodied AI. The work provides a practical, architecture-agnostic method to push VLA performance toward an empirical upper bound without retraining.

Abstract

Vision-Language Action (VLA) models have shown remarkable progress in robotic manipulation by leveraging the powerful perception abilities of Vision-Language Models (VLMs) to understand environments and directly output actions. However, by default, VLA models may overly attend to image tokens in the task-irrelevant region, which we describe as 'distracting tokens'. This behavior can disturb the model from the generation of the desired action tokens in each step, affecting the success rate of tasks. In this paper, we introduce a simple yet effective plug-and-play Distracting Token Pruning (DTP) framework, which dynamically detects and prunes these distracting image tokens. By correcting the model's visual attention patterns, we aim to improve the task success rate, as well as exploring the performance upper boundaries of the model without altering its original architecture or adding additional inputs. Experiments on the SIMPLER Benchmark (Li et al., 2024) show that our method consistently achieving relative improvements in task success rates across different types of novel VLA models, demonstrating generalizability to transformer-based VLAs. Further analysis reveals a negative correlation between the task success rate and the amount of attentions in the task-irrelevant region for all models tested, highlighting a common phenomenon of VLA models that could guide future research. We also publish our code at: https://anonymous.4open.science/r/CBD3.

Figures (9)

• Figure 1: Overview of our proposed Distracting Token Pruning (DTP) method for improving Vision-Language-Action (VLA) models in robotic manipulation. Left: Input consisting of the current visual observation and natural language instruction. Middle: Comparison between the original VLA model (top) may focus on task-irrelevant regions that lead to task failure, while our DTP-enhanced approach (bottom) creates more focused attention on task-critical areas, leading to the improvements in the task success rate.
• Figure 2: Detailed architecture of our Distracting Token Pruning (DTP) method for improving visual attention in Vision-Language-Action models. The method consists of three main components: (a) Important Region Construction: Using selected transformer layers $C$ to calculate the relevance score between image and prompt tokens, which forms the task-related important region $G$. (b) Visual Attention Pattern Construction: Creating the output to image token attention heatmap $A$ from all attention layers, weighted by the visual attention proportion. It captures where the model focuses when generating actions. (c) Distracting Token Pruning: For any image token in the unimportant region, if its attention value is greater than $\tau\cdot a_m$ , it will be treated as distracting tokens, and will be pruned.
• Figure 3: Visualization of the distracting token pruning process across different VLA models and tasks. The figure shows the Relevance Heatmap (first row), Important Region (second row), Visual Attention Patterns (third row), and Final Pruning Mask (last row) for three VLA models (SpatialVLA spatialvla, Nora nora, UniVLA UniVLA) on two representative tasks ('Put The Spoon...' and 'Stack The Block...'). The comparison reveals how our method identifies and prunes distracting visual tokens across different model architectures. Please refer to Appendix \ref{['more_visualizations']} for more visualization cases.
• Figure 4: Exploration of performance upper bounds under varying tolerance parameter $\tau$ across different VLA models. Dashed lines denote the baseline success rates without DTP, while solid lines illustrate performance under DTP with different $\tau$ values. The success rate is annotated at peak to indicate $\hat{\tau}$. The results confirm that tuning $\tau$ enables the selection of optimal tolerance levels, thereby maximizing potential performance gains.
• Figure 5: Unimportant attention in VLA models. (Top-left) Distribution of unimportant attention values between success and failure episodes across three models. (Top-right, Bottom-left, Bottom-right) Temporal evolution of unimportant attention for each model, with episode steps normalized to $[0,1]$ to provide a unified view of the dynamic changes.