Résumé Technique : Modulation Sensible à la Qualité pour les Transformers de Diffusion

Énoncé du Problème

Les modèles de diffusion de texte-en-image modernes, en particulier les Diffusion Transformers (DiT) comme Stable Diffusion 3.5 (SD3.5), s'appuient sur des signaux de conditionnement — spécifiquement les embeddings de pas de temps (timestep embeddings) et les embeddings de prompts agrégés (pooled prompt embeddings) — pour moduler le processus de débruitage via l'Adaptive Layer Normalization (AdaLN). Bien que ces signaux communiquent efficacement le niveau de bruit actuel et l'intention sémantique, ils manquent d'informations sensibles à la qualité. Par conséquent, le mécanisme de modulation ne peut pas réagir directement aux modes de défaillance spécifiques à l'échantillon, tels que l'accumulation d'erreurs, l'interprétation incorrecte du prompt ou les caractéristiques hallucinées qui émergent à mesure que l'image latente évolue. Le conditionnement de base est indépendant de l'état latent actuel ($\partial M_{base} / \partial x_t = 0$), ce qui empêche le modèle de corriger les problèmes de fidélité basés sur la qualité réelle de l'image intermédiaire.

Méthodologie

Les auteurs proposent le Quality Representation Module (QRM), un adaptateur léger basé sur un transformer conçu pour injecter des signaux sensibles à la qualité dans le processus de diffusion sans modifier le backbone pré-entraîné.

Architecture

Le QRM est implémenté comme un décodeur transformer qui prédit des décalages de modulation additifs ($M_{qrm}$) aux paramètres AdaLN de base ($M_{base}$).

• Entrées : Le module reçoit l'image latente actuelle ($x_t$), l'embedding de prompt agrégé ($p$), l'embedding de pas de temps ($t$), et les paramètres de modulation de base ($M_{base}$).
• Traitement :
  1. L'image latente $x_t$ est projetée en tokens latents via une convolution $1\times1$ apprise et un pooling adaptatif.
  2. Les paramètres de modulation de base sont normalisés et projetés pour former des tokens de requête ($Z_0$).
  3. Une mémoire clé-valeur ($KV$) est construite à partir des tokens latents, du prompt et des embeddings de pas de temps.
  4. À travers des couches de cross-attention et de self-attention, les tokens de requête sont attentifs à la mémoire $KV$ pour génire une représentation raffinée.
  5. La sortie est projetée à nouveau vers les largeurs natives de l'AdaLN pour produire le vecteur de décalage $M_{qrm}$.
• Application : La modulation finale appliquée aux blocs du transformer est $M_t = M_{base,t} + M_{qrm,t}$.

Stratégie d'Entraînement (ReFL)

Le QRM est entraîné à l'aide d'une stratégie de Reward Feedback Learning (ReFL) :

• Composants Gelés : Le backbone de diffusion (débruiteur, encodeurs de texte, VAE), l'échantillonneur (sampler) et le modèle de récompense restent gelés. Seuls les paramètres du QRM sont mis à jour.
• Objectif : Le module apprend à maximiser un score de récompense scalaire dérivé d'un modèle de récompense gelé (par exemple, HPSv2.1 ou CLIPScore) appliqué à une image décodée à partir d'une étape de débruitage intermédiaire.
• Fonctions de Perte : Le papier utilise soit une perte de maximisation de la récompense, soit une perte de recherche de marge (margin-seeking loss), qui pousse l'image modulée par le QRM à surpasser la base par une marge spécifique.
• Sélection des Pas de Temps : Pour améliorer l'efficacité et se concentrer sur la formation structurelle, le QRM n'est activé que pendant un sous-ensemble prédéfini de pas de temps à haut bruit (par exemple, $t \in \{50, \dots, 26\}$), plutôt que pendant l'ensemble du calendrier de diffusion.

Inférence

Pendant l'inférence, le module QRM entraîné agit comme un adaptateur plug-in. Il prédit des décalages de modulation basés sur l'état latent actuel et les ajoute aux paramètres de base. Cela permet au modèle de diriger la trajectoire de débruitage vers des états de plus haute qualité sans réentraîner le backbone ou modifier le calendrier d'échantillonnage.

Contributions Clés

1. Architecture QRM : Introduction d'un module transformer léger qui apprend des représentations sensibles à la qualité à partir des entrées existantes du modèle (latent, prompt, pas de temps) pour affiner les paramètres de modulation.
2. Intégration Non-Invasive : Un mécanisme pour augmenter la modulation du transformer de diffusion avec des signaux de qualité d'image latente sans modifier le backbone de diffusion, le VAE ou le calendrier d'échantillonnage sous-jacent.
3. Modulation Guidée par la Récompense : Démonstration que les mises à jour de modulation prédites ($M_{qrm}$) peuvent efficacement affiner le processus de débruitage à chaque pas de temps, améliorant l'adhérence au prompt et la fidélité de l'image.
4. Évaluation Complète : Études d'ablation approfondies analysant les effets de la taille de l'architecture QRM, des plages de pas de temps actifs et des métriques de récompense d'entraînement sur la qualité de l'image.

Résultats Expérimentaux

Les auteurs ont évalué le QRM en utilisant Stable Diffusion 3.5 (SD3.5) sur le benchmark Parti-Prompts.

• Performance Quantitative :
  • Comparé à la base SD3.5, le QRM entraîné avec HPSv2.1 comme signal de récompense a obtenu les scores les plus élevés parmi toutes les méthodes évaluées (y compris l'affinement de prompt, le fine-tuning et les approches ReNeg).
  • Améliorations : +0,20 en score esthétique (Aesthetic score) et +0,82 en score HPSv2.1.
  • Lorsqu'il est entraîné avec CLIPScore, le QRM améliore l'alignement avec les préférences humaines (HPSv2.1) tout en maintenant une qualité perceptuelle comparable.
• Analyse Qualitative :
  • Les visualisations des trajectoires de débruitage montrent que l'application du QRM pendant les premiers pas de temps, à haut bruit, conduit à l'émergence de structures sémantiques reconnaissables nettement plus tôt que dans la version de base.
  • L'application du QRM uniquement lors des étapes intermédiaires ou tardives produit des améliorations moindres, confirmant que la structure globale est établie tôt dans le processus de diffusion.
• Résultats d'Ablation :
  • Capacité : Les modèles QRM plus larges optimisent plus fortement la récompense spécifique d'entraînement (par exemple, HPSv2.1) mais présentent de légers compromis dans d'autres métriques (par exemple, CLIPScore).
  • Plage de Pas de Temps : La modulation pendant les premiers pas de temps ($50 \to 26$) était systématiquement plus efficace pour l'alignement que la modulation dans les étapes ultérieures.
  • Métrique de Récompense : Le choix du modèle de récompense (CLIPScore vs HPSv2.1) a directement façonné les biais perceptuels appris par le QRM.

Signification et Revendications

Le papier affirme que les interventions au niveau de la modulation fournissent une voie évolutive et efficace pour aligner les grands modèles de diffusion avec les signaux de préférence humaine. Contrairement aux méthodes qui nécessitent un fine-tuning coûteux de l'ensemble du backbone (par exemple, LoRA ou le fine-tuning complet), le QRM atteint des améliorations de qualité constantes en opérant entièrement via des mises à jour de modulation additives conditionnées sur la représentation latente évolutive.

Les auteurs positionnent le QRM comme un adaptateur flexible et sensible à la qualité, qui peut être piloté par différents signaux de récompense et appliqué à diverses architectures de diffusion transformers. Ils reconnaissent que le QRM hérite des biais des modèles de récompense d'entraînement et suggèrent que les travaux futurs pourraient explorer l'entraînement multi-récompense et la sélection de pas de temps adaptative pour améliorer davantage la robustesse et la généralisation. Le travail démontre que des gains significatifs en fidélité d'image et en alignement de prompt peuvent être obtenus sans modifier les poids pré-entraînés des modèles de pointe.