Control Consistency Losses for Diffusion Bridges

Probleemstelling

Het simuleren van de dynamiek van een diffuusproces (diffusion process) dat is voorwaardelijk gesteld op een specifieke eindtoestand $X_T = x_T$, staat bekend als het probleem van het diffuusbrug (diffusion bridge). Dit is een fundamenteel probleem in wetenschappen zoals chemie (bijv. overgangspad-sampling), financiën en vormevolutie.

De grootste uitdaging ontstaat bij zeldzame gebeurtenissen (rare events). Wanneer de onvoorwaardelijke dynamiek (de basisprocessen) zelden de eindtoestand bereikt, falen traditionele methoden (zoals MCMC of Sequential Monte Carlo) of recente neurale benaderingen die op onvoorwaardelijke steekproeven trainen. Deze methoden krijgen onvoldoende data uit de regio rond de eindtoestand om een nauwkeurige geleide dynamiek te leren. Bestaande neurale methoden die wel proberen dit op te lossen (zoals Neural Guided Diffusion Bridges van Yang et al., 2025), vereisen vaak backpropagatie door gesimuleerde trajecten, wat de schaalbaarheid en rekentijd aanzienlijk beperkt.

Methodologie: Self-Consistency Losses

De auteurs stellen een nieuwe aanpak voor, genaamd Control Consistency Diffusion Bridge (CCDB), die gebaseerd is op een zelfconsistentie-eigenschap (self-consistency property) van de optimale besturingsdrift $u^*(t, x)$.

1. Theoretische Grondslag

In plaats van de verwachting over de onvoorwaardelijke dynamiek te benaderen (wat problematisch is bij zeldzame gebeurtenissen), leiden de auteurs een relatie af die de besturing op een tijdstip $s$ relateert aan de besturing op een later tijdstip $t$ langs de gesimuleerde trajecten.
De kern van de methode is Stelling 2.1 (Zelfconsistentie-eigenschap I):
Voor $0 < s < t < T$ voldoet de optimale besturing $u^*$ aan:
$$u^*(s, x) = \mathbb{E}[J_{t|s}^\top u^*(t, X_t) \mid X_s = x]$$
Waarbij $J_{t|s} = \nabla_{X_s} X_t$ het Jacobian-proces is dat de gevoeligheid meet van de toestand op tijdstip $t$ voor verstoringen op tijdstip $s$.

Deze eigenschap stelt dat de besturing op een vroeg tijdstip de verwachte (geprojecteerde) waarde van de besturing op een later tijdstip moet zijn, waarbij de projectie wordt uitgevoerd via de Jacobiaan van het traject.

2. Het Leringsalgoritme (CCDB)

Het algoritme leert een neurale besturing $u_\theta$ door een fixed-point loss te minimaliseren die deze zelfconsistentie afgedwongen.

• Loss Functie: De loss meet de kwadratische fout tussen de voorspelde besturing $u_\theta(s, X_s)$ en het doelwit, berekend als een verwachte projectie van de besturing op een later tijdstip:
  $$L(\theta) = \mathbb{E} \left[ \| u_\theta(s, X_s) - U_s \|^2 \right]$$
  Waarbij $U_s$ een doelpunt is dat wordt berekend door terugwaarts te integreren langs het gesimuleerde traject, gebruikmakend van de Jacobiaan en een wegingsfunctie $\alpha_t$.
• Online Training: Het algoritme werkt in een iteratieve, online modus. Het simuleert trajecten met de huidige besturing, berekent de doelpunten (targets) en update de parameters.
• Geen Backpropagatie door Trajecten: Een cruciaal voordeel is dat de loss niet vereist dat gradients worden gepropageerd door de gesimuleerde trajecten (geen backpropagation through time). De doelpunten worden als "stop-gradient" behandeld tijdens de update. Dit maakt de methode veel schaalbaarder dan methoden die KL-divergentie minimaliseren via backpropagatie.

3. Implementatie Details

• Randvoorwaarden: Om te garanderen dat het proces daadwerkelijk de eindtoestand $x_T$ bereikt, wordt de drift in de laatste tijdstap direct naar $x_T$ "gejumped". Dit forceert de brug-eigenschap per constructie.
• Gradient-vorm: De auteurs dwingen de besturing om een gradient-vorm te hebben ($u = \nabla \phi$), wat theoretisch vereist is voor de optimaliteit.
• Generalisatie (SC2): Voor systemen met onstabiele dynamiek (waar de Jacobiaan $J_{t|s}$ kan exploderen, bijvoorbeeld bij hoge energiebarrières), introduceren de auteurs een gegeneraliseerde zelfconsistentie-eigenschap (SC2). Deze gebruikt een 'auxiliary process' met een aangepaste drift om de groei van de Jacobiaan te onderdrukken, wat numerieke stabiliteit biedt.

Kernbijdragen

1. Nieuwe Zelfconsistentie Eigenschap: De afleiding van een fundamentele eigenschap van de optimale besturing die de relatie tussen besturingen op verschillende tijdstippen beschrijft via het Jacobian-proces.
2. Schaalbaar Leringsalgoritme (CCDB): Een algoritme dat diffusion bridges leert zonder backpropagatie door gesimuleerde trajecten, wat leidt tot een aanzienlijke snelheidswinst.
3. Generalisatie voor Stabiliteit: De introductie van een geavanceerde variant (SC2) die werkt in settings met grote Jacobiaan-groei, waar eerdere methoden zouden falen.
4. Theoretische Connectie: Het leggen van een brug tussen diffusion bridges en bredere Stochastic Optimal Control (SOC) problemen, waarbij recente methoden zoals "Adjoint Matching" als speciale gevallen uit hun raamwerk voortkomen.

Resultaten

De auteurs evalueren CCDB op diverse benchmarks, variërend van lineaire processen tot complexe moleculaire dynamica:

• Ornstein-Uhlenbeck & CIR Modellen: CCDB convergeert sneller en nauwkeuriger dan bestaande methoden, zelfs bij zeldzame gebeurtenissen. Het presteert vergelijkbaar met de state-of-the-art NGDB-methode, maar is ongeveer 3 keer sneller in trainingstijd.
• Double-Well Potentiaal: In dit scenario met hoge energiebarrières (waar de Jacobiaan explodeert) faalt de standaard methode (SC1). De generalisatie (SC2) slaagt echter en leert de juiste dynamiek, terwijl methoden die trainen op onvoorwaardelijke data (SDB, FB) volledig falen.
• Moleculaire Dynamica (Alanine Dipeptide & Müller-Brown):
  • De methode slaagt erin overgangspaden te vinden tussen metastabiele toestanden in complexe potentiaalvelden.
  • In 100 dimensies behoudt de methode zijn prestaties, wat aantoont dat de schaalbaarheid goed is en de moeilijkheid voornamelijk wordt bepaald door de complexiteit van het pad en niet door de dimensie.
  • De trainingstijd voor complexe moleculaire systemen is aanzienlijk lager dan bij NGDB (15 minuten vs. 3 uur op krachtige hardware).

Betekenis en Impact

Deze paper biedt een belangrijke doorbraak in het veld van het simuleren van voorwaardelijke stochastische processen:

• Efficiëntie: Door backpropagatie door trajecten te vermijden, wordt het mogelijk om diffusion bridges te trainen op complexe, hoge-dimensionale problemen die eerder onhandelbaar waren voor neurale methoden.
• Robuustheid: De introductie van de generalisatie (SC2) lost het probleem van numerieke instabiliteit op bij zeldzame gebeurtenissen met hoge barrières, een veelvoorkomend probleem in chemische en fysieke simulaties.
• Unificatie: Het werk verenigt het domein van diffusion bridges met recente ontwikkelingen in Stochastic Optimal Control, wat nieuwe inzichten biedt voor het ontwerpen van trainingsdoelen voor generatieve modellen en besturingsproblemen.

Kortom, CCDB biedt een sneller, stabieler en theoretisch onderbouwd alternatief voor bestaande methoden om diffusion bridges te leren, met name waardevol voor toepassingen in de computationele chemie en het bestuderen van zeldzame gebeurtenissen.