Training-free inference acceleration for World Action Models
MARS: Systematic Multi-Level Acceleration for
World Action Models
1School of Intelligence Science and Technology, Peking University
2State Key Laboratory of General Artificial Intelligence, Peking University
3University of Technology Sydney
4University of Liverpool
Acceleration at a glance
One policy, multiple efficient operating points
By adjusting acceleration strength across three levels, MARS enables flexible efficiency–performance trade-offs.
Abstract
World Action Models (WAMs) have emerged as a promising alternative for embodied intelligence by jointly predicting future world states and robot actions. However, their heavy visual encoding and iterative generative inference incur substantial latency, limiting their applicability to high-frequency closed-loop control. Existing training-free acceleration methods typically optimize a single stage, yielding limited speedups or degraded control under aggressive reuse. In this work, we systematically profile the inference overhead of WAMs and observe abundant redundancy across Transformer computation, diffusion denoising, and consecutive policy queries. However, the safe removal of such redundancy depends strongly on its impact on action prediction and control sensitivity. Based on these observations, we introduce MARS, a training-free and plug-and-play framework for multi-level WAM acceleration. MARS selectively removes low-impact attention computation, adaptively accelerates diffusion inference according to motion granularity, and speculatively reuses predicted future visual representations across consecutive queries with online verification. Together, these complementary mechanisms provide remarkable acceleration while limiting approximation error. Extensive experiments across LIBERO, LIBERO-PRO, and RoboCasa show that MARS achieves up to 4.00×, 3.94×, and 2.10× end-to-end speedups, respectively, while maintaining task success rate close to the original policy. Moreover, MARS generalizes across different model backbones and across simulation and real-robot experiments, substantially improving the efficiency–performance trade-off of WAM inference.
Motivation
Profiling WAM inference reveals redundancy at the Transformer, diffusion, and cross-query levels.
Redundancy across multiple levels of WAM inference. (a) Non-uniform attention importance. (b) Stable denoising outputs and DiT residuals. (c) Observed and predicted visual maps with the corresponding error map. Diagnostics use Cosmos Policy in LIBERO.
Transformer-Level Redundancy
Non-Uniform Attention
Many cross-attention branches and self-attention heads have low importance, while a smaller subset remains influential. This non-uniform importance motivates selective attention skipping.
Diffusion-Level Redundancy
Stable Output States and Residuals
Denoising inputs vary substantially across steps, whereas denoising outputs and DiT residuals remain comparatively stable. This motivates one-step generation and the multi-step solver with residual reuse.
Cross-Query Redundancy
Speculative Visual Representation
Predicted visual maps closely match observations at the next policy query, with small errors in the corresponding error map. This motivates verified reuse of predicted future visual tokens.
Method
MARS reduces redundancy at the Transformer, diffusion, and cross-query levels.
Within a DiT forward pass
Transformer-Level Acceleration
Skips low-impact cross-attention branches and prunes low-importance self-attention heads while retaining components critical to action prediction.
Within a policy query
Diffusion-Level Acceleration
Uses the provisional action to estimate motion granularity: coarse-grained motions use one-step generation; fine-grained stages retain the original multi-step solver with DiT residual reuse.
Across consecutive policy queries
Cross-Query Acceleration
Uses shallow VAE features to verify consistency between predicted and observed visual changes. When verification passes, it reuses predicted future visual tokens and skips the remaining VAE encoding.
Quantitative Results
End-to-end acceleration on Cosmos Policy, measured with matched three-seed evaluation.
LIBERO
| Method | Success Rate (%, ↑) / Speedup (↑) | Success Rate (↑) |
Speedup (↑) | Latency (↓) | FLOPs (↓) | |||
|---|---|---|---|---|---|---|---|---|
| Goal | Object | Spatial | Long | |||||
| Cosmos Policy | 98.20 / 1.00 | 99.87 / 1.00 | 97.47 / 1.00 | 96.67 / 1.00 | 98.05% | 1.00× | 1.4001 s | 100.00% |
| + DiCache | 98.07 / 1.70 | 99.87 / 1.66 | 97.80 / 1.67 | 97.27 / 1.72 | 98.25% | 1.69× | 0.8273 s | 48.12% |
| + C3ache* | 97.80 / 2.19 | 99.00 / 2.16 | 97.53 / 2.25 | 94.47 / 2.30 | 97.20% | 2.24× | 0.6256 s | 33.18% |
| + X-Cache* | 97.87 / 1.11 | 99.67 / 1.17 | 97.47 / 1.12 | 97.00 / 1.21 | 98.00% | 1.17× | 1.2017 s | 80.20% |
| + Ours (Speed) | 96.73 / 3.89 | 99.07 / 4.08 | 97.60 / 3.87 | 96.27 / 4.07 | 97.42% | 4.00× | 0.3498 s | 18.39% |
| + Ours (Acc.) | 98.47 / 3.06 | 99.60 / 3.29 | 97.80 / 3.10 | 96.60 / 3.19 | 98.12% | 3.17× | 0.4414 s | 19.19% |
LIBERO-PRO
| Method | Success Rate (%, ↑) / Speedup (↑) | Success Rate (↑) |
Speedup (↑) | Latency (↓) | FLOPs (↓) | |||
|---|---|---|---|---|---|---|---|---|
| Goal | Object | Spatial | Long | |||||
| Cosmos Policy | 43.67 / 1.00 | 46.13 / 1.00 | 47.40 / 1.00 | 45.33 / 1.00 | 45.63% | 1.00× | 1.4044 s | 100.00% |
| + DiCache | 44.13 / 1.67 | 46.27 / 1.66 | 47.73 / 1.63 | 45.07 / 1.73 | 45.80% | 1.68× | 0.8347 s | 50.32% |
| + C3ache* | 42.80 / 2.22 | 45.53 / 2.20 | 47.53 / 2.28 | 42.60 / 2.30 | 44.62% | 2.26× | 0.6222 s | 32.86% |
| + X-Cache* | 43.73 / 1.17 | 46.20 / 1.22 | 47.80 / 1.17 | 44.33 / 1.26 | 45.52% | 1.22× | 1.1549 s | 75.55% |
| + Ours (Speed) | 44.20 / 3.91 | 45.67 / 3.86 | 47.33 / 3.96 | 44.93 / 3.99 | 45.53% | 3.94× | 0.3564 s | 18.80% |
| + Ours (Acc.) | 43.67 / 3.05 | 47.00 / 3.08 | 47.60 / 3.06 | 44.60 / 3.11 | 45.72% | 3.08× | 0.4558 s | 19.05% |
RoboCasa
| Method | S.R. (↑) | Speedup (↑) | Latency (↓) | FLOPs (↓) |
|---|---|---|---|---|
| Cosmos Policy | 66.69% | 1.00× | 2.2958 s | 100.00% |
| + DiCache | 66.47% | 1.50× | 1.5305 s | 49.99% |
| + C3ache* | 65.00% | 1.81× | 1.2653 s | 33.64% |
| + X-Cache* | 66.37% | 1.20× | 1.9202 s | 72.86% |
| + Falcon | 66.47% | 1.21× | 1.8968 s | 74.56% |
| + Ours | 66.94% | 2.10× | 1.0911 s | 22.50% |
With Different Base Models
One framework, diverse policies
The compatible Transformer-level and diffusion-level components transfer beyond Cosmos Policy to additional WAM and VLA backbones without retraining.
Results with Different Base Models
LingBot-VA can only be evaluated on Long because only its Long checkpoint is publicly available.
| Method | Success Rate (%, ↑) / Speedup (↑) | Success Rate (↑) | Speedup (↑) | Latency (↓) | |||
|---|---|---|---|---|---|---|---|
| Goal | Object | Spatial | Long | ||||
| Fast-WAM | 95.87 / 1.00 | 99.40 / 1.00 | 97.13 / 1.00 | 93.93 / 1.00 | 96.58% | 1.00× | 0.5365 s |
| + Ours | 96.07 / 3.43 | 99.00 / 3.25 | 96.93 / 3.43 | 93.20 / 2.93 | 96.30% | 3.17× | 0.1693 s |
| LingBot-VA | — | — | — | 97.40 / 1.00 | 97.40% | 1.00× | 13.4105 s |
| + Ours | — | — | — | 98.13 / 2.88 | 98.13% | 2.88× | 4.6558 s |
| π0.5 | 97.80 / 1.00 | 98.67 / 1.00 | 98.40 / 1.00 | 92.93 / 1.00 | 96.95% | 1.00× | 0.1462 s |
| + Ours | 97.87 / 1.29 | 99.07 / 1.30 | 98.47 / 1.33 | 92.27 / 1.24 | 96.92% | 1.28× | 0.1145 s |
Qualitative Results
Citation
Read the paper on arXiv. BibTeX will be added soon.
Follow the repository for updates →