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

Résumé Technique : Distillation Collaborative à Peu d'Étapes et Quantification à Bas Débit pour les Modèles de Diffusion Vidéo à Double Expert Wan2.2

1. Énoncé du Problème

Les grands modèles de diffusion vidéo, tels que le Wan2.2-T2V-A14B, atteignent une haute qualité visuelle mais font face à des barrières de déploiement significatives dues à deux goulots d'étranglement primaires :

1. Coût Computationnel : Le processus de débruitage itératif nécessite un grand nombre d'étapes d'échantillonnage, ce qui rend l'inférence coûteuse en termes de latence.
2. Empreinte Mémoire : Les modèles possèdent un nombre élevé de paramètres résidents (environ 27B au total, avec environ 14B actifs par étape dans l'architecture à double expert), dépassant les budgets mémoire de nombreux scénarios de déploiement.

Les stratégies de compression existantes échouent souvent lorsqu'elles sont appliquées directement à ces modèles. La simple réduction des étapes d'échantillonnage sans réentraînement conduit à une faible qualité, tandis que la quantification de la trajectoire originale à étapes longues introduit souvent des décalages de distribution d'activation lors de l'inférence, provoquant une propagation d'erreurs temporelles et une dégradation de la fidélité vidéo. De plus, Wan2.2 utilise une architecture de mélange d'experts (MoE - Mixture of Experts) à double expert (expert de haut bruit et expert de bas bruit), ce qui complique les approches de quantification standard qui traitent le modèle comme un tronc dense.

2. Méthodologie

Le document propose un pipeline de compression orienté vers le déploiement qui intègre la distillation à peu d'étapes et la quantification à bas débit de manière co-conçue. Le pipeline suit une séquence spécifique : adaptation du modèle $\rightarrow$ distillation à peu d'étapes $\rightarrow$ quantification à bas débit $\rightarrow$ construction de la calibration $\rightarrow$ génération $\rightarrow$ évaluation.

Composantes Techniques Clés :

• Gestion de l'Architecture Sensible aux Experts Duaux :

  • Le pipeline respecte le mécanisme de routage spécifique de Wan2.2, où un expert de haut bruit gère la disposition globale/mouvement et un expert de bas bruit affine les textures/détails.
  • Calibration par Branche : Au lieu de regrouper les activations, la méthode capture et calibre les statistiques séparément pour les branches de haut bruit et de bas bruit. Cela empêche le quantificateur de mélanger des distributions d'activation incompatibles.
  • Protection des Couches d'Entrée : Les couches précoces et les blocs sensibles au routage sont conservés en haute précision ou calibrés plus agressivement. Ceci est critique car les erreurs dans la phase initiale de haut bruit (initialisation de la disposition/mouvement) se propagent à travers toute la séquence, alors que les erreurs ultérieures sont souvent locales.

• Distillation à Peu d'Étapes via l'Appariement de Distribution :

  • La méthode emploie la Distillation par Appariement de Distribution (DMD) pour apprendre une trajectoire de débruitage plus courte plutôt que de simplement tronquer l'échantillonneur.
  • Le modèle étudiant est optimisé pour correspondre à la distribution de l'enseignant sur le nombre réduit d'étapes, avec un objectif sensible aux étapes ($L_{distill} = \lambda_1 L_{high} + \lambda_2 L_{low}$) qui respecte les rôles distincts des phases de haut bruit et de bas bruit.

• Quantification Co-conçue :

  • Calibration Alignée sur la Trajectoire : Une innovation centrale est d'effectuer la calibration de la quantification sur l'étudiant distillé à peu d'étapes plutôt que sur le modèle original à étapes longues. Cela garantit que le quantificateur observe les statistiques d'activation qu'il rencontrera lors de l'inférence réelle, réduisant ainsi le décalage de distribution.
  • Représentation HiF4 (HiFloat4) : Le pipeline utilise un format de virgule flottante à blocs de 4 bits (HiF4) conçu pour l'inférence à bas débit. Ce format présente une hiérarchie d'échelle à trois niveaux (Global E6M2, Niveau-2 E1, Niveau-3 E1) et une amplitude signée S1P2. Cette conception offre une large plage dynamique globale (environ $2^{-50}$ à $2^{18}$) pour gérer les plages d'activation hétérogènes entre les phases de débruitage, tout en préservant les détails locaux grâce aux micro-exposants.
  • Granularité : Les poids sont quantifiés par canal, et les activations sont quantifiées par jeton (token) pour suivre les variations au niveau du prompt et de la trame.

• Ingénierie et Évaluation :

  • Gestion de la Mémoire : Pour éviter les erreurs de type Out-Of-Memory (OOM) sur des GPU uniques (ex: H100 80Go), la calibration et le remplacement sont effectués couche par couche, avec des tampons auxiliaires chargés à la demande.
  • Évaluation en Boucle Fermée : Le système utilise VBench pour évaluer cinq dimensions spécifiques : Cohérence du Sujet (SC), Qualité Esthétique (AQ), Qualité d'Imagerie (IQ), Cohérence Globale (OC) et Fluidité du Mouvement (MS). Le protocole d'évaluation garantit une comparaison équitable en utilisant des prompts, des résolutions et des calendriers d'échantillonnage identiques pour les bases de référence.

3. Contributions Clés

1. Co-conception de la Distillation et de la Quantification : Le document démontre que calibrer la quantification sur la trajectoire distillée à peu d'étapes (plutôt que sur le modèle original à étapes longues) réduit considérablement le décalage de distribution d'activation, permettant au modèle compressé de maintenir une qualité proche de l'étudiant en pleine précision.
2. Stratégie de Compression Sensible au MoE : La méthode introduit un schéma de calibration par branche et de protection des couches d'entrée spécifiquement adapté à la structure à double expert de Wan2.2, empêchant le mélange des statistiques d'activation de haut et de bas bruit.
3. Intégration de HiF4 pour la Diffusion Vidéo : L'application de l'échelle hiérarchique de HiF4 s'avère efficace pour les modèles de diffusion vidéo, équilibrant le besoin d'une large plage globale (pour éviter le dépassement lors des phases de haut bruit) avec la précision locale (pour préserver la texture lors des phases de bas bruit).
4. Validation Empirique du Nombre d'Étapes : L'étude évalue systématiquement les configurations de 4, 8, 20 et 40 étapes, identifiant le point d'équilibre optimal entre qualité et efficacité.

4. Résultats Expérimentaux

Les expériences ont été menées sur le modèle Wan2.2-T2V-A14B en utilisant le jeu de données OpenS2V pour la calibration et VBench pour l'évaluation.

• Performance vs Baselines :

  • À 4 étapes, le modèle quantifié maintient une erreur mineure par rapport au modèle pleine précision du même nombre d'étapes, mais reste derrière la base de référence originale à étapes longues.
  • À 8 étapes, le modèle quantifié (Score Moyen : 0,7002) commence à dépasser la base de référence pleine précision originale (0,6976).
  • À 20 étapes, le modèle quantifié atteint la meilleure performance (Score Moyen : 0,7074), surpassant à la fois la base de référence pleine précision originale et le modèle distillé pleine précision du même nombre d'étapes (0,7051).
  • À 40 étapes, la performance plafonne, indiquant des rendements décroissants à mesure que la longueur de la trajectoire augmente.

• Analyse des Métriques (Configuration à 20 Étapes) :

  • Le modèle quantifié a amélioré la base de référence originale en Cohérence du Sujet (0,9629 contre 0,9495), en Qualité Esthétique, en Qualité d'Imagerie et en Fluidité du Mouvement.
  • La Cohérence Globale est restée comparable à celle du modèle du même nombre d'étapes.
  • Les résultats indiquent que la calibration alignée sur la distribution réussit à récupérer les pertes liées au bas débit sans sacrifier la stabilité temporelle.

• Sensibilité à l'Erreur :

  • L'analyse révèle que les perturbations sur une seule trame affectent principalement la Qualité d'Imagerie et la Qualité Esthétique, tandis que les erreurs accumulées sur les étapes se manifestent par une Dérive du Sujet et un Tremblement du Mouvement. La protection des couches d'entrée proposée atténue efficacement ces derniers points.

5. Signification et Revendications

Le document affirme que le pipeline proposé offre une voie pratique pour déployer des modèles de diffusion vidéo de haute qualité sous des contraintes de mémoire GPU et de latence limitées.

• Compromis Qualité-Efficacité : Les auteurs identifient la configuration à 20 étapes comme étant la configuration optimale, fournissant une capacité de débruitage suffisante pour récupérer la qualité temporelle et visuelle tout en nécessant nettement moins d'étapes que les échantillonneurs conventionnels.
• Changement Stratégique : Le travail soutient que la quantification ne doit pas être un "ajout ultérieur de perte", mais une partie intégrante de la stratégie d'accélération. En co-concevant la distillation et la quantification, la méthode transforme la conversion à bas débit en un composant stable de la pipeline d'inférence.
• Prêt pour le Déploiement : L'utilisation de la protection des couches d'entrée et de la calibration par branche répond aux modes de défaillance spécifiques des modèles vidéo à double expert (ex: dérive du sujet et instabilité du mouvement), rendant le modèle compressé adapté à la génération interactive et au traitement par lots sous ressources contraintes.

Le document conclut que bien que le pipeline actuel se concentre sur la quantification post-entraînement, le cadre établit une boucle de vérification où les échements spécifiques aux métriques (ex: chutes de Cohérence du Sujet) fournissent des signaux de débogage directs pour affiner la calibration et le nombre d'étapes.