Learning to Accelerate Vision-Language-Action Models through Adaptive Visual Token Caching

TL;DR

Vision-Language-Action models offer strong multimodal capabilities for embodied robotics but incur high inference cost from reprocessing full visual inputs each frame. This work proposes Learnable Adaptive Caching (LAC), a framework that couples a Cached Token Selector with a Cache Ratio Predictor to learn task-aware token reuse on top of a frozen VLA backbone, guided by motion signals via optical flow and trained with differentiable discrete decision learning. Through a two-stage training procedure—attention-alignment initialization followed by joint optimization with Gumbel-Softmax—LAC achieves substantial speedups (up to 1.76x) while improving task success on LIBERO and SIMPLER benchmarks and on real robots. The results demonstrate that task-driven adaptive computation can dramatically enhance efficiency and robustness of VLA systems, enabling real-world deployment of powerful visuomotor policies.

Abstract

Vision-Language-Action (VLA) models have demonstrated remarkable generalization capabilities in robotic manipulation tasks, yet their substantial computational overhead remains a critical obstacle to real-world deployment. Improving inference efficiency is therefore essential for practical robotic applications. Existing acceleration methods often rely on heuristic or static strategies--such as rule-based token caching or pruning--that are decoupled from task objectives and fail to adapt to dynamic scene changes. In this work, we reformulate inference acceleration as a learnable policy optimization problem and propose a novel framework that integrates a dynamic, task-aware decision-making process directly into the VLA model. At its core are two lightweight, cooperative modules: a Cached Token Selector, which determines which tokens should be reused, and a Cache Ratio Predictor, which controls how many tokens to reuse. Training these modules is non-trivial due to their discrete decisions. We address this by adopting a differentiable relaxation that allows gradient-based end-to-end optimization. Extensive experiments on the LIBERO and SIMPLER benchmarks, as well as real-robot evaluations, show that our method achieves a 1.76x wall-clock inference speedup while simultaneously improving the average success rate by 1.9 percentage points (from 75.0% to 76.9%) on LIBERO and by 5.0 percentage points on real-world tasks, significantly outperforming existing baselines. This work highlights the potential of learning task-aware computational allocation policies, paving the way for VLA models that are both powerful and efficient.

Figures (8)

• Figure 1: Our learnable caching (LAC) outperforms rule-based baselines in both efficiency and performance. (a) Rule-based selectors are decoupled from the task objective. They often rely on proxies, but where the model attends is not necessarily what the task requires. This leads to suboptimal, task-agnostic caching (20% cache ratio). (b) Our selector is optimized with direct task gradients ($\nabla\mathcal{L}_{\text{VLA}}$), ensuring its decisions align with actual task requirements. This enables a more aggressive adaptive caching strategy, leading to a much higher cache ratio (40%) and greater efficiency, while simultaneously improving the success rate.
• Figure 2: Overview of the two-stage training framework for LAC.Stage I: Initialization. The Cached Token Selector is initially pretrained to align with the VLA's internal attention map, providing a stable and informative initialization. Stage II: Joint Optimization. Following this, the Cached Token Selector and Cache Ratio Predictor are jointly optimized with the VLA backbone held frozen. The selector learns token-level saliency for dynamic reuse, while the predictor estimates the optimal cache ratio from a discrete set. Gradients from the task loss $\mathcal{L}_{\text{VLA}}$ are backpropagated through both modules via the differentiable selection mechanism, enabling end-to-end policy learning for adaptive computation.
• Figure 3: Qualitative visualization of learned vs. rule-based caching. LAC's policy (bottom) uses a learned saliency map to focus computation on the moving gripper while caching the static background (orange). In contrast, the rule-based method (top) detrimentally caches the static target basket, a task-agnostic error that could lead to task failure. This demonstrates the superiority of LAC's learned approach, which is further stabilized by recovered tokens (green).
• Figure 4: Effect of Token Reusing/Pruning Ratio on LIBERO-Spatial. Comparison of different methods under varying token reusing (or pruning) ratios. Left: Task success rate. Right: CUDA inference time. Our LAC maintains consistently high performance as the reuse ratio increases, while pruning-based methods such as SparseVLM degrade rapidly. In addition, LAC achieves the lowest CUDA latency across all settings, confirming its genuine wall-clock acceleration.
• Figure 5: Qualitative visualization of our learned policy (LAC) on a real-world manipulation task. The Token Saliency map (middle row) demonstrates that the model correctly identifies the moving end-effector and its immediate interaction space as the most critical regions. Consequently, the Cache Mask (bottom row) shows an efficient policy where static background elements are cached, focusing computation only on task-relevant dynamics. Visualizations for two other tasks can be found in the Appendix.