Complementary Future Token Pruning for World Action Models

Baekseung Kim; Junhoo Lee; Sohyeon Hong; Suhyun Jeon; Minkyu Kim; Seungyeon Kim; Donghoon Han; Nojun Kwak

cs.ROOct 14, 2026

Complementary Future Token Pruning for World Action Models

Authors: Baekseung Kim, Junhoo Lee, Sohyeon Hong, Suhyun Jeon, Minkyu Kim, Seungyeon Kim, Donghoon Han, Nojun Kwak

Organizations: Seoul National University · Korea Advanced Institute of Science & Technology · Dnotitia

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.27x while improving task success rates. At a 50% pruning ratio, it achieves 1.29-1.33x average speedups on both simulation benchmarks, with overall success rates within 1.0 percentage point of the baseline.

Figures & tables

Figure 1: Rethinking future-video generation for action prediction.

Figure 2: Overview of CF-Prune. $k$ denotes the current denoising step, and $K$ is the total number of denoising steps.

Panel of Figure 3 — Figure 3: Real-world experimental setup and the three manipulation tasks used for evaluation.

Method	Variant	Success Rate (%)				Speedup		FLOPs
Method	Variant	Simple	Moderate	Complex	Overall	Coarse	Average	FLOPs
Dense	–	40.3	34.1	27.6	36.5	–	$1.00\times$	100%
SpecPrune-VLA	Obs-only (70%)	33.4	33.8	24.7	32.3	$1.07\times$	$1.06\times$	93.1%
SpecPrune-VLA	Obs + future (25%)	36.9	28.2	23.5	32.2	$1.26\times$	$1.23\times$	77.5%
SP-VLA	Obs-only (70%)	34.4	24.1	21.2	29.2	$1.05\times$	$1.05\times$	93.1%
SP-VLA	Obs + future (25%)	37.5	32.3	23.5	33.8	$1.24\times$	$1.24\times$	77.5%

Table 1: Results on RoboLab. Success rates are averaged over tasks. Coarse phase latency excludes precision phases, while average latency and FLOPs are averaged over all inference calls.

Method	Variant	Success Rate (%)					Speedup		FLOPs
Method	Variant	Artic.	Transfer	Control	Nav.	Overall	Coarse	Average	FLOPs
Dense	–	95.0	55.0	72.0	38.0	65.0	–	$1.00\times$	100%
SpecPrune-VLA	Obs-only (70%)	65.0	17.0	72.0	44.0	49.5	$1.06\times$	$1.05\times$	95.4%
SpecPrune-VLA	Obs + future (25%)	92.0	51.0	78.0	36.0	64.3	$1.21\times$	$1.17\times$	82.8%
SP-VLA	Obs-only (70%)	82.0	27.0	75.0	42.0	56.5	$1.07\times$	$1.07\times$	93.7%
SP-VLA	Obs + future (25%)	94.0	38.0	90.0	37.0	64.8	$1.18\times$	$1.17\times$	82.8%

Table 2: Results on four representative RoboCasa tasks, with 100 rollouts per task. The columns represent articulation (OpenStandMixerHead), object transfer (CoffeeSetupMug), device control (TurnOnElectricKettle), and navigation (NavigateKitchen).

Method	Success Rate (%)				Speedup
Method	Emergency Button	Block Pick-and-Place	Block Stacking	Overall	Coarse	Average
Dense	84.0	92.0	80.0	85.3	–	$1.00\times$
CF-Prune ( $\alpha=35\%$ )	88.0	96.0	88.0	90.7	$1.37\times$	$1.27\times$

Table 3: Real-world evaluation results. Success rates are measured over 25 rollouts per task, and Overall denotes the average success rate across the three tasks.

Selection Strategy	Success Rate (%)
Instruction-only	34.0
Action-only	30.8
Instruction-first	34.3
Action-first (Ours)	35.3

Table 4: Token selection strategy ablation on RoboLab at 35% pruning.

Panel of Figure 4 — Figure 4: Layer-wise normalized attention entropy at the first denoising step across sampled RoboLab tasks. In each subfigure, the left heatmap shows entropy over rollout progress averaged across tasks, while the right shows task-wise entropy averaged over sampled rollout stages.

Panel of Figure 5 — Figure 5: Effect of Gaussian smoothing on token pruning. Navy cells mark pruned positions in the first future latent frame, overlaid on the input observation.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Camera views used in our real-world experiments. The right and front views provide third-person observations, while the wrist-mounted camera captures local interactions near the gripper.

Panel of Figure 7 — Figure 7: Comparison of the token selection strategies in Table 4 on the same RoboLab scene at 35% pruning. All variants use the same attention scores and $5\times 5$ Gaussian smoothing. Navy cells indicate pruned positions.

Panel of Figure 8 — Figure 8: Examples of future-token pruning at 35% and 50% on RoboLab, RoboCasa, and recorded real-world inputs. Navy cells indicate pruned positions.

Figure 9: Representative real-world rollouts of the dense baseline and CF-Prune on emergency-button pressing, block pick-and-place, and block stacking. Each row shows eight selected frames illustrating task progression from left to right.

Efficient Inference for World Action Models

Efficient inference for world action models cuts the cost of predicting action-conditioned futures for responsive control and planning.

See all 80 papers

Sparse-WAM: Accelerating World Action Models via Action-Guided Sparse Imagination
Sep 30, 2026Xinling Xie, Haodong Wang, Jiazhi Mi +8Efficient Inference for World Action Models Token Pruning
NJU · HKUST · HIT +1
Efficient-WAM: A 1B-Parameter World-Action Model with Low-Cost Future Imagination
Jun 8, 2026Jiajun Li, Tiecheng Guo, Yifan Ye +9World Models for Robotics Efficient Inference for World Action Models
1The University of Hong Kong · 2Peking University · 3Muka Robotics +2
What Makes World Action Models Generalize? An Empirical Study of Test-Time Future Modeling
Sep 28, 2026Renping Zhou, Zanlin Ni, Zihao Fan +8Action-Conditioned World Models Efficient VLA Model Inference
Leap Lab, Tsinghua University · University of Science and Technology of China · Beijing Institute of Technology

Sep 30, 2026J/K move · Enter open · S save

Sparse-WAM: Accelerating World Action Models via Action-Guided Sparse Imagination

Xinling Xie, Haodong Wang, Jiazhi Mi, Zhiming Liu+7

NJU · HKUST · HIT+1

Efficient Inference for World Action Models Token Pruning World Action Models Video World Models

Open paper AskPDF arXiv

World-action models (WAMs) leverage pretrained video models to improve generalization in robot control by jointly predicting future visual states and actions. This capability comes at a substantial inference cost, as dense future-frame tokens are repeatedly processed during denoising. Prior methods address this by token pruning that prioritizes visual fidelity to reduce denoising costs in video diffusion models. However, these methods do not use action relevance to determine which future-frame tokens to retain during joint denoising in WAMs. In this paper, we propose Sparse-WAM, a training-free framework for action-guided sparse imagination that selectively processes future-frame tokens to accelerate WAM inference. We observe substantial overlap in the spatial distribution of attention from action tokens to future-frame tokens (action-to-future attention) between consecutive denoising steps, despite continued updates to the future representations. Motivated by this, we develop Action-Guided Token Selection to retain frame-specific action-relevant regions together with cross-frame context. However, a naive implementation can incur attention-scoring and token-packing overhead that offsets the computational savings from pruning. We therefore introduce Pilot, an efficient engine that reduces sparse inference overhead through lightweight scoring and cross-step reuse of token selections. On LIBERO with FastWAM-Joint and RoboLab-120 with Cosmos 3 Edge, Sparse-WAM achieves inference speedups of approximately

2.0\times

and

1.8\times

, respectively, over dense eager inference on an NVIDIA RTX 4090, while largely preserving task performance.

Abstract

Efficient Inference for World Action Models

Explore similar work

Sparse-WAM: Accelerating World Action Models via Action-Guided Sparse Imagination

Efficient-WAM: A 1B-Parameter World-Action Model with Low-Cost Future Imagination

What Makes World Action Models Generalize? An Empirical Study of Test-Time Future Modeling