Flash-WAM: Modality-Aware Distillation for World Action Models

Technische Samenvatting: Flash-WAM

Probleemstelling

World Action Models (WAMs) vormen een veelbelovende richting voor embodied AI, waarbij toekomstige videoframes en robotacties gezamenlijk worden gegenereerd via iteratieve diffusieprocessen. Hoewel deze modellen sterke prestaties leveren op manipulatie-benchmarks, creëert hun afhankelijkheid van tientallen denoising-stappen voor zowel de video- als de actiestroom een prohibitieve computationele kost. Zo vereist het state-of-the-art LingBot-VA model bijvoorbeeld 8,1 seconden per chunk op een NVIDIA L40S GPU, wat real-time closed-loop controle onmogelijk maakt.

Hoewel stap-distillatie een natuurlijke oplossing biedt voor het versnellen van diffusiemodellen, schieten bestaande "off-the-shelf" methoden tekort wanneer ze worden toegepast op de gezamenlijke video-actie setting. De kern van het probleem ligt in de asymmetrische ruis-schedules die door WAMs worden gehanteerd:

• Videostromen maken gebruik van hoge Signal-to-Noise Ratio (SNR) verschuivingen, waardoor de trainingsmassa geconcentreerd is bij hoge ruisniveaus ($\sigma$).
• Actiestromen maken gebruik van lagere SNR-verschuivingen, waardoor de trainingsmassa over het volledige bereik wordt verspreid met een significante dichtheid bij lage $\sigma$.

Standaard consistency distillation-methoden (bijv. LCM) vertrouwen op een enkele consistentiefunctie. De auteurs tonen aan dat deze functies doorgaans kwadratische gradiëntverdwijning ($|b(\sigma)| \propto \sigma^2$) vertonen naarmate $\sigma \to 0$. Bij een uniforme toepassing op een gezamenlijk model leidt dit ertoe dat het gradiëntsignaal voor de actiestroom (die traint bij lage $\sigma$) verdwijnt, wat resulteert in een instorting van de nauwkeurigheid van de acties en de taaksuccesratio (dalend van >90% naar ~24% in experimenten), zelfs als de videoreconstructie intact blijft.

Methodologie: Flash-WAM

De auteurs stellen Flash-WAM voor, een modaliteit-bewust stap-distillatieframework dat ontworpen is om deze structurele incompatibiliteit op te lossen. In plaats van een enkele consistentiefunctie op beide modaliteiten toe te passen, selecteert Flash-WAM verschillende leden van de consistentiefunctie-familie die specifiek zijn afgestemd op het ruisregime van elke stroom.

1. Structurele Analyse van Gradiëntschaling

Het paper stelt vast dat de sterkte van het gradiëntsignaal in consistency distillation proportioneel is aan $|b(\sigma)|$, waarbij de consistentiefunctie wordt gedefinieerd als $f(x_\sigma, \sigma) = a(\sigma)x_\sigma + b(\sigma)v_\theta$.

• Propositie 1: Voor elke geldige consistentiefunctie is de optimale gradiëntschaling nabij $\sigma=0$ lineair ($O(\sigma)$), wat alleen haalbaar is als $b'(0) \neq 0$.
• Falen van Standaard LCM: Standaard LCM-parametrisaties voldoen aan $b'(0)=0$, wat resulteert in kwadratische schaling ($O(\sigma^2)$). Dit is voldoende voor video-training bij hoge $\sigma$, maar catastrofaal voor actie-training bij lage $\sigma$.

2. Modaliteit-bewuste Parametrisatie

Flash-WAM lost dit op door verschillende parametrisaties toe te wijzen:

• Actiestroom (Low-$\sigma$ Regime): De auteurs leiden een lineaire gradiëntschalings-parametrisatie af waarbij $a(\sigma) = 1$ en $b(\sigma) = -\sigma$. Dit zorgt ervoor dat $|b(\sigma)| = \sigma$, wat een sterk, lineair gradiëntsignaal biedt over de gehele low-noise trainingsdistributie van de actiestroom.
• Videostroom (High-$\sigma$ Regime): Voor de videostroom, die geconcentreerd is bij hoge $\sigma$, behouden de auteurs de standaard variance-preserving Karras-parametrisatie (gebruikt in LCM). Deze keuze is gemotiveerd door de noodzaak van variantiebehoud en begrensde outputbereiken bij hoge ruisniveaus, eigenschappen die minder kritisch zijn voor de laag-dimensionale, begrensde actie-targets.

3. Gezamenlijk Trainingsdoel

Het student-model wordt getraind om de consistentie-eigenschap voor beide modaliteiten simultaan te voldoen met behulp van een gecombineerde loss:
$$\mathcal{L} = \mathcal{L}_v + \lambda_a \mathcal{L}_a$$
waarbij $\mathcal{L}_v$ en $\mathcal{L}_a$ hun respectievelijke modaliteit-specifieke consistentiefuncties gebruiken. De architectuur blijft een gedeelde transformer-backbone, waarbij de modaliteit-bewuste parametrisatie alleen de loss-heads beïnvloedt, waardoor de per-stap rekenkosten gelijk blijven aan die van de teacher.

Belangrijkste Resultaten

Het framework werd geïmplementeerd op LingBot-VA, waarbij de inferentie werd gecomprimeerd naar één stap voor beide modaliteiten (1v/1a) of één videostap en twee actiestappen (1v/2a).

• Versnelling: Flash-WAM bereikt een 23× versnelling, waardoor de per-chunk latentie wordt verminderd van 8,1 seconden naar 348 ms op een NVIDIA L40S, wat real-time inferentie mogelijk maakt (onder de 500 ms budget).
• Simulatie Benchmarks:
  • RoboTwin 2.0: Bij 1v/2a bereikt Flash-WAM 85,5% succes (tegenover 91,2% van de teacher), vergeleken met 23,9% voor naive joint LCM. Bij 1v/1a behoudt het 81,4% succes.
  • LIBERO: Bij 1v/2a bereikt het een gemiddeld succes van 95,7% (tegenover 98,6% van de teacher), waarmee het video-only distillatie baselines overtreft.
• Real-world Deployment: Geëvalueerd op een Unitree G1 humanoid robot over drie manipulatietaken.
  • De ongeaccelereerde teacher (3v/10a) behaalde 66,7%.
  • Het verminderen van stappen zonder distillatie (1v/2a) liet de prestaties dalen naar 40,0%.
  • Naive Video-only LCM (1v/2a) behaalde 43,3%.
  • Flash-WAM (1v/2a) behaalde 60,0%, waarmee het de ongeaccelereerde prestaties substantieel herstelde.
• Kwalitatieve Analyse: Visualisaties laten zien dat terwijl naive distillatie zorgt voor het verdwijnen of vervagen van objecten, Flash-WAM de scène-structuur en objectidentiteit behoudt.

Betekenis en Claims

Het paper claimt dat Flash-WAM het eerste framework is dat succesvol real-time inferentie voor gezamenlijke video-actie diffusiemodellen mogelijk maakt door een structurele foutmodus aan te pakken die inherent is aan het toepassen van single-modality distillatie op asymmetrische multi-modale systemen.

• Diagnose: Het werk identificeert dat de incompatibiliteit tussen standaard consistency distillation en gezamenlijke diffusie geen hyperparameter-afstemmingsprobleem is, maar een fundamentele mismatch in de beschikbaarheid van het gradiëntsignaal veroorzaakt door asymmetrische ruis-schedules.
• Oplossing: Door de selectie van consistentiefuncties te funderen in een structurele analyse van gradiëntschaling, biedt Flash-WAM een geprincipeerde methode om de taaksucces te behouden terwijl de denoising-stappen worden teruggebracht naar één per modaliteit.
• Impact: De resultaten demonstreren dat real-time robotbesturing met WAMs haalbaar is op commodity edge-hardware zonder de generatieve capaciteiten op te offeren die WAMs superieur maken aan statische Vision-Language-Action (VLA) policies. De auteurs benadrukken dat hun aanpak de oorspronkelijke WAM-inferentiestructuur behoudt, in tegenstelling tot methoden die video-generatie volledig verwijderen.