Technische Samenvatting: Efficiënte Adjoint Matching voor het Fine-tunen van Diffusiemodellen

Probleemstelling

Beloningsgebaseerd fine-tunen is de standaardbenadering geworden om vooraf getrainde diffusie- en flow-modellen af te stemmen op menselijke voorkeuren bij tekst-naar-beeld (T2I) generatie. Onder deze methoden biedt Adjoint Matching (AM) een principiële formulering door beloningsfine-tunen te beschouwen als een Stochastisch Optimale Controle (SOC) probleem. AM beoogt een vooraf getraind model te transformeren naar een doeldistributie $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$, waarbij $r(x)$ een beloningsfunctie is.

Echter, AM kampt met aanzienlijke computerefficiëntieproblemen die de praktische toepassing belemmeren:

1. Knelpunt bij Forward Simulatie: AM vereist het simuleren van volledige generatieve trajecten met een stochastische differentiaalvergelijking (SDE) oplosser onder "geheugenloze" dynamica. Dit vereist een groot aantal functiewaardebepalingen (NFE's) en het opslaan van het volledige traject in het geheugen.
2. Knelpunt bij Backward Simulatie: Om het controlesignaal te berekenen, vereist AM een backward Ordinary Differential Equation (ODE) simulatie van de adjoint toestand langs het gesampelde traject. Dit houdt in dat Jacobiaan-Vectoren Producten (JVP's) door het vooraf getrainde netwerk bij elke integratiestap worden berekend, wat computereconomisch duur is.

De auteurs observeren dat deze knelpunten intrinsiek verbonden zijn met de niet-triviale basisdrift die wordt overgenomen van het vooraf getrainde model in standaard AM-formuleringen.

Methodologie: Efficiënte Adjoint Matching (EAM)

Het artikel stelt Efficiënte Adjoint Matching (EAM) voor, wat het SOC-probleem herformuleert om beide computerefficiëntieknelpunten te elimineren, terwijl de theoretische garantie om de belonings-georiënteerde doeldistributie te bereiken behouden blijft.

1. Herontwerp van Basisdynamica

EAM introduceert twee structurele voorwaarden voor de basisdynamica:

• Lineariteit: De basisdrift $b(x,t)$ wordt herontworpen om lineair te zijn in $x$ (d.w.z. $b(x,t) = D(t)x$).
• Geheugenloosheid: De basisdynamica blijft geheugenloos, waardoor tussenliggende toestanden direct kunnen worden gesampeld vanuit het eindpunt.

Door de diffusiecoëfficiënt $\sigma(t)$ vast te leggen om overeen te komen met de standaardopstelling, leiden de auteurs een specifieke familie van toelaatbare lineaire drifts af (Vergelijking 15) die aan deze voorwaarden voldoen. Deze lineaire drift maakt een gesloten-vorm oplossing voor de adjoint toestand mogelijk, waardoor de noodzaak voor backward ODE-simulatie en JVP-berekeningen wordt weggenomen.

2. Correctie van Eindkosten

Het veranderen van de basisdrift verandert de einddistributie van het ongecontroleerde systeem. Om ervoor te zorgen dat de optimale gecontroleerde distributie nog steeds overeenkomt met de gewenste belonings-georiënteerde distributie $p^\star_1$, leidt EAM een gecorrigeerde eindkosten $g(x)$ af. De gradiënt van deze kosten, $\nabla g(x)$, bevat termen voor de vooraf getrainde score, de nieuwe basiseindscore en de beloningsgradiënt (Vergelijking 16).

3. Trainingsalgoritme

De EAM-trainingslus (Algoritme 1) verloopt als volgt:

1. Eindpuntgeneratie: In plaats van een volledige SDE-uitrol, wordt het eindpuntbeeld $X_1$ gegenereerd met behulp van een deterministische ODE-oplosser met enkele stappen (bijv. 10 NFE's).
2. Constructie van Tussenliggende Toestanden: Tussenliggende toestanden $X_t$ worden direct gesampeld vanuit de oorspronkelijke ruis-kern $q_t(\cdot | X_1)$, geconditioneerd op het gegenereerde eindpunt, waardoor de noodzaak voor forward SDE-simulatie wordt omzeild.
3. Gesloten-vorm Adjoint: De adjoint toestand wordt berekend via een gesloten-vorm schaling van de eindpuntgradiënt $\nabla g(X_1)$, waardoor de backward simulatie volledig wordt verwijderd.
4. Parameterisering van Controle: De controle $u(x,t)$ wordt geparameteriseerd via een trainbaar snelheidsmodel (bijv. LoRA) dat is geïnitieerd vanuit het vooraf getrainde model, waardoor fine-tunen kan starten vanuit de vooraf getrainde dynamica.

Belangrijkste Bijdragen

• Algoritmevoorstel: De introductie van EAM, een op beloningsgradiënten gebaseerd fine-tuning-algoritme dat dezelfde theoretische doelen bereikt als AM, maar met aanzienlijk lagere computerkosten.
• Theoretische Karakterisering: Een karakterisering van lineaire basisdrifts die voldoen aan de vereisten voor efficiënte trajectconstructie (geheugenloosheid) en gesloten-vorm adjoint-berekening.
• Empirische Validatie: Aantonen dat EAM tot 4$\times$ sneller convergeert dan AM op standaard T2I-benchmarks, terwijl het de prestaties over meerdere metrieken behoudt of overtreft.

Resultaten

De auteurs hebben EAM geëvalueerd op Stable Diffusion 3.5-Medium (SD3.5-M) met behulp van de Pick-a-Pic dataset, en dit vergeleken met AM en de baseline SD3.5-M.

• Efficiëntie: EAM verlaagt de trainingskosten per iteratie met tot 4$\times$. Dit wordt bereikt door 40-staps SDE-simulaties te vervangen door 10-staps ODE-eindpunt-sampling en backward adjoint-simulaties te elimineren.
• Prestaties: Op benchmarks waaronder PickScore, ImageReward, HPSv2.1, CLIPScore en Aesthetics, komt EAM consistent overeen met of overtreft AM.
  • In single-reward (PickScore) instellingen behaalde EAM een PickScore van 23,16 tegenover 23,00 voor AM.
  • In multi-reward instellingen (PickScore + HPSv2.1 + Aesthetics) behaalde EAM een PickScore van 22,93 en een Aesthetics-score van 6,56, beter dan AM (6,21).
• Kwalitatieve Verbeteringen: Visuele vergelijkingen geven aan dat EAM fijnmazige details en complexe composities beter behoudt dan AM, dat soms details gladstrijkt. Dit wordt toegeschreven aan de consistentie van EAM tussen training (geconditioneerde ruis op het eindpunt) en inferentie-trajecten.

Betekenis en Claims

Het artikel beweert dat de primaire beperkingen van Adjoint Matching niet inherent zijn aan de SOC-formulering zelf, maar artefacten zijn van de specifieke keuze van basisdrift. Door de basisdrift te herontwerpen om lineair en geheugenloos te zijn, en dienovereenkomstig de eindkosten aan te passen, verwijdert EAM de twee grote computerefficiëntieknelpunten (forward SDE-simulatie en backward adjoint-simulatie) zonder in te leveren op de theoretische afstemmingsgaranties.

De auteurs positioneren EAM als een praktische, schaalbare oplossing voor beloningsfine-tunen die de steekproefficiëntie van op gradiënten gebaseerde methoden behoudt, terwijl de computereconomische last drastisch wordt verminderd, waardoor hoogwaardige afstemming van diffusiemodellen toegankelijker wordt. Het werk claimt niet alle optimalisatieuitdagingen op te lossen (bijv. buffer replay of adaptieve schema's worden aangegeven als toekomstig werk), maar vestigt een nieuwe benchmark voor efficiëntie in dit domein.