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

2026-09-30

AI Summary

World action models suffer from prohibitive inference costs due to the dense processing of future frames and lack guidance on action relevance. This work proposes a training-free acceleration framework that introduces, for the first time, an action-guided token selection mechanism alongside a Pilot engine. By leveraging a lightweight scoring algorithm to identify critical regions and exploiting attention overlap for cross-step cache reuse, the method achieves sparse inference with minimal overhead. Integrating diffusion models with token pruning techniques, this approach yields approximately 2.0× and 1.8× inference speedups on the LIBERO and RoboLab benchmarks, respectively, while preserving task performance without degradation.

Abstract

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.

Problem

• World Action Models
• Inference Acceleration
• Token Pruning
• Robot Control
• Video Diffusion

Innovation

• World Action Models
• Action-Guided Token Selection
• Sparse Imagination
• Training-free Acceleration
• Pilot Engine