Mixed neural posterior estimation for simulators with discrete and continuous parameters

Technische Samenvatting: Gemengde Neuronale Posterior Schatting voor Simulatoren met Discrete en Continue Parameters

Probleemstelling
Neuronale Posterior Schatting (NPE) is een standaardinstrument geworden voor Inference op Basis van Simulatie (SBI), wat snelle, geamortiseerde parameterschatting mogelijk maakt voor complexe simulatoren met onberekenbare likelihoods. Echter, standaard NPE-methoden gaan uit van continue parameterruimten, waarbij de posterior doorgaans wordt gemodelleerd als een continue kansdichtheid met behulp van normaliserende flows of diffusiemodellen. Deze aanname faalt voor een brede klasse van wetenschappelijke simulatoren waarbij de parameterruimte gemengd is, met zowel continue dimensies (bijvoorbeeld snelheden, geleidingsvermogens) als discrete dimensies (bijvoorbeeld categorische variabelen, indexen voor modelselectie, locaties van veranderingpunten).

Bestaande benaderingen voor gemengde inferentie ondervinden aanzienlijke beperkingen. Standaard Markov Chain Monte Carlo (MCMC)-methodes vereisen zorgvuldige afstelling, amortiseren de inferentiekosten niet over observaties heen en hebben vaak moeite met discrete variabelen (bijvoorbeeld door analytische marginalisatie of gespecialiseerde discontinuïteiten in Hamilton-dynamica te vereisen). Recente op diffusie gebaseerde SBI-methoden proberen discrete parameters in continue ruimten in te bedden of ze te behandelen als modelindexen, maar deze kunnen computationally zwaar zijn of complexe Riemanniaanse embeddings vereisen. Bijgevolg worden simulatoren met gemengde parameterruimten door huidige SBI-toolboxes onvoldoende ondersteund.

Methodologie: Gemengde Neuronale Posterior Schatting (MNPE)
De auteurs stellen Gemengde Neuronale Posterior Schatting (MNPE) voor, een uitbreiding van NPE die discrete en continue parameters gezamenlijk behandelt via een gefactoriseerde inferentienetwerk.

1. Factorisatiestrategie:
   De gezamenlijke posterior $p(\theta_d, \theta_c | x)$ over discrete parameters $\theta_d$ en continue parameters $\theta_c$ wordt gefactoriseerd als:
   $$p(\theta_d, \theta_c | x) = p(\theta_d | x) \cdot p(\theta_c | \theta_d, x)$$
   Dit splitst het probleem op in twee sub-taken: classificatie over de discrete ruimte en conditionele dichtheidsschatting in de continue ruimte, geconditioneerd op de discrete toestand.

2. Netwerkarchitectuur:

   • Discrete Component: De term $p(\theta_d | x)$ wordt gemodelleerd met een Masked Autoregressive Density Estimator (MADE). Dit autoregressieve netwerk geeft kansen voor elke discrete dimensie uit, geconditioneerd op de voorafgaande discrete dimensies en de observatie $x$. Dit maakt het hanteren van willekeurige aantallen discrete parameters met variërende aantal klassen mogelijk.
   • Continue Component: De term $p(\theta_c | \theta_d, x)$ wordt gemodelleerd met een standaard conditioneel generatief model (bijvoorbeeld een neuronale spline flow of diffusiemodel) dat geconditioneerd is op zowel de observatie $x$ als de discrete parameters $\theta_d$.
   • Trainingsdoel: Het netwerk wordt gezamenlijk getraind om de negatieve log-kans van de gemengde parameters te minimaliseren:
     $$\mathcal{L} = -\log q(\theta_d, \theta_c | x) = -\log q(\theta_c | \theta_d, x) - \log q(\theta_d | x)$$
     Het continue deel gebruikt de standaard flow-likelihood, terwijl het discrete deel reduceert tot een cross-entropy-verlies.

3. Calibratiediagnostiek:
   Om de kwaliteit van gemengde posteriors te beoordelen zonder referentieoplossingen, stellen de auteurs een gecombineerd diagnostisch kader voor:

   • Continue Dimensies: Gebruik van Simulation-Based Calibration (SBC) via rangstatistieken om uniformiteit in de marginale rangverdelingen te controleren.
   • Discrete Dimensies: Gebruik van Expected Calibration Error (ECE) en betrouwbaarheidsdiagrammen, waarbij de discrete inferentie wordt behandeld als een classificatieprobleem.
   • Baselines: De auteurs leiden empirische baselines voor eindige steekproeven af voor beide metrieken om werkelijke miscalibratie te onderscheiden van statistische ruis, met erkenning dat perfecte calibratie met eindige steekproeven onbereikbaar is.

Belangrijkste Resultaten
De auteurs valideren MNPE op drie simulatoren met toenemende complexiteit:

1. Berekenbaar Gaussisch Speelvoorbeeld:

   • Een simulator met één continue en één discrete latente variabele die gekoppeld zijn via een gedeelde observatie.
   • MNPE herstelt de analytische grondwaarheid-posterior nauwkeurig, inclusief bimodale regimes waar de discrete toestand onzeker is.
   • De Classifier Two-Sample Test (C2ST)-score convergeert naar het toevalsniveau (~0,5) met ongeveer 1.000 simulaties, wat aangeeft dat de posterior niet te onderscheiden is van de ware posterior.

2. Tandem-wachtrij Simulator:

   • Een simulator met drie continue snelheden en twee discrete aantallen servers ($|D|=25$).
   • MNPE wordt vergeleken met MCMC-referenties (PyMC en NumPyro) waarbij discrete dimensies analytisch gemarginaliseerd zijn.
   • Met 100.000 simulaties komen de MNPE-posteriors overeen met de MCMC-referentie in zowel continue als discrete dimensies.
   • Calibratiediagnostiek (SBC-ranghistogrammen en ECE-betrouwbaarheidsdiagrammen) bevestigt dat de posterior goed gekalibreerd is, waarbij de continue dimensies de baseline voor eindige steekproeven benaderen en de discrete dimensies overeenkomen met de verwachte fout van een perfect gekalibreerde classifier.

3. Hodgkin–Huxley Biofysische Simulator:

   • Een "black-box" simulator voor neurale activiteit met onberekenbare likelihoods, waarbij continue geleidingsvermogens en discrete kanaalidentiteiten (bijvoorbeeld aanwezigheid van specifieke ionkanalen) betrokken zijn.
   • MNPE slaagt erin gezamenlijke posteriors af te leiden voor continue geleidingsvermogens en categorische kanaaltypes.
   • De methode produceert goed gekalibreerde posteriors over verschillende trainingsbudgetten, wat de analyse van compensatiemechanismen tussen ionkanalen mogelijk maakt (bijvoorbeeld de relatie tussen L-type calciumkanalen en kaliumgeleidingsvermogen).

Betekenis en Beweringen
Het artikel beweert dat MNPE een kritieke kloof in SBI opvult door een unificerend, geamortiseerd kader te bieden voor simulatoren met gemengde parameterruimten. De primaire bijdragen zijn:

• Generaliseerbaarheid: Het ondersteunt willekeurige aantallen discrete en continue parameters en is compatibel met elk continu generatief netwerk (flows, diffusie).
• Efficiëntie: In tegenstelling tot MCMC amortiseert MNPE de kosten van simulatie en training, wat snelle inferentie op nieuwe observaties mogelijk maakt via een enkele forward pass.
• Calibratie: Het biedt een principiële methode voor het beoordelen van calibratie in gemengde settings, waarbij rang-gebaseerde en classificatie-gebaseerde metrieken worden gecombineerd met empirische baselines.
• Implementatie: De methode is geïmplementeerd in het publiek beschikbare Python-pakket sbi, waardoor het toegankelijk is voor domeinwetenschappers zonder uitgebreide afstelling.

De auteurs wijzen beperkingen inherent aan de NPE-familie aan, zoals de behoefte aan grote trainingsdatasets voor complexe posteriors en gevoeligheid voor modelmisspecificatie. Zij verduidelijken ook dat hun calibratieprocedure marginale calibratie per dimensie beoordeelt in plaats van volledige gezamenlijke calibratie, wat voor gemengde ruimten nog steeds uitdagend blijft.