Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

Dit paper introduceert prompt reinjectie, een trainingsvrije methode die het vergeten van prompts in multimodale Diffusion Transformers tegengaat door vroege tekstrepresentaties later in het proces opnieuw in te brengen, wat leidt tot verbeterde instructievolging en generatiekwaliteit.

De kern

Stel je voor dat je een zeer getalenteerde schilder hebt die fantastische plaatjes kan maken op basis van wat je hem vertelt. Je zegt: "Teken een blauwe kat die op een rode bal zit." De schilder begint te werken, maar naarmate hij dieper in het proces duikt (hij voegt steeds meer details toe), begint hij zijn eigen instructies te vergeten.

Op het eind is de kat misschien nog wel blauw, maar hij zit op een groene auto in plaats van op een rode bal. Of misschien is de kat wel weg, en heb je alleen nog maar een rode bal.

Dit is precies het probleem dat de auteurs van dit paper hebben ontdekt bij moderne AI-tekst-naar-beeld modellen (zoals SD3, FLUX en Qwen-Image). Ze noemen dit "Prompt Vergeten" (Prompt Forgetting).

Hier is de uitleg in simpele taal, met een paar creatieve vergelijkingen:

1. Het Probleem: De "Fluisterende Boodschapper"

In deze AI-modellen werken tekst en beelden samen in één groot brein (een Transformer).

• Hoe het zou moeten werken: De tekst (je instructie) is als een hoofdingenieur die de hele bouw van het huis (het beeld) leidt. Hij staat aan het begin en geeft de orders.
• Wat er eigenlijk gebeurt: Naarmate de AI dieper in de lagen van haar eigen brein duikt om het beeld te maken, verandert de tekst steeds meer. Het is alsof de hoofdingenieur zijn instructies fluistert aan een collega, die het fluistert aan een volgende, en die weer aan een volgende.
• Het resultaat: Na 20 of 30 "fluisteringen" (laagjes) is de originele boodschap verdraaid. De AI vergeet details als "vier honden" (in plaats van drie) of "een kat rechts van een hond". De tekst-informatie verdwijnt langzaam in de nevel.

2. De Oplossing: "Prompt Herinjectie"

De auteurs hebben een slimme, gratis truc bedacht om dit op te lossen. Ze noemen het Prompt Reinjection (Prompt Herinjectie).

Stel je voor dat je een lange, moeilijke wandeling maakt met een vriend die de weg vergeet.

• De oude manier: Je hoopt dat hij het zich herinnert, maar hij loopt steeds meer de verkeerde kant op.
• De nieuwe manier (Prompt Reinjection): Je geeft je vriend elke paar minuten een frisse kopie van de originele kaart terug. Je zegt: "Hé, kijk even naar dit stukje papier van het begin, daar stond 'rechtsaf bij de boom'."

In de AI betekent dit:

1. De AI haalt de tekst-informatie uit het begin van het proces (waar het nog heel helder en fris is).
2. Tijdens het maken van het plaatje, in de diepere lagen waar de verwarring begint, injecteert de AI deze frisse tekst-informatie opnieuw in het systeem.
3. Dit zorgt ervoor dat de AI zich blijft herinneren wat je eigenlijk wilde, zelfs als diep in het proces de verwarring toeslaat.

3. Waarom werkt dit zo goed?

De onderzoekers hebben ontdekt dat de AI vooral moeite heeft met ruimtelijke details (links/rechts, boven/onder) en aantallen (vier, vijf).

• Zonder de truc: "Teken een blauwe auto links van een rode fiets." -> De AI tekent een blauwe auto en een rode fiets, maar ze staan door elkaar.
• Met de truc: De AI krijgt steeds weer de herinnering: "Nee, de auto moet links staan!" en corrigeert zichzelf direct.

Het resultaat is dat de AI veel beter doet wat je vraagt, zonder dat de kwaliteit van het plaatje verslechtert. Het plaatje blijft mooi, maar de instructies worden nu perfect gevolgd.

Samenvatting in één zin

Deze paper zegt: "Moderne AI-verfkwasten vergeten hun instructies naarmate ze dieper in het schilderij duiken, maar als we ze af en toe een 'herinneringsbriefje' van het begin teruggeven, maken ze precies wat we willen."

Het is een slimme, goedkope manier om AI slimmer te maken zonder dat je de hele AI opnieuw hoeft te trainen. Je geeft haar gewoon een beetje extra geheugensteun op het moment dat ze het nodig heeft.

Technische samenvatting

Titel: Prompt Reinjection: Verlichting van Prompt Vergeten in Multimodale Diffusie-Transformers

1. Het Probleem: Prompt Vergeten (Prompt Forgetting)

De auteurs identificeren een fundamenteel probleem in moderne Multimodale Diffusie-Transformers (MMDiTs), zoals Stable Diffusion 3 (SD3), SD3.5, FLUX en Qwen-Image. Hoewel deze modellen tekst- en visuele tokens gezamenlijk verwerken in een unified transformer-stack (in tegenstelling tot oudere U-Net-architecturen met statische cross-attention), vertonen ze een fenomeen dat "Prompt Forgetting" wordt genoemd.

• Oorzaak: Er bestaat een supervisie-asymmetrie. De trainingsdoelstelling (bijv. $\epsilon$-voorspelling) is louter gericht op het reconstrueren van de visuele latenten. Visuele tokens ontvangen directe supervisie, terwijl tekstuele tokens alleen indirect worden bijgewerkt via hun invloed op de visuele reconstructie via gezamenlijke attention-mechanismen.
• Gevolg: Omdat er geen expliciete semantische constraints zijn voor de tekst-tokens, ondergaan deze representaties een significante drift naarmate de diepte van het netwerk toeneemt. De fijne details van de prompt (zoals attributen, aantallen en ruimtelijke relaties) worden geleidelijk "vergeten" of onherstelbaar in de diepere lagen.
• Observatie: Empirische analyses tonen aan dat de recoverability (herstelbaarheid) van token-niveau attributen monotoon afneemt met de diepte van het model. Ruimtelijke relaties lijken het meest kwetsbaar te zijn.

2. Methodologie: Prompt Reinjection

Om dit probleem op te lossen, stellen de auteurs Prompt Reinjection voor. Dit is een training-vrije (training-free) interventie die tijdens de inferentie wordt toegepast. Het doel is om de hoge kwaliteit van de semantische informatie uit de vroege lagen opnieuw te injecteren in de diepere blokken van het model.

De methode bestaat uit twee hoofdcomponenten om de distributie- en geometrische discrepanties tussen lagen op te lossen:

1. Distributie Ankeren en Herstel (Distribution Anchoring & Restoration):

   • Omdat features in verschillende lagen verschillende schalen en offsets hebben, worden de features eerst genormaliseerd met Layer Normalization (LN).
   • Na de injectie worden de features weer geprojecteerd naar de oorspronkelijke statistische verdeling (mean en standaarddeviatie) van de doel-laag om de stabiliteit van de generatie te waarborgen.

2. Geometrie Alignering via Orthogonale Procrustes:

   • Normalisatie lost alleen schaal- en verschuivingsproblemen op, maar niet de rotatie van het coördinatenstelsel in de latent space.
   • Tijdens een eenmalige kalibratiefase (met een dataset zoals COCO-5K) wordt een orthogonale rotatiematrix $R$ berekend die de features van de oorsprongslaag (shallow) zo goed mogelijk aligneert met de doel-laag (deep) via Singular Value Decomposition (SVD).
   • Tijdens de inferentie worden de oorspronkelijke features eerst gealigneerd met $R$ en vervolgens als een residual toegevoegd aan de diepere features:
     $$T_{added} = \hat{T}_{tgt} + w \cdot \hat{T}_{ori}R$$
     Waarbij $w$ een hyperparameter is die de injectiestrength regelt.

Selectie van Lagen:

• Oorsprongslaag ($l_{ori}$): Meestal de eerste laag na de initiële scherpe distributietransitie (bijv. laag 1 of 2), waar de tekstfeatures nog de meeste semantische fideliteit behouden.
• Doellaag: Alle diepere blokken na de oorsprongslaag.

3. Belangrijkste Bijdragen

• Identificatie van Prompt Forgetting: Het paper biedt rigoureuze bewijzen (via Conditional K-Nearest Neighbor Alignment en layer-wise probing) dat token-niveau prompt-informatie progressief verloren gaat in de diepere lagen van MMDiTs.
• Training-Free Oplossing: De introductie van Prompt Reinjection, een methode die geen hertraining van het model vereist en direct toepasbaar is op bestaande inferentie-pipelines.
• Geometrische Alignering: Het toepassen van Orthogonale Procrustes-transformatie om cross-layer feature mismatch effectief op te lossen, wat cruciaal is voor stabiele injectie.
• Uitgebreide Validatie: De methode wordt getest op vier verschillende state-of-the-art modellen (SD3, SD3.5, FLUX, Qwen-Image) en diverse benchmarks.

4. Resultaten

De experimenten tonen consistente verbeteringen aan op meerdere fronten:

• Instruction Following: Op de GenEval benchmark verbeterde Prompt Reinjection de algehele scores aanzienlijk:
  • SD3.5: +6.48%
  • FLUX: +5.64%
  • De verbetering is het meest opvallend bij taken die complexe redenering vereisen, zoals spatial relations (ruimtelijke relaties), counting (tellen) en multi-object composition.
• Kwaliteit en Aesthetiek: De verbetering in instructie-opvolging gaat niet ten koste van de beeldkwaliteit. Metrieken zoals HPSv2 (Human Preference Score), ImageReward, PickScore en CLIP-score blijven gelijk of verbeteren licht.
• Kwalitatieve Verbetering: Visuele vergelijkingen tonen aan dat het model beter in staat is om complexe prompts te volgen, zoals "een zwarte kat en een witte snor" of "vier handtassen", waar de basismodellen vaak fouten maken in aantallen of attributen.
• Efficiëntie: De methode voegt slechts een minimale rekentijd en geheugenn overhead toe (ongeveer 8% extra FLOPs per blok voor de geometrische alignering), wat een uitstekende afweging is voor de winst in prestaties.

5. Betekenis en Impact

Dit paper is significant omdat het een fundamenteel inzicht biedt in het gedrag van Multimodale Diffusie-Transformers: dat de gezamenlijke verwerking van tekst en beeld, hoewel krachtig, leidt tot semantische drift van de tekstrepresentaties.

• Paradigmaverschuiving: Het suggereert dat tekst in MMDiTs niet als een statische conditie kan worden behandeld, maar dat de dynamiek van de tekstrepresentatie actief moet worden beheerd.
• Praktische Toepasbaarheid: Omdat Prompt Reinjection training-vrij is, kunnen ontwikkelaars en onderzoekers direct bestaande, geavanceerde modellen (zoals FLUX of SD3.5) verbeteren zonder de enorme kosten van hertraining.
• Toekomstperspectief: Het paper opent de deur voor toekomstig werk dat zich richt op het toevoegen van directe supervisie aan de teksttak tijdens training of het leren van dynamische injectie-weights per timestep.

Kortom, Prompt Reinjection biedt een elegante en effectieve oplossing voor het "vergeten" van prompts in de nieuwste generatie beeldgeneratiemodellen, waardoor de betrouwbaarheid en precisie van tekst-naar-beeld generatie aanzienlijk wordt verhoogd.