Flash-WAM: Modality-Aware Distillation for World Action Models
Flash-WAM uses modality-aware step distillation to reduce inference from 50 steps to 1, achieving 23× speedup while maintaining high task success rates.
Key Findings
Methodology
This paper introduces Flash-WAM, a modality-aware step distillation framework that addresses the asymmetry in noise schedules between video and action streams. It employs two distinct consistency functions: a linear-gradient scaling for low-noise action data and a variance-preserving approach for high-noise video data. Structural analysis of the consistency function family guides the design, ensuring effective gradient signals across different noise regimes. The framework enables single-step inference for both modalities, significantly accelerating joint diffusion models while preserving task success. The training involves separate consistency losses for each modality, combined into a unified objective, leveraging shared transformer architecture.
Key Results
- On LingBot-VA, Flash-WAM reduces inference steps from 50 to 1, decreasing latency from 8100ms to 415ms, a 23× speedup. It maintains high success rates: RoboTwin 2.0 at 85.5%, LIBERO at 95.7%, and achieves 60% average success on real-world tasks, outperforming naive distillation at 24%.
- Experimental results across benchmarks show that the modality-aware approach effectively mitigates the gradient vanishing problem in low-noise action regimes, enabling high-fidelity joint generation with minimal steps.
- The method demonstrates practical viability for real-time robotic control, with significant improvements in inference speed and task success, validating the theoretical analysis.
Significance
This work addresses the critical bottleneck of inference speed in multimodal diffusion models, enabling real-time robotic applications. By tailoring consistency functions to each modality’s noise characteristics, it overcomes fundamental training failures caused by asymmetric noise schedules. The approach bridges the gap between high-quality generative modeling and practical deployment, fostering advances in autonomous systems, virtual environments, and interactive AI. Long-term, it paves the way for end-to-end, efficient multimodal generative systems that can operate on edge devices, reducing hardware costs and expanding application scenarios.
Technical Contribution
The key technical contributions include: 1) identifying the structural failure mode caused by uniform consistency functions under asymmetric noise schedules; 2) deriving a principled framework for modality-specific consistency functions based on structural analysis; 3) proving that linear gradient scaling is optimal in low-noise regimes, and designing corresponding functions for actions and videos; 4) achieving single-step joint inference with shared transformer architecture, significantly reducing latency. These innovations provide theoretical guarantees and practical improvements over existing methods, enabling fast, high-fidelity multimodal generation.
Novelty
This research is the first to explicitly address the challenge of asymmetric noise schedules in joint video-action diffusion models. By introducing modality-aware consistency functions grounded in structural analysis, it overcomes the gradient vanishing problem that plagues naive distillation approaches. Unlike prior work focused on single-modality acceleration, this framework dynamically matches the consistency function to each modality’s noise regime, representing a fundamental shift in how multimodal diffusion models are trained and deployed. This novelty opens new avenues for efficient, scalable multimodal generative modeling.
Limitations
- The approach relies on pre-trained model structures and assumptions about noise schedule separability, which may limit generalization to more complex or unstructured scenarios.
- Tuning the modality-specific parameters requires careful analysis and may not transfer seamlessly across different architectures or tasks.
- Performance may degrade under extreme noise conditions or in highly dynamic environments, necessitating further robustness enhancements.
Future Work
Future research will explore adaptive modulation of consistency functions to handle more diverse noise regimes, integrate reinforcement learning for dynamic schedule adjustment, and extend the framework to multi-modal interactions involving language, tactile, and other sensory inputs. Additionally, efforts will focus on improving robustness and generalization, enabling deployment in more complex, real-world settings.
AI Executive Summary
The rapid advancement of diffusion models has revolutionized generative AI, producing high-quality images, videos, and complex multimodal outputs. However, their iterative inference process remains computationally intensive, limiting real-time applications such as robotic control. Existing acceleration techniques like distillation often struggle in multimodal settings due to the asymmetric noise schedules used for different data streams. This paper introduces Flash-WAM, a novel framework that employs modality-aware consistency functions to address this challenge. By analyzing the structural properties of the consistency function family, the authors design tailored functions for video and action streams, ensuring effective gradient signals across diverse noise regimes. This approach enables collapsing multiple denoising steps into a single step per modality, dramatically reducing latency. Experiments on LingBot-VA demonstrate a reduction from 8 seconds to 0.35 seconds per chunk, a 23× speedup, while maintaining high task success rates (85.5% RoboTwin 2.0, 95.7% LIBERO). The method also achieves a 60% success rate on real-world robot tasks, outperforming naive distillation. These results signify a major step toward real-time, high-fidelity multimodal generative models for robotics and beyond. The framework’s theoretical grounding and practical validation suggest broad applicability, with future directions including adaptive schedule tuning, robustness improvements, and extension to more complex multi-sensory systems. Overall, Flash-WAM bridges the gap between high-quality generative modeling and real-time deployment, opening new horizons for autonomous AI systems.
Deep Analysis
Background
近年来，深度生成模型，尤其是扩散模型，在图像、视频生成中取得巨大成功。代表性工作如DDPM（Denoising Diffusion Probabilistic Models）和其变体，已广泛应用于高质量内容生成。然而，这些模型的推理过程依赖大量迭代，导致速度缓慢，难以满足实时需求。随着机器人自主控制和虚拟环境的兴起，快速生成和决策成为关键。近年来，逐步蒸馏技术被引入以加速推理，但在多模态任务中，视频和动作的噪声调度差异带来了新挑战。现有方法多关注单一模态，难以应对多模态的噪声调度不匹配，导致训练崩溃或性能下降。
Core Problem
多模态扩散模型（WAMs）在视频和动作的联合生成中面临噪声调度不匹配的问题。视频模态采用高噪声调度以捕获空间结构，动作模态则偏向低噪声以保证精度。这种差异导致一致性蒸馏在训练中出现梯度信号消失，模型难以同时优化两个模态，严重制约推理速度和任务性能。解决这一问题对于实现机器人端到端的实时控制具有重要意义。
Innovation
本文的创新点包括：1）提出模态感知的逐步蒸馏策略，根据噪声调度的不同，设计匹配的参数化一致性函数；2）通过结构分析，证明低噪声区域应采用线性梯度缩放策略，而高噪声区域采用方差保持策略，有效缓解梯度信号消失问题；3）在LingBot-VA模型上实现单步推理，将推理时间从8秒缩短到0.35秒，满足实时控制需求。这一突破为多模态扩散模型的快速推理提供了新思路。
Methodology
- �� 识别多模态扩散中噪声调度不匹配导致的梯度信号失衡问题；• 设计模态感知的匹配一致性函数，动作模态采用线性梯度缩放（b(σ) = −σ），视频模态采用方差保持（cskip和cout参数化）；• 结合结构分析，确保在不同噪声域中一致性函数的梯度缩放满足理论最优；• 训练过程中，分别对视频和动作模态引入不同的损失函数，利用单一模型实现多模态的联合蒸馏；• 在LingBot-VA模型上进行大规模实验，验证推理速度和任务成功率。
Experiments
采用RoboTwin 2.0和LIBERO两个基准，评估模型在多任务中的表现。对比Naive蒸馏和其他加速技术，重点关注推理时间和成功率。硬件方面，使用NVIDIA L40S GPU，设定500ms为实时控制阈值。通过不同的噪声调度参数，验证模态感知策略的有效性。还在实际机器人平台测试，确保方法的实用性。实验还包括消融分析，验证不同一致性函数的影响。
Results
Flash-WAM将推理步骤从50步缩减到1步，推理时间从8秒降至0.35秒，速度提升23倍，任务成功率在RoboTwin 2.0达85.5%，在LIBERO达95.7%，在实际机器人上达60%，优于Naive蒸馏的24%。这些结果验证了模态感知策略在缓解梯度消失和提升速度方面的有效性。
Applications
该技术适用于自主机器人、虚拟仿真和交互式AI系统，特别是在需要快速响应和高精度的场景。通过单步推理，降低硬件成本，提升系统反应速度，为工业自动化、智能制造提供技术支持。未来结合强化学习和多模态交互，有望推动自主系统的智能化升级。
Limitations & Outlook
目前方法依赖预训练模型结构，可能在极端噪声或复杂场景中表现不足。模态感知策略需调优，泛化能力有限。在极端噪声环境下性能仍有待提升，未来需增强模型鲁棒性和适应性。
Plain Language Accessible to non-experts
想象你在厨房做饭，有两个任务：一个是切菜，一个是煮汤。切菜需要细心和精准，不能太快也不能太慢；煮汤则需要长时间的火候控制。以前用一样的刀和火调节，效率不高。现在，有了聪明的厨具，能根据任务自动调整刀锋和火力，让切菜和煮汤都变得又快又好。这个厨具就像Flash-WAM，用不同策略对待视频和动作两个“任务”，让机器人在做事时既快又准，像专业厨师一样高效。
ELI14 Explained like you're 14
想象你在学校参加比赛，有两个项目：拼图和跑步。拼图需要耐心和观察，跑步要速度。以前用一样的方法训练，效果不理想。现在，有个聪明的教练，他会根据每个项目的特点，给你不同的训练方案：拼图时慢慢观察，跑步时冲刺。这样一来，你在比赛中就能表现得更棒，赢得更多奖牌。这个教练就像Flash-WAM，用不同的方法帮机器人更快、更准地完成任务。
Glossary
Diffusion Model (扩散模型)
一种生成模型，通过逐步添加和去除噪声，逐渐生成高质量数据。
用于视频和动作的联合生成。
Consistency Distillation (一致性蒸馏)
一种加速扩散模型推理的技术，通过训练模型在不同噪声水平保持一致性。
核心技术，用于单步推理。
Noise Schedule (噪声调度)
控制扩散过程中噪声强度变化的策略，影响模型训练和生成效果。
不同模态采用不同调度策略。
Modality-aware (模态感知)
根据不同模态的特性设计匹配策略，提升多模态模型效率。
本文的关键思想。
Gradient Scaling (梯度缩放)
调整训练中梯度大小，改善学习信号传递。
分析一致性函数族的基础。
Open Questions Unanswered questions from this research
- 1 多模态扩散模型在极端噪声环境下的性能仍需验证，尤其在复杂场景中模型的鲁棒性和泛化能力不足。
- 2 如何进一步优化模态感知策略，使其适应不同模型架构和任务需求，是未来研究的重要方向。
Abstract
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.