Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation

Technische Samenvatting: Null-Text Test-Time Alignment van Text-to-Image Diffusiemodellen

Probleemstelling
Pre-getrainde text-to-image diffusiemodellen zijn vaak getraind op grootschalige, ongecureerde webdatasets die misalignede of ongewenste inhoud bevatten. Het afstemmen (alignment) van deze modellen op specifieke menselijke waarden of doelstellingen (bijv. esthetische kwaliteit, specifieke voorkeuren) tijdens de inferentie is cruciaal voor betrouwbare inzetbaarheid. Bestaande alignment-methoden vallen over het algemeen uiteen in twee categorieën, die beide kampen met significante beperkingen:

1. Fine-tuning methoden: Deze updaten de modelparameters om een beloningsfunctie te optimaliseren. Hoewel effectief, zijn ze computationeel duur en gevoelig voor belonings-overoptimalisatie (reward over-optimization of reward hacking), waarbij het model overfit op proxy-beloningen, waardoor het generalisatie over andere beloningen en outputdiversiteit verliest.
2. Test-Time Alignment (TTA) methoden: Deze optimaliseren latente variabelen of ruisvariabelen, samplen uit beloningsgeconditioneerde posteriors, of gebruiken discrete zoekalgoritmen zonder de modelgewichten bij te werken. Echter, deze methoden lijden vaak aan belonings-onder- of overoptimalisatie vanwege de zeer ongestructureerde aard van latente en ruisruimtes. Optimalisatie in deze ruimtes kan leiden tot "reward hacking", waarbij het model niet-semantische ruispatronen exploiteert om de doelbeloning te maximaliseren, terwijl de visuele kwaliteit, semantische coherentie en prestaties op andere metrieken verslechteren.

Methodologie: Null-Text Test-Time Alignment (Null-TTA)
De auteurs stellen Null-TTA voor, een training-vrij framework dat diffusiemodellen afstemt door de onvoorwaardelijke (null-text) embedding binnen Classifier-Free Guidance (CFG) te optimaliseren, in plaats van latente of ruisvariabelen te manipuleren.

• Kerninzicht: In CFG dient de onvoorwaardelijke embedding ($\phi$) als een geometrisch anker voor de generatieve distributie van het model. In tegenstelling tot ongestructureerde ruisruimtes, bezit de tekst-embeddingruimte een gestructureerd semantisch karakter. Optimalisatie binnen deze ruimte zorgt ervoor dat de alignment plaatsvindt op een semantisch coherente manifold, wat voorkomt dat het model niet-semantische ruisartefacten exploiteert.
• Doelfunctie: De methode formuleert een geregulariseerde doelstelling die de generatieve distributie van het model naar een doelbeloning stuurt, terwijl de distributieve drift van het pre-getrainde model wordt bestraft. De doelstelling maximaliseert de verwachte beloning $R(x_0)$ terwijl de KL-divergentie tussen de door de geoptimaliseerde embedding $\phi'$ geïnduceerde distributie en de oorspronkelijke distributie geïnduceerd door $\phi$ wordt geminimaliseerd:
  $$\max_{\phi'} \left( \lambda_1 \mathbb{E}_{p(x_0|\phi')}[R(x_0)] - \lambda_2 \text{KL}(p(x_{0:T}, \phi') || p(x_{0:T}, \phi)) \right)$$
  De KL-term wordt gedecomposeerd in lokale divergenties tussen opeenvolgende denoising-stappen en een divergentie tussen de embedding-distributies. Dit wordt vereenvoudigd tot een gesloten vorm die de verschillen in voorspelde ruis ($\tilde{\epsilon}$) en de verschillen in embeddings ($\phi' - \phi$) omvat.
• Optimalisatiestrategie:
  • Adaptieve Regularisatie: Het regularisatiegewicht $\lambda_2$ wordt adaptief geëvanneerd (geanneald) naarmate het denoising-proces vordert ($t \to 0$). Sterke regularisatie in de vroege stappen stabiliseert de optimalisatie, terwijl zwakkere regularisatie in latere stappen fijnere alignment mogelijk maakt.
  • Greedy Search: In elke denoising-stap hanteert de methode een lichtgewicht greedy search. Het samplet $K$ kandidaat-latents uit de transitiekernel, schat de schone sample in via de Tweedie-formule, scoort deze met het beloningsmodel, en selecteert deterministisch de hoogst scorende kandidaat voor de volgende stap.
  • Niet-differentieerbare Beloningen: Voor scenario's waar de beloningsfunctie een black box is (niet-differentieerbaar), gebruikt de methode zeroth-order gradiëntschatting met behulp van willekeurige perturbaties van de null-text embedding.

Belangrijkste Bijdragen

1. Null-TTA Framework: Een nieuw, training-vrij TTA-framework dat alignment uitvoert in de gestructureerde semantische ruimte van tekst-conditioning, waardoor de valkuilen van ongestructureerde latente/ruis-optimalisatie worden vermeden.
2. Geprincipeerde Doelstelling: Een KL-geregulariseerde doelstelling die de generatieve distributie van het model zelf heroriënteert naar de doelbeloning, wat zorgt voor stabiele en efficiënte alignment zonder parameterupdates. Dit voorkomt reward hacking en behoudt de semantische coherentie.
3. State-of-the-Art Prestaties: Uitgebreide experimenten tonen aan dat Null-TTA superieure doelbeloningsprestaties bereikt terwijl het sterke cross-reward generalisatie en perceptuele kwaliteit behoudt.

Experimentele Resultaten
De auteurs evalueerden Null-TTA op Stable Diffusion v1.5 en SDXL met diverse doelbeloningen (HPSv2, PickScore, Aesthetic Score, ImageReward) en niet-differentieerbare doelstellingen (JPEG-compressibiliteit).

• Beloningsoptimalisatie vs. Generalisatie: Null-TTA verschuift de Pareto-front consequent naar buiten vergeleken met baselines (inclusief DNO, DAS, DSearch en TDPO). Het bereikt hogere doelbeloningen terwijl het de heldengehouden (held-out) beloningen aanzienlijk minder degradeert dan concurrenten, wat duidt op succesvolle alignment zonder overoptimalisatie.
• Kwalitatieve Prestaties: In uitdagende categorieën (tellen, compositionaliteit, ruimtelijk redeneren, ongebruikelijke kleuren, stijltransfer, onmogelijke scènes) produceert Null-TTA afbeeldingen die de prompt-restricties trouwer voldoen terwijl de globale coherentie behouden blijft, waarbij foutmodi zoals onjuiste objectaantallen of verstrengelde relaties worden vermeden.
• Gebruikersstudie: Een blind gebruikersonderzoek met 16 deelnemers toonde een statistisch significante voorkeur voor Null-TTA boven baselines, wat bevestigt dat beloningswinsten vertalen naar perceptuele kwaliteit.
• Computationele Efficiëntie: Null-TTA bereikt sterkere doelbeloningsprestaties met een lager GPU-geheugengebruik en vergelijkbare of snellere wall-clock tijden dan baselines. Deze efficiëntie komt voort uit het terugpropageren door enkel de cross-attention lagen (het updaten van de embedding) in plaats van de volledige U-Net.
• Niet-differentieerbare Beloningen: De methode optimaliseerde succesvol voor JPEG-bestemingsgrootte (een black-box doelstelling) met behulp van zeroth-order gradiënten, waarbij competitieve compressiecijfers werden behaald terwijl de visuele kwaliteit beter werd behouden dan bij gradiënt-gebaseerde baselines.

Significantie en Claims
Het artikel vestigt semantische-ruimte optimalisatie als een effectief en nieuw paradigma voor Test-Time Alignment. Door gebruik te maken van het gestructureerde karakter van de tekst-embeddingruimte, biedt Null-TTA inherente regularisatie die reward hacking voorkomt en de pre-getrainde kennis van het model behoudt. De auteurs beweren dat deze aanpak een eenvoudige maar krachtige alternatief biedt voor bestaande TTA-methoden, wat betrouwbare, efficiënte en generaliseerbare alignment van diffusiemodellen tijdens de inferentie mogelijk maakt zonder de computationele kosten van fine-tuning.