Dtp Framework Achieves Higher Vision-Language Action Success Rates by Pruning Tokens

Vision-Language Action (VLA) models are increasingly applied to robotic manipulation but often underperform due to unnecessary attention to irrelevant visual details, which researchers Chenyang Li, Jieyuan Liu (University of California, San Diego), Bin Li, and colleagues refer to as “distracting tokens.” To address this challenge, the team introduces Distracting Token Pruning (DTP), a framework that dynamically identifies and removes these disruptive image tokens, effectively refocusing the model’s attention and improving task success. This simple, plug-and-play approach consistently enhances performance on the SIMPLER Benchmark, revealing a common flaw in VLA models and providing a pathway to unlock their full potential without modifying model architecture or adding additional inputs.

The study shows that VLA models frequently attend to task-irrelevant regions, which disrupts the generation of desired action tokens and reduces task success. The DTP framework operates in three main stages: (1) constructing important regions based on prompt-image interactions, (2) analyzing action attention heatmaps to identify areas of focus, and (3) selectively pruning tokens in unimportant regions using an intersection-based strategy with a tolerance parameter τ. Experiments on the SIMPLER Benchmark demonstrate that DTP consistently improves task success rates across various transformer-based VLA architectures. Further analysis establishes a negative correlation between task success and attention allocated to task-irrelevant regions, highlighting a pervasive weakness in existing models.

The DTP process begins with relevance-based construction of important regions, leveraging both the user prompt and the visual observation. Attention scores are then calculated for each image token, quantifying the model’s focus. Tokens with attention weights below a threshold defined by τ are dynamically pruned before action token generation, ensuring that the model prioritizes task-relevant visual information. This approach refines visual attention patterns without modifying the underlying VLA architecture or introducing additional inputs.

Experiments were conducted on the SIMPLER Benchmark using multiple VLA models, including SpatialVLA, Nora, and UniVLA. SpatialVLA’s task success increased from 37.5% without DTP to 68.7% with DTP. Nora improved from 29.2% to 74.0%, and UniVLA rose from 6.2% to 68.7%, demonstrating consistent and substantial gains across all tested models. Analysis revealed a strong negative correlation between attention to irrelevant regions and task success, confirming the effectiveness of DTP in mitigating this issue. Extended evaluations on WidowX and Google Robot tasks further validated the generalizability of DTP, with relative improvements observed across different robots and architectures. On the LIBERO Benchmark, Nora achieved a +6.6% absolute gain on the challenging LIBERO-10 suite, along with 1.4–2.6% improvements on other benchmarks.

The researchers also investigated the effect of varying the tolerance parameter τ. Smaller τ values prune more tokens, while larger values converge toward the original model’s performance. Notably, UniVLA required a larger τ, suggesting it initially allocates higher attention to distracting tokens. This analysis indicates that refining visual attention patterns can enhance task success without altering model architecture or adding inputs. Ablation studies confirmed the superiority of targeted pruning compared to random or simplified methods, emphasizing that reducing attention to irrelevant areas—what the authors term “attention leakage”—is key to improving VLA reliability and efficiency.

Overall, DTP establishes a novel and effective method for correcting visual attention patterns in VLA models, boosting task success rates, and enabling more reliable and efficient robotic manipulation. The framework’s plug-and-play design allows easy adaptation to existing VLA systems, and the authors provide publicly available code to facilitate further research in embodied AI and robotic manipulation. Future work may explore adaptive methods for automatically selecting the optimal τ value, further enhancing the framework’s usability and performance. These findings underscore the importance of controlling visual attention to achieve more generalizable and robust robotic manipulation capabilities.