GenSBI: Generative Methods for Simulation-Based Inference in JAX

GenSBI is een nieuwe open-source JAX-bibliotheek die flow matching, score matching en denoising diffusion-modellen implementeert met op transformatoren gebaseerde architecturen om een native, end-to-end simulatie-gebaseerd inferentiekader te bieden voor onderzoekers die JAX gebruiken, met het behalen van hoge nauwkeurigheid en goed gekalibreerde posterieuren op standaard benchmarks.

De kern

Het Grote Plaatje: Het "Black Box"-probleem

Stel je een ongelooflijk complexe machine voor, zoals een gigantische, futuristische koffiezetapparaat. Je kunt de knoppen (parameters) op verschillende instellingen draaien, en de machine spitst een kop koffie (data) uit. Je kunt dit een miljoen keer doen: draai de knoppen naar instelling A, krijg koffie A; draai ze naar instelling B, krijg koffie B.

Stel nu dat iemand je een specifieke kop koffie geeft en vraagt: "Welke knopinstellingen heb je gebruikt om dit te maken?"

Dit is het Simulation-Based Inference (SBI)-probleem. In de wetenschap zijn deze "koffiezetapparaten" complexe simulaties van het universum, het menselijk brein of deeltjesbotsingen. Het probleem is dat de machine, hoewel hij uitstekend koffie maakt, vreselijk is in het vertellen hoe hij een specifieke kop heeft gemaakt. De wiskunde om het proces terug te rekenen is te moeilijk om direct op te lossen.

De Oude Manier versus de Nieuwe Manier

De Oude Manier (De Verwerpingsmethode):
Lange tijd probeerden wetenschappers dit op te lossen door te gokken. Ze draaiden willekeurig aan de knoppen, maakten een kop koffie en keken of het smaakte naar de doelkoffie. Als het dichtbij zat, hielden ze de gok; zo niet, dan gooiden ze het weg.

• De Tekortkoming: Als de koffiemachine 100 knoppen heeft, is dit als proberen een specifiek zandkorreltje op een strand te vinden door blind te gokken. Het duurt eeuwen en verspilt veel koffie.

De Nieuwe Manier (Neurale SBI):
In plaats van te gokken en weg te gooien, begonnen wetenschappers een "slimme assistent" (een neurale netwerken) te trainen. Ze tonen de assistent miljoenen voorbeelden van "Knopinstellingen → Koffiekop"-paren. De assistent leert het patroon. Zodra deze getraind is, weet hij, als je hem een nieuwe kop koffie toont, direct de knopinstellingen.

• Het Voordeel: Dit heet amortisatie. Je betaalt de kosten voor het trainen van de assistent één keer. Daarna is het achterhalen van de instellingen voor elke nieuwe kop koffie direct.

De Kloof: Het "JAX"-probleem

Tot nu toe waren de beste "slimme assistenten" voor deze taak gebouwd met een specifiek programmeer toolkit genaamd PyTorch.
Echter, een groeiend aantal wetenschappers en ingenieurs schakelt over op een ander toolkit genaamd JAX. JAX is als een hoogpresterende sportwagen: het is sneller, verwerkt meerdere motoren (GPU's/TPU's) beter en is geweldig voor complexe wiskunde.

• Het Probleem: Als je je koffiemachine in JAX bouwt, kon je de beste "slimme assistenten" niet gebruiken omdat ze alleen in PyTorch werkten. Je zat vast aan oudere, langzamere tools of moest je hele project vertalen, wat een pijnlijke zaak is.

De Oplossing: GenSBI

De auteurs presenteren GenSBI, een nieuwe open-source bibliotheek die de beste "slimme assistenten" naar de JAX-wereld brengt. Denk hierbij aan een universele adapter die je toelaat de meest geavanceerde AI-tools aan te sluiten op je op JAX gebaseerde koffiemachine.

Hier is wat GenSBI speciaal maakt, met eenvoudige analogieën:

1. Drie Verschillende "Leerstijlen" (Generatieve Methoden)

Net als studenten op verschillende manieren leren, leren deze AI-modellen het "Knop-naar-Koffie"-patroon op drie verschillende manieren. GenSBI ondersteunt alle drie, zodat je de beste kunt kiezen voor je taak:

• Flow Matching: Stel je voor dat je een rechte lijn tekent van een leeg canvas naar een voltooid schilderij. Deze methode leert die rechte lijn te tekenen. Het is snel, efficiënt en zeer stabiel.
• Denoising Diffusion (EDM): Stel je voor dat je begint met een TV-scherm vol ruis en het langzaam opkuisst totdat het beeld verschijnt. Deze methode leert hoe je het "ruis" moet "reinigen". Het is zeer krachtig maar kan wat meer stappen vereisen.
• Score Matching: Stel je een wandelaar voor die probeert de top van een berg te vinden door altijd bergop te lopen. Deze methode leert de "helling" van de data om het zoeken te sturen.

2. De "Transformer"-hersenen

Het artikel introduceert drie specifieke soorten "hersenen" (neurale netwerkarchitecturen) voor deze assistenten:

• SimFormer: Een "Zwitsers zakmes"-hersen. Het kan naar de knoppen en de koffie tegelijk kijken en elke relatie tussen hen achterhalen.
• Flux1: Een hersen aangepast van een beroemde beeldgenerator. Het is geweldig in het bekijken van een specifieke koffiekop en direct de knoppen raden.
• Flux1Joint: Een nieuwe, super-hersen die het beste van beide combineert. Het leert de hele relatie tussen knoppen en koffie in één keer. Dit is krachtig omdat het vragen kan beantwoorden zoals "Welke koffie zou deze knopinstelling maken?" en "Welke knoppen maakten deze koffie?" zonder opnieuw getraind te hoeven worden.

3. De "Veiligheidscontrole" (Kalibratie)

In de wetenschap kun je de AI niet zomaar vertrouwen; je moet weten of het liegt. Als de AI zegt dat er 90% kans is dat de knoppen op "Hoog" stonden, is het dan echt 90% van de tijd juist?
GenSBI wordt geleverd met ingebouwde Veiligheidscontroles (zoals SBC, TARP en LC2ST). Dit zijn als stress-tests. Ze voeren duizenden simulaties uit om ervoor te zorgen dat het vertrouwen van de AI overeenkomt met de realiteit. Als de AI te zelfverzekerd of verward is, markeren deze tools dit direct.

De Resultaten: Werkt het?

De auteurs testten GenSBI op standaard "koffiemachine"-puzzels (benchmarks) die wereldwijd door wetenschappers worden gebruikt.

• Nauwkeurigheid: De AI leerde de instellingen bijna perfect te raden. Op een schaal waarbij 0,5 "perfect niet te onderscheiden van de waarheid" is, scoorde GenSBI tussen 0,50 en 0,56. Dit is bijna ideaal.
• Snelheid: Omdat het op JAX draait, is het snel. Het kan trainen op miljoenen voorbeelden en vervolgens het antwoord raden voor een nieuwe kop koffie in milliseconden.
• Veelzijdigheid: Het werkte goed, of de data nu simpele getallen waren of complexe afbeeldingen (zoals foto's van zwaartekrachtslenzen of geluidsgolven van zwarte gaten).

Samenvatting

GenSBI is een nieuwe toolkit die wetenschappers die de programmeertaal JAX gebruiken, toelaat de meest geavanceerde, moderne AI-methoden te gebruiken voor het oplossen van "reverse-engineering"-problemen. Het biedt drie verschillende leerverbintenissen, krachtige nieuwe AI-architecturen en ingebouwde veiligheidscontroles, die allemaal samenwerken om wetenschappers te helpen de verborgen oorzaken achter complexe data te achterhalen—of het nu de geboorte van het universum is of de verspreiding van een virus.

Waar te vinden: De code is gratis en open-source op GitHub, klaar voor iedereen om te gebruiken.

Technische samenvatting

Technische Samenvatting: GenSBI – Generatieve Methoden voor Simulatiegebaseerde Inferentie in JAX

1. Het Probleem

Simulatiegebaseerde Inferentie (SBI) lost het inverse probleem op van het afleiden van parameters $\theta$ uit observaties $x$ wanneer de likelihood-functie $p(x|\theta)$ onberekenbaar is. Dit scenario komt veel voor in de moderne wetenschap, variërend van kosmologische N-lichaamsimulaties en deeltjesfysica-eventgeneratoren tot epidemiologische modellen en astronomie van zwaartekrachtgolven. In deze gevallen definieert de simulator de likelihood slechts impliciet via een steekproefprocedure; deze kan niet analytisch worden geëvalueerd.

Traditionele likelihood-vrije methoden, zoals Approximate Bayesian Computation (ABC) en benaderingsmethoden voor likelihood, lijden onder de vloek van de dimensionaliteit, afhankelijkheid van handgemaakte samenvattingsstatistieken of rigide parametrische aannames. Neuronale SBI is opgekomen als een superieur paradigma, waarbij flexibele neurale dichtheidschatters worden getraind op gesimuleerde paren $\{(\theta^{(i)}, x^{(i)})\}$ om de doeldistributie direct te leren.

Er bestaat echter een aanzienlijke kloof in het software-ecosysteem. Hoewel de dominante SBI-bibliotheek, sbi, is gebouwd op PyTorch en normaliserende flows, flow matching en diffusiemodellen ondersteunt, ontbreekt het onderzoekers die voormodellen en analysepijplijnen ontwikkelen in het JAX-ecosysteem aan een native optie voor moderne generatieve SBI. Bestaande JAX-gebaseerde tools (bijv. sbijax) ondersteunen niet de nieuwste op transformers gebaseerde architecturen of het volledige scala aan continu-tijd generatieve formuleringen (flow matching en diffusie) die bewezen effectief zijn voor dichtheidsschatting.

2. Methodologie

GenSBI is een open-source, JAX-native bibliotheek die drie distincte generatieve frameworks voor dichtheidsschatting implementeert, allemaal verenigd onder één modulaire interface:

2.1 Generatieve Formuleringen

De bibliotheek implementeert drie uitwisselbare methoden, waarbij de wiskundige formulering wordt ontkoppeld van de neurale architectuur:

1. Flow Matching: Lert een snelheidsveld $v_\theta(x, t)$ dat een eenvoudige prior (doorgaans Gaussiaans) transporteert naar de data-distributie langs rechte, optimale-transportpaden. Deze methode maakt gebruik van een conditionele optimale-transport (CondOT)-scheduler, wat resulteert in bijna lineaire ODE-trajectoria die numerieke integratiefouten verminderen en efficiënt steekproeven trekken met minder solver-stappen mogelijk maken.
2. Denoising Diffusion (EDM): Implementeert het Elucidating the Design Space of Diffusion-Based Generative Models (EDM)-framework. Het traint een gepreconditioneerde denoiser om een ruisvervuilingsproces om te draaien dat wordt gedefinieerd door een stochastische differentiaalvergelijking (SDE) of een deterministische probability flow ODE. Deze aanpak biedt formele statistische garanties via de variatiele ondergrens wanneer getraind met likelihood-gewichting.
3. Score Matching: Implementeert score-based generatieve modellering via Variance Preserving (VP) of Variance Exploding (VE) SDE's. Het traint een netwerk om de scorefunctie $\nabla_x \log p_t(x)$ te schatten, waarmee een reverse-time SDE van ruis naar data wordt geleid.

2.2 Neuronale Architecturen

GenSBI introduceert drie op transformers gebaseerde backbones, die verder gaan dan de traditionele masked autoregressive flows (MAF's) en neural spline flows (NSF's) die in eerdere SBI-werkzaamheden werden gebruikt:

• SimFormer: Aangepast van Gloeckler et al. [26], verwerkt deze architectuur een enkele token-sequentie die de gezamenlijke vector $z = (\theta, x)$ voorstelt. Het gebruikt een conditionele mask om conditionele, gezamenlijke en unconditional inferentiemodi dynamisch te hanteren binnen één model.
• Flux1: Aangepast van het FLUX.1-imagegeneratiemodel [44], hanteert deze architectuur een twee-stroomontwerp (afzonderlijke observatie- en conditioneringsstromen) met double-stream blokken en adaptieve layer-normalisatie (adaLN-Zero). Het is geoptimaliseerd voor conditionele dichtheidsschatting.
• Flux1Joint: Een nieuwe architectuur geïntroduceerd in GenSBI die de expressieve single-stream transformer-blokken van Flux1 combineert met het masking-mechanisme van SimFormer. Het maakt gezamenlijke dichtheidsschatting mogelijk met de voordelen van moderne transformer-gating en self-attention, waardoor conditionering na training op elk subset van variabelen mogelijk is.

2.3 Software-architectuur

De bibliotheek is ontworpen rondom een strategiepatroon dat drie assen ontkoppelt:

1. Generatieve Methode: Het wiskundige kader (Flow Matching, EDM, Score Matching).
2. Inferentiemodus: Het pijplijntype (Conditioneel, Gezamenlijk, Unconditioneel).
3. Neurale Backbone: De specifieke transformer-architectuur.

Dit ontwerp stelt gebruikers in staat om elk component te vervangen (bijv. van Flow Matching naar EDM, of van Flux1 naar SimFormer) via configuratie zonder de trainingslus of inferentiepijplijn te herschrijven. De bibliotheek integreert diep met het JAX-ecosysteem, waarbij Flax wordt gebruikt voor neurale netwerken, diffrax voor ODE/SDE-oplossers, NumPyro voor probabilistische programmering en Orbax voor checkpointing.

2.4 Kalibratie en Validatie

Omdat goed gekalibreerde posteriors niet-onderhandelbaar zijn voor wetenschappelijke toepassingen, integreert GenSBI vier diagnostische hulpmiddelen als eersteklas componenten:

• Simulation-Based Calibration (SBC): Controleert op uniformiteit in de rang van ware parameters binnen posterior-steekproeven.
• TARP (Tests of Accuracy with Random Points): Evalueert verwachte dekkingkansen met Jeffreys-vertrouwensintervallen.
• LC2ST (Local Classifier Two-Sample Test): Biedt observatie-specifieke correctheidsevaluatie.
• Marginal Coverage: Empirische checks op per-dimensie geloofwaardige intervallen.

3. Belangrijkste Bijdragen

Het artikel presenteert de volgende specifieke bijdragen:

1. Drie Generatieve Formuleringen: Een verenigde implementatie van flow matching, score matching en denoising diffusion (EDM) in JAX, waardoor ze uitwisselbaar kunnen worden gebruikt met dezelfde neurale backbone.
2. State-of-the-Art Architecturen: De levering van drie op transformers gebaseerde modellen (SimFormer, Flux1, Flux1Joint), inclusief de nieuwe Flux1Joint die gate-gemoduleerde transformer-blokken uitbreidt naar gezamenlijke dichtheidsschatting.
3. Ingebouwde Kalibratie-diagnostiek: Integratie van SBC-, TARP-, LC2ST- en marginal coverage-tests direct in de bibliotheekworkflow.
4. Benchmark-validatie: Uitgebreide validatie op standaard SBI-benchmarks (SBIBM) en geavanceerde wetenschappelijke toepassingen (zwaartekrachtgolven, sterke lensing), wat concurrerende prestaties en goed gekalibreerde posteriors aantoont.

4. Resultaten

De auteurs valideren GenSBI op zeven benchmarktaken, waaronder vijf uit de SBIBM-suite en twee geavanceerde toepassingen met hoogdimensionale gestructureerde data (tijdreeksen van zwaartekrachtgolven en afbeeldingen van sterke lensing).

• Kwaliteit van de Posterior: Op SBIBM-taken bereikt GenSBI bijna ideale Classifier Two-Sample Test (C2ST)-scores (0,50–0,56, waarbij 0,50 ideaal is). Bijvoorbeeld, op de uitdagende SLCP-taak behaalde het Flux1Joint-model met flow matching een C2ST van 0,534, wat beter presteerde dan SimFormer (0,566) en standaard NPE (0,742).
• Kalibratie: De TARP-diagnostiekcurves voor alle geteste configuraties liggen op de diagonaal, wat wijst op goed gekalibreerde posterior-dekking over verschillende dimensionaliteiten en posterior-geometrieën.
• Efficiëntie en Schaalbaarheid: De bibliotheek toont aan dat flow matching en score matching convergeren naar vergelijkbare C2ST-scores naarmate het simulatiebudget toeneemt. Hoewel EDM soms grotere budgetten vereist om de prestaties van flow matching op gezamenlijke schattings taken te evenaren, blijken alle methoden effectief.
• Geavanceerde Toepassingen: GenSBI verwerkt succesvol hoogdimensionale observaties (bijv. $2 \times 8192$ tijdreeksen voor zwaartekrachtgolven, $64 \times 64$ afbeeldingen voor lensing) met behulp van geleerde embedding-netwerken (CNN's) gekoppeld aan de transformer-backends, en bereikt goed gekalibreerde posteriors zonder referentie-grondwaarheid.
• Vergelijking: GenSBI komt overeen met of overtreft bestaande baselines (OneFlowSBI, SimFormer, sbi NPE) op alle taken en simulatiebudgetten, en bereikt deze resultaten met een bijna uniforme trainingsconfiguratie, waardoor de behoefte aan uitgebreide per-taak hyperparameter-tuning wordt vermeden.

5. Betekenis en Claims

Het artikel positioneert GenSBI als een cruciale toevoeging aan het landschap van wetenschappelijke software, specifiek om de kloof te vullen voor JAX-native simulatiegebaseerde inferentie met moderne generatieve modellen.

• Domeinonafhankelijkheid: Het framework is ontworpen voor elk veld waar een stochastische simulator een impliciete likelihood definieert, van fysica tot epidemiologie en neurowetenschappen.
• Modulariteit en Composeerbaarheid: Door de generatieve methode, architectuur en inferentiemodus te ontkoppelen, stelt GenSBI onderzoekers in staat om te experimenteren met verschillende combinaties zonder architecturale beperkingen, een flexibiliteit die in veel bestaande PyTorch-gebaseerde tools ontbreekt.
• Wetenschappelijke Strenge: De auteurs benadrukken dat de opname van rigoureuze kalibratie-diagnostiek een kernfunctie is, waardoor wordt gegarandeerd dat de resulterende posteriors niet alleen expressief zijn, maar ook wetenschappelijk betrouwbaar.
• Toekomstperspectief: Het artikel noteert bescheiden de huidige beperkingen, zoals het ontbreken van normaliserende flows (die efficiënter zijn voor likelihood-evaluatie in MCMC-lussen) en de noodzaak van verder testen op zeer hoogdimensionale parameterruimten ($\dim(\theta) > 10$). Het modulaire ontwerp is echter bedoeld om de toevoeging van deze functies in toekomstige releases te faciliteren.

Kortom, GenSBI biedt een robuust, uitbreidbaar en gekalibreerd framework voor neurale posterior-schatting in JAX, waarbij het gebruik maakt van de expressiviteit van flow matching en diffusiemodellen in combinatie met transformer-architecturen om complexe inferentieproblemen in de natuurwetenschappen op te lossen.