Flash-WAM: Дистилляция с учетом модальности для мировых моделей действий

Аннотация

Модели «мир-действие» (World-action models, WAMs) совместно генерируют будущее видео и действия робота через итеративную диффузию, демонстрируя высокую производительность на эталонных тестах манипуляций, но требуя десятков шагов шумоподавления — цена, исключающая управление в реальном времени. Дистилляция шагов стала естественным решением, однако стандартные методы нарушаются в совместной постановке видео-действие, поскольку видеопоток и поток действий используют разные графики шума со сдвигом по SNR и на этапе обучения характеризуются существенно различными маргинальными распределениями шума — асимметрией, которую не могут учесть методы дистилляции для одной модальности. Мы представляем Flash-WAM — фреймворк дистилляции шагов с учётом модальности, вдохновлённый дистилляцией согласованности, который выбирает функцию согласованности для каждой модальности в соответствии с её режимом шума: параметризацию с линейным масштабированием градиента для низкошумового режима потока действий в сочетании с параметризацией с сохранением дисперсии для высокошумового режима видеопотока, что обосновано структурным анализом семейства функций согласованности, характеризующим достижимое масштабирование градиента при граничном условии согласованности. Реализованный на LingBot-VA, Flash-WAM сжимает вывод до одного шага в каждой модальности. На RoboTwin 2.0 это снижает задержку на фрагмент с 8,1 секунды до 348 мс на NVIDIA L40S — 23-кратное ускорение, обеспечивающее вывод в реальном времени. Flash-WAM сохраняет успешность выполнения задач на эталонных тестах моделирования (85,5% на RoboTwin 2.0, 95,7% на LIBERO) и существенно восстанавливает производительность в реальном мире (в среднем 60% на гуманоидном роботе Unitree G1), тогда как наивная дистилляция согласованности при том же бюджете шагов падает до 24%.

English

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 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.