Flash-WAM: Modality-Aware Distillation for World Action Models

Arman Akbari, Ci Zhang, Arash Akbari, Lin Zhao, Yixiao Chen, Weiwei Chen, Xuan Zhang, Geng Yuan, Yanzhi Wang · N/A · 2026

Summary

Abstract Summary

World-action models (WAMs) jointly generate future video and robot actions through iterative diffusion, achieving strong performance on manipulation benchmarks but requiring tens of denoising steps, a cost that precludes real-time control. Step distillation has emerged as the natural remedy, but off-the-shelf methods break down in the joint video-action setting because video and action streams use different SNR-shifted noise schedules and reach training with substantially different marginal noise distributions, an asymmetry that single-modality distillation methods cannot accommodate. We introduce Flash-WAM, a modality-aware step-distillation framework inspired by consistency distillation that selects the consistency function for each modality to match its noise regime: a linear-gradient-scaling parametrization for the action stream's low-noise regime, paired with a variance-preserving parametrization for the video stream's high-noise regime, grounded in a structural analysis of the consistency-function family that characterizes the achievable gradient scaling under the consistency boundary condition. Instantiated on LingBot-VA, Flash-WAM compresses inference to a single step in each modality. On RoboTwin 2.0, this reduces per-chunk latency from 8.1 seconds to 348 ms on NVIDIA L40S, a 23x speedup that enables real-time inference. Flash-WAM preserves task success on simulation benchmarks (85.5% RoboTwin 2.0, 95.7% LIBERO) and substantially recovers real-world performance (60% average on a Unitree G1 humanoid robot), while naive consistency distillation drops to 24% at the same step budget.

Key Points

Demonstrates humanoid robot control on real hardware
Addresses dexterous manipulation challenges
Uses simulation or synthetic data for training
Uses knowledge distillation to compress models

Preliminaries

Modern world-action models generate future visual states and the corresponding action sequences using flow matching, a continuous-time generative process that transforms noise into data through iterative denoising. The cost of this iterative process motivates step distillation, in which a student model $\theta_{S}$ is trained to reproduce the output of a pretrained teacher model $\theta_{T}$ in fewer denoising steps.

Methodology

We introduce Flash-WAM , a step-distillation framework for joint video-action diffusion models. Under the asymmetric noise schedules of Section 3.3 ($s^{v}>s^{a}$), the two streams reach the consistency-distillation loss in structurally different noise regimes, and the consistency function applied to each stream must be selected to match its regime.

Section 4.1 establishes why off-the-shelf consistency distillation fails when applied uniformly to both modalities.

Methodology — Flash-WAM: Modality-Aware Distillation for World Action Models — Figure 3 : Real-world evaluation suite on Unitree G1 humanoid robot.

Experiments

We apply our method to the released LingBot-VA model [ 15 ] (shared backbone version), a state-of-the-art, open-source world-action model whose parameter count is small enough for commodity edge deployment, where step distillation has the most practical impact.

Other recent WAMs (Motus [ 1 ] , DreamZero [ 35 ] ) adopt different architectural formulations or integrate their own inference-optimization stacks at the architecture level, and fall outside the scope of our analysis. Per-chunk latency is measured on a single NVIDIA L40S GPU.

Experiments — Flash-WAM: Modality-Aware Distillation for World Action Models — Figure 4 : Qualitative comparison on RoboTwin task “pick_diverse_bottles”, generated with open-loop setting.

Sources and Demonstrations

Figures are reproduced from Akbari et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.