Technische Samenvatting: Stitched Value Model for Diffusion Alignment (StitchVM)

1. Probleemstelling

Diffusie- en stromingsgebaseerde generatieve modellen blinken uit in het vastleggen van data-distributies, maar vereisen vaak afstemming op taakspecifieke beloningen (bijv. prompt-trouw, esthetische voorkeur, menselijke voorkeur) om praktisch bruikbaar te zijn. De kernuitdaging bij deze afstemming is dat beloningen doorgaans worden gedefinieerd voor schone uitvoerafbeeldingen ($x_0$), terwijl afstemmingsprocedures (zowel tijdens inferentie als training) de verwachte beloning moeten schatten van schone samples die worden gegenereerd door ruisige tussenliggende latente variabelen ($z_t$). Deze schatting staat bekend als de waardefunctie.

Bestaande methoden voor het benaderen van deze waardefunctie lijden onder aanzienlijke afwegingen:

• Tweedie-benadering: Schat de posterior-middellijn van het schone sample en berekent de beloning op die proxy. Hoewel efficiënt, introduceert dit aanzienlijke bias in regimes met hoge ruis en vereist het dure extra evaluaties van de denoiser en de VAE-decoder.
• Monte Carlo (MC)-benadering: Voert meerdere denoising-trajecten uit vanaf $z_t$ en middelt de beloningen. Hoewel nauwkeuriger, leidt dit tot verbodsbepalende rekentkosten vanwege de noodzaak van meerdere volledige rollouts.
• Gestudeerde waarde-modellen: Het direct trainen van een waardefunctie voor ruisige latente variabelen is een open probleem. Bestaande pogingen vertrouwen vaak op training op kleine schaal of smalle voorkeursdatasets, wat resulteert in inferieure nauwkeurigheid en generalisatie vergeleken met reward-modellen op basisniveaus in pixelruimte. Bovendien is het trainen van een nieuw waarde-model voor elke nieuwe diffusie-ruggengraat of reward-model computatief niet volhoudbaar.

2. Methodologie: StitchVM

De auteurs stellen StitchVM (Stitched Value Model) voor, een raamwerk dat de capaciteiten van voorgeprente, hoogwaardige reward-modellen in pixelruimte overbrengt naar het regime van ruisige latente variabelen met minimale rekentkosten. De aanpak maakt gebruik van model stitching, een techniek die lagen van verschillende neurale netwerken opnieuw combineert.

Kernarchitectuur

StitchVM construeert een hybride model door een afgeknotte, bevroren diffusie-ruggengraat (het "hoofd") te koppelen aan een gesneden, voorgeprent reward-model (de "staart") via een lichtgewicht stitching-laag.

• Hoofd: Een diffusie-ruggengraat (bijv. SD 3.5, FLUX) afgeknot bij laag $i$. Dit component verwerkt van nature ruisige latente variabelen en extrahert ruisbewuste kenmerken.
• Staart: Een voorgeprent reward-model (bijv. CLIP, HPSv2, Aesthetic Predictor) beginnend bij laag $j$. Dit component behoudt de robuuste, taakspecifieke reward-capaciteit die tijdens grootschalige pretraining is geleerd.
• Stitching-laag ($s_\psi$): Een lichtgewicht module die de kenmerkruimte van het diffusie-hoofd afstemt op de reward-staart.

Trainingsprocedure

De training van StitchVM vindt plaats in twee fasen:

1. Interface-selectie (gesloten vorm): De auteurs zoeken naar de optimale stitching-indices $(i, j)$ waarbij de diffusie-kenmerken (van ruisige latente variabelen) en reward-kenmerken (van schone afbeeldingen) het meest compatibel zijn. Dit wordt bereikt door een lineaire mapping $W$ te fitten tussen gekoppelde kenmerken die zijn geëxtraheerd uit een kleine probe-set van schone afbeeldingen en hun corresponderende ruisige latente variabelen. Het paar dat de feature-matching-verlies minimaliseert, wordt geselecteerd.
2. Lichtgewicht finetuning: Zodra de interface is vastgesteld, worden de stitching-laag en de afgeknotte reward-staart gefinetuned met behulp van ongelabelde schone afbeeldingen. Het doel is om de waardefunctie $V(z_t)$ te regresseren naar de oorspronkelijke rewardscore $r(z_0)$ van de schone afbeelding. Deze off-policy doelstelling stelt het model in staat om de verwachte beloning van schone samples te leren die zijn geconditioneerd op ruisige latente variabelen, zonder volledige rollouts te vereisen.

Toepassing op Afstemming

StitchVM dient als directe waardefunctieschatting, die benaderingen vervangt bij zowel inferentie-tijd als training-tijd afstemming:

• Inferentie-tijd:
  • DPS (Diffusion Posterior Sampling): Vervangt de op Tweedie gebaseerde gradiënt door de directe gradiënt van StitchVM ($\nabla_{z_t} V(z_t)$), waardoor backpropagatie door de denoiser en VAE overbodig wordt.
  • FK Steering: Maakt "voorstelschaal" mogelijk, waarbij elke deeltje meerdere lokale voorstellen voortbrengt. Deze worden goedkoop gescoord door StitchVM om het beste kandidaat te selecteren, een proces dat te kostbaar zou zijn met Tweedie-benaderingen.
• Training-tijd:
  • Directe Reward Finetuning (bijv. AlignProp, DRaFT): Staat toe dat training stopt bij tussenliggende ruisige latente variabelen $z_\tau$ en de StitchVM-predictie gebruikt als toezichtdoel, waardoor volledige trajectrollouts en lange backpropagatieketens worden vermeden.
  • DiffusionNFT: Vervangt terminale schone beloningen door tussentijdse waarde-schattingen, waardoor efficiënte online RL op het voorwaartse proces mogelijk wordt.

3. Belangrijkste Bijdragen

• Een nieuw paradigma voor waarde-modellering: StitchVM toont aan dat krachtige reward-modellen in pixelruimte effectief kunnen worden overgebracht naar het regime van ruisige latente variabelen zonder grootschalige hertraining, simpelweg door ze te stitchen aan een diffusie-ruggengraat.
• Efficiëntie en nauwkeurigheid: De methode bereikt state-of-the-art prestaties in waardeschatting terwijl het uitzonderlijk lichtgewicht is. Bijvoorbeeld, het stitchen van CLIP ViT-L in een SD 3.5 Medium waarde-model kost slechts $\approx 10$ GPU-uren.
• Prestatiewinst in afstemming:
  • Inferentie: DPS wordt 3,2$\times$ sneller met een 50% reductie in piek GPU-geheugen. FK steering met StitchVM presteert aanzienlijk beter dan standaard FK steering en Best-of-$N$ sampling, met name door efficiënte selectie van lokale voorstellen mogelijk te maken.
  • Training: DiffusionNFT wordt 2,3$\times$ sneller, en directe reward finetuning (DRaFT) wordt 1,3$\times$ sneller terwijl het hogere generatiekwaliteit bereikt (bijv. +30% GenEval) door effectief toezicht bij stappen met hoge ruis.
• Generalisatie: De aanpak werkt over diverse diffusie-ruggengraten (SD 3.5 Medium/Large, FLUX) en reward-modellen (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor).

4. Experimentele Resultaten

Het artikel evalueert StitchVM over meerdere benchmarks:

• Kwaliteit van de waardefunctie: Op ruisige latente variabelen met variërende ruisniveaus ($\sigma \in [0.1, 0.9]$) komt StitchVM dicht in de buurt van de prestaties van de oorspronkelijke schone reward-modellen in zero-shot retrieval (MSCOCO, Flickr30K), voorkeurvoorspelling (HPDv2, ImageReward) en esthetische scoring (AVA). Het presteert aanzienlijk beter dan baselines zoals VAE-stitching (dat instort bij hoge ruis) en modellen die vanaf nul zijn getraind op ruisige data (bijv. NoisyCLIP, DiNa-LRM).
• Inferentie-afstemming: In DPS- en FK-steering-experimenten op DrawBench verbeteren StitchVM-versterkte methoden consistent metrieken zoals ImageReward, HPSv2 en GenEval, terwijl ze de rekentkosten drastisch reduceren. Opvallend is dat het verhogen van het aantal lokale voorstellen ($M$) in FK steering met StitchVM efficiënter blijkt dan het verhogen van het aantal deeltjes ($N$).
• Training-afstemming: Trainingscurven tonen aan dat StitchVM-varianten hogere eindscores bereiken met minder GPU-uren in vergelijking met standaard DRaFT en DiffusionNFT. Kwalitatieve resultaten geven aan dat met StitchVM aangevulde modellen sneller convergeren en eerder in de training scherper, hogere kwaliteit afbeeldingen produceren.

5. Betekenis en Claims

Het artikel beweert dat StitchVM een generieke sjabloon biedt voor het combineren van latente diffusie-ruggengraten met feedforward-modellen in pixelruimte zonder afbreuk te doen aan de uitgebreide pretraining van een van beide. Door krachtige reward-modellen naar de latente ruimte te tillen, maakt het een nieuwe stijl van diffusie-afstemming mogelijk waarbij de juiste waardefunctie voor ruisige latente variabelen één keer wordt geconstrueerd en wordt afgeschreven over vele samples en iteraties.

De auteurs positioneren StitchVM niet als vervanging voor pretraining van basismodellen, maar als een praktische, laagkostbare mechanisme om deze modellen aan te passen aan specifieke taken. Ze benadrukken dat deze aanpak waarde-modelgebaseerde afstemming toegankelijk en efficiënt maakt, wat onderzoek naar veiligere en betrouwbaardere generatieve modellering mogelijk versnelt. Het werk suggereert dat het "stitchen" van voorgeprente componenten een levensvatbaar alternatief is voor de verbodsbepalende kosten van het trainen van waardefuncties vanaf nul of het vertrouwen op bevooroordeelde benaderingen.