Riepilogo Tecnico: MobileWan

Problematica

I recenti progressi nella diffusione video sono stati guidati dalla scalabilità di architetture basate su transformer (DiT) a miliardi di parametri, migliorando significativamente la fedeltà visiva e la coerenza del movimento. Tuttavia, la generazione video su dispositivi mobili rimane sensibilmente indietro rispetto ai sistemi su scala server a causa dei severi vincoli sulla memoria di picco, sulla complessità quadratica dell'auto-attenzione (self-attention), l'overhead del decoder e la latenza end-to-end. I modelli mobili esistenti sono tipicamente limitati a backbone di piccole dimensioni (0,4–1,8 miliardi di parametri), il che inevitabilmente limita la qualità della generazione e crea un divario marcato tra le prestazioni mobile e quelle su scala server. La sfida principale è che i colli di bottiglia non sono confinati a un singolo modulo; semplicemente sostituire l'attenzione softmax con l'attenzione lineare spesso richiede un riaddestramento da zero o comporta un gap di espressività, mentre gli approcci ibridi non hanno ancora convertito pienamente il backbone in un'architettura ricorrente a bassa memoria adatta al deployment su dispositivi mobili.

Metodologia

Gli autori propongono MobileWan, un sistema che dimostra come un video diffusion transformer da 5 miliardi di parametri su scala server (basato su Wan2.2-5B) possa essere distribuito efficientemente su hardware mobile con memoria limitata attraverso una combinzione di riformulazione ricorrente e compressione strutturata. La metodologia consiste in quattro fasi principali di ottimizzazione:

1. Pruning dei Teste Apprendibile (Learnable Head Pruning)

Invece del pruning euristico dei blocchi, a cui gli autori hanno riscontrato che Wan2.2-5B è altamente sensibile, introducono un metodo di pruning delle teste apprendibile.

• Meccanismo: Viene introdotto un gate scalare apprendibile per ogni testa di auto-attenzione, modulandone il contributo. Questi gate sono ottimizzati end-to-end utilizzando un regolarizzatore di sparsità e una strategia di apprendimento con bias di rumore.
• Strategia di Addestramento: Il modello viene addestrato con una distribuzione ad alto bias di rumore (estraendo rumore $\sigma \sim \text{LogitNormal}(1.5, 1)$) durante la fase di apprendimento del gate. Questo enfatizza la struttura globale e il movimento, rendendo le decisioni di pruning più robuste. Il processo prevede uno schema in tre fasi: esplorazione del gate morbido (soft), annealing della temperatura verso valori binari e fine-tuning finale con un'architettura con pruning fisso.
• Risultato: Questo approccio supera il pruning euristico, particolarmente nei rapporti di pruning aggressivi, producendo riduzioni misurabili dei FLOP con una perdita di qualità limitata.

2. Distillazione della Ricorrenza (Recurrence Distillation)

Per affrontare la complessità quadratica della memoria dell'auto-attenzione, gli autori convertono il transformer pre-addestrato in un modello autoregressivo chunk-wise.

• Attenzione Ibrida: Utilizzano un meccanismo di attenzione ibrida che combina l'attenzione lineare causale con l'attenzione softmax locale. Il modello presta attenzione a un piccolo set di token locali tramite softmax per il contesto immediato e utilizza l'attenzione lineare per il contesto storico.
• Riformulazione RNN: Durante l'inferenza, l'auto-attenzione lineare causale permette al modello di operare come una Rete Neurale Ricorrente (RNN). Mantiene aggregati correnti compatti (vettori di stato $s_t$ e $z_t$) per i chunk passati invece di memorizzare l'intera matrice di attenzione.
• Beneficio: Ciò riduce l'impronta di memoria affinché sia costante indipendentemente dalla lunghezza della sequenza, consentendo la generazione di video lunghi su dispositivi con RAM limitata.

3. Distillazione dei Passaggi (Step Distillation)

Per ridurre la latenza, gli autori applicano la Distribution Matching Distillation (DMD) e la sua variante D-DMD.

• Approccio: Distillano il modello teacher multi-step in un modello student a pochi step.
• Ottimizzazione: Gli esperimenti hanno confrontato la DMD standard, la D-DMD (che divide la perdita in termini di Distribution Matching e di augmentazione del Classifier-Free Guidance) e la DMD2 (che incorpora la perdita avversaria).
• Selezione: Sebbene la DMD2 con 2 step abbia raggiunto il grado dinamico e i punteggi VBench più elevati, gli autori hanno notato che la D-DMD produce campioni visivamente accattivanti con un grado dinamico leggermente inferiore. Il sistema finale utilizza una configurazione ottimizzata per 3 passi di denoising (NFE) per bilanciare qualità e velocità.

4. Ottimizzazione del Decoder

I classici Video VAE utilizzano convoluzioni 3D causali che sono computazionalmente costose.

• Modifica: Gli autori si basano sull'efficiente decoder LightX2V ma ne modificano l'architettura per estendere la finestra di look-back causale.
• Meccanismo: Invece di condizionarsi sui frame passati immediati, la predizione del frame $F_t$ è condizionata su un contesto temporale più lungo (fino a $F_{t-4}$).
• Risultato: Ciò riduce gli artefatti temporali e il flickering pur mantenendo l'efficienza delle convoluzioni 2D, migliorando significativamente la qualità della ricostruzione (PSNR e VBench) rispetto ad altri decoder efficienti.

5. Pipeline di Quantizzazione

Per consentire il deployment sul NPU di Snapdragon® 8 Gen 5, il modello subisce un processo di quantizzazione in due fasi:

1. Post-Training Quantization (PTQ): Comprime i pesi a 8 bit.
2. Quantization-Aware Training (QAT): Un processo di distillazione in due fasi (prima per blocco e poi end-to-end) allinea il modello quantizzato con il teacher a precisione intera, recuperando l'accuratezza persa durante la PTQ.

Risultati Chiave

MobileWan raggiunge i seguenti parametri di prestazione su un dispositivo mobile (NPU Snapdragon® 8 Gen 5):

• Specifiche di Generazione: Genera video di 5 secondi a risoluzione 480×832, 16 FPS.
• Latenza: Latenza end-to-end di 20 secondi (inclusi 6,65 secondi per lo step di denoising del DiT).
• Qualità: Raggiunge un punteggio VBench di 83,79, che è paragonabile al modello Wan2.2 su scala server (83,12) e supera significativamente lo stato dell'arte mobile precedente (Neodragon: 81,61).
• Memoria: La riformulazione ricorrente riduce il requisito di memoria di picco da "Out of Memory" (OOM) su mobile a 9,9 GB per il DiT e 1,9 GB per il decoder.
• Studio Utente: In uno studio cieco sugli utenti, i partecipanti hanno preferito MobileWan rispetto a Neodragon l'80% delle volte. Nei confronti con il server Wan2.2 finetunato, MobileWan non è stato preferito nel 53% dei casi, ma è stato preferito nel 22%, indicando che è competitivo con il modello server in metà dei casi.

Significato e Rivendicazioni

Il paper sostiene di colmare il divario di qualità tra la generazione video mobile e quella su scala server sfidando l'assunto che la generazione video mobile richieda modelli piccoli.

• Scalabilità: Dimostra che un video DiT da 5 miliardi di parametri può essere distribuito nativamente su hardware mobile commerciale attraverso una compressione fine e una riformulazione ricorrente.
• Efficienza vs Qualità: Il lavoro prova che ottimizzazioni aggressive (pruning, ricorrenza, distillazione) possono essere applicate senza una perdita estrema di qualità, raggiungendo un punteggio VBench (83,79) che rivaleggia con i modelli server.
• Integrazione di Sistema: MobileWan è presentato come il primo modello di diffusione video da 5B distribuibile su un dispositivo mobile commerciale, stabilendo un nuovo stato dell'arte nella generazione video mobile.

Gli autori riconoscono i limiti, tra cui una qualità subottimale per certi tipi di scene (es. volti umani a media distanza), occasionali sovrasaturazioni o riduzione del movimento a seconda del metodo di distillazione, e potenziali discontinuità temporali o flickering in alcuni video generati. Intendono affrontare questi aspetti nel lavoro futuro.