Collaborative Few-Step Distillation and Low-Bit Quantization for Wan2.2 Dual-Expert Video Diffusion Models

Riepilogo Tecnico: Distillazione Collaborativa a Pochi Step e Quantizzazione a Basso Bit per i Modelli di Diffusione Video Dual-Expert Wan2.2

1. Definizione del Problema

I grandi modelli di diffusione video, come il Wan2.2-T2V-A14B, raggiungono un'elevata qualità visiva ma affrontano barriere significative alla distribuzione a causa di due colli di bottiglia primari:

1. Costo Computazionale: Il processo di denoising iterativo richiede un numero elevato di step di campionamento, rendendo l'inferenza costosa in termini di latenza.
2. Impronta di Memoria: I modelli possiedono un numero elevato di parametri residenti (circa 27B totali, con circa 14B attivi per step nell'architettura dual-expert), superando i budget di memoria di molti scenari di distribuzione.

Le strategie di compressione esistenti spesso falliscono quando applicate direttamente a questi modelli. Ridurre semplicemente gli step di campionamento senza riaddestramento porta a una scarsa qualità, mentre quantizzare la traiettoria originale a molti step introduce spesso disallineamenti nella distribuzione delle attivazioni durante l'inferenza, causando la propagazione dell'errore temporale e la degradazione della fedeltà del video. Inoltre, Wan2.2 utilizza un'architettura Mixture of Experts (MoE) dual-expert (esperto ad alto rumore e esperto a basso rumore), che complica gli approcci di quantizzazione standard che trattano il modello come un tronco denso.

2. Metodologia

Il documento propone una pipeline di compressione orientata alla distribuzione che integra la distillazione a pochi step e la quantizzazione a basso bit in modo co-progettato. La pipeline segue una sequenza specifica: adattamento del modello $\rightarrow$ distillazione a pochi step $\rightarrow$ quantizzazione a basso bit $\rightarrow$ costruzione della calibrazione $\rightarrow$ generazione $\rightarrow$ valutazione.

Componenti Tecniche Chiave:

• Gestione dell'Architettura Consapevole del Dual-Expert:

  • La pipeline rispetta il meccanico di routing specifico di Wan2.2, dove un esperto ad alto rumore gestisce il layout globale/movimento e un esperto a basso rumore rifinisce texture/dettagli.
  • Calibrazione per Branch: Invece di raggruppare le attivazioni, il metodo cattura e calibra le statistiche separatamente per i rami ad alto e basso rumore. Ciò evita che il quantizzatore mescoli distribuzioni di attivazione incompatibili.
  • Protezione degli Strati di Ingresso (Entrance-Layer Protection): Gli strati iniziali e i blocchi sensibili al routing sono mantenuti a una precisione superiore o calibrati più aggressivamente. Ciò è critico perché gli errori nella fase iniziale ad alto rumore (inizializzazione di layout/movimento) si propagano attraverso l'intera sequenza, mentre gli errori successivi sono spesso locali.

• Distillazione a Pochi Step tramite Matching della Distribuzione:

  • Il metodo impiega la Distribution-Matching Distillation (DMD) per apprendere una traiettoria di denoising più breve invece di limitarsi a troncare il campionatore.
  • Il modello studente è ottimizzato per corrispondere alla distribuzione del modello insegnante attraverso il numero ridotto di step, con un obiettivo consapevole della fase ($L_{distill} = \lambda_1 L_{high} + \lambda_2 L_{low}$) che rispetta i distinti ruoli delle fasi ad alto e basso rumore.

• Quantizzazione Co-Progettata:

  • Calibrazione Allineata alla Traiettoria: Un'innovazione centrale è eseguire la calibrazione della quantizzazione sul modello studente distillato a pochi step piuttosto che sul modello originale a molti step. Ciò assicura che il quantizzatore osservi le statistiche di attivazione che incontrerà durante l'inferenza effettiva, riducendo il disallineamento della distribuzione.
  • Rappresentazione HiF4 (HiFloat4): La pipeline utilizza un formato a virgola mobile a blocchi a 4 bit (HiF4) progettato per l'inferenza a basso bit. Questo formato presenta una gerarchia di scaling a tre livelli (Global E6M2, Level-2 E1, Level-3 E1) e un'ampiezza con segno S1P2. Questo design fornisce un ampio intervallo dinamico globale (circa $2^{-50}$ a $2^{18}$) per gestire intervalli di attivazione eterogenei tra le fasi di denoising, preservando al contempo il dettaglio locale attraverso micro-esponenti.
  • Granularità: I pesi sono quantizzati per canale, e le attivazioni sono quantizzate per token per tracciare le variazioni a livello di prompt e di frame.

• Ingegneria e Valutazione:

  • Gestione della Memoria: Per evitare errori di Out-Of-Memory (OOM) su singole GPU (es. H100 da 80GB), la calibrazione e la sostituzione vengono eseguite layer-per-layer, con buffer ausiliari caricati su richiesta.
  • Valutazione a Ciclo Chiuso: Il sistema utilizza VBench per valutare cinque dimensioni specifiche: Coerenza del Soggetto (SC), Qualità Estetica (AQ), Qualità dell'Immagine (IQ), Coerenza Complessiva (OC) e Fluidità del Movimento (MS). Il protocollo di valutazione garantisce un confronto equo utilizzando prompt, risoluzioni e programmi di campionamento identici per i baseline.

3. Contributi Chiave

1. Co-Design di Distillazione e Quantizzazione: Il documento dimostra che calibrare la quantizzazione sulla traiettoria distillata a pochi step (anziché sul modello originale a molti step) riduce significativamente il disallineamento della distribuzione delle attivazioni, permettendo al modello compresso di mantenere una qualità vicina a quella dello studente a precisione intera.
2. Strategia di Compressione Consapevole del MoE: Il metodo introduce uno schema di calibrazione per branch e protezione degli strati di ingresso specificamente adattato alla struttura dual-expert di Wan2.2, prevenendo la miscelazione delle statistiche di attivazione ad alto e basso rumore.
3. Integrazione di HiF4 per la Diffusione Video: L'applicazione dello scaling gerarchico di HiF4 si dimostra efficace per i modelli di diffusione video, bilanciando la necessità di un ampio intervallo globale (per evitare l'overflow nelle fasi ad alto rumore) con la precisione locale (per preservare la texture nelle fasi a basso rumore).
4. Validazione Empirica del Numero di Step: Lo studio valuta sistematicamente le impostazioni di 4, 8, 20 e 40 step, identificando il punto di equilibrio ottimale tra qualità ed efficienza.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul modello Wan2.2-T2V-A14B utilizzando il dataset OpenS2V per la calibrazione e VBench per la valutazione.

• Performance rispetto ai Baseline:

  • A 4 step, il modello quantizzato mantiene un piccolo errore rispetto allo stesso modello a precisione intera, ma resta indietro rispetto al baseline originale a molti step.
  • A 8 step, il modello quantizzato (Punteggio Medio: 0.7002) inizia a superare il baseline originale a precisione intera (0.6976).
  • A 20 step, il modello quantizzato raggiunge la migliore performance (Punteggio Medio: 0.7074), superando sia il baseline originale a precisione intera che il modello studente a precisione intera con lo stesso numero di step (0.7051).
  • A 40 step, la performance si stabilizza, indicando rendimenti decrescenti dall'aumento della lunghezza della traiettoria.

• Analisi delle Metriche (Impostazione a 20 Step):

  • Il modello quantizzato ha migliorato il baseline originale in Coerenza del Soggetto (0.9629 vs 0.9495), Qualità Estetica, Qualità dell'Immagine e Fluidità del Movimento.
  • La Coerenza Complessiva è rimasta comparabile al modello con lo stesso numero di step.
  • I risultati indicano che la calibrazione allineata alla distribuzione ha recuperato con successo le perdite del basso bit senza sacrificare la stabilità temporale.

• Sensibilità all'Errore:

  • L'analisi rivela che le perturbazioni del singolo frame influenzano principalmente la Qualità dell'Immagine e la Qualità Estetica, mentre gli errori accumulati attraverso gli step si manifestano come deriva del soggetto (Subject Drift) e tremolio del movimento (Motion Jitter). La protezione degli strati di ingresso proposta mitiga efficacemente questi ultimi.

5. Significato e Rivendicazioni

Il documento sostiene che la pipeline proposta offre un percorso pratico per distribuire modelli di diffusione video di alta qualità sotto budget limitati di memoria GPU e latenza.

• Compromesso Qualità-Efficienza: Gli autori identificano l'impostazione a 20 step come la configurazione ottimale, fornendo una capacità di denoising sufficiente per recuperare la qualità temporale e visiva pur richiedendo significativamente meno step rispetto ai campionatori convenzionali.
• Svolta Strategica: Il lavoro sostiene che la quantizzazione non debba essere un "aggiunta postuma lossy", ma una parte integrante della strategia di accelerazione. Co-progettando distillazione e quantizzazione, il metodo trasforma la conversione a basso bit in una componente stabile della pipeline di inferenza.
• Prontezza alla Distribuzione: L'uso della protezione degli strati di ingresso e della calibrazione per branch affronta i modi di fallimento specifici dei modelli video dual-expert (es. deriva del soggetto e instabilità del movimento), rendendo il modello compresso adatto alla generazione interattiva e all'elaborazione batch con risorse limitate.

Il documento conclude che, sebbene l'attuale pipeline si concentri sulla quantizzazione post-training, il framework stabilisce un ciclo di verifica in cui i fallimenti specifici delle metriche (es. cali nella Coerenza del Soggetto) fornisano segnali diretti di debug per perfezionare la calibrazione e il numero di step.