Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

KV-Efficient VLA: A Method to Speed up Vision Language Models with RNN-Gated Chunked KV Cache

Wanshun Xu, Long Zhuang, Lianlei Shan

TL;DR

This work tackles the real-time inference bottlenecks of Vision-Language-Action models caused by growing KV caches and long-context attention. It introduces KV-Efficient, a model-agnostic memory compression mechanism that partitions the KV cache into fixed-size chunks, aggregates each chunk with an MLP, and uses an LSTM-based recurrent gate to selectively retain high-utility history while preserving recent detail. The approach preserves a fixed recent window and compresses older context, resulting in substantial efficiency gains: approximately $24.6\%$ reduction in FLOPs, $1.34\times$ faster inference, and $1.87\times$ memory savings for KV caches, with seamless integration into existing VLA stacks. The method is validated theoretically via a cost-model and empirically on diverse VLA backbones and datasets (OpenX-Embodiment, RLBench), demonstrating scalable inference for long-horizon robotic tasks while maintaining flexibility for future real-world deployment.

Abstract

Vision-Language-Action (VLA) models offer a unified framework for robotic perception and control, but their ability to scale to real-world, long-horizon tasks is limited by the high computational cost of attention and the large memory required for storing key-value (KV) pairs during inference, particularly when retaining historical image tokens as context. Recent methods have focused on scaling backbone architectures to improve generalization, with less emphasis on addressing inference inefficiencies essential for real-time use. In this work, we present KV-Efficient VLA, a model-agnostic memory compression approach designed to address these limitations by introducing a lightweight mechanism to selectively retain high-utility context. Our method partitions the KV cache into fixed-size chunks and employs a recurrent gating module to summarize and filter the historical context according to learned utility scores. This design aims to preserve recent fine-grained detail while aggressively pruning stale, low-relevance memory. Based on experiments, our approach can yield an average of 24.6% FLOPs savings, 1.34x inference speedup, and 1.87x reduction in KV memory. Our method integrates seamlessly into recent VLA stacks, enabling scalable inference without modifying downstream control logic.

KV-Efficient VLA: A Method to Speed up Vision Language Models with RNN-Gated Chunked KV Cache

TL;DR

This work tackles the real-time inference bottlenecks of Vision-Language-Action models caused by growing KV caches and long-context attention. It introduces KV-Efficient, a model-agnostic memory compression mechanism that partitions the KV cache into fixed-size chunks, aggregates each chunk with an MLP, and uses an LSTM-based recurrent gate to selectively retain high-utility history while preserving recent detail. The approach preserves a fixed recent window and compresses older context, resulting in substantial efficiency gains: approximately reduction in FLOPs, faster inference, and memory savings for KV caches, with seamless integration into existing VLA stacks. The method is validated theoretically via a cost-model and empirically on diverse VLA backbones and datasets (OpenX-Embodiment, RLBench), demonstrating scalable inference for long-horizon robotic tasks while maintaining flexibility for future real-world deployment.

Abstract

Vision-Language-Action (VLA) models offer a unified framework for robotic perception and control, but their ability to scale to real-world, long-horizon tasks is limited by the high computational cost of attention and the large memory required for storing key-value (KV) pairs during inference, particularly when retaining historical image tokens as context. Recent methods have focused on scaling backbone architectures to improve generalization, with less emphasis on addressing inference inefficiencies essential for real-time use. In this work, we present KV-Efficient VLA, a model-agnostic memory compression approach designed to address these limitations by introducing a lightweight mechanism to selectively retain high-utility context. Our method partitions the KV cache into fixed-size chunks and employs a recurrent gating module to summarize and filter the historical context according to learned utility scores. This design aims to preserve recent fine-grained detail while aggressively pruning stale, low-relevance memory. Based on experiments, our approach can yield an average of 24.6% FLOPs savings, 1.34x inference speedup, and 1.87x reduction in KV memory. Our method integrates seamlessly into recent VLA stacks, enabling scalable inference without modifying downstream control logic.

Paper Structure

This paper contains 10 sections, 11 equations, 4 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Works
  3. Methods
  4. Overall Framework
  5. KV Cache Chunking with Recurrent Gating
  6. Computation Cost Analysis
  7. Experiments
  8. Theoretical Results
  9. Experiment:
  10. Summary

Figures (4)

  • Figure 1: VLA is pretrained on diverse robotic datasets, incorporating the KV Efficient Module. The input data consists of prompts, images and robotic state, while the output comprises robotic action vectors.
  • Figure 2: KV-Efficient Framework. The framework consists of two main steps: (1) Chunking the KV cache into fixed-size segments and aggregating them, and (2) Using an RNN module to update and select which chunks to retain or discard, enabling efficient and selective compression for attention mechanisms.
  • Figure 3: Training Loss for KV-Efficient module
  • Figure :