RNE: plug-and-play diffusion inference-time control and energy-based training

Deze paper introduceert de Radon-Nikodym Estimator (RNE), een plug-and-play framework dat een fundamenteel verband legt tussen marginaal dichtheden en overgangskernen, waardoor het zowel schatting van diffusiedichtheden, besturing tijdens de inferentie als energiegebaseerde training voor zowel continue als discrete diffusiemodellen verenigt.

De kern

De Radon-Nikodym Schattinger (RNE): Een Gids voor het Besturen van AI-Generatie

Stel je voor dat je een kunstenaar bent die een schilderij maakt door eerst een canvas vol met willekeurige vlekken (ruis) te nemen en die vlekken langzaam weg te werken tot een prachtig landschap verschijnt. Dit is hoe moderne Diffusiemodellen werken. Ze zijn de sterren van de AI-wereld, van het maken van foto's tot het ontwerpen van nieuwe medicijnen.

Maar er is een probleem: deze kunstenaars weten precies hoe ze van een vlek naar een landschap gaan, maar ze weten niet precies hoe waarschijnlijk het is dat ze op een bepaald punt in het proces een specifieke vlek hebben. Het is alsof je een routebeschrijving hebt, maar geen kaart van de kans dat je op een bepaalde hoek staat.

Zonder die "kanskaart" (de dichtheid) kun je de kunstenaar niet goed sturen. Je kunt hem niet zeggen: "Maak het landschap iets warmer" of "Zorg dat het een auto is die op een eiland staat".

In dit paper introduceren de auteurs RNE (Radon-Nikodym Estimator). Laten we dit uitleggen met een paar creatieve metaforen.

1. De Twee Wegen: Voorwaarts en Achteruit

Stel je een rivier voor.

• De voorwaartse weg: Een bootje stroomt mee met de stroming (van een schoon landschap naar een modderige rivier). Dit is het "ruisproces".
• De achterwaartse weg: Je probeert de boot terug te sturen stroomopwaarts naar het schoon landschap. Dit is het "denoising-proces" dat de AI gebruikt om afbeeldingen te maken.

De auteurs zeggen: "Wist je dat deze twee wegen, als je ze precies tegenover elkaar bekijkt, eigenlijk dezelfde route zijn, alleen in omgekeerde richting?"

In de wiskunde heet dit dat de Radon-Nikodym-afgeleide (een soort verhouding tussen de twee wegen) altijd gelijk is aan 1. Dit klinkt saai, maar het is goud waard. Het betekent dat als je de "ruis" (de voorwaartse weg) kent, je automatisch de "kans" (de achterwaartse weg) kunt berekenen.

RNE is de tool die deze twee wegen aan elkaar plakt. Het vertelt je: "Als je hier bent in het landschap, en je kijkt terug naar de rivier, dan is de kans dat je hier belandde precies zo."

2. Waarom is dit zo handig? (De "Plug-and-Play" Magie)

Voorheen was het moeilijk om deze AI-modellen te sturen. Mensen gebruikten "heuristic" methoden (probeer-en-fout), wat vaak leidde tot rare resultaten of vertekende beelden.

Met RNE kunnen we nu drie dingen doen, alsof we een universele afstandsbediening hebben:

A. Het Besturen tijdens het Maken (Inference-Time Control)

Stel je voor dat je een AI vraagt om een foto te maken van een "gele ambulance".

• Zonder RNE: De AI probeert het, maar misschien wordt het een gele bus of een rode ambulance.
• Met RNE: We kunnen de AI tijdens het proces "sturen". We zeggen: "Hé, op dit punt in het proces, is de kans dat dit een ambulance is nog niet hoog genoeg. Laten we de boot een beetje meer sturen richting 'ambulance'."

RNE berekent precies hoeveel je moet sturen, zonder dat je de AI opnieuw hoeft te trainen. Het werkt als een GPS die je route corrigeert terwijl je rijdt, zodat je precies op de bestemming aankomt, zelfs als je afwijkt.

B. Het Maken van "Fusion" Kunst (Model Composition)

Stel je hebt twee AI's: één die goed is in het tekenen van katten, en één die goed is in het tekenen van ruimteschepen.
Je wilt een ruimtekat.
Vroeger was dit heel lastig. Met RNE kunnen we de twee modellen "samenvoegen" alsof we twee muzieknummers mixen. De tool berekent hoe je de katten-route en de ruimtevaart-route combineert om een nieuwe, perfecte route te vinden die beide eigenschappen heeft.

C. Het Trainen van Betere Modellen (Energy-Based Training)

Soms willen we dat de AI niet alleen plaatjes maakt, maar ook begrijpt waarom iets mooi is (de "energie" of kwaliteit).
RNE fungeert als een slimme leraar. Tijdens het trainen kijkt de leraar naar de route die de AI neemt en zegt: "Je bent hier een beetje onzeker, je hebt de kansberekening niet goed gedaan. Probeer het nog eens." Dit helpt de AI om veel sneller en accurater te leren, zonder dat het veel extra rekenkracht kost.

3. De "Referentie" (De Anker)

Er is nog een klein probleem: als je de rivier te snel afstroomt (te grote stappen), kun je de boot kwijtraken. De berekeningen worden onstabiel.

De auteurs lossen dit op door een anker te gebruiken. Ze vergelijken de moeilijke, echte rivier met een simpele, voorspelbare rivier (een standaard rivier die ze zelf hebben ontworpen).

• Ze meten het verschil tussen de echte rivier en de simpele rivier.
• Omdat de simpele rivier makkelijk te berekenen is, wordt de hele berekening veel stabieler en nauwkeuriger.

Het is alsof je een zware last draagt, maar je hebt een stevige stok (het anker) om je evenwicht te bewaren.

Samenvatting in één zin

RNE is een slimme wiskundige tool die het onzichtbare "kanskaartje" van een AI-kunstenaar zichtbaar maakt, zodat we die kunstenaar tijdens het werk perfect kunnen sturen, combineren en trainen, zonder dat we de hele kunstenaar opnieuw hoeven te bouwen.

Het is de sleutel om van "AI die willekeurig plaatjes maakt" naar "AI die precies doet wat we willen" te gaan.

Technische samenvatting

Probleemstelling

Diffusiemodellen genereren data door geleidelijk ruis te verwijderen, wat overeenkomt met de tijdsomkering van een verstoringsproces. Hoewel deze modellen uitstekend presteren in het genereren van hoge kwaliteit samples, is er een fundamentele beperking in veel praktische toepassingen:

1. Ontbrekende randdichtheden: Voor taken zoals inferentie-tijd controle (inference-time control), posterior sampling, of het combineren van meerdere modellen, is kennis nodig van de marginal densities ($p_t$) langs het generatie-traject. Voor score-based diffusiemodellen is deze dichtheid echter doorgaans onberekenbaar (intractable).
2. Afhankelijkheid van heuristieken: Bestaande methoden voor controle (zoals guidance) zijn vaak heuristisch en introduceren bias. Probabilistische inferentietechnieken (zoals Sequential Monte Carlo - SMC) kunnen bias elimineren, maar vereisen een nauwkeurige schatting van de dichtheid of de verhouding tussen dichtheden, wat traditioneel rekenkundig zeer duur is (bijv. door divergentie-berekeningen).
3. Gebrek aan eenheid: Bestaande benaderingen voor dichtheidsschatting, SMC-gewichten en energy-based training zijn versnipperd en missen een gezamenlijk theoretisch raamwerk.

Methodologie: De Radon-Nikodym Estimator (RNE)

De kern van dit paper is de introductie van de Radon-Nikodym Estimator (RNE). Dit is een unificerend framework dat gebaseerd is op het concept dat de Radon-Nikodym-afgeleide (de dichtheidsverhouding) tussen een stochastisch proces en zijn tijdsomkering altijd gelijk is aan 1.

Kernconcepten:

• Tijdsomkering: Voor een voorwaarts proces (verstorend) met drift $\mu_t$ en een achterwaarts proces (denoising) met drift $\nu_t$, die elkaars tijdsomkering zijn, geldt dat de padmaat (path measure) identiek is.
• De Identiteit: De verhouding tussen de marginale dichtheid op tijdstip $\tau$ en $\tau'$ kan worden uitgedrukt als de verhouding van de overgangskernen (transition kernels) van het voorwaartse en achterwaartse proces:
  $$\frac{p_\tau(Y_\tau)}{p_{\tau'}(Y_{\tau'})} = R^\nu_\mu(Y_{[\tau, \tau']})$$
  Waarbij $R$ de Radon-Nikodym-estimator is, berekend als een product van Gaussische kernen over het discretisatie-traject.
• Plug-and-Play Aard: In plaats van de complexe divergentie van het score-netwerk te berekenen, gebruikt RNE alleen de bekende overgangskernen (noising en denoising kernels) van het getrainde model.

Stabilisatie:
Om instabiliteit en cumulatieve fouten bij discretisatie te voorkomen, introduceert het paper een analytische referentie (vaak een lineair proces met een Gaussische startverdeling). Door de RNE te normaliseren ten opzichte van deze referentie, worden de varianties van de kernen uitgelijnd, wat leidt tot een stabielere en nauwkeurigere schatting zonder extra rekentijd.

Belangrijkste Bijdragen

1. Unificatie van Methoden: RNE biedt een enkel perspectief dat diverse bestaande methoden verenigt, waaronder:

   • Twisted Diffusion Sampler (TDS)
   • Feynman-Kac correctors en steering
   • Itô dichtheidsschatters
   • Fokker-Planck regularisatie
     Het paper toont aan dat deze methoden special cases zijn van het RNE-framework.

2. Inferentie-tijd Controle (RNC):
   De auteurs introduceren de Radon-Nikodym Corrector (RNC) voor Sequential Monte Carlo (SMC). RNC maakt het mogelijk om importance weights te berekenen voor:

   • Annealing: Het veranderen van de temperatuur van het model.
   • Reward-tilting: Het genereren van samples die voldoen aan een beloningsfunctie (posterior sampling).
   • Model Composition: Het combineren van meerdere diffusiemodellen (bijv. voor multi-target drug design).
   • Flexibiliteit: In tegenstelling tot eerdere methoden (zoals Feynman-Kac correctors) die specifieke drifts vereisen om divergentietermen te cancelen, staat RNC toe dat gebruikers vrij kiezen tussen verschillende voorwaartse en achterwaartse processen om de variantie te minimaliseren.

3. Energy-Based Training Regularisatie:
   RNE wordt gebruikt als een eenvoudige en efficiënte regularisatieterm voor het trainen van energy-based diffusiemodellen. Deze regularisatie dwingt het model om de relatie tussen de marginale dichtheid en de overgangskernen te respecteren, wat leidt tot nauwkeurigere energiestructuren zonder de noodzaak om de divergentie van het score-netwerk te schatten.

4. Modality-Agnostisch:
   Het framework is niet beperkt tot continue diffusiemodellen (SDEs). Het is ook toepasbaar op discrete diffusiemodellen en Continuous Time Markov Chains (CTMC), wat het breed inzetbaar maakt voor tekst, discrete data en andere modaliteiten.

Resultaten

De auteurs evalueren RNE op diverse taken en datasets:

• Inferentie-tijd Annealing (ALDP & LJ-systemen):

  • RNC presteert aanzienlijk beter dan baselines (zoals "anneal score" zonder SMC) en is vergelijkbaar of superieur aan de state-of-the-art Feynman-Kac Corrector (FKC).
  • Door de keuze van drift-parameters ($c_a, c_b$) kan de effectieve steekproefgrootte (ESS) en de diversiteit van de samples worden geoptimaliseerd.
  • RNE toont betere schaalbaarheid bij het verhogen van het aantal deeltjes (particles) tijdens inferentie.

• Multi-Target Small Molecule Design:

  • Bij het combineren van modellen voor twee verschillende eiwitdoelen, bereikt RNC een hogere succesrate in het genereren van liganden met betere docking scores dan heuristische score-sommatie en FKC.

• Traject Stitching (Maze Navigation):

  • RNC slaagt erin om korte trajecten naadloos te verbinden naar een doelwit met een succesrate van 100%, terwijl methoden zonder SMC (alleen guidance) faalden.

• Training van Energy-Based Modellen:

  • Op 2D en 100D Gaussian Mixture Models (GMM) en ALDP data, leidt RNE-regularisatie tot een significant nauwkeurigere schatting van de energie (en dus de dichtheid) vergeleken met standaard Denoising Score Matching (DSM).
  • Bij het schatten van vrije energie (via Thermodynamic Integration) verbetert RNE de nauwkeurigheid aanzienlijk (van ~27.30 naar ~29.28, dichter bij de referentiewaarde van 29.43).

• CTMC Toepassing:

  • Succesvolle toepassing op discrete diffusie (MaskGIT) voor tekst-naar-beeld generatie met beloningsgebaseerde sturing, wat de effectiviteit op discrete ruimtes bewijst.

Significantie

Dit paper biedt een fundamentele doorbraak in het begrijpen en toepassen van diffusiemodellen:

1. Theoretische Unificatie: Het lost de versnippering op tussen verschillende probabilistische inferentiemethoden door ze te reduceren tot één wiskundige identiteit (RNE).
2. Praktische Toepasbaarheid: Het "plug-and-play" karakter maakt het mogelijk om geavanceerde controlemechanismen (zoals SMC) toe te passen zonder de onderliggende modelarchitectuur te moeten herontwerpen of zware divergentie-berekeningen uit te voeren.
3. Efficiëntie: Het elimineert de noodzaak van dure divergentie-schattingen, wat het toepasbaar maakt op grotere en complexere modellen.
4. Brede Reikwijdte: Door de toepasbaarheid op zowel continue als discrete processen, opent het de deur voor geavanceerde controle in een breed scala aan domeinen, van moleculair ontwerp tot robotica en beeldgeneratie.

Kortom, RNE transformeert diffusiemodellen van puur generatieve modellen naar krachtige, controleerbare probabilistische systemen met een solide theoretisch fundament.