Résumé Technique : NanoFLUX

Énoncé du Problème

Les modèles de diffusion Text-to-Image (T2I) à grande échelle, tels que la famille FLUX (montant jusqu'à 17 milliards de paramètres) et Qwen-Image (20 milliards de paramètres), ont atteint une qualité visuelle de pointe. Cependant, leur échelle massive crée une barrière significative pour le déploiement sur les appareils, nécessitant des infrastructures GPU coûteuses ou une inférence basée sur le cloud. Les solutions existantes pour les appareils mobiles reposent souvent sur des modèles plus petits qui manquent de qualité comparable ou utilisent des techniques de compression (élagage, quantification, réduction d'étapes) qui ne traitent pas pleinement la redondance architecturale des Transformers de Diffusion (DiTs) à grande échelle et de leurs encodeurs de texte associés. Le défi consiste à compresser ces modèles massifs pour les appareils mobiles tout en préservant la qualité de génération et en minimisant la latence d'inférence.

Méthodologie : NanoFLUX

Les auteurs introduisent NanoFLUX, un modèle T2I de 2,4 milliards de paramètres distillé à partir du modèle enseignant FLUX.1-Schnell de 17 milliards. La compression est réalisée via un pipeline de distillation progressive ciblant les deux composants les plus lourds : le Transformer de Diffusion (DiT) de 12 milliards et l'encodeur de texte T5-XXL de 5 milliards.

1. Compression du Transformer (12B $\to$ 2B)

Les auteurs emploient une stratégie de compression progressive en quatre étapes pour réduire la taille du DiT :

• Étape C1 : Élagage des têtes d'attention (12B $\to$ 5B) :

  • Analyse : La décomposition en valeurs singulières (SVD) sur les sorties des têtes d'attention a révélé que toutes les têtes ne sont pas indépendantes ; beaucoup contiennent des informations redondantes.
  • Action : Le nombre de têtes d'attention ($H$) a été réduit uniformément de 24 à 16.
  • Distillation : Une perte au niveau des caractéristiques a été introduite pour garantir que les têtes restantes préservent collectivement l'information des têtes élaguées. L'étudiant a été initialisé avec les 16 premières têtes de l'enseignant.

• Étape C2 : Réduction de la dimension des têtes (5B $\to$ 3B) :

  • Analyse : La SVD a été appliquée à nouveau aux caractéristiques du modèle de 5B.
  • Action : La dimension de la tête ($d_H$) a été réduite de 128 à 96, conservant les principaux composants singuliers. Cela a produit un modèle de 3 milliards de paramètres.

• Étape C3 : Élagage de la profondeur par fusion de blocs (3B $\to$ 2,5B) :

  • Analyse : La similitude cosinus entrée-sortie a été calculée pour les blocs du transformer. Les blocs Single-Stream (SS) présentaient une forte similitude (>0,85), indiquant une redondance fonctionnelle.
  • Action : Au lieu de supprimer ces blocs, les auteurs ont fusionné une chaîne de 15 blocs SS redondants en un seul bloc en faisant la moyenne de leurs paramètres. Cela a réduit la profondeur du transformer de 38 à 24 blocs.

• Étape C4 : Normalisation de couche statique (2,5B $\to$ 1,8B) :

  • Analyse : Les coefficients de l'Adaptive Layer Normalization (AdaLN) prédits par le modèle présentaient une faible variance entre les échantillons, suggérant que le conditionnement dynamique issu des projections CLIP était redondant.
  • Action : Les coefficients dynamiques d'AdaLN ont été remplacés par des moyennes statiques précalculées (Static-LN). Cela a supprimé environ 0,7 milliard de paramètres. Le modèle a été affiné brièvement avec des coefficients gelés pour récupérer la qualité.

2. Sous-échantillonnage progressif des tokens

Pour traiter la complexité quadratique de l'auto-attention sans sacrifier les détails :

• Mécanisme : Un module de sous-échantillonnage ($D$) basé sur ResNet et un module de sur-échantillonnage ($U$) ont été insérés dans le transformer.
• Stratégie hybride : Le modèle opère sur un schéma de résolution hybride. Les premières étapes de diffusion (capturant la structure globale) traitent des tokens sous-échantillonnés (réduisant le nombre de tokens par 4), tandis que les étapes ultérieures (affinant les détails) traitent des tokens à pleine résolution.
• Entraînement : Une stratégie de "Sous-échantillonnage progressif des tokens" (PTD) a été utilisée pour stabiliser l'entraînement. Les modules ont été insérés et entraînés de manière incrémentielle, déplaçant le point de sous-échantillonnage plus tôt dans le réseau au fil des itérations.
• Résultat : 23 blocs sur 43 transformer opèrent sur des tokens sous-échantillonnés, réduisant considérablement la latence tout en maintenant un traitement haute résolution là où cela est nécessaire.

3. Distillation de l'encodeur de texte (5B $\to$ 330M)

• Cible : Distiller l'encodeur T5-XXL de 5 milliards en un modèle T5-Large de 330 millions.
• Méthode : Un processus de distillation en deux étapes :
  1. Correspondance d'embedding : Un MLP est ajouté à l'étudiant pour correspondre à la dimension de sortie de l'enseignant, minimisant l'erreur quadratique moyenne (MSE) entre les embeddings de prompts.
  2. Distillation par bloc : L'étudiant est entraîné pour correspondre aux états cachés du prompt de l'enseignant à des couches spécifiques du transformer. Crucialement, la méthode exploite des signaux visuels provenant des premières couches du débruiteur (via l'attention conjointe) pour guider l'encodeur de texte, atténuant ainsi l'accumulation d'erreurs.
  3. Stabilité : Les gradients sont détachés pour les étapes de bruit élevé ($t < \hat{t}$) afin d'éviter une supervision instable, se concentrant sur la distillation des étapes ultérieures et des trois premières couches du transformer.

Résultats Clés

• Taille du modèle : Le modèle final NanoFLUX possède 2,4 milliards de paramètres (2B pour le DiT + 330M pour l'encodeur de texte), soit une réduction de 7$\times$ par rapport au modèle enseignant FLUX.1-Schnell de 17 milliards.
• Performance :
  • Sur les appareils mobiles, NanoFLUX génère des images de 512 $\times$ 512 en environ 2,5 secondes.
  • Les benchmarks quantitatifs (One-IG, DPG, GenEval, HPDv3) montrent que NanoFLUX maintient des performances comparables au modèle enseignant, avec seulement une légère dégradation sur certaines métriques (par exemple, le HPSv3 passe de 11,45 à 10,41).
  • Les résultats qualitatifs indiquent que le modèle préserve la qualité d'image et l'adhérence au prompt de manière comparable au modèle original de 17 milliards.
• Efficacité : La stratégie de sous-échantillonnage progressif des tokens équilibre avec succès la latence d'inférence et la qualité de génération, permettant au modèle de fonctionner efficacement sur du matériel aux ressources limitées.

Signification et Revendications

L'article affirme que NanoFLUX est le premier framework à étudier de manière exhaustive la compression de modèles à grande échelle pour la génération T2I, ciblant spécifiquement des modèles comme FLUX. Sa signification réside dans :

1. Faisabilité de la génération sur appareil : Démontrer qu'une génération T2I de haute qualité est réalisable sur les appareils mobiles sans dépendre d'une infrastructure cloud.
2. Aperçu architectural : Fournir une voie pratique pour réduire l'échelle des grands modèles génératifs en identifiant et en exploitant des redondances architecturales spécifiques (têtes d'attention, profondeur de blocs, dynamique de normalisation) plutôt qu'en s'appuyant uniquement sur une quantification générique ou une réduction d'étapes.
3. Généralisabilité : Offrir des perspectives et un pipeline de distillation progressive qui peuvent se généraliser aux futurs modèles génératifs à grande échelle.

Les auteurs positionnent ce travail comme une étape vers l'élargissement de l'accessibilité à l'IA générative de haute qualité pour les utilisateurs disposant de ressources de calcul limitées.