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

* Equal contribution   † Corresponding author

4.00×LIBERO speedup
3.94×LIBERO-PRO speedup
2.10×RoboCasa speedup
0additional training

Acceleration at a glance

One policy, multiple efficient operating points

By adjusting acceleration strength across three levels, MARS enables flexible efficiency–performance trade-offs.

LIBERO-PRO efficiency–performance trade-off
LIBERO-PRO efficiency–performance trade-off Success rate versus average latency for Cosmos Policy, DiCache, C³ache, X-Cache, and five MARS operating points. Bubble area represents FLOPs. 47.0 46.5 46.0 45.5 45.0 44.5 44.0 0.3 0.5 0.7 0.9 1.1 1.3 1.5 Average Latency (s) Success Rate (%) Method (Size ∝ FLOPs) Cosmos-Policy + DiCache + C³ache + X-Cache + Ours Better efficiency–performance trade-off
Pre-trained WAM
Original inference
Baseline
100.00%relative performance
1.00×speedup
Pre-trained WAM
+ MARS
Performance-oriented
100.20%relative performance
3.08×speedup
Speed-oriented
98.55%relative performance
4.50×speedup
01

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.

02

Motivation

Profiling WAM inference reveals redundancy at the Transformer, diffusion, and cross-query levels.

WAM redundancy analysis: attention importance, denoising input and output differences, DiT residual stability, and observed versus predicted visual maps with an error map

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.

(a)

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.

(b)

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.

(c)

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.

03

Method

MARS reduces redundancy at the Transformer, diffusion, and cross-query levels.

MARS framework: Transformer-level attention skipping, diffusion-level acceleration, and cross-query visual token reuse with a shallow VAE verifier
01

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.

02

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.

03

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.

04

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 Policy98.20 / 1.0099.87 / 1.0097.47 / 1.0096.67 / 1.0098.05%1.00×1.4001 s100.00%
+ DiCache98.07 / 1.7099.87 / 1.6697.80 / 1.6797.27 / 1.7298.25%1.69×0.8273 s48.12%
+ C3ache*97.80 / 2.1999.00 / 2.1697.53 / 2.2594.47 / 2.3097.20%2.24×0.6256 s33.18%
+ X-Cache*97.87 / 1.1199.67 / 1.1797.47 / 1.1297.00 / 1.2198.00%1.17×1.2017 s80.20%
+ Ours (Speed)96.73 / 3.8999.07 / 4.0897.60 / 3.8796.27 / 4.0797.42%4.00×0.3498 s18.39%
+ Ours (Acc.)98.47 / 3.0699.60 / 3.2997.80 / 3.1096.60 / 3.1998.12%3.17×0.4414 s19.19%

LIBERO-PRO

Method Success Rate (%, ↑) / Speedup (↑) Success
Rate (↑)
Speedup (↑) Latency (↓) FLOPs (↓)
Goal Object Spatial Long
Cosmos Policy43.67 / 1.0046.13 / 1.0047.40 / 1.0045.33 / 1.0045.63%1.00×1.4044 s100.00%
+ DiCache44.13 / 1.6746.27 / 1.6647.73 / 1.6345.07 / 1.7345.80%1.68×0.8347 s50.32%
+ C3ache*42.80 / 2.2245.53 / 2.2047.53 / 2.2842.60 / 2.3044.62%2.26×0.6222 s32.86%
+ X-Cache*43.73 / 1.1746.20 / 1.2247.80 / 1.1744.33 / 1.2645.52%1.22×1.1549 s75.55%
+ Ours (Speed)44.20 / 3.9145.67 / 3.8647.33 / 3.9644.93 / 3.9945.53%3.94×0.3564 s18.80%
+ Ours (Acc.)43.67 / 3.0547.00 / 3.0847.60 / 3.0644.60 / 3.1145.72%3.08×0.4558 s19.05%

RoboCasa

Method S.R. (↑) Speedup (↑) Latency (↓) FLOPs (↓)
Cosmos Policy66.69%1.00×2.2958 s100.00%
+ DiCache66.47%1.50×1.5305 s49.99%
+ C3ache*65.00%1.81×1.2653 s33.64%
+ X-Cache*66.37%1.20×1.9202 s72.86%
+ Falcon66.47%1.21×1.8968 s74.56%
+ Ours66.94%2.10×1.0911 s22.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.

Fast-WAM3.17×
LingBot-VA2.88×
π0.51.28×

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-WAM95.87 / 1.0099.40 / 1.0097.13 / 1.0093.93 / 1.0096.58%1.00×0.5365 s
+ Ours96.07 / 3.4399.00 / 3.2596.93 / 3.4393.20 / 2.9396.30%3.17×0.1693 s
LingBot-VA———97.40 / 1.0097.40%1.00×13.4105 s
+ Ours———98.13 / 2.8898.13%2.88×4.6558 s
π0.597.80 / 1.0098.67 / 1.0098.40 / 1.0092.93 / 1.0096.95%1.00×0.1462 s
+ Ours97.87 / 1.2999.07 / 1.3098.47 / 1.3392.27 / 1.2496.92%1.28×0.1145 s
05

Qualitative Results

MARS qualitative results figure one
MARS qualitative results figure two
06

Citation

Read the paper on arXiv. BibTeX will be added soon.

Follow the repository for updates →