🤖 machine learning

Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization

Dit artikel introduceert off-policy log-dispersie-regularisatie (LDR), een nieuw raamwerk dat de data-efficiëntie en prestaties van Boltzmann-generatoren verbetert door gebruik te maken van doelenergie-labels om het energielandschap te regulariseren zonder dat daarvoor extra on-policy monsters nodig zijn.

Oorspronkelijke auteurs: Henrik Schopmans, Christopher von Klitzing, Pascal Friederich

Gepubliceerd 2026-08-06

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Henrik Schopmans, Christopher von Klitzing, Pascal Friederich

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je het weer probeert te voorspellen, maar in plaats van naar wolken te kijken, probeer je uit te vogelen hoe een minuscuul, onzichtbaar molecuul wiebelt en draait in een druppel water. In de wereld van de computationele wetenschap is dit een enorme puzzel. Moleculen zijn als complexe, hoogdimensionale doolhoven waar ze in diepe dalen (stabiele vormen) kunnen vast komen te zitten of hoge heuvels (onstabiele vormen) kunnen beklimmen. Om te begrijpen hoe medicijnen werken of hoe eiwitten vouwen, moeten wetenschappers elke mogelijke vorm in kaart brengen die een molecuul kan aannemen, gewogen naar hoe waarschijnlijk het is dat die vorm voorkomt. Dit wordt het bemonsteren van een "Boltzmann-verdeling" genoemd.

Het probleem is dat moleculen zo ingewikkeld zijn dat traditionele methoden aanvoelen als het proberen in kaart te brengen van een heel continent door telkens slechts één stap te zetten. Je zou jarenlang alleen maar rondjes kunnen lopen om één berg, zonder ooit de andere kant te zien. Om dit te versnellen, gebruiken wetenschappers "Boltzmann-generatoren". Denk aan deze als superintelligente, door AI aangedreven teleporters. In plaats van stap voor stap te lopen, leren ze de kaart kennen en kunnen ze direct "teleporteren" naar elke geldige vorm die een molecuul kan aannemen. Maar er is een addertje onder het gras: het trainen van deze teleporters is ongelooflijk duur. Elke keer dat de AI een vorm raadt, moet een computer de energie berekenen, wat is als het vragen aan een zeer trage, zeer dure orakel voor de prijs van een ticket. Als de AI miljoenen keren de orakel moet vragen om de kaart te leren, wordt het hele proces te traag om nuttig te zijn.

Dit artikel introduceert een slimme nieuwe truc om deze AI-teleporters veel sneller te leren, met behruik van een methode genaamd Off-Policy Log-Dispersion Regularization (LDR). De auteurs, Henrik Schopmans, Christopher von Klitzing en Pascal Friederich, realiseerden zich dat ze tijdens het trainen van de AI een cruciale aanwijzing negeerden: de energiewaarden zelf. Meestal krijgt de AI alleen te horen: "Hier zijn enkele vormen die het molecuul aannam; leer om deze na te maken." Maar het artikel suggereert een betere aanpak: "Hier zijn enkele vormen, en hier is precies hoeveel energie elke vorm kost. Zorg ervoor dat je gissingen overeenkomen met die kosten."

Ze stellen een nieuwe "regularizer" voor, die werkt als een coach die de AI extra feedback geeft. Stel je voor dat je leert een landschap te tekenen. Een standaard leraar zegt misschien: "Teken hier een boom." Maar deze nieuwe coach zegt: "Teken een boom hier, en zorg ervoor dat de schaduw precies overeenkomt met de hoogte van de heuvel." Als je tekening te donker of te licht is vergeleken met de werkelijke energie van de heuvel, geeft de coach je een zachte duw om de vorm van je tekening aan te passen. Deze "duw" is de LDR. Het vereist niet dat de AI de orakel vraagt om nieuwe, dure energieberekeningen; het gebruikt simpelweg de energielabels die al beschikbaar waren van de data.

De resultaten zijn behoorlijk indrukwekkend. In hun tests ontdekten de auteurs dat het gebruik van deze nieuwe coachingmethode de AI in staat stelde de kaart met veel minder dure energieberekeningen te leren. Bijvoorbeeld, bij het trainen op een molecuul genaamd alanine dipeptide, behaalde de nieuwe methode kwalitatief hoogwaardige resultaten met slechts 1 miljoen samples, terwijl de oude methoden 5 miljoen nodig hadden voor een vergelijkbare kwaliteit. In zelfs uitdagendere tests met een groter molecuul (alanine hexapeptide) was de verbetering zo significant dat de AI de kaart kon leren met 10 keer minder energie-evaluaties dan de vorige beste methoden.

Het artikel laat ook zien dat deze truc werkt wanneer de data "biased" is — wat betekent dat de AI slechts een klein, vreemd hoekje van de wereld van het molecuul te zien heeft gekregen. Normaal gesproken zou dit de AI in de war brengen, maar de LDR-coach hielp het hem toch om de rest van de kaart te begrijpen. De auteurs testten dit op zowel eenvoudige 2D-vormen als complexe 3D-moleculen, en in elk geval leerde de AI sneller en maakte hij minder fouten. Ze toonden zelfs aan dat het werkt wanneer de AI moet leren zonder dat hij überhaupt echte molecuulvormen ziet, waarbij hij enkel vertrouwt op energieberekeningen.

Kortom, dit artikel vindt geen nieuw type molecuul of een nieuwe natuurwet uit. In plaats daarvan verzint het een betere manier om de AI te leren de bestaande kaart te lezen. Door de energielabels als een "vorm-regularizer" te gebruiken, suggereren de auteurs dat we het simuleren van moleculen tot tien keer efficiënter kunnen maken. Dit betekent dat wetenschappers in de toekomst mogelijk nieuwe medicijnen kunnen ontwerpen of complexe biologische processen sneller kunnen begrijpen, omdat de computers die het zware werk doen, de dure orakel niet zo vaak om hulp hoeven te vragen. De methode is flexibel, werkt met verschillende soorten data en verschillende soorten AI-modellen, wat suggereert dat het een standaardinstrument kan worden voor iedereen die probeert de microscopische wereld te simuleren.

Technische Samenvatting: Efficiënt Trainen van Boltzmann-generatoren met Off-Policy Log-dispersie Regularisatie

1. Probleemstelling

Het samplen uit ongenormaliseerde waarschijnlijkheidsdensiteiten, specifiek de Boltzmann-verdeling van fysische systemen, is een fundamentele uitdaging in de computationele fysica en chemie. De doelverdeling wordt gedefinieerd als $p_X(x) = \tilde{p}_X(x)/Z$ , waarbij $\tilde{p}_X(x) = \exp(-E(x)/k_B T)$ de ongenormaliseerde densiteit is, $E(x)$ de energiefunctie is en $Z$ een onhandelbare normalisatieconstante is.

Hoewel klassieke methoden zoals Moleculaire Dynamica (MD) en Markov Chain Monte Carlo (MCMC) asymptotisch correct zijn, lijden ze vaak aan een trage menging en hoge computationele kosten door lange, gecorreleerde trajecten, met name in hoogdimensionale of energetisch complexe systemen. Boltzmann-generatoren (BGs) bieden een generatief modelleringsalternatief, wat onafhankelijk samplen uit evenwichtsverdelingen mogelijk maakt. Hun praktische bruikbaarheid wordt echter beperkt door data-efficiëntie. Standaard trainingsdoelstellingen vertrouwen doorgaans op:

Evenwichtssamples uit de doelverdeling (vaak duur om te genereren via MD).
Uitgebreide energie-evaluaties in variationele instellingen.

Bestaande benaderingen maken er vaak geen volledig gebruik van dat de beschikbare target-energielabels tijdens de datasetgeneratie worden benut, of ze vereisen aanvullende on-policy samples en energie-evaluaties om deze informatie effectief te incorporeren.

2. Methodologie: Off-Policy Log-dispersie Regularisatie (LDR)

De auteurs stellen Off-Policy Log-dispersie Regularisatie (LDR) voor, een framework dat target-energie informatie integreert in de training van Boltzmann-generatoren zonder dat daarvoor extra on-policy samples of extra energie-evaluaties nodig zijn.

Kernconcept

De methode bouwt voort op de log-variantie doelstelling, die traditioneel wordt gebruikt voor on-policy variationele training. De auteurs definiëren een functie $f_\theta(x)$ die het verschil vertegenwoordigt tussen de negatieve log-modeldensiteit en de negatieve log-targetdensiteit (ongenormaliseerde belangweights):
$f_\theta(x) = -\log q_\theta^X(x) - (-\log \tilde{p}_X(x)) = -\log q_\theta^X(x) - \frac{E(x)}{k_B T}$
Bij het optimum, waar $q_\theta^X = p_X$ , moet $f_\theta(x)$ constant zijn ( $\log Z$ ).

De Log-dispersie Doelstelling

De auteurs generaliseren de log-variantie doelstelling naar een familie van log-dispersie doelstellingen:
$\mathcal{L}_\theta^{LD} = \mathbb{E}_{r_X} \left[ \rho(f_\theta(x) - c_\theta) \right], \quad c_\theta = \mathbb{E}_{r_X}[f_\theta(x)]$
waarbij $\rho(u) \geq 0$ een dispersiemaat is (bijv. $|u|^p$ ) en $r_X$ een referentieverdeling is.

LDR-L2 ( $p=2$ ): Minimaliseert de variantie (tweede centrale moment), wat de standaard log-variantie doelstelling herstelt.
LDR-L1 ( $p=1$ ): Minimaliseert het eerste absolute centrale moment. De auteurs hypothetiseren dat deze variant minder gevoelig is voor energie-outliers en zorgt voor een stabielere optimalisatie, wat cruciaal is gezien de heavy-tailed aard van de belangweights in hoogdimensionale energielandschappen.

Off-Policy Regularisatiestrategie

Een cruciaal theoretisch inzicht (Appendix A) is dat als de referentieverdeling $r_X$ (bijv. een vaste dataset) niet over de volledige configuratieruimte beschikt, het minimaliseren van de log-dispersie doelstelling alleen geen divergentie is; het faalt in het beperken van het model buiten de data-manifold, wat leidt tot een incorrecte normalisatie.

Om dit op te lossen, wordt LDR toegepast als een regularisator naast een standaard data-gebaseerde doelstelling (bijv. Forward KL, Flow Matching of Score Matching):
$\mathcal{L}_\theta = \lambda_{data}\mathcal{L}_\theta^{data} + \lambda_{LD}\mathcal{L}_\theta^{LD}$

De data-gebaseerde term zorgt voor globale normalisatie en convergentie naar de correcte doelverdeling.
De LDR-term fungeert als een vorm-regularisator, die de voorstelverdeling afstemt op het target-energielandschap door de dispersie van de log belangweights op de data-manifold te minimaliseren.

Deze aanpak maakt het mogelijk om willekeurige referentieverdelingen te gebruiken, inclusief vertekende simulatiedata of hulp-energie-gelabelde datasets, zonder dat hiervoor extra target-energie-evaluaties tijdens de training nodig zijn.

3. Belangrijkste Bijdragen

Adaptatie naar de Off-Policy Setting: De auteurs passen de log-variantie doelstelling van on-policy variationele training aan naar een off-policy setting, waarbij het wordt gebruikt als een regularisator op vaste, energie-gelabelde datasets.
Generalisatie naar Log-dispersie: Ze generaliseren de doelstelling naar een familie van log-dispersie verliezen, waarbij ze een L1-variant introduceren (LDR-L1) die een verbeterde stabiliteit biedt tegen outliers vergeleken met de standaard L2-variant.
Brede Toepasbaarheid: Het framework wordt effectief gedemonstreerd over drie verschillende trainingsregimes:
- Training op onbevooroordeelde evenwichtsdata.
- Training op vertekende (biased) simulatiedata (gebruikmakend van importance sampling en gemengde datasets).
- Zuiver variationele training zonder toegang tot target samples (alleen gebruikmakend van target energie-evaluaties).
Integratie met State-of-the-Art: De methode is succesvol geïntegreerd met Constrained Mass Transport (CMT), de huidige state-of-the-art voor variationele training, wat de efficiëntie aanzienlijk verbetert.

4. Experimentele Resultaten

De auteurs evalueerden LDR op verschillende benchmarks, waaronder een 2D Gaussian Mixture Model (GMM), Alanine Dipeptide ( $d=60$ ), en Alanine Hexapeptide ( $d=180$ ). Prestaties werden gemeten middels Negative Log-Likelihood (NLL) en Effective Sample Size (ESS).

Belangrijkste Bevindingen:

Onbevooroordeelde Data: Op onbevooroordeelde datasets verbeterde LDR zowel de NLL als de ESS aanzienlijk. Voor Alanine Dipeptide bereikte LDR een prestatie met $1 \times 10^6$ samples die superieur was aan niet-geregulariseerde training met $5 \times 10^6$ samples, wat wijst op een 5x verbetering in data-efficiëntie.
Vertekende Data: Wanneer getraind op vertekende datasets (gegenereerd via korte trajecten), leverde LDR, toegepast op de ge-resamplede evenwichtsdataset ( $D_{IS}$ ) en de oorspronkelijke vertekende dataset, substantiële winst op. LDR maakte het mogelijk om hoogwaardige modellen te trainen met slechts $1 \times 10^5$ importance samples plus de initiële constructie van de vertekende dataset, terwijl niet-geregulariseerde baselines meer dan 10 keer zoveel importance samples vereisten.
Variationele Training (Geen Target Samples): In de zuiver variationele setting met CMT, verminderde LDR het aantal benodigde target energie-evaluaties met ongeveer een orde van grootte. Voor Alanine Dipeptide bereikte CMT + LDR een vergelijkbare prestatie als standaard CMT met slechts $1 \times 10^7$ evaluaties in plaats van $1 \times 10^8$ . Voor Alanine Hexapeptide voorkwam LDR mode collapse in regimes waar vanilla CMT faalde (bereikte >30% ESS versus <6% voor vanilla CMT).
Cartesiaanse Coördinaten: De winst werd gereproduceerd bij het direct trainen op Cartesiaanse coördinaten met TarFlow, wat bevestigt dat de voordelen niet gebonden zijn aan interne coördinatenrepresentaties.
Stabiliteit: LDR-L1 bood over het algemeen stabielere trainingsdynamiek (lagere gradiëntnorm-variantie) dan LDR-L2, met name in hoogdimensionale systemen met heavy-tailed belangweights.

5. Betekenis en Claims

Het artikel claimt dat LDR een eenvoudig maar krachtig framework biedt voor het integreren van target-energie informatie in de training van Boltzmann-generatoren. De primaire betekenis ligt in:

Data-efficiëntie: Het verbetert consistent de uiteindelijke prestaties en data-efficiëntie over diverse settings, waarbij het vaak de benodigde hoeveelheid samples of energie-evaluaties met een orde van grootte vermindert.
Flexibiliteit: Het is breed toepasbaar en ondersteunt onbevooroordeelde, vertekende en zuiver variationele training zonder dat daarvoor extra on-policy samples of extra energie-evaluaties nodig zijn.
Complementariteit: De auteurs merken op dat LDR complementair is aan bestaande variantiereductie-technieken zoals path gradients (die target gradiënten gebruiken), aangezien LDR gebruikmaakt van target energie-waarden (zero-order informatie).

De auteurs benadrukken dat hoewel LDR de afhankelijkheid van de datasetgrootte niet elimineert, het de noodzaak voor massale datasets aanzienlijk vermindert door de hoeveelheid informatie die uit elke sample wordt geëxtraheerd te vergroten. Het framework is ontworpen om gemakkelijk te worden geïntegreerd in bestaande trainingspipelines voor normalizing flows en potentieel voor diffusie-gebaseerde generatoren.