ImageRAGTurbo: Towards One-step Text-to-Image Generation with Retrieval-Augmented Diffusion Models

Probleemstelling

Diffusiemodellen zijn momenteel de leidende technologie voor tekst-naar-beeldgeneratie, maar ze kampen met een significant nadeel: het iteratieve steekproefproces (sampling) dat nodig is om ruis om te zetten in een coherent beeld, veroorzaakt hoge latentie.

• Huidige beperkingen: Standaard modellen vereisen 25 tot 50 stappen. Snellere "few-step" modellen (1-4 stappen), zoals die gebruikmaken van distillatie, verminderen de latentie aanzienlijk, maar gaan vaak ten koste van de beeldkwaliteit en de nauwkeurigheid van de prompt-uitlijning (prompt fidelity).
• Specifiek uitdaging: In het extreme geval van één-stap generatie (one-step) hebben bestaande modellen moeite om complexe visuele concepten correct weer te geven (bijv. een boot met een vuurtoren op de achtergrond).
• Bestaande oplossingen: Methoden zoals modelquantisatie of caching versnellen het proces maar vereisen nog steeds veel stappen. Het trainen van nieuwe few-step modellen via distillatie is computatief zwaar en kostbaar.

Methodologie: ImageRAGTurbo

De auteurs stellen ImageRAGTurbo voor, een raamwerk dat Retrieval-Augmented Generation (RAG) combineert met few-step diffusiemodellen. Het doel is om de kwaliteit en prompt-uitlijning te verbeteren zonder de snelheid te offeren.

Het proces verloopt als volgt:

1. Retrieval: Gegeven een doeltekstprompt ($p_{tgt}$), worden relevante tekst-beeldparen ($p_{retr}, x_{retr}$) opgehaald uit een database.
2. H-space Injectie: In plaats van het hele model opnieuw te trainen, wordt de informatie uit de opgehaalde beelden geïnjecteerd in de H-space van de UNet-denoiser. De H-space verwijst naar de feature maps van de diepste laag van de UNet, die verantwoordelijk zijn voor hoogwaardige concepten (zoals objecten en attributen).
3. Efficiënte Finetuning met H-Adapter:
   • De auteurs tonen eerst aan dat directe injectie zonder finetuning al helpt, maar dat het optimaliseren van de mengsterkte (blending strength) per prompt een dure zoektocht vereist.
   • Om dit op te lossen, introduceren ze een trainbare adapter in de H-space. Deze adapter gebruikt een cross-attention mechanisme om automatisch de correlatie tussen de opgehaalde content en de doelprompt te bepalen en deze te mengen.
   • De adapter is lichtgewicht en wordt gefinetuned terwijl de rest van de denoiser (de UNet) vaststaat (frozen).
4. Trainingsprocedure:
   • Het model wordt getraind via Latent Adversarial Training. Een discriminator onderscheidt tussen de latente representaties van de student (few-step) en de leraar (meerdere stappen).
   • Er worden extra loss-functies gebruikt voor stabiliteit: Score Distillation Loss en Latent LPIPS Loss.
   • De decoder van de student wordt ook geoptimaliseerd met Low-Rank Adaptation (LoRA).

Belangrijkste Bijdragen

1. Nieuw Framework: Een retrieval-verrijkt raamwerk voor few-step beeldgeneratie dat de prompt-uitlijning en beeldkwaliteit aanzienlijk verbetert zonder de snelheid te verminderen.
2. Efficiënte Finetuning: Een methode die gebruikmaakt van een lichtgewicht adapter in de H-space, wat de rekenvereisten verlaagt ten opzichte van traditionele trainingsmethoden voor few-step modellen.
3. Superieure Prestaties: Uitgebreide experimenten tonen aan dat de methode bestaande few-step benaderingen overtreft op zowel kwaliteit als prompt-uitlijning, terwijl de inferentiesnelheid vergelijkbaar blijft.

Resultaten

De auteurs hebben hun methode getest op de MS-COCO en TIFA benchmarks en vergeleken met modellen zoals Stable Diffusion Turbo, Latent Consistency Models en RDM (Retrieval-Augmented Diffusion Model).

• Kwaliteit vs. Snelheid: ImageRAGTurbo genereert beelden in één stap (1 NFE - Number of Function Evaluations) met een kwaliteit die vergelijkbaar is met een 50-staps Stable Diffusion model.
• Metrics:
  • MS-COCO: Bereikte een CLIP-score van 0.323 (vs. 0.319 voor SD Turbo) en een lagere FID-score (betere beeldkwaliteit).
  • TIFA (Prompt Fidelity): Behaalde een TIFA-score van 0.801 (vs. 0.779 voor SD Turbo zonder retrieval). Dit is een verbetering van 2.2% ten opzichte van de snelste concurrenten.
  • Vergelijking met 50-staps: Hoewel het iets onder de 50-staps teacher ligt (ongeveer 1.2% lager in TIFA), is het 25x sneller.
• Efficiëntie: De totale inferentietijd is 116.7 ms per beeld, wat bijna gelijk is aan SD Turbo (113.8 ms). De extra overhead door de retrieval en de adapter is slechts 2.5%.
• Kwalitatieve resultaten: Het model slaagt er beter in dan andere few-step modellen om specifieke objecten en ruimtelijke relaties correct weer te geven (bijv. "een boot op het water met een vuurtoren op de achtergrond").

Betekenis en Toekomstperspectief

ImageRAGTurbo is een doorbraak omdat het het compromis tussen snelheid en kwaliteit in text-to-image generatie effectief doorbreekt.

• Het bewijst dat het gebruik van externe kennis (retrieved data) de taak voor het diffusiemodel vereenvoudigt, waardoor minder stappen nodig zijn om een hoogwaardig resultaat te bereiken.
• Het biedt een praktische oplossing voor real-time en interactieve toepassingen waar lage latentie cruciaal is.
• Toekomst: De auteurs wijzen erop dat het huidige systeem gebaseerd is op CLIP-retrieval. Toekomstig werk zou kunnen focussen op fijnkorreliger (fine-grained) en compositie-gebaseerde retrieval om nog nauwkeurigere resultaten te behalen.

Kortom, ImageRAGTurbo stelt een nieuwe standaard neer voor snelle, hoogwaardige beeldgeneratie door slimme integratie van retrieval-technieken in de latent space van diffusiemodellen.