PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kai Ding, Yang He, Ruijie Quan, Yi Yang

arXiv:2610.11401v1cs.ROcs.CVcs.LG

TL;DR

Closed-loop World Action Models repeatedly prefill a video DiT, a major computational cost that existing training-free accelerations leave dense. WAM-Cache selectively reuses layerwise keys and values, combining action attention, visual latent surprise, and an age bound. It reduces prefill computation across simulation and real-world evaluations while staying close to dense-policy performance.

  • Problem

    Closed-loop WAMs repeatedly run the video-DiT prefill, while existing training-free accelerations leave this major per-chunk cost dense.

    Source: §1 · p. 1

  • Method

    WAM-Cache retains layerwise key–value representations across chunks, refreshing tokens selected by action-expert attention and latent surprise, with an age bound on reuse.

    Source: §03 · pp. 2–3

  • Results

    WAM-Cache reduces prefill FLOPs across RoboTwin 2.0, LIBERO, and real-world manipulation while staying close to dense-policy performance.

    Source: §5.2 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    Refresh decisions based on action-expert attention address downstream sensitivity that visual drift alone does not capture, while the age bound limits staleness.

    Source: §03 · pp. 2–3, §5.3 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

World Action Models pair a video-generation DiT with an action expert, grounding manipulation in physical dynamics learned from video. In closed-loop control, repeatedly prefilling the video DiT makes computation a central bottleneck.

1 INTRODUCTION

  • A WAM pairs a video-generation DiT with an action expert that conditions on its representations, grounding control in dynamics learned from video.The video DiT runs at every closed-loop control cycle as the robot observes and replans action chunks.Source: §1 · p. 1
  • The video DiT prefill encodes observations into layerwise key–value pairs queried by the action expert, and dominates per-chunk cost.Existing training-free accelerations primarily cache intermediate denoiser representations, leaving this forward-pass bottleneck untouched.Source: §1 · p. 1

03 Age-bounded refresh ⋯

WAM-Cache reuses layerwise key–value representations across inference chunks, refreshing tokens using action attention and visual latent surprise with an age bound. It reduces prefill computation across simulation and real-world settings while staying close to the dense policy.

03 Age-bounded refresh ⋯

  • WAM-Cache retains layerwise key–value representations across chunks and recomputes only a sparse set of visual tokens at each chunk.The video DiT prefills that refresh set; the action expert then uses fresh and retained keys and values together.Source: §03 · p. 2, Figure 1
  • Action-conditioned selection refreshes tokens heavily attended to by the action expert, while latent surprise tracks feature drift since a token’s last refresh.The two content signals select largely disjoint subsets and are combined by union.Source: §03 · pp. 2–3
  • An age bound forces refresh after a fixed number of consecutive reuses, preventing rank-based content signals from leaving a token stale indefinitely.WAM-Cache requires no retraining and can compose with diffusion-step reduction, model distillation, and denoiser caching.Source: §03 · p. 3
  • 40.2% fewer prefill FLOPs on real-world tasks accompanies success comparable to the dense model, 82.5% against 85.0%.Across RoboTwin 2.0, LIBERO, and real-world manipulation, reported prefill reductions range from 32% to 42%.Source: §03 · p. 3

2 RELATED WORK

Prior work applies video models to robotics and accelerates policy inference through retraining, caching, or token reuse. WAM-Cache addresses the video prefill in WAMs, where visual prefill and action decoding use decoupled modules.

World Action Models

  • Video models have supported robotics planning, pretraining, data generation, and policy backbones; recent WAMs pair video DiTs with action decoders.Fast-WAM and GigaWorld-Policy show that future imagination can be dropped with little loss in task success, leaving a single video DiT prefill.Source: §2 · p. 3

Efficient Policy Inference and Token Reuse

  • Prior acceleration methods distill or compress models, reduce denoising steps, cache denoiser features, or reuse tokens by change or attention.Unlike methods that select reuse inside a unified backbone, WAM-Cache uses downstream policy attention and strictly caps token staleness across decoupled visual-prefill and action-decoding modules.Source: §2 · p. 3

3 PRELIMINARIES

A World Action Model processes an observation and language instruction with a video DiT and an action-expert DiT. The video DiT produces layerwise key–value pairs that the action expert queries to predict action chunks.

World Action Models

  • A WAM maps an observation and text instruction to an action chunk, with the robot executing a prefix before observing again.Each observe-and-act cycle is an inference chunk; the video DiT processes the observation’s visual tokens in one forward pass.Source: §3 · p. 4
  • At each of its denoising steps, the action expert cross-attends to the video prefill’s keys and values to predict the action chunk.A VAE and patch-embedding projection convert the observation into visual tokens for the video DiT.Source: §3 · p. 4

Key–Value Caching in DiTs

  • A dense video-DiT prefill produces one key–value pair per visual token at every layer; caching stores these pairs for later reuse.Each video-DiT self-attention query attends over all visual-token keys and values.Source: §3 · p. 4

4 METHOD

WAM-Cache retains layerwise key–value representations across inference chunks and recomputes a sparse token set. Its refresh rule combines latent surprise, action attention, and age-bounded renewal.

4.1 OVERVIEW: STALENESS-BOUNDED KV REUSE

  • After dense initialization, each chunk refreshes the union of surprise-, attention-, and age-selected tokens, overwriting their cached key–value pairs while retaining the others.Refreshed-token queries attend over all visual keys and values; no token is reused for more than G consecutive chunks.Source: §4.1 · p. 4

4.2 LATENT SURPRISE

  • 53.1% success for surprise alone on clean RoboTwin 2.0 is 37.0 points below the dense model, without the age bound.The latent-displacement ranking can assign low priority to regions that change little but matter for contact or alignment.Source: §4.2 · p. 5, Figure 2 · p. 5

4.3 ACTION-CONDITIONED SELECTION

  • 80.7% success for surprise plus action attention is 27.6 points above surprise alone but 9.5 points below the dense model.Action attention adds visually stable, policy-critical regions by selecting tokens from preceding-chunk attention, available before prefill at no extra cost.Source: §4.3 · pp. 5–6

4.4 AGE-BOUNDED REFRESH

  • The age bound leaves a 1.8-point gap to the dense model while removing 41.6% of prefill FLOPs.It forces refreshes for tokens that neither content score selected for G chunks, while limiting staleness rather than the size of accumulated error.Source: Figure 2 · p. 5, §4.4 · p. 6

5 EXPERIMENTS

WAM-Cache is evaluated on simulation benchmarks and real-robot tasks, with experiments testing its compute savings, refresh rules, and hyperparameter choices. It stays close to dense-prefill success while reducing video DiT prefill cost.

5.1 IMPLEMENTATION DETAILS

  • The evaluation applies WAM-Cache to Fast-WAM on RoboTwin 2.0 and LIBERO in simulation and two AIRBOT Play tasks; default action denoising uses N = 2.All models run in bfloat16 on one NVIDIA RTX PRO 6000 Blackwell GPU; the two-step dense model stays within 1.5 points of ten steps on RoboTwin 2.0, while one step loses more than 20.Source: §5 · p. 6, §5.1 · p. 6, §5.2 · pp. 6–7

5.2 SIMULATION EXPERIMENTS

  • 41.6% and 39.3% fewer prefill FLOPs on clean and randomized RoboTwin 2.0 come with 1.8- and 1.3-point success gaps; LIBERO uses 32.3% fewer FLOPs at 95.85% versus 96.50%.At a larger RoboTwin budget, success comes within 0.1 points of dense prefill in both conditions; LIBERO also matches the dense average at a 26.1% reduction.Source: Table 1 · p. 7, §5.2 · p. 7, Table 2 · p. 8
  • At visual contexts of 120, 480, and 1080 tokens, clean-observation prefill speedups are 1.23×, 1.47×, and 1.67×, respectively.At 1080 tokens, this reaches 98% of the 1.71× bound; the randomized configuration reaches 1.58×, or 96% of its bound.Source: Figure 3 · p. 7, §5.2 · p. 8

5.3 ABLATION STUDY

  • 60.4–65.3% success across pixel-difference, latent-surprise, and K/V-drift rules leaves even the drift oracle 25 points below dense prefill.The rules were compared under the same content budget, isolating which tokens they refresh.Source: §5.3 · p. 8
  • Action attention alone reaches 81.5% success; adding 0.1 surprise budget raises it to 88.4% because the signals select largely disjoint tokens.Their union is the cheapest informed rule, with ρ̄ = 0.584 versus 0.643 and 0.615.Source: §5.3 · p. 8
  • Without an age bound, success remains 9.5 points below dense prefill; setting G = 2 recovers 7.7 points and leaves a 1.8-point gap.The bound prevents low-ranked tokens from being skipped indefinitely while their cached error compounds.Source: §5.3 · p. 8
  • Moving from G = 1 to G = 2 saves 10.2 points of FLOPs for 0.6 points of success; success rises steeply with attention budget b_a to 0.4, then saturates.At b_a = 0.6, success is within 0.1 points of dense prefill while removing 28.2% of prefill FLOPs; G = 2 is used across benchmarks and on the real robot without retuning.Source: §5.3 · pp. 8–9, Figure 4 · p. 9

5.4 REAL-WORLD EXPERIMENTS

  • 40.2% fewer prefill FLOPs on AIRBOT Play matches dense success on carrot placement and loses one trial on cube stacking.The simulation budgets were used without retuning; each task had 20 evaluation trials.Source: Table 4 · p. 9, §5.4 · p. 9

6 CONCLUSION

WAM-Cache reuses video DiT key–value representations with sparse refreshes, combining policy attention, visual surprise, and an age bound. It reduces prefill FLOPs across simulation and real-world manipulation while remaining close to the dense policy.

6 CONCLUSION

  • 32–42% fewer video DiT prefill FLOPs kept WAM-Cache within 0.7–1.8 points of the dense policy in simulation and 2.5 points on a real robot.The conclusion reports CUDA-latency speedups that grow as visual context scales.Source: §6 · p. 9
  • Visual change or representation drift alone is insufficient; WAM-Cache combines action-expert cross-attention with visual latent surprise and bounds token age.The conclusion says downstream action accuracy is governed by where the policy attends rather than what moved.Source: §6 · p. 9

A WAM-CACHE ALGORITHM

  • Each chunk selects surprise, attention, and age-triggered tokens, unions them, recomputes their video DiT representations, and updates their ages.The action expert then runs over the prefilled key–value pairs.Source: A WAM-CACHE ALGORITHM · p. 16

B COST ANALYSIS

The cost analysis relates prefill computation to the refreshed-token ratio and accounts for selection overhead and age-bound refreshes. The default measured steady-state ratio is below the stated bounds.

B COST ANALYSIS

  • The prefill FLOP accounting includes self-attention projections and attention products, cross-attention projections and products, and the MLP’s two linear layers.Text keys and values are projected once per episode and reused by both dense Fast-WAM and WAM-Cache.Source: B COST ANALYSIS · p. 16
  • The refreshed-token ratio scales every prefill FLOP term, so the relative saving 1 − ρt depends on the refresh policy, not Sv, d, or df.Refreshed queries still attend over all visual keys.Source: B COST ANALYSIS · p. 16
  • Selection overhead comprises latent-surprise distances, reused action-attention maps, age checks, and top-k operations, each below one module’s projection cost.These overheads are not visible in measured CUDA latency.Source: B COST ANALYSIS · p. 16
  • The content-based recompute ratio is at most bs + ba, while age-triggered refreshes can exceed that bound in an individual chunk.Repeated forced refreshes of one token are separated by at least G + 1 chunks.Source: B COST ANALYSIS · p. 16
  • 0.584 is the measured steady-state recompute ratio for the default configuration, below the stated bounds of 0.83 and 0.67.The ratio can also fall below bs + ba when the surprise and attention selections overlap.Source: B COST ANALYSIS · p. 16

C IMPLEMENTATION DETAILS

The implementation uses Fast-WAM with a shared-attention video DiT and action expert, and varies refresh budgets across benchmarks. Evaluation compares against adapted caching and uses task-level episode protocols.

C IMPLEMENTATION DETAILS

  • Fast-WAM pairs the Wan2.2-5B video DiT with a 1B-parameter action expert, while concatenated camera images yield 120 visual tokens on RoboTwin 2.0 and 98 on LIBERO.The text encoder and video VAE remain unchanged.Source: C IMPLEMENTATION DETAILS · p. 17
  • RoboTwin 2.0 uses larger surprise budgets under distractor randomization, while LIBERO uses larger attention budgets because its wrist view emphasizes the gripper and manipulated object.The reported choices are bs = 0.2 and 0.3 for RoboTwin 2.0; on LIBERO, ba = 0.5 and 0.6 at bs = 0.1.Source: C IMPLEMENTATION DETAILS · p. 17
  • VLA-Cache is adapted to the video DiT prefill, reusing static image-patch tokens unless they receive high video-to-text attention.Its static threshold is calibrated to match WAM-Cache’s recompute ratio.Source: C IMPLEMENTATION DETAILS · p. 17
  • Evaluation runs 100 episodes per RoboTwin 2.0 task and 50 per LIBERO task, with success averaged over tasks.Each chunk executes 24 actions before replanning on RoboTwin 2.0 and 10 on LIBERO.Source: C IMPLEMENTATION DETAILS · p. 17

D DENOISING STEPS AND PREFILL SHARE

WAM-Cache retains performance at Fast-WAM’s original ten action-denoising steps, while the step sweep supports using two steps by default. The video DiT prefill remains a major per-chunk cost, especially at two steps.

D DENOISING STEPS AND PREFILL SHARE

  • A single denoising step can average across valid coordination modes in bimanual tasks, whereas it suffices for nearly unimodal single-arm tasks.The paper uses N = 2 by default and compares every caching success rate with the dense model at the same step count.Source: D DENOISING STEPS AND PREFILL SHARE · p. 17
  • N = 2 is 1.3 and 1.5 points below N = 10 on RoboTwin 2.0’s clean and randomized scenes, and 0.6 points below on LIBERO average.The dense step sweep is nearly flat on LIBERO, and N = 1 loses more than 20 points on RoboTwin 2.0.Source: Table 5 · p. 17, D DENOISING STEPS AND PREFILL SHARE · p. 17
  • 41.6% and 39.5% of prefill FLOPs were removed at N = 10 under clean and randomized scenes, with dense-policy gaps of 0.6 and 1.1 points, respectively.At the larger budget, WAM-Cache matches the dense model under clean scenes and comes within 0.2 points under randomized scenes, without retuning.Source: D DENOISING STEPS AND PREFILL SHARE · pp. 17–18, Table 6 · p. 18
  • At N = 2, prefill accounts for 83.5% of the two DiTs’ FLOPs and 57.5% of per-chunk FLOPs including the VAE; WAM-Cache’s 41.6% prefill reduction cuts two-DiT FLOPs by 35%.The two-DiT total falls from 1.262 to 0.824 TFLOPs per chunk; the reduction is 24% including the VAE.Source: Table 7 · p. 18, D DENOISING STEPS AND PREFILL SHARE · p. 18

E REAL-WORLD SETUP AND ROLLOUTS

The real-world evaluation tests Carrot in Bowl and Stack Cubes on AIRBOT Play, using teleoperated demonstrations and binary-success trials. Observations combine external and wrist camera views.

E REAL-WORLD SETUP AND ROLLOUTS

  • AIRBOT Play observations combine synchronized external and wrist RGB views at 20 Hz, and actions are streamed to the arm at 20 Hz.The two views are concatenated into one frame before VAE encoding.Source: E REAL-WORLD SETUP AND ROLLOUTS · p. 18
  • Each task used 100 teleoperated demonstrations, 30k training steps, and 20 evaluation trials with object placements drawn from the demonstration ranges.Trials were scored as binary success without human intervention.Source: E REAL-WORLD SETUP AND ROLLOUTS · p. 18
  • Figure 5 shows the arm approaching, grasping, transporting, and releasing an object while the rest of the scene stays fixed.The frames run from the initial scene through grasping and transport to placement for both tasks.Source: E REAL-WORLD SETUP AND ROLLOUTS · p. 18, Figure 5 · p. 19

Related papers

Loading 2610.11401v1…