Complementary Future Token Pruning for World Action Models

Abstract

World action models (WAMs) jointly predict future visual states and robot actions, but processing dense future-frame tokens incurs substantial inference latency. We propose CF-Prune, a training-free framework for pruning future-frame tokens in pretrained WAMs. Our key insight is that future visual predictions support action generation and therefore need not preserve every visual detail required for high-fidelity video synthesis. CF-Prune combines instruction and action attention to identify task-relevant tokens during the initial denoising forward pass, even when future-frame and action representations remain highly noisy. An entropy-based phase identification mechanism further retains dense computation during precision phases to mitigate pruning-induced performance degradation. Experiments on RoboLab, RoboCasa, and real-world manipulation tasks demonstrate consistent inference speedups. At a 35% pruning ratio, CF-Prune achieves average speedups of 1.19–1.27× while improving task success rates. At a 50% pruning ratio, it achieves 1.29–1.33× average speedups on both simulation benchmarks, with overall success rates within 1.0 percentage point of the baseline.

Key numbers

1. The problem: future frames dominate the compute

World action models build on video diffusion backbones and denoise future video and robot actions together. After VAE compression, a future video is still a spatiotemporal grid of T_v × H_v × W_v tokens, far outnumbering the H_a action tokens and dominating every denoising step.

But a WAM uses future frames as an intermediate reference for action generation, not as a video to show anyone. They must keep task-relevant semantics and dynamics, not every visual detail — so a dense grid is redundant for action prediction.

Pruning these tokens is not straightforward: they are the model's outputs, denoised from Gaussian noise, so their relevance is hard to judge early, and judging it late forfeits the speedup. Existing WAM acceleration reduces denoising steps or redesigns the model, which requires retraining.

2. Method

CF-Prune makes its token selection once, inside the initial denoising forward pass, and applies it for the rest of denoising.

• Instruction-based spatial prior. Observation tokens carry no diffusion noise, so observation-to-instruction attention marks instruction-relevant regions of the scene. The scoring layer is the one whose attention is most concentrated (lowest normalized entropy).
• Action-based visual relevance. Action-to-future attention measures which regions the action tokens rely on — per frame for the moving wrist camera, shared across frames for fixed third-person views. A 5 × 5 Gaussian blur keeps the selection spatially contiguous.
• Complementary selection. For pruning ratio α, keep the top (1 − α)/2 tokens by action relevance, then another (1 − α)/2 by instruction relevance among the rest.
• Phase identification. An entropy contrast across camera views and early vs. late actions detects precision phases near object interaction; pruning is switched off there and applied only in coarse phases.

3. Results

Success (%) on 120 RoboLab tasks with Cosmos3-Nano-Policy-DROID; speedup and FLOPs are averaged over all inference calls. Latency on an RTX A6000. The paper also reports observation-only variants of the VLA pruning baselines.

RoboLab. At α = 35%, CF-Prune is 1.23× faster with 26.8% fewer FLOPs and raises overall success from 36.5% to 37.8%. At α = 50% it reaches 1.33× with only a 0.8 pp drop. Extending VLA pruning masks to future frames costs more success at similar speed, and Sparse VideoGen's profiling overhead outweighs its savings.

RoboCasa. At α = 35%, CF-Prune is 1.19× faster with 15.6% fewer FLOPs and lifts overall success from 65.0% to 69.0%. At α = 50% (1.29×) overall success is 64.0%, but navigation drops from 38.0% to 18.0%, so the effect of aggressive pruning varies by task.

Success (%) over 25 rollouts per task on a Franka Emika Panda with wrist, right and front cameras.

4. Ablations

Both signals are needed. At 35% pruning without phase identification, instruction-only selection reaches 34.0% on RoboLab and action-only 30.8%; the complementary action-first selection reaches 35.3%, above instruction-first ordering (34.3%).

Phase identification helps. Keeping dense computation in detected precision phases lifts overall success from 35.3% to 37.8% while flagging only 31.4% of inference calls, so the average speedup stays at 1.23×.

Takeaway

Treating future frames as support for action prediction, not as video to be rendered faithfully, exposes redundancy that a pretrained WAM can shed at inference time. The attention available in the very first, still-noisy denoising pass is enough to decide which future tokens action generation needs.

Citation

@article{kim2026cfprune,
  title   = {Complementary Future Token Pruning for World Action Models},
  author  = {Kim, Baekseung and Lee, Junhoo and Hong, Sohyeon and Jeon, Suhyun and Kim, Minkyu and Kim, Seungyeon and Han, Donghoon and Kwak, Nojun},
  journal = {arXiv preprint arXiv:2610.19957},
  year    = {2026}
}