Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

Ce papier propose une méthode sans entraînement appelée « réinjection de prompt » pour atténuer l'oubli progressif des consignes textuelles dans les Transformers de diffusion multimodaux, améliorant ainsi la capacité de suivi des instructions et la qualité de génération d'images.

L'essentiel

Imaginez que vous demandez à un artiste très talentueux de dessiner quelque chose de très précis, comme : "Un chat noir avec une moustache blanche, assis sur un fauteuil vert, à côté d'un vase rouge."

1. Le Problème : L'Artiste qui Oublie (Le "Prompt Forgetting")

Dans les modèles d'IA modernes (comme SD3, FLUX ou Qwen-Image), l'artiste ne regarde pas votre description une seule fois au début et l'oublie ensuite. Au contraire, il lit votre texte et dessine en même temps, couche par couche, comme si il construisait un gratte-ciel étage par étage.

Le problème découvert par les chercheurs :
À mesure que l'artiste construit les étages supérieurs de son dessin (les couches profondes de l'IA), il commence à oublier les détails fins de votre description initiale.

• Au début, il se souvient parfaitement qu'il y a un "chat".
• Mais vers la fin, alors qu'il ajoute les derniers détails, il a l'impression que le "vase rouge" est devenu un "vase bleu", ou qu'il a oublié de mettre le chat "à côté" du fauteuil.

C'est ce qu'ils appellent l'"Oubli du Prompt". L'information textuelle se dilue et se perd au fur et à mesure que l'image se précise. C'est comme si vous chuchotiez une recette à un cuisinier, et qu'à chaque fois qu'il ajoute un ingrédient, il oublie un peu plus ce que vous aviez dit au début.

2. La Solution : La "Réinjection" (Le Mémo-Post-it)

Pour résoudre ce problème, les chercheurs proposent une astuce géniale appelée "Prompt Reinjection" (Réinjection du Prompt).

L'analogie du Post-it :
Imaginez que vous êtes cet artiste. Vous avez une feuille de papier avec votre recette complète.

• Sans la solution : Vous lisez la recette, puis vous la posez sur la table. Plus vous avancez dans la cuisine, plus vous vous éloignez du papier, et plus vous risquez d'oublier les détails.
• Avec la solution (Prompt Reinjection) : À chaque étape de la cuisson (chaque couche de l'IA), quelqu'un vient vous rappeler ce que vous aviez écrit sur le papier au tout début. Il ne vous donne pas la recette entière, mais il vous "réinjecte" le souvenir frais de ce que vous deviez faire.

Techniquement, le modèle prend les informations textuelles "fraîches" (celles du début, qui sont encore très claires) et les réinjecte dans les couches profondes où l'IA commence à se tromper. C'est comme si on collait un post-it frais sur chaque étage du gratte-ciel en construction pour rappeler : "N'oublie pas : le vase est ROUGE, pas bleu !"

3. Comment ça marche sans réapprendre ?

La chose la plus cool, c'est que cette méthode est "sans entraînement" (training-free).

• On n'a pas besoin de rééduquer l'artiste (ce qui prendrait des mois et des millions de dollars).
• On ne change pas le cerveau de l'IA.
• On se contente de lui donner un petit coup de pouce au moment où elle dessine, en lui rappelant gentiment ce qu'elle a oublié.

C'est comme si vous donniez un petit coup de coude à votre ami pour lui rappeler de ne pas oublier le sel, sans avoir besoin de lui apprendre à cuisiner à nouveau.

4. Les Résultats : Des Dessins Plus Fidèles

Grâce à cette astuce, les résultats sont bluffants :

• Mieux pour les nombres : Si vous demandez "4 chiens", l'IA dessine vraiment 4 chiens, pas 3 ou 5.
• Mieux pour les couleurs : Le "chat noir" reste noir, même si le fond est complexe.
• Mieux pour la position : Le "vase à droite du chat" reste bien à droite.

Les chercheurs ont testé cela sur plusieurs modèles (SD3, FLUX, etc.) et ont vu que l'IA obéit beaucoup mieux aux instructions, sans que la qualité de l'image ne baisse. L'image reste belle, mais elle respecte enfin toutes les règles que vous aviez fixées.

En résumé

Ce papier dit essentiellement : "Les IA d'aujourd'hui oublient vos instructions compliquées quand elles en ont fini avec le début de la tâche. Notre solution est de leur rappeler constamment vos instructions tout au long du processus, comme un mémo qui ne s'efface jamais, pour obtenir un résultat parfait."

C'est une astuce simple, intelligente et gratuite pour rendre nos robots artistes beaucoup plus obéissants !

Résumé technique

1. Le Problème : L'Oubli de Prompt (Prompt Forgetting)

Les auteurs identifient un phénomène critique dans les Transformers de Diffusion Multimodaux (MMDiTs), tels que SD3, SD3.5, FLUX et Qwen-Image. Contrairement aux architectures précédentes (U-Net) où le texte agit comme une condition statique externe, les MMDiTs traitent les tokens textuels et les latents visuels de manière unifiée au sein d'une même pile de transformeurs, permettant une interaction bidirectionnelle.

Cependant, cette évolution conjointe crée un déséquilibre de supervision :

• Supervision asymétrique : L'objectif de perte (denoising) est défini uniquement dans l'espace latent visuel. Les tokens visuels reçoivent une supervision directe, tandis que les représentations textuelles ne sont mises à jour que de manière indirecte via l'attention conjointe.
• Conséquence : À mesure que la profondeur du réseau augmente, les caractéristiques textuelles subissent une dérive significative. Les informations sémantiques fines du prompt (attributs, relations spatiales, nombres) sont progressivement « oubliées » ou deviennent irrécupérables dans les couches profondes.
• Preuve empirique : Les auteurs montrent, via des analyses de dérive de distribution (CKNNA, PCA) et des sondages (probing) de récupération d'attributs, que la capacité à décoder les attributs linguistiques diminue de manière monotone avec la profondeur des couches.

2. Méthodologie : Prompt Reinjection (Réinjection de Prompt)

Pour résoudre ce problème sans réentraînement (training-free), les auteurs proposent Prompt Reinjection, une intervention au moment de l'inférence.

Principe de base :
La méthode consiste à réinjecter les caractéristiques textuelles issues des couches peu profondes (où l'information sémantique est encore intacte) dans les blocs transformeurs plus profonds, via des connexions résiduelles.

Mécanisme technique :
La réinjection naïve échoue souvent en raison de décalages de distribution (échelle, décalage) et géométriques (système de coordonnées) entre les couches. Pour y remédier, la méthode intègre deux étapes d'alignement :

1. Ancrage et Restauration de la Distribution (Distribution Anchoring) :

   • Normalisation des caractéristiques d'origine (couches peu profondes) et de cible (couches profondes) via la Normalisation par Couche (Layer Normalization).
   • Après la fusion, les caractéristiques sont projetées pour restaurer la moyenne et l'écart-type originaux de la couche cible, assurant la stabilité numérique.

2. Alignement Géométrique (Geometry Alignment) :

   • Utilisation d'une transformation de Procrustes Orthogonale pour aligner les systèmes de coordonnées latents entre les couches.
   • Une matrice de rotation orthogonale $R$ est calculée (une seule fois lors d'une phase de calibration sur un ensemble de données comme COCO-5K) pour minimiser l'erreur de reconstruction entre les caractéristiques normalisées de la couche source et de la couche cible.
   • Lors de l'inférence, les caractéristiques sources sont alignées via cette rotation avant d'être ajoutées.

Formule de réinjection :
Pour une couche cible $l$, les nouvelles caractéristiques textuelles $T^{(l)}_{new}$ sont obtenues en ajoutant une version alignée et pondérée des caractéristiques de la couche source $l_{ori}$ :
$$T^{(l)}_{new} = \hat{T}^{(l)} + w \cdot (\hat{T}^{(l_{ori})} \cdot R)$$
Où $w$ est un poids d'injection faible (ex: 0.025) et $\hat{T}$ désigne les caractéristiques normalisées.

3. Contributions Clés

1. Identification et Quantification de l'Oubli : Première caractérisation systématique du phénomène d'oubli de prompt dans les MMDiTs, démontré par une dégradation monotone de la récupération d'attributs linguistiques en profondeur.
2. Solution sans Réentraînement : Développement d'une méthode d'inférence (Prompt Reinjection) qui améliore la fidélité aux instructions sans modifier les poids du modèle ni nécessiter de LoRA ou de fine-tuning.
3. Alignement Inter-couches : Introduction d'un pipeline d'alignement (statistique + géométrique) permettant la fusion efficace de caractéristiques provenant de différentes profondeurs du réseau, résolvant les problèmes de compatibilité de distribution.
4. Amélioration Généralisée : Démonstration que la méthode fonctionne sur une variété d'architectures MMDiTs (SD3, SD3.5, FLUX, Qwen-Image) et pour divers types de contraintes (couleurs, comptage, relations spatiales).

4. Résultats Expérimentaux

Les évaluations ont été menées sur plusieurs benchmarks de référence : GenEval, DPG-Bench et T2I-CompBench++.

• Suivi des Instructions (Instruction Following) :

  • Sur GenEval, Prompt Reinjection améliore les scores globaux de SD3.5 (+6,48 %) et de FLUX (+5,64 %).
  • Les gains sont particulièrement marqués pour les tâches de relations spatiales (Position) et de comptage, qui sont les plus touchées par l'oubli de prompt.
  • Le modèle Qwen-Image (20B paramètres) montre également des améliorations significatives (+3% en comptage, +6,3% en position), prouvant que même les grands modèles souffrent de cette dégradation.

• Qualité de l'Image :

  • Les métriques de préférence humaine (HPSv2, ImageReward, PickScore) et d'alignement global (CLIP Score) restent stables ou s'améliorent légèrement.
  • Cela confirme que la méthode corrige la dérive sémantique sans introduire d'artefacts visuels ni dégrader la fidélité de l'image.

• Analyse Qualitative :

  • Les comparaisons montrent que le modèle avec réinjection respecte beaucoup mieux les contraintes complexes (ex: « quatre chiens », « un oiseau au-dessus d'un ballon », « une clé sur un poulet ») par rapport aux modèles de base qui échouent souvent sur le nombre ou la position.

• Coût Computations :

  • L'ajout de la réinjection entraîne une surcharge minime (environ 8,83 % de FLOPs supplémentaires par bloc pour l'alignement complet), ce qui est négligeable par rapport au gain de performance.

5. Signification et Impact

Ce travail est significatif car il remet en question l'hypothèse implicite selon laquelle les MMDiTs préservent naturellement l'information textuelle tout au long du processus de débruitage. En identifiant l'asymétrie de supervision comme cause racine de l'oubli, les auteurs proposent une solution élégante et efficace.

Implications futures :

• Architecture : Cela suggère que les futures architectures MMDiTs devraient peut-être intégrer une supervision directe sur la branche textuelle ou des mécanismes de rétention sémantique.
• Optimisation : La méthode offre un moyen simple d'améliorer les modèles existants déployés sans coût de réentraînement, ce qui est crucial pour l'industrie.
• Compréhension : L'étude fournit des outils d'analyse (CKNNA, sondage de couches) essentiels pour comprendre la dynamique interne des modèles de diffusion multimodaux.

En résumé, Prompt Reinjection est une avancée majeure pour la fiabilité des générateurs d'images par IA, garantissant que les instructions complexes sont respectées du début à la fin du processus de génération.