Flash-WAM: Destilación Consciente de la Modalidad para Modelos de Acción del Mundo

Resumen

Los modelos mundo-acción (WAMs, por sus siglas en inglés) generan conjuntamente video futuro y acciones robóticas mediante difusión iterativa, logrando un alto rendimiento en benchmarks de manipulación, pero requiriendo decenas de pasos de eliminación de ruido, un costo que impide el control en tiempo real. La destilación de pasos ha surgido como el remedio natural, pero los métodos preexistentes fallan en el contexto conjunto video-acción porque los flujos de video y acción utilizan diferentes programaciones de ruido desplazadas por SNR y llegan al entrenamiento con distribuciones de ruido marginales sustancialmente diferentes, una asimetría que los métodos de destilación de una sola modalidad no pueden manejar. Presentamos Flash-WAM, un marco de destilación de pasos consciente de la modalidad inspirado en la destilación de consistencia que selecciona la función de consistencia para cada modalidad para que coincida con su régimen de ruido: una parametrización de escalado de gradiente lineal para el régimen de bajo ruido del flujo de acciones, junto con una parametrización de preservación de varianza para el régimen de alto ruido del flujo de video, basada en un análisis estructural de la familia de funciones de consistencia que caracteriza el escalado de gradiente alcanzable bajo la condición de frontera de consistencia. Implementado en LingBot-VA, Flash-WAM comprime la inferencia a un solo paso en cada modalidad. En RoboTwin 2.0, esto reduce la latencia por fragmento de 8.1 segundos a 348 ms en NVIDIA L40S, una aceleración de 23 veces que permite la inferencia en tiempo real. Flash-WAM preserva el éxito de las tareas en benchmarks de simulación (85.5% RoboTwin 2.0, 95.7% LIBERO) y recupera sustancialmente el rendimiento en el mundo real (60% promedio en un robot humanoide Unitree G1), mientras que la destilación de consistencia ingenua cae al 24% con el mismo presupuesto de pasos.

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.