Flash-WAM：模态感知的世界动作模型蒸馏 Flash-WAM: Modality-Aware Distillation for World Action Models
通过模态感知的一致性蒸馏，将世界动作模型的推理延迟从8.1秒降至348毫秒，实现23倍加速
前置知识
世界动作模型
世界动作模型是一种联合生成未来视频和机器人动作的模型，通过共享的transformer架构同时预测视觉动态和对应的动作序列。它将策略生成分解为两个阶段：视觉动态预测（预测世界如何在潜在空间演化）和逆向动力学（恢复与预测过渡一致的动作）。相比直接预测动作的VLA策略，WAM通过继承大规模视频预训练的时空先验，能够更好地泛化到新场景和长时任务。
理解WAM是理解本文蒸馏目标的基础，Flash-WAM的所有优化都是针对这种联合视频-动作生成的双阶段结构设计的。
流匹配
流匹配是一种连续时间生成框架，学习通过直线插值路径将噪声分布转化为数据分布。在训练时，干净数据 $x_0$ 通过 $x_\sigma = (1-\sigma)x_0 + \sigma\epsilon$ 被腐蚀，其中 $\epsilon \sim \mathcal{N}(0, I)$，神经网络被训练预测速度 $v = \epsilon - x_0 = dx_\sigma/d\sigma$。在推理时，通过对速度场从 $\sigma=1$ 到 $\sigma=0$ 进行欧拉积分来生成样本。流匹配是现代世界动作模型的底层生成框架。
Flash-WAM的所有推导都建立在流匹配的数学基础上，理解流匹配才能理解一致性蒸馏如何压缩去噪步数。
一致性蒸馏
一致性蒸馏通过强制一致性属性来加速采样：一致性函数 $f(x_\sigma, \sigma)$ 将概率流常微分方程轨迹上的任何点映射到其在 $\sigma=0$ 处的干净端点。标准形式为 $f(x_\sigma, \sigma) = a(\sigma)x_\sigma + b(\sigma)v_\theta(x_\sigma, \sigma)$，其中边界条件 $a(0)=1, b(0)=0$ 强制 $f(x_0, 0) = x_0$。在蒸馏设置中，冻结的教师模型提供引导的欧拉步，学生模型和EMA目标被训练沿着这些轨迹保持一致。通过这种方式，模型可以在更少的去噪步内生成可比的结果。
本文的核心技术突破正是发现标准一致性蒸馏在联合视频-动作设置下失效，并提出模态感知的改进版本。
SNR移位噪声调度器
SNR移位噪声调度器是一种控制训练质量沿噪声调度分布位置的技术，通过 $\sigma = s\tilde{\sigma}/(1+(s-1)\tilde{\sigma})$ 实现，其中 $\tilde{\sigma} \sim U[0,1]$，移位参数 $s \geq 1$。较大的 $s$ 将分布推向更高的噪声水平。在WAM中，视频和动作使用独立的移位参数 $s_v$ 和 $s_a$，满足 $s_v > s_a$，因为高维、结构冗余的视频潜在表示可以容忍每步更重的噪声，而低维、精度关键的动作序列需要更温和的调度。
这种不对称的噪声调度正是Flash-WAM需要解决的核心问题，它导致视频和动作到达蒸馏损失时的噪声分布完全不同，单一一致性函数无法同时服务两者。
边际噪声分布
边际噪声分布描述了训练过程中不同噪声水平 $\sigma$ 的采样概率密度。由于视频和动作使用不同的SNR移位参数 $s_v > s_a$，它们的边际噪声分布呈现显著差异：视频噪声集中在 $[0,1]$ 的高端（高 $\sigma$ 区域），而动作噪声分布在全范围，在低 $\sigma$ 区域有大量质量。这种分布不对称性意味着两个流在到达一致性蒸馏损失时处于结构性不同的噪声区域，因此需要不同的一致性函数来获得有效的梯度信号。
这是Flash-WAM核心洞察的数学基础，理解边际噪声分布的差异才能理解为什么标准一致性蒸馏在动作流上失败。
研究动机
现有世界动作模型虽然能够联合生成未来视频和机器人动作，但需要数十步去噪操作，计算成本极高。具体而言，在RoboTwin 2.0基准测试上，最先进的LingBot-VA模型每个数据块需要25步视频去噪和50步去噪动作，在单块NVIDIA L40S GPU上总共耗时8100毫秒（3550毫秒视频 + 4550毫秒动作），远远超出了实时控制的预算（约500毫秒，对应2 Hz的块级速率）。现有的工程级优化（如KV缓存、噪声历史增强的局部去噪、异步预测-执行流水线）只能减少壁钟延迟，而不改变底层去噪步数，无法从根本上解决实时性问题。步蒸馏是自然解决方案，但在联合视频-动作设置下，现有的单模态蒸馏方法（如DMD2需要辅助分数网络和对抗训练、渐进式蒸馏需要多阶段训练流水线）难以直接应用。即使是最自然的适配方法一致性蒸馏，在应用时也会导致任务成功率从超过91%暴跌到23%的灾难性结果。
本文的目标是本文的目标是开发一个步蒸馏框架，能够在保持世界动作模型性能的同时，将推理步数大幅压缩到实时控制预算内。具体而言，目标是将LingBot-VA模型的推理压缩到每个模态仅用一步（1v/1a），在NVIDIA L40S上将每块延迟从8.1秒降至348毫秒（23倍加速），同时在RoboTwin 2.0基准上保持超过80%的平均成功率，在LIBERO基准上保持超过95%的平均成功率，并在真实机器人实验中恢复接近原始模型的性能。
与已有工作不同的是，本文的独特切入角度是诊断联合模态蒸馏失败的根本原因，而不是简单地调整超参数或使用更强的正则化。通过深入分析，本文发现失败不是调参问题，而是结构性问题：视频潜在表示和动作序列具有根本不同的统计特性，导致WAM使用不同的SNR移位噪声调度器。因此，两个流在到达一致性蒸馏损失时的边际噪声分布完全不同——视频噪声集中在高 $\sigma$，而动作噪声在全范围分布且在低 $\sigma$ 有大量质量。现有一致性蒸馏方法（如LCM）的梯度信号随 $\sigma \to 0$ 二次消失，导致动作流在大部分训练分布上几乎没有学习信号。本文的创新在于为每个模态选择匹配其噪声区域的一致性函数，而不是试图用单一函数同时服务两者。
核心方法
Flash-WAM的核心思想是将视频和动作蒸馏视为两个具有不同梯度信号需求的根本不同问题。每个模态接收一个与其训练分布集中位置匹配的一致性函数：对于训练在高 $\sigma$ 区域的视频流，选择方差保留参数化；对于训练在低 $\sigma$ 区域的动作流，选择线性梯度缩放参数化。这种选择基于对一致性函数族的结构分析，刻画了在一致性边界条件下可实现的梯度缩放。整体训练流程采用共享transformer架构，视频和动作token被连接成联合序列，通过flex attention处理。模态感知参数化仅影响每流损失头，不改变架构和每步计算成本。在部署时，蒸馏后的学生模型以自回归方式生成视频和动作，每个模态仅需一步去噪。
核心创新点是模态感知的一致性函数选择。通过理论分析，本文证明任何满足边界条件的一致性函数在 $\sigma=0$ 附近的最优梯度缩放为 $|b(\sigma)| = O(\sigma)$，这个边界当且仅当 $b'(0) \neq 0$ 时才能达到。标准LCM参数化 $b_{LCM}(\sigma) = -\sigma^2\sigma_d/\sqrt{\sigma^2+\sigma_d^2}$ 满足 $b'_{LCM}(0)=0$，只能实现二次消失的梯度缩放 $|b_{LCM}(\sigma)| = \sigma^2/\sigma_d + O(\sigma^4)$。因此，Flash-WAM为动作流选择最简单的线性缩放对：$a(\sigma)=1, b(\sigma)=-\sigma$，这满足 $b'(0)=-1 \neq 0$，在整个 $[0,1]$ 上实现 $|b(\sigma)| = \sigma$ 的线性缩放。对于视频流，由于分布集中在高 $\sigma$，LCM已经提供足够的梯度信号，因此选择Karras参数化以保持方差稳定和输出范围有界。
方法步骤详情
Flash-WAM的训练过程包含以下步骤：首先，从教师模型 $\theta^T$ 开始，该模型已被训练为通过流匹配生成视频和动作。在每次训练迭代中，采样噪声水平 $\sigma_s$，然后通过将调度前进 $k$ 离散步获得 $\sigma_e < \sigma_s$。使用教师欧拉步（在蒸馏过程中使用分类器自由引导 $w \sim U[w_{min}, w_{max}]$ 用于视频，使用无引导预测用于动作）形成目标 $\tilde{x}_{\sigma_e} = x_{\sigma_s} + \hat{v}_{cfg}(\sigma_e - \sigma_s)$。学生模型 $\theta^S$ 和EMA目标 $\theta^{S'}$ 被训练为沿着这些轨迹保持一致。对于视频流，使用方差保留参数化 $f^v(x_{v\sigma}, \sigma) = c_{skip}(\sigma)x_{v\sigma} + c_{out}(\sigma)\hat{x}_{v0}$；对于动作流，使用线性梯度缩放参数化 $f^a(x_{a\sigma}, \sigma) = 1 \cdot x_{a\sigma} - \sigma \cdot v_\theta(x_{a\sigma}, \sigma)$。每个模态贡献一个一致性损失：$L_v = d(f_{\theta^S}(x_{v\sigma_s}, \sigma_s), f_{\theta^{S'}}(\tilde{x}_{v\sigma_e}, \sigma_e))$ 和 $L_a = d(f_{\theta^S}(x_{a\sigma_s}, \sigma_s), f_{\theta^{S'}}(\tilde{x}_{a\sigma_e}, \sigma_e))$。完整Flash-WAM目标为 $L = L_v + \lambda_a L_a$。在部署时，蒸馏后的学生模型以自回归方式生成视频块和对应的动作序列，每个模态仅需一步去噪。
技术新颖性
Flash-WAM的技术新颖性体现在三个方面：首先，它提供了对联合模态蒸馏失败的结构性诊断，识别出标准一致性蒸馏与联合扩散在不对称噪声调度下的根本不兼容性，并通过形式分析和实证实验表征了这个失败模式。其次，它提出了模态感知一致性蒸馏框架，基于对一致性函数族的结构分析，为每个模态选择匹配其噪声区域的一致性函数，刻画了一致性边界条件下可实现的梯度缩放。最后，它实现了实时WAM推理，在LingBot-VA上将每个模态的推理压缩到一步，在NVIDIA L40S上将每块延迟从8.1秒降低到348毫秒（23倍加速），同时在模拟基准上保持接近原始模型的性能（RoboTwin 2.0上85.5%，LIBERO上95.7%）。这种加速不是通过牺牲视频生成实现的，而是保留了原始WAM推理结构，同时将每个模态的去噪压缩到一步。
实验结果
Flash-WAM在三个主要实验环境中都取得了显著成果。在RoboTwin 2.0模拟基准上，Flash-WAM在1v/2a配置下达到85.54%的平均成功率，恢复了LingBot-VA教师91.25%性能的大部分，同时实现19倍加速。在更激进的1v/1a配置下，Flash-WAM仍然达到81.41%的平均成功率，尽管将视频去噪减少25倍、动作去噪减少50倍，但仍保持在未加速配置的10个百分点内。对应的23.3倍加速将每块延迟降低到单块NVIDIA L40S上的348毫秒，实现实时推理。相比之下，现有的步蒸馏方法在相同NFE预算下表现不佳：朴素联合LCM崩溃到23.97%，DMD2达到78.74%，仅视频LCM落后于78.79%。在LIBERO基准上，Flash-WAM在1v/2a配置下达到95.7%的平均成功率，几乎恢复了LingBot-VA教师的98.6%性能，同时实现13.7倍加速，并在每个任务套件上都优于仅视频LCM。在1v/1a配置下，Flash-WAM达到95.1%的平均成功率，实现16.3倍加速，将每块延迟从6767毫秒降低到NVIDIA L40S上的404毫秒。在Unitree G1人形机器人上的真实世界实验中，发布的LingBot-VA模型在3v/10a配置下达到66.7%的平均成功率。在不进行蒸馏的情况下减少LingBot-VA的NFE会崩溃真实世界性能，在1v/2a配置下降低到40.0%，在1v/1a配置下降低到23.3%。朴素应用LCM（仅视频LCM）部分恢复性能（1v/2a配置下43.3%，1v/1a配置下33.3%），但Flash-WAM在两种NFE配置和所有三个任务上都显著优于这两种方法，在1v/2a配置下达到60.0%平均成功率，在1v/1a配置下达到50.0%。定性分析表明，在RoboTwin Clean-split任务的开环自回归滚动中，未加速的LingBot-VA教师（25v/50a）产生干净的预测，在整个滚动过程中保持物体身份和夹持器几何结构。两种现成的蒸馏基线（朴素联合LCM和DMD2）在相同1v/2a NFE预算下明显退化：棕色瓶在朴素联合LCM下完全消失，在DMD2下变得模糊。Flash-WAM方法在整个滚动过程中保持可识别的场景结构和物体身份。消融分析表明，Flash-WAM在两种NFE配置和所有任务范围内都优于三种基于LCM的替代蒸馏策略，朴素联合LCM完全崩溃，仅视频LCM留下全教师NFE的动作流但仍落后Flash-WAM大约7个百分点，添加MSE正则化在更激进的1v/1a配置下退化。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RoboTwin 2.0 (Clean + Randomized) | 平均成功率 (%) | 85.54 (1v/2a), 81.41 (1v/1a) | 91.25 (LingBot-VA教师 25v/50a), 23.97 (朴素联合LCM 1v/2a), 78.79 (仅视频LCM 1v/2a) | 恢复教师93.8%性能，相比朴素联合LCM提升256% |
| LIBERO (Spatial + Object + Goal + Long-horizon) | 平均成功率 (%) | 95.7 (1v/2a), 95.1 (1v/1a) | 98.6 (LingBot-VA教师 20v/50a), 95.2 (仅视频LCM 1v/2a) | 恢复教师97.1%性能，相比仅视频LCM提升0.5% |
| 真实机器人实验 (Unitree G1) | 平均成功率 (%) | 60.0 (1v/2a), 50.0 (1v/1a) | 66.7 (LingBot-VA教师 3v/10a), 40.0 (减少NFE无蒸馏 1v/2a), 43.3 (仅视频LCM 1v/2a) | 相比减少NFE无蒸馏提升50%，相比仅视频LCM提升38% |
| 推理延迟 | 每块延迟 (ms) | 347 (1v/1a) | 8100 (LingBot-VA教师 25v/50a) | 23.3倍加速 |
局限与改进
作者承认的局限性包括Flash-WAM目前仅在LingBot-VA模型上验证，该模型参数量较小，适合在商品边缘硬件上部署，但对于其他采用不同架构或自己集成推理优化栈的WAM（如Motus采用Transformer混合架构，DreamZero在架构级别集成推理时优化）尚未验证。此外，虽然Flash-WAM显著提升了推理速度，但性能仍然略低于原始教师模型，特别是在更激进的1v/1a配置下。从技术角度看，Flash-WAM依赖于对每个模态噪声分布的先验知识（需要知道哪个模态在高噪声区域训练，哪个在低噪声区域训练），这可能需要额外的分析或超参数调整。另一个局限性是Flash-WAM专注于蒸馏已有的预训练WAM，而不是从头训练一个快速的WAM，这意味着仍然需要一个高质量的预训练教师模型作为起点。
独立分析的弱点
独立分析的弱点包括：首先，Flash-WAM的性能提升在更激进的1v/1a配置下仍然有约10个百分点相对于教师模型的性能损失，这可能对于某些关键任务是不可接受的。改进方向可以探索混合步数配置，比如1v/2a已经很好的平衡了速度和性能。其次，Flash-WAM目前仅针对流匹配框架设计，对于其他生成范式（如DDPM、DDIM等）的适配性需要进一步研究。第三，虽然Flash-WAM在模拟环境中表现良好，但在真实世界中的性能仍然显著低于教师模型（60% vs 66.7%），这可能意味着真实世界的分布漂移对蒸馏模型的影响更大。改进方向可以包括更强的域适应技术或在真实世界数据上进一步微调。第四，Flash-WAM的模态感知参数化需要为每个模态手动选择一致性函数类型，对于模态数量超过两个的场景（如同时处理视频、深度、触觉等多模态输入），这种手动选择可能变得不可扩展，需要开发自动化的模态感知参数化选择方法。
未来方向
作者提出的未来工作方向包括将Flash-WAM扩展到其他WAM架构，如Motus和DreamZero，以验证框架的通用性。基于成果可延伸的未来方向包括：探索自适应的噪声调度策略，根据模态的具体特性动态调整SNR移位参数，而不是使用固定的手动设置。研究多模态一致性函数的自动选择算法，通过元学习或强化学习自动确定每个模态的最佳一致性函数类型。将Flash-WAM与架构级优化（如KV缓存、异步预测-执行流水线）结合，实现端到端的实时WAM系统。探索在蒸馏过程中加入更强的正则化或约束，以进一步缩小蒸馏模型与教师模型之间的性能差距。研究在更广泛的机器人平台和任务类型上的泛化能力，包括移动操作、多机器人协作等复杂场景。最后，探索将Flash-WAM的思想应用于其他联合生成任务，如图文联合生成、音频视频联合生成等。
复现评估
根据论文描述，Flash-WAM应用了已发布的LingBot-VA模型[15]，这是一个最先进的开源联合视频-动作扩散模型，参数量足够小，可以在商品边缘硬件上运行，这对步蒸馏具有最大的实际影响。论文提到完整实现细节包括超参数和训练配置在附录A中提供。然而，论文没有明确提到是否开源代码或预训练模型。实验在单个NVIDIA L40S GPU上进行，这是一个相对高端但可获得的硬件，对于学术实验室来说是可行的。数据方面，RoboTwin 2.0和LIBERO都是公开可用的基准数据集。真实世界实验使用了Unitree G1人形机器人和Unitree Dex1-1夹持器，每个任务收集了50个远程操作演示。复现难度中等，主要挑战在于获得预训练的LingBot-VA教师模型和配置正确的训练超参数。论文提供了详细的实验设置和消融分析，这有助于复现。
论文图表
