Technische Samenvatting: MobileWan

Probleemstelling

Recente vooruitgang in video-diffusie is gedreven door het schalen van transformer-gebaseerde architecturen (DiT) naar miljarden parameters, wat de visuele getrouwheid en bewegingscoherentie aanzienlijk heeft verbeteren. Echter, mobiele videogeneratie loopt substantieel achter op server-schaal systemen vanwege strikte beperkingen op de piekgeheugenbehoefte, de kwadratische complexiteit van self-attention, decoder-overhead en end-to-end latentie. Bestaande mobiele modellen zijn doorgaans beperkt tot kleine backbones (0,4–1,8 miljard parameters), wat onvermijdelijk de generatiekwaliteit beperkt en een duidelijke kloof creëert tussen mobiele en server-schaal prestaties. De kernuitdaging is dat knelpunten niet beperkt zijn tot één enkele module; het simpelweg vervangen van softmax attention door lineaire attention vereist vaak het volledig opnieuw trainen of veroorzaakt een expressiviteitskloof, terwijl hybride benaderingen de backbone nog niet volledig hebben omgezet naar een laag-geheugen recurrente architectuur die geschikt is voor mobiele implementatie.

Methodologie

De auteurs stellen MobileWan voor, een systeem dat aantoont dat een server-schaal 5B-parameter video-diffusietransformer (gebaseerd op Wan2.2-5B) efficiënt kan worden ingezet op geheugenbeperkte mobiele hardware door een combinatie van recurrente herformulering en gestructureerde compressie. De methodologie bestaat uit vier primaire optimalisatiestadia:

1. Leerbare Head Pruning

In plaats van heuristische block pruning, waar de auteurs ontdekten dat Wan2.2-5B zeer gevoelig voor is, introduceren zij een methode voor learnable head pruning.

• Mechanisme: Een leerbare scalaire gate wordt geïntroduceerd voor elke self-attention head, die de bijdrage ervan moduleert. Deze gates worden end-to-end geoptimaliseerd met een sparsiteitsregularisator en een noise-biased leerstrategie.
• Trainingsstrategie: Het model wordt getraind met een hoog-ruis-gebiaste distributie (waarbij ruis $\sigma \sim \text{LogitNormal}(1.5, 1)$ wordt getrokken) tijdens de gate-leern fase. Dit legt de nadruk op globale structuur en beweging, waardoor de pruning-beslissingen robuuster worden. Het proces omvat een drie-fasen schema: zachte gate-exploratie, temperatuur-annealing naar binaire waarden, en finale fine-tuning met een gefixeerde geprunte architectuur.
• Resultaat: Deze aanpak presteert beter dan heuristische pruning, met name bij agressieve pruning-ratio's, wat resulteert in meetbare FLOP-reducties met beperkt kwaliteitsverlies.

2. Recurrente Distillatie

Om de kwadratische geheugencomplexiteit van self-attention aan te pakken, zetten de auteurs de vooraf getrainde transformer om in een chunk-wise autoregressief model.

• Hybride Attention: Zij maken gebruik van een hybride attention-mechanisme dat causale lineaire attention combineert met lokale softmax attention. Het model let op een kleine lokale token-set via softmax voor onmiddellijke context en gebruikt lineaire attention voor historische context.
• RNN Herformulering: Tijdens de inferentie staat de causale lineaire attention het model toe om als een Recurrent Neural Network (RNN) te opereren. Het behoudt compacte lopende aggregaten (state vectoren $s_t$ en $z_t$) voor voorgaande chunks in plaats van de volledige attention-matrix op te slaan.
• Voordeel: Dit vermindert de geheugenvoetafdruk zodat deze constant blijft, ongeacht de sequentielengte, wat de generatie van lange video's mogelijk maakt op apparaten met beperkt RAM.

3. Stap-distillatie

Om de latentie te verlagen, passen de auteurs Distribution Matching Distillation (DMD) en de variant D-DMD toe.

• Aanpak: Zij distilleren het multi-step teacher-model naar een few-step student-model.
• Optimalisatie: Experimenten vergeleken standaard DMD, D-DMD (het splitsen van de loss in Distribution Matching en Classifier-Free Guidance augmentatie termen) en DMD2 (dat een adversariële loss incorporeert).
• Selectie: Hoewel DMD2 met 2 stappen de hoogste dynamische graad en VBench-scores behaalde, merkten de auteurs op dat D-DMD visueel aantrekkelijke samples produceerde met een iets lagere dynamische graad. Het uiteindelijke systeem maakt gebruik van een configuratie die geoptimaliseerd is voor 3 denoising stappen (NFEs) om een balans te vinden tussen kwaliteit en snelheid.

4. Decoder Optimalisatie

Standaard Video VAE's gebruiken causale 3D-convoluties die computationeel duur zijn.

• Modificatie: De auteurs bouwen voort op de efficiënte LightX2V decoder, maar modificeren de architectuur om het causale look-back venster te verlengen.
• Mechanisme: In plaats van te conditioneren op onmiddellijke voorgaande frames, wordt de voorspelling van frame $F_t$ geconditioneerd op een langere temporele context (tot $F_{t-4}$).
• Resultaat: Dit vermindert temporele artefacten en flikkeringen, terwijl de efficiëntie van 2D-convoluties behouden blijft, wat de reconstructiekwaliteit (PSNR en VBench) aanzienlijk verbetert ten opzichte van andere efficiënte decoders.

5. Kwantisatie Pipeline

Om implementatie op de Snapdragon® 8 Gen 5 NPU mogelijk te maken, ondergaat het model een twee-fasen kwantisatieproces:

1. Post-Training Quantization (PTQ): Comprimeert gewichten naar 8-bit.
2. Quantization-Aware Training (QAT): Een twee-fasen distillatieproces (eerst blok-gewijs, daarna end-to-end) zorgt ervoor dat het gekwantiseerde model overeenkomt met het full-precision teacher-model, waardoor de nauwkeurigheid die verloren ging tijdens PTQ wordt hersteld.

Belangrijkste Resultaten

MobileWan behaalt de volgende prestatie-indicatoren op een mobiel apparaat (Snapdragon® 8 Gen 5 NPU):

• Generatie Specificaties: Genereert 5-seconden video's op 480×832 resolutie, 16 FPS.
• Latentie: End-to-end latentie van 20 seconden (inclusief 6,65 seconden voor de DiT denoising stap).
• Kwaliteit: Behaalt een VBench score van 83,79, wat vergelijkbaar is met het server-schaal Wan2.2 model (83,12) en de vorige mobiele state-of-the-art (Neodragon: 81,61) significant overtreft.
• Geheugen: De recurrente herformulering vermindert de piekgeheugenvereiste van "Out of Memory" (OOM) op mobiel naar 9,9 GB voor de DiT en 1,9 GB voor de decoder.
• Gebruikersstudie: In een blind gebruikersonderzoek gaven deelnemers in 80% van de gevallen de voorkeur aan MobileWan boven Neodragon. In vergelijkingen met het gefinetunede server Wan2.2 model gaf de helft van de gevallen de voorkeur aan MobileWan niet, maar werd het in 22% van de gevallen wel verkozen, wat aangeeft dat het competitief is met het server-model in de helft van de gevallen.

Betekenis en Claims

Het artikel beweert de kwaliteitskloof tussen mobiele en server-schaal videogeneratie te overbruggen door de aanname uit te dagen dat mobiele videogeneratie kleine modellen vereist.

• Schaalbaarheid: Het demonstreert dat een 5B-parameter video DiT lokaal kan worden ingezet op commerciële mobiele hardware door middel van fijnmazige compressie en recurrente herformulering.
• Efficiëntie vs. Kwaliteit: Het werk bewijst dat agressieve optimalisaties (pruning, recessie, distillatie) kunnen worden toegepast zonder extreem kwaliteitsverlies, waarbij een VBench score (83,79) wordt behaald die wedijvert met server-modellen.
• Systeemintegratie: MobileWan wordt gepresenteerd als de eerste 5B-schaal video-diffusiemodel die direct op een commercieel mobiel apparaat kan worden ingezet, waarmee een nieuwe state-of-the-art wordt gevestigd in mobiele videogeneratie.

De auteurs erkennen beperkingen, waaronder suboptimale kwaliteit voor bepaalde soorten scènes (bijv. menselijke gezichten op afstand), incidentele oververzadiging of verminderde beweging afhankelijk van de distillatiemethode, en mogelijke temporele discontinuïteiten of flikkeringen in sommige gegenereerde video's. Zij zijn van plan dit in toekomstig werk aan te pakken.