Driving · 2026-10-09

WAM-Cache: Staleness-Bounded KV Reuse for Efficient World Action Models

Kai Ding, Yang He, Ruijie Quan, Yi Yang

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

WAM-Cache is a training-free KV-reuse method for World Action Models that keeps video DiT layerwise key–value pairs across inference chunks and recomputes only a sparse refresh set selected by (i) latent surprise, (ii) action-conditioned attention, and (iii) an age bound. On Fast-WAM, it cuts video DiT prefill FLOPs by 32–42% while staying close to dense policy success.

Why it matters

In closed-loop control, WAMs must run the video DiT every chunk, and the prefill dominates compute. The paper argues that simply refreshing “what changed” is not sufficient because stale KV pairs affect accuracy based on where the action expert attends. WAM-Cache provides a staleness-bounded, training-free way to reduce KV prefill cost while controlling accuracy degradation.

Method

Maintain cached layerwise KV pairs across chunks and refresh only a token set ℛ_t each chunk; unrefreshed tokens keep stored KV while refreshed tokens are recomputed in-place.
Select ℛ_t as a union of latent-surprise TopK, action-attention TopK from the preceding chunk, and an age-bound set A_t that forces refresh when reused for G consecutive chunks.
No retraining; the action expert queries fresh KV for ℛ_t and cached KV for unselected tokens, composing with other orthogonal speedups (e.g., step reduction, distillation, denoiser caching).

Limitation

The paper does not report any explicit limitation section in the provided text.

Abstract (from arXiv)

World Action Models (WAMs) enable generalist robot manipulation by conditioning an action expert on representations from a pretrained video Diffusion Transformer (DiT). In closed-loop control, the video DiT runs at every chunk to encode the current observation into layerwise key-value (KV) pairs that the action expert queries. This prefill dominates the per-chunk computational cost, yet existing training-free accelerations leave it fully dense. We present WAM-Cache, a training-free framework that retains layerwise key-value representations across chunks and recomputes only a sparse refresh set of tokens. Crucially, we find that the intuitive heuristic of refreshing visually drifted tokens plateaus far below the dense baseline, even with an oracle predicting ground-truth KV drift. Downstream action accuracy is instead governed by where the action expert attends, not by what moved. WAM-Cache therefore selects the refresh set by uniting the action expert's cross-attention with visual latent surprise, complemented by a strict age bound that suppresses compounding error. On Fast-WAM, WAM-Cache cuts video DiT prefill FLOPs by 32-42% across RoboTwin 2.0, LIBERO, and real-world experiments, while staying within 0.7-1.8 percentage points of the dense policy in simulation and 2.5 points on a real robot.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-08 · Indexed by Pub-AI: 2026-10-09 UTC

preprint