Technische Samenvatting: Kwaliteitsbewuste Modulatie voor Diffusie-Transformers

Probleemstelling

Moderne tekst-naar-afbeelding diffusiemodellen, in het bijzonder Diffusion Transformers (DiT) zoals Stable Diffusion 3.5 (SD3.5), vertrouwen op conditioneringssignalen—specifiek tijdstap-embeddings en gepoolde prompt-embeddings—om het denoisingsproces te moduleren via Adaptive Layer Normalization (AdaLN). Hoewel deze signalen effectief de huidige ruisniveau en semantische intentie communiceren, missen ze kwaliteitsbewuste informatie. Bijgevolg kan het modulatiemechanisme niet direct reageren op de specifieke foutmodi van een sample, zoals foutaccumulatie, incorrecte interpretatie van de prompt of gehallucineerde kenmerken die verschijnen naarmate de latente afbeelding evolueert. De baseline-conditionering is onafhankelijk van de huidige latente staat ($\partial M_{base} / \partial x_t = 0$), wat voorkomt dat het model fideliteitsproblemen kan corrigeren op basis van de werkelijke kwaliteit van de tussenliggende afbeelding.

Methodologie

De auteurs stellen de Quality Representation Module (QRM) voor, een lichtgewicht transformer-gebaseerde adapter die is ontworpen om kwaliteitsbewuste signalen in het diffusieproces te injecteren zonder de vooraf getrainde backbone aan te passen.

Architectuur

QRM is geïmplementeerd als een transformer-decoder die additieve modulatie-offsets ($M_{qrm}$) voorspelt voor de baseline AdaLN-parameters ($M_{base}$).

• Inputs: De module ontvangt de huidige latente afbeelding ($x_t$), de gepoolde prompt-embedding ($p$), de tijdstap-embedding ($t$) en de baseline modulatieparameters ($M_{base}$).
• Verwerking:
  1. De latente afbeelding $x_t$ wordt geprojecteerd in latente tokens via een geleerde $1\times1$ convolutie en adaptieve pooling.
  2. Baseline modulatieparameters worden genormaliseerd en geprojecteerd om query-tokens ($Z_0$) te vormen.
  3. Een key-value geheugen ($KV$) wordt geconstrueerd uit de latente tokens, de prompt en de tijdstap-embeddings.
  4. Door middel van cross-attention en self-attention lagen attenderen de query-tokens op het $KV$-geheugen om een verfijnde representatie te genereren.
  5. De output wordt teruggeprojecteerd naar de natuurlijke AdaLN-breedtes om de offset-vector $M_{qrm}$ te produceren.
• Toepassing: De uiteindelijke modulatie die op de transformer-blokken wordt toegepast is $M_t = M_{base,t} + M_{qrm,t}$.

Trainingsstrategie (ReFL)

QRM wordt getraind met een Reward Feedback Learning (ReFL) strategie:

• Bevroren Componenten: De diffusie-backbone (denoiser, tekstencoders, VAE), de sampler en het beloningsmodel blijven bevroren. Alleen de QRM-parameters worden bijgewerkt.
• Doelstelling: De module leert een scalaire beloningsscore te maximaliseren die afkomstig is van een bevroren beloningsmodel (bijv. HPSv2.1 of CLIPScore) toegepast op een afbeelding die is gedecodeerd uit een tussenliggende denoisingsstap.
• Verliesfuncties: Het paper maakt gebruik van ofwel een beloningsmaximalisatieverlies of een margin-seeking verlies, wat de QRM-gemoduleerde afbeelding ertoe aanzet de baseline met een specifieke marge te overtreffen.
• Tijdstapselectie: Om de efficiëntie en de focus op structurele vorming te verbeteren, wordt QRM alleen geactiveerd tijdens een vooraf gedefinieerde subset van tijdstappen met hoge ruis (bijv. $t \in \{50, \dots, 26\}$), in plaats van tijdens het gehele diffusieschema.

Inferentie

Tijdens de inferentie fungeert de getrainde QRM-module als een plug-in adapter. Het voorspelt modulatie-offsets op basis van de huidige latente staat en voegt deze toe aan de baseline-parameters. Dit stelt het model in staat om het denoisingspad te sturen naar staten met een hogere kwaliteit zonder de backbone te hertrainen of het sampling-schema te wijzigen.

Belangrijkste Bijdragen

1. QRM Architectuur: Introductie van een lichtgewicht transformer-module die kwaliteitsbewuste representaties leert van bestaande modelinputs (latent, prompt, tijdstap) om modulatieparameters te verfijnen.
2. Niet-invasieve Integratie: Een mechanisme om de modulatie van de diffusion transformer aan te vullen met latente afbeeldingkwaliteitssignalen zonder de onderliggende diffusion backbone, VAE of het sampling-schema te wijzigen.
3. Beloningsgestuurde Modulatie: Demonstratie dat voorspelde modulatie-updates ($M_{qrm}$) effectief het denoisingsproces kunnen verfijnen bij elk tijdstap, wat de prompt-adherentie en beeldgetrouwheid verbetert.
4. Uitgebreide Evaluatie: Uitgebreide ablatie-studies die de effecten analyseren van de QRM-architectuur grootte, actieve tijdstapsbereiken en trainingsbeloningsmetrieken op de beeldkwaliteit.

Experimentele Resultaten

De auteurs hebben QRM geëvalueerd met behulp van Stable Diffusion 3.5 (SD3.5) op de Parti-Prompts benchmark.

• Kwantitatieve Prestaties:
  • Vergeleken met de SD3.5 baseline, behaalde QRM getraind met HPSv2.1 als het beloningssignaal de hoogste scores onder alle geëvalueerde methoden (inclus inclusief prompt-verfijning, fine-tuning en ReNeg-benaderingen).
  • Verbeteringen: +0,20 in Aesthetic score en +0,82 in HPSv2.1 score.
  • Wanneer getraind met CLIPScore, verbeterde QRM de menselijke voorkeurs-alignement (HPSv2.1) terwijl de vergelijkbare perceptuele kwaliteit behouden bleef.
• Kwalitatieve Analyse:
  • Visualisaties van denoisings-trajecten laten zien dat het toepassen van QRM tijdens de vroege, hoog-ruis tijdstappen leidt tot herkenbare semantische structuren die aanzienlijk eerder ontstaan dan in de baseline.
  • Het toepassen van QRM alleen in de midden- of latere stadia leverde kleinere verbeteringen op, wat bevestigt dat de globale structuur vroeg in het diffusieproces wordt vastgesteld.
• Ablatie-bevindingen:
  • Capaciteit: Grotere QRM-modellen die sterker zijn geoptimaliseerd voor de specifieke trainingsbeloning (bijv. HPSv2.1) vertoonden een lichte trade-off in andere metrieken (bijv. CLIPScore).
  • Tijdstapbereik: Modulatie tijdens vroege tijdstappen ($50 \to 26$) was consequent effectiever voor alignement dan modulatie in latere stadia.
  • Beloningsmetriek: De keuze van het beloningsmodel (CLIPScore vs. HPSv2.1) vormde direct de perceptuele biases die door QRM werden geleerd.

Betekenis en Claims

Het paper claimt dat modulatie-niveau interventies een schaalbare en effectieve weg bieden voor het afstemmen van grote diffusiemodellen op menselijke voorkeurssignalen. In tegen tegenstelling tot methoden die dure fine-tuning van de volledige backbone vereisen (zoals LoRA of volledige fine-tuning), bereikt QRM consistente kwaliteitsverbeteringen door volledig te opereren via additieve modulatie-updates geconditioneerd op de evoluerende latente representatie.

De auteurs positioneren QRM als een flexibele, kwaliteitsbewuste adapter die gestuurd kan worden door verschillende beloningssignalen en toegepast kan worden op diverse diffusion transformer-architecturen. Ze erkennen dat QRM de biases van de trainingsbeloningsmodellen erft en suggeren dat toekomstig werk de mogelijkheden kan verkennen van multi-reward training en adaptieve tijdstapselectie om de robuustheid en generalisatie verder te versterken. Het werk demonstreert dat aanzienlijke winsten in beeldgetrouwheid en prompt-alignement kunnen worden behaald zonder de vooraf getrainde gewichten van state-of-the-art modellen te wijzigen.