PriorGuide: Test-Time Prior Adaptation for Simulation-Based Inference

Titel: PriorGuide: Test-Time Prior Adaptation voor Simulatie-gebaseerde Inference

1. Het Probleem

Simulatie-gebaseerde inferentie (SBI) is een krachtige methode voor Bayesiaanse inferentie in complexe systemen (zoals engineering en neurowetenschappen) waar de waarschijnlijkheidsfunctie (likelihood) $p(x|\theta)$ niet analytisch te berekenen is, maar simulaties $x \sim p(x|\theta)$ wel mogelijk zijn.

Moderne SBI-methoden gebruiken vaak geamortiseerde inferentie (bijv. met diffusiemodellen), waarbij een neurale net wordt getraind op een enorme dataset van gesimuleerde parameter-data paren $(\theta, x)$. Een cruciale beperking van deze aanpak is dat het getrainde model onlosmakelijk verbonden is met de trainingsprior $p_{train}(\theta)$ die werd gebruikt om de data te genereren.

In de praktijk willen onderzoekers echter vaak:

1. Expertkennis integreren: Een specifiekere prior gebruiken dan de brede uniforme prior die tijdens training werd gebruikt.
2. Sensitiviteitsanalyse uitvoeren: Controleren hoe robuust conclusies zijn bij verschillende aannames over de prior.

Huidige oplossingen voor het aanpassen van de prior zijn problematisch:

• Opnieuw trainen: Het trainen van een nieuw model voor elke nieuwe prior is computatief zwaar en onpraktisch, vooral bij dure simulatoren.
• Meta-priors: Bestaande methoden (zoals ACE of Distribution Transformers) trainen op een "meta-prior" (een verzameling van mogelijke priors). Dit beperkt de flexibiliteit tot vooraf gedefinieerde prior-families en kan niet omgaan met complexe, niet-factoriserende priors die niet in de trainingsset stonden.

2. Methodologie: PriorGuide

PriorGuide introduceert een methode om een reeds getraind diffusiemodel aan te passen aan een nieuwe doel-prior $q(\theta)$ op het moment van inferentie (test-time), zonder het model opnieuw te hoeven trainen.

Kernprincipes:

• Prior Ratio als Gids: Het doel is om te samplingen uit de nieuwe posterior $q(\theta|x) \propto q(\theta)p(x|\theta)$. Dit kan worden herschreven als:
  $$q(\theta|x) \propto \underbrace{\frac{q(\theta)}{p_{train}(\theta)}}_{r(\theta)} \cdot p(\theta|x)$$
  Hierbij is $r(\theta)$ de prior ratio. PriorGuide gebruikt deze ratio als een guidance term in het diffusieproces om de generatie te sturen van de oorspronkelijke posterior $p(\theta|x)$ naar de nieuwe posterior $q(\theta|x)$.

• Analytische Benadering: De exacte berekening van de guidance term vereist een verwachting over de reverse transition kernel, wat intractabel is. PriorGuide lost dit op met twee benaderingen:

  1. Reverse Kernel: De overgang $p(\theta_0 | \theta_t, x)$ wordt benaderd als een multivariate Gaussische verdeling (gebaseerd op Tweedie's formule).
  2. Prior Ratio (GMM): De ratio $r(\theta) = q(\theta)/p_{train}(\theta)$ wordt benaderd als een Gaussische Mixture Model (GMM). Omdat zowel $q$ als $p_{train}$ analytisch bekend zijn, kan deze GMM worden gefit (via gradient descent) om de ratio te representeren.

• De Update-regel: Door de GMM-benadering van de ratio en de Gaussische kernel te combineren, kan de guidance term analytisch worden berekend. Dit resulteert in een correctie van de voorspelling van het getrainde model:
  $$\mu_{new} = \mu_{old} + \text{correctie term gebaseerd op GMM-componenten}$$

• Langevin Dynamics voor Nauwkeurigheid: Omdat de Gaussische benadering van de kernel een approximatie is, introduceert PriorGuide Langevin dynamics-stappen na elke diffusiestap. Dit fungeert als een correctiemechanisme (annealed MCMC) dat de steekproefkwaliteit verbetert, ten koste van extra rekentijd op het testmoment.

3. Belangrijkste Bijdragen

1. Een principieel raamwerk voor test-time prior adaptatie: PriorGuide maakt het mogelijk om pre-getrainde diffusiemodellen voor SBI flexibel aan te passen aan willekeurige nieuwe priors zonder retraining.
2. Novel Guidance Approximation: De introductie van een GMM-benadering voor de prior ratio, wat een analytische oplossing biedt voor de guidance term en Monte Carlo schattingen (met bias en variance) vermijdt.
3. Trade-off tussen Kosten en Nauwkeurigheid: Het framework biedt een mechanisme om de nauwkeurigheid te verhogen door meer compute-investering (meer diffusiestappen en Langevin-correcties) op het testmoment, wat het concept van "test-time compute" voor Bayesiaanse inferentie belichaamt.
4. Uitgebreide Validatie: Empirisch bewijs op diverse SBI-problemen (van 2D benchmarks tot complexe tijdreeksmodellen in neurowetenschappen en engineering).

4. Resultaten

De auteurs evalueren PriorGuide op zes verschillende SBI-problemen, waaronder Two Moons, Ornstein-Uhlenbeck Process (OUP), het Turin-model (radio-propagatie), en Bayesian Causal Inference (BCI).

• Posterior Inference: PriorGuide overtreft de basisdiffusiemodel (Simformer) aanzienlijk in alle scenario's, vooral bij sterke of complexe (mixture) priors. Het presteert vaak beter dan of gelijk aan bestaande methoden voor prior-adaptatie zoals ACE (Amortized Conditioning Engine), vooral bij niet-factoriserende priors.
• Posterior Predictive Inference: De methode werkt ook effectief voor het voorspellen van toekomstige data onder nieuwe priors, met resultaten die dicht bij de ground truth liggen.
• Robuustheid: De methode is robuust tegen matige verschuivingen tussen de trainingsprior en de test-prior. Hoewel prestaties afnemen bij extreme "out-of-distribution" priors, degradeert de methode geleidelijk in plaats van catastrofaal te falen.
• Kosten-Nauwkeurigheid: De Pareto-frontier analyse toont aan dat het toevoegen van Langevin-stappen (NL) de nauwkeurigheid (gemeten in MMTV) significant verbetert ten opzichte van puur diffusiestappen, binnen een redelijke rekentijd.

5. Betekenis en Toekomstperspectief

PriorGuide lost een fundamenteel probleem op in het veld van simulatie-gebaseerde inferentie: de koppeling tussen een getraind model en de specifieke prior die bij de training werd gebruikt.

• Praktische Impact: Het stelt onderzoekers in staat om post-hoc expertkennis toe te passen en uitgebreide sensitiviteitsanalyses uit te voeren zonder de enorme kosten van het opnieuw trainen van modellen met dure simulatoren.
• Paradigmaverschuiving: Het ondersteunt de trend naar "foundation models" voor SBI, waarbij één model dient als basis en specifieke eisen (zoals een prior) dynamisch worden ingebracht tijdens de inferentie.
• Beperkingen: De methode vereist dat de nieuwe prior $q(\theta)$ voldoende overlap heeft met de trainingsprior $p_{train}(\theta)$ om te voorkomen dat het model in "out-of-distribution" gebieden terechtkomt waar de score-functie slecht is getraind. Ook kan de berekening van de GMM-guidance in zeer hoge dimensies (>20) computatie-intensief worden.

Kortom, PriorGuide biedt een flexibele, efficiënte en nauwkeurige oplossing voor het aanpassen van Bayesiaanse inferentie aan veranderende aannames, wat essentieel is voor robuuste wetenschappelijke modellering.