Résumé Technique : MobileWan

Énoncé du Problème

Les progrès récents dans la diffusion vidéo ont été portés par le passage à l'échelle d'architectures basées sur les transformers (DiT) atteignant des milliards de paramètres, améliorant considérablement la fidélité visuelle et la cohérence du mouvement. Cependant, la génération de vidéos sur mobile accuse un retard substantiel par rapport aux systèmes de classe serveur en raison des contraintes strictes sur la mémoire de pointe, de la complexité quadratique de l'auto-attention, de la surcharge du décodeur et de la latence de bout en bout. Les modèles mobiles existants sont généralement limités à de petits backbones (0,4 à 1,8 milliard de paramètres), ce qui limite inévitablement la qualité de génération et crée un écart marqué entre les performances mobiles et celles des serveurs. Le défi central est que les goulots d'étranglement ne se limitent pas à un seul module ; simplement remplacer l'attention softmax par une attention linéaire nécessite souvent de réentraîner à partir de zéro ou entraîne un déficit d'expressivité, tandis que les approches hybrides n'ont pas encore pleinement converti le backbone en une architecture récurrente à faible mémoire adaptée au déploiement mobile.

Méthodologie

Les auteurs proposent MobileWan, un système qui démontre qu'un transformer de diffusion vidéo de 5 milliards de paramètres (basé sur Wan2.2-5B) peut être déployé efficacement sur du matériel mobile à mémoire contrainte grâce à une combinaison de reformulation récurrente et de compression structurée. La méthodologie se compose de quatre étapes d'optimisation principales :

1. Élagage de Têtes Apprenable (Learnable Head Pruning)

Au lieu d'un élagage de blocs heuristique, auquel les auteurs ont trouvé que Wan2.2-5B est très sensible, ils introduisent une méthode d'élagage de têtes apprenable.

• Mécanisme : Une porte scalaire apprenable est introduite pour chaque tête d'auto-attention, modulant sa contribution. Ces portes sont optimisées de bout en bout à l'aide d'un régularisateur de parcimonie et d'une stratégie d'apprentissage biaisée par le bruit.
• Stratégie d'entraînement : Le modèle est entraîné avec une distribution fortement biaisée par le bruit (tirant le bruit $\sigma \sim \text{LogitNormal}(1.5, 1)$) pendant la phase d'apprentissage des portes. Cela met l'accent sur la structure globale et le mouvement, rendant les décisions d'élagage plus robustes. Le processus implique un programme en trois phases : exploration de porte douce (soft gate), recuit de la température vers des valeurs binaires, et ajustement final avec une architecture élaguée fixe.
• Résultat : Cette approche surpasse l'élagage heuristique, particulièrement aux ratios d'élagage agressifs, produisant des réductions mesurables de FLOPs avec une perte de qualité limitée.

2. Distillation Récurrente

Pour traiter la complexité mémoire quadratique de l'auto-attention, les auteurs convertissent le transformer pré-entraîné en un modèle autorégressif par morceaux (chunk-wise).

• Attention Hybride : Ils emploient un mécanisme d'attention hybride combinant une attention linéaire causale avec une attention softmax locale. Le modèle prête attention à un petit ensemble de tokens locaux via softmax pour le contexte immédiat et utilise l'attention linéaire pour le contexte historique.
• Reformulation RNN : Lors de l'inférence, l'attention linéaire causale permet au modèle de fonctionner comme un réseau de neurones récurrents (RNN). Il maintient des agrégats de fonctionnement compacts (vecteurs d'état $s_t$ et $z_t$) pour les morceaux passés plutôt que de stocker la matrice d'attention complète.
• Bénéfice : Cela réduit l'empreinte mémoire pour qu'elle soit constante quelle que soit la longueur de la séquence, permettant la génération de vidéos longues sur des appareils dotés d'une RAM limitée.

3. Distillation d'Étapes (Step Distillation)

Pour réduire la latence, les auteurs appliquent la Distillation de Correspondance de Distribution (DMD) et sa variante D-DMD.

• Approche : Ils distillent le modèle enseignant multi-étapes vers un modèle étudiant à quelques étapes.
• Optimisation : Les expériences ont comparé le DMD standard, le D-DMD (divisant la perte en termes de correspondance de distribution et d'augmentation de guidage sans classificateur/Classifier-Free Guidance) et le DMD2 (incorporant une perte adversaire).
• Sélection : Bien que le DMD2 avec 2 étapes ait atteint le degré dynamique et les scores VBench les plus élevés, les auteurs ont noté que le D-DMD produisait des échantillons visuellement attrayants avec des degrés dynamiques légèrement inférieurs. Le système final utilise une configuration optimisée pour 3 étapes de débruitage (NFEs) afin d'équilibrer qualité et vitesse.

4. Optimisation du Décodeur

Les VAE vidéo standards utilisent des convolutions 3D causales qui sont coûteuses en termes de calcul.

• Modification : Les auteurs s'appuient sur le décodeur efficace LightX2V mais modifient son architecture pour étendre la fenêtre de regard en arrière (look-back window) causale.
• Mécanisme : Au lieu de se conditionner sur les trames passées immédiates, la prédiction de la trame $F_t$ est conditionnée sur un contexte temporel plus long (jusqu'à $F_{t-4}$).
• Résultat : Cela réduit les artefacts temporels et le scintillement tout en maintenant l'efficacité des convolutions 2D, améliorant considérablement la qualité de reconstruction (PSNR et VBench) par rapport aux autres décodeurs efficaces.

5. Pipeline de Quantification

Pour permettre le déploiement sur le NPU Snapdragon® 8 Gen 5, le modèle subit un processus de quantification en deux étapes :

1. Quantification Post-Entraînement (PTQ) : Compresse les poids en 8 bits.
2. Entraînement Sensible à la Quantification (QAT) : Un processus de distillation en deux étapes (par bloc puis de bout en bout) aligne le modèle quantifié avec l'enseignant en pleine précision, récupérant la précision perdue lors de la PTQ.

Résultats Clés

MobileWan atteint les métriques de performance suivantes sur un appareil mobile (NPU Snapdragon® 8 Gen 5) :

• Spécifications de génération : Génère des vidéos de 5 secondes en résolution 480×832, 16 FPS.
• Latence : Latence de bout en bout de 20 secondes (incluant 6,65 secondes pour l'étape de débruitage DiT).
• Qualité : Atteint un score VBench de 83,79, ce qui est comparable au modèle serveur Wan2.2 (83,12) et surpasse nettement l'état de l'art mobile précédent (Neodragon : 81,61).
• Mémoire : La reformulation récurrente réduit l'exigence de mémoire de pointe de "Out of Memory" (OOM) sur mobile à 9,9 Go pour le DiT et 1,9 Go pour le décodeur.
• Étude Utilisateur : Dans une étude utilisateur en aveugle, les participants ont préféré MobileWan à Neodragon 80 % du temps. Dans les comparaisons avec le modèle serveur Wan2.2 affiné, MobileWan n'a pas été préféré dans 53 % des cas, mais a été préféré dans 22 % des cas, indiquant qu'il est compétitif avec le modèle serveur dans la moitié des cas.

Signification et Revendications

Le papier prétend combler l'écart de qualité entre la génération vidéo mobile et celle de classe serveur en remettant en question l'hypothèse selon laquelle la génération vidéo mobile nécessite de petits modèles.

• Scalabilité : Il démontre qu'un DiT vidéo de 5 milliards de paramètres peut être déployé nativement sur du matériel mobile commercial grâce à une compression fine et une reformulation récurrente.
• Efficacité vs Qualité : Ce travail prouve que des optimisations agressives (élagage, récurrence, distillation) peuvent être appliquées sans une perte de qualité extrême, atteignant un score VBench (83,79) qui rivalise avec les modèles serveurs.
• Intégration Système : MobileWan est présenté comme le premier modèle de diffusion vidéo de l'échelle 5B déployable sur un appareil mobile commercial, établissant un nouveau état de l'art pour la génération vidéo mobile.

Les auteurs reconnaissent des limites, notamment une qualité sous-optimale pour certains types de scènes (ex: visages humains à mi-distance), un certain oversaturation ou une réduction du mouvement selon la méthode de distillation, ainsi que de potentiels discontinuités temporelles ou scintillements dans certaines vidéos générées. Ils prévoient d'aborder ces points dans des travaux futurs.