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

TL;DR

This paper introduces a learnable, task-aware token caching framework that significantly accelerates vision-language-action models in robotic tasks while improving success rates, addressing the inefficiency of prior static methods.

Contribution

It proposes a novel adaptive token caching method with differentiable training, enabling dynamic inference acceleration tailored to task demands in VLA models.

Findings

Achieves 1.76x inference speedup on benchmarks.

Improves success rate by up to 5 percentage points.

Outperforms existing static caching baselines.

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.