PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2606.05254v1cs.LGcs.CVcs.RO

TL;DR

Joint video-action world models are too slow for real-time control, and existing distillation methods do not accommodate their asymmetric noise regimes. Flash-WAM uses modality-matched consistency distillation to achieve single-step inference while preserving benchmark performance and substantially recovering real-world performance.

  • Problem

    Joint video-action diffusion models use asymmetric noise schedules, so a single consistency function cannot effectively distill both modalities for real-time control.

  • Method

    Flash-WAM selects modality-specific consistency functions matched to each stream’s noise regime, supported by structural analysis of achievable gradient scaling.

  • Results

    348 ms per chunk and 23× speedup reduce LingBot-VA inference to single steps while achieving 85.5% RoboTwin 2.0 and 95.7% LIBERO success.

  • Takeaways & Limitations

    Flash-WAM enables real-time LingBot-VA inference while retaining near-original simulation success and recovering substantially stronger real-world performance than naive distillation.

  • Takeaways & Limitations

    Experiments are primarily in simulation, while broader deployment beyond shared-backbone WAMs and analysis of high-noise regimes remain future work.

Abstract

from arXiv · show

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 \textbf{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 $23{\times}$ 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.

1 Introduction

Flash-WAM addresses the failure of off-the-shelf consistency distillation in joint video-action diffusion by matching each modality’s consistency function to its asymmetric noise regime. On LingBot-VA, it enables single-step inference with substantially lower latency while preserving strong simulation and real-world performance.

  • Motivation: Joint video-action WAMs require iterative denoising for both future-video latents and conditioned action sequences, making both stages major contributors to per-chunk latency.The model first denoises future video latents, then decodes the next action sequence conditioned on predicted frames.
  • Motivation: Off-the-shelf consistency distillation fails because video and action streams have different statistical properties and asymmetric SNR-shifted noise schedules.Video is high-dimensional and structurally redundant, whereas actions are low-dimensional and precision-critical.
  • Method: Flash-WAM selects modality-specific consistency functions: variance-preserving for video’s high-σ regime and linear-gradient-scaling for action’s low-σ regime.The framework is grounded in structural analysis of the consistency-function family and achievable gradient scaling under the consistency boundary condition.
  • Results: 348 ms per-chunk latency replaces 8.1 seconds on NVIDIA L40S, delivering up to 23× speedup with single-step inference in each modality on RoboTwin 2.0.This compression brings WAM inference below the real-time control budget.
  • Results: 85.5% RoboTwin 2.0 and 95.7% LIBERO are recovered at one video step and two action steps, while Unitree G1 performance reaches 60% average success across three manipulation tasks.At one video step and one action step, success remains 81.4% on RoboTwin 2.0 and 95.1% on LIBERO.

2 Related Works

Recent unified world-action models couple video and action generation within shared architectures, while step-distillation methods reduce diffusion inference to a small number of steps. Step distillation primarily follows either teacher trajectories or matches denoising distributions.

  • Unified World Action Models: Unified world-action models jointly generate video and actions through shared or coupled architectures.LingBot-VA uses a shared transformer backbone; Motus couples vision-language, video, and action generators via cross-attention.
  • Unified World Action Models: LingBot-VA uses autoregressive video-action diffusion, Motus uses Mixture-of-Transformers, and DreamZero reduces denoising steps through inference-time architectural optimization.
  • Step Distillation: Step-distillation methods compress iterative diffusion denoising into a small number of inference steps through trajectory-following or distribution-matching approaches.Trajectory-following methods include progressive distillation and consistency models, which respectively halve sampling steps iteratively or map trajectory points to a shared clean endpoint.

3 Preliminaries

World-action models use flow matching to jointly generate future video latents and corresponding actions, but iterative denoising makes real-time control costly. Consistency distillation reduces this cost, yet modality-specific SNR-shifted schedules place video and action streams in different noise regimes that a shared consistency function cannot serve uniformly.

  • 3.2 Consistency Distillation: Consistency models map points along a probability-flow ODE trajectory to the clean endpoint, while distillation trains a student and EMA target against teacher-guided Euler steps.The consistency function f(xσ, σ) = a(σ)xσ + b(σ)vθ obeys a(0) = 1 and b(0) = 0.
  • 3.3 World-Action Models: World-action models factor generation into visual dynamics and inverse dynamics, sampling future video latents before the corresponding action sequence.Both stages share transformer parameters and are realized through flow matching.
  • 3.3 World-Action Models: Nv + Na sequential transformer passes dominate per-chunk latency, preventing real-time control; 25v/50a denotes 25 video and 50 action denoising steps.The notation specifies separate numbers of function evaluations for the two modalities.
  • 3.4 Modality-Aware Distillation: Independent video and action SNR-shifted schedulers concentrate training mass in different noise regimes, so one shared consistency function cannot serve both modalities uniformly.Video tolerates heavier per-step noise, whereas precision-critical actions require a gentler schedule.

4 Methodology

Flash-WAM addresses joint video-action distillation by matching each modality’s consistency function to its distinct noise regime. Its analysis identifies a structural low-noise gradient-scaling limitation and motivates modality-specific parametrization choices for action and video streams.

  • Motivation: Uniform consistency distillation fails because video concentrates near high σ while action spreads across the full range with substantial low-σ training mass.This asymmetry is presented as a structural failure mode rather than a tunable inefficiency.
  • Gradient-scaling analysis: The consistency-loss gradient scales pointwise as |b(σ)|, so small |b(σ)| yields little learning signal regardless of prediction quality.Valid consistency functions satisfy f(xσ, σ) = a(σ)xσ + b(σ)vθ with a(0) = 1 and b(0) = 0.
  • Gradient-scaling analysis: |b(σ)| = O(σ) is the best achievable low-noise scaling, attained exactly when b′(0) ≠ 0; LCM instead has quadratic vanishing.LCM’s gradient-scale factor at σ = 0.1 is roughly 36× smaller than in the high-σ video regime.
  • Action stream: The action stream uses a consistency function with b exactly linear in σ, no tunable hyperparameter, and |b(σ)| = σ throughout [0, 1].This construction satisfies the consistency boundary condition and realizes the optimal linear low-σ scaling.
  • Video stream: For high-σ video, Flash-WAM prioritizes Karras-style variance preservation and bounded output range because these stabilize high-dimensional video latents.The method notes that these properties are largely irrelevant for low-dimensional, bounded action targets.
  • Framework design: Flash-WAM’s novelty is principled per-modality selection: it explains which consistency-function family members to use in each regime and why.The modality-aware choices affect only per-stream loss heads, leaving the shared architecture and per-step compute cost unchanged from the teacher.

5 Experiments

Experiments on RoboTwin 2.0, LIBERO, and a Unitree G1 humanoid show that Flash-WAM enables single-step real-time inference while retaining strong simulation and real-world performance. Across benchmarks and ablations, it consistently outperforms off-the-shelf and alternative LCM-based distillation strategies.

  • Benchmarks: Flash-WAM is evaluated on RoboTwin 2.0, LIBERO, and a Unitree G1 humanoid robot, covering simulation robustness, task-suite performance, and real-world control.RoboTwin 2.0 includes Clean and Randomized splits, while LIBERO contains Spatial, Object, Goal, and Long-horizon suites.
  • RoboTwin 2.0: 85.54% average success: Flash-WAM at 1v/2a recovers most of LingBot-VA’s 91.25% on RoboTwin 2.0 at a 19× speedup.At 1v/1a, Flash-WAM achieves 81.41% average success.
  • RoboTwin 2.0: 23.3× speedup reduces RoboTwin per-chunk latency to 348 ms on a single NVIDIA L40S, enabling real-time inference despite 25× and 50× denoising reductions.At 1v/2a, naive joint LCM reaches 23.97%, DMD2 reaches 78.74%, and video-only LCM reaches 78.79%.
  • LIBERO: 95.7% average success: Flash-WAM at 1v/2a nearly matches the LIBERO teacher’s 98.6% while achieving a 13.7× speedup.At 1v/1a, it achieves 95.1% average success at 16.3× speedup, reducing latency from 6,767 ms to 404 ms.
  • Ablation Analysis: Across RoboTwin splits, NFE configurations, and horizons, Flash-WAM outperforms alternative LCM strategies, while naive joint LCM collapses and action regularization cannot replace action distillation.Naive joint LCM drops to 25.88% on the Clean split and falls 24 points below plain video-only LCM.

6 Conclusion

Flash-WAM is introduced as a step-distillation framework for joint video-action diffusion models, motivated by a structural failure mode caused by asymmetric per-modality noise schedules. On a Unitree G1 humanoid robot, it achieves 60% across three manipulation tasks and outperforms non-distilled and video-only baselines at the same step budget.

  • Conclusion: Flash-WAM is introduced as a step-distillation framework for joint video-action diffusion models.The framework addresses a structural failure mode in off-the-shelf consistency distillation.
  • Conclusion: 60% across three manipulation tasks is achieved by Flash-WAM in real-world experiments on a Unitree G1 humanoid robot.This result substantially outperforms both reduced-NFE inference without distillation and Video-only LCM at the same step budget.
  • Conclusion: 40% for reduced-NFE inference without distillation and 43.3% for Video-only LCM are surpassed by Flash-WAM at the same step budget.These are the reported comparison results for the Unitree G1 real-world experiments.

A Implementation Details · A.1 Libero Finetuning · A.2 Flash-WAM Distillation Hyperparameters

The implementation adapts LingBot-VA separately to each LIBERO suite before distillation, then trains Flash-WAM with shared procedures and reported hyperparameters. Fine-tuning and distillation each require approximately 24 hours on four H100 GPUs per suite.

  • A Implementation Details: The appendix provides implementation details for Flash-WAM and all baselines across fine-tuning, distillation, and baseline-specific configurations.The fine-tuning and distillation procedures are covered in Sections A.1 and A.2, while baseline choices are specified separately in Section A.3.
  • A.1 Libero Finetuning: The released LingBot-VA base checkpoint is first adapted separately to each LIBERO suite before step distillation.The checkpoint was pretrained on multi-task data, and the adaptation follows the LingBot-VA protocol.
  • A.1 Libero Finetuning: 4,000 training steps are used to fine-tune the LingBot-VA base checkpoint on each LIBERO suite.The fine-tuning hyperparameters are reported in Table 5.
  • A.1 Libero Finetuning: Approximately 24 hours on 4 H100s are required to fine-tune each LIBERO suite.This is the reported per-suite fine-tuning cost.
  • A.2 Flash-WAM Distillation Hyperparameters: 2,000 steps are used to distill Flash-WAM from the fine-tuned LingBot-VA teacher on each LIBERO suite.The distillation hyperparameters are reported in Table 6.
  • A.2 Flash-WAM Distillation Hyperparameters: Approximately 24 hours on 4× H100 GPUs are required to distill Flash-WAM on each LIBERO suite.The same training duration is reported for each suite.
  • A.2 Flash-WAM Distillation Hyperparameters: The distillation procedure is also applied to Naive joint LCM, Video-only LCM, and Video-only LCM + reg. on RoboTwin for fair comparison.These are the LCM-based baselines receiving the same training procedure.

A.3 Baseline Implementations … A.3.3 Video-only LCM + reg

The baselines share Flash-WAM’s data, checkpoint, and training iterations, while differing in distillation objectives. They progressively modify video and action streams through uniform LCM, video-only distillation, or action flow-matching regularization.

  • A.3 Baseline Implementations: All baselines use the same training data, base checkpoint, and number of training iterations as Flash-WAM, differing only in their distillation objective.
  • A.3.1 Naive Joint LCM: Naive Joint LCM applies the standard LCM consistency function uniformly to video and action streams.
  • A.3.1 Naive Joint LCM: Naive Joint LCM computes modality-specific consistency losses and combines them as L = Lv + λaLa with λa = 1.0.
  • A.3.2 Video-only LCM: Video-only LCM distills only video, leaving action unchanged during training and at the teacher’s full 50-step inference schedule.
  • A.3.3 Video-only LCM + reg: Video-only LCM + reg adds an action-stream flow-matching regularizer, enabling both video and action streams to run at reduced NFE during inference.
  • A.3.3 Video-only LCM + reg: The video stream retains standard LCM consistency supervision, while the action stream uses an MSE flow-matching loss anchored to the demonstration distribution.
  • A.3.3 Video-only LCM + reg: The action regularizer constructs noisy actions from clean demonstrations and sampled schedule noise, supervises velocity prediction, and applies a per-channel validity mask.
  • A.3.3 Video-only LCM + reg: Video-only LCM + reg combines the video consistency loss with the action regularizer, while otherwise matching Flash-WAM’s hyperparameter configuration.

A.3.4 DMD2 Baseline Implementations

The paper implements two reproducible DMD2 baselines for joint video-action diffusion: Video-only DMD2 + reg for the main results and Joint DMD2 for appendix evaluation. Both share the backbone, scoring procedure, and critic objective while differing in rollout modalities and student losses.

  • Variant overview: The two baselines are Video-only DMD2 + reg, used in the main results, and Joint DMD2, evaluated as the fully-joint appendix variant.Their shared networks, scoring procedure, and critic objective are retained across variants.
  • Networks: Three full-copy joint video-action backbones define the frozen reference model θT, trainable student θS, and trainable critic θC.The reference defines the real score, the student is a K-step generator initialized from θT, and the critic tracks the student distribution.
  • Student rollout: K = 4 student denoising steps use a uniform noise-band schedule, with Joint DMD2 denoising both modalities and Video-only DMD2 + reg denoising only video.For the video-only variant, action inputs are directly perturbed ground-truth actions at corresponding noise levels.
  • Losses: Joint DMD2 applies distribution matching to video and action, whereas Video-only DMD2 + reg applies it only to video and adds an action flow-matching regularizer.The action regularizer supervises the student’s final-step action prediction against the target velocity from the perturbed ground-truth action.

B Additional Experimental Results

At the aggressive 1v/1a single-step setting on RoboTwin 2.0, Table 7 compares distillation strategies and shows Joint DMD2 substantially underperforms Video-only DMD2 + reg and Flash-WAM.

  • RoboTwin 2.0: 1v/1a is the most aggressive single-step configuration evaluated for the RoboTwin 2.0 comparison.Table 7 reports all distillation strategies at this setting.
  • RoboTwin 2.0: Table 7 includes Joint DMD2, the fully-joint DMD2 variant excluded from the main paper.Joint DMD2 is described in Section A.3.4.
  • RoboTwin 2.0: 52.7% Clean and 48.5% Randomized are Joint DMD2’s average RoboTwin 2.0 results.These results are reported at the 1v/1a setting.
  • RoboTwin 2.0: Roughly 14 points below Video-only DMD2 + reg and 30 points below Flash-WAM, Joint DMD2 performs substantially worse on average.The comparison uses Video-only DMD2 + reg from the main results.

C Limitations and future work · D LLM Usage

The paper identifies limitations around simulation-only evaluation, architectural scope, and incomplete gradient-scaling analysis, while describing a bounded use of LLMs for writing and terminology checks. Future work includes physical-robot deployment and broader architectural and analytical extensions.

  • C Limitations and future work: Experiments were conducted in simulation, leaving deployment on physical robots for future work.The passage explicitly identifies real-world deployment as an unresolved next step.
  • C Limitations and future work: Flash-WAM currently targets shared-backbone WAMs rather than multi-model architectures with separate per-modality sub-models.Extending the framework to multi-model architectures is identified as a natural next step.
  • C Limitations and future work: The analysis characterizes optimal gradient scaling in the low-σ regime where actions train.The paper notes that this low-noise analysis does not yet cover the corresponding high-σ regime.
  • C Limitations and future work: A corresponding analysis for the high-σ regime is needed to complete the analytical picture.This is presented as an explicit future analytical direction.
  • D LLM Usage: LLMs were used in a limited and clearly bounded role during paper preparation.Their stated uses included grammar improvement, sentence rephrasing, alternative wording, and tightening verbose passages.
  • D LLM Usage: LLMs were occasionally used to verify that technical phrasing followed standard diffusion and step-distillation conventions.The passage distinguishes this terminology check from the paper’s technical claims and results.
Loading 2606.05254v1…