Nested Sampling with Slice-within-Gibbs: Efficient Evidence Calculation for Hierarchical Bayesian Models

Hieronder volgt een gedetailleerde technische samenvatting van het paper "Nested Sampling with Slice-within-Gibbs: Efficient Evidence Calculation for Hierarchical Bayesian Models" van David Yallup, geschreven in het Nederlands.

Probleemstelling

Bayesiaanse inferentie in hoge dimensies, met name bij hiërarchische modellen, staat voor twee grote uitdagingen:

1. Schaalbaarheid van Nested Sampling: Nested Sampling is een krachtige methode voor het schatten van de Bayesiaanse bewijskracht (evidence, $Z$), wat essentieel is voor modelvergelijking. Echter, de standaardimplementatie vereist het trekken van steekproeven uit een voorwaartse verdeling die beperkt is door een likelihood-drempel. Bij hiërarchische modellen met $J$ groepen en lokale parameters $\theta_j$ en gedeelde hyperparameters $\psi$, moet bij elke MCMC-proposal de globale likelihood-beperking ($\sum \ell_j > \ell^*$) worden gecontroleerd. Een naïeve implementatie herberekent de totale likelihood bij elke stap, wat leidt tot een rekentijd van $O(J)$ per proposal. Omdat er $O(J)$ parameters zijn die updates nodig hebben, resulteert dit in een totale complexiteit van $O(J^2)$ per vervanging van een "live point". In de praktijk leidt dit tot een totale algoritme-complexiteit van $O(J^3)$, wat onhoudbaar wordt voor modellen met duizenden groepen.
2. Beperkingen van Gradiënt-gebaseerde Methoden: Hoewel methoden zoals Hamiltonian Monte Carlo (HMC) en NUTS goed schalen door gebruik te maken van gradiënten, worstelen ze vaak met de complexe geometrieën van hiërarchische modellen (zoals de "funnel"-problematiek). Bovendien kunnen deze methoden de bewijskracht ($Z$) niet direct schatten zonder extra, vaak kostbare, technieken.

Methodologie: NS-SwiG

Het paper introduceert Nested Sampling with Slice-within-Gibbs (NS-SwiG), een algoritme dat de structurele onafhankelijkheid van hiërarchische modellen benut om de rekentijd drastisch te verminderen.

1. Likelihood-Budget Decompositie:
De kern van de innovatie is het decomponeren van de globale likelihood-beperking in per-blok budgetten.

• In plaats van bij elke lokale update van een groep $k$ de som van alle $J$ likelihoods te herberekenen, wordt een "budget" $B_k$ gedefinieerd:
  $$B_k = \ell^* - S + \ell_k(\theta_k, \psi)$$
  waarbij $S$ de gecachete totale log-likelihood is.
• De globale voorwaarde $\sum \ell_j > \ell^*$ wordt hiermee omgezet in een lokale voorwaarde voor groep $k$: $\ell_k(\theta'_k, \psi) > B_k$.
• Omdat $S$ en de individuele $\ell_j$ worden gecachet, kan de budgetcontrole in $O(1)$ tijd worden uitgevoerd in plaats van $O(J)$.

2. Constrained Slice-within-Gibbs Kernel:
Het algoritme gebruikt een binnenste lus (inner loop) gebaseerd op een Slice-within-Gibbs kernel om de beperkte voorwaartse verdeling te bemonsteren:

• Buitenste update (Hyperparameters): Er wordt een slice-update uitgevoerd voor de hyperparameters $\psi$.
• Binnenste sweep (Lokale parameters): Er wordt een sweep uitgevoerd over alle $J$ lokale blokken $\theta_j$. Voor elk blok wordt een slice-update uitgevoerd die voldoet aan het lokale budget $B_j$.
• Door deze blokbewegingen te combineren met caching, daalt de kosten per volledige sweep van $O(J^2)$ naar $O(J)$.

3. Implementatie:
Het algoritme is geïmplementeerd in JAX binnen het BlackJAX framework, wat vectorisatie en parallelle verwerking van de "live points" mogelijk maakt. Het ondersteunt ook Markov-structuren (zoals in stochastische volatiliteitsmodellen) door de budgetten aan te passen aan lokale buren in plaats van alleen onafhankelijke groepen.

Belangrijkste Bijdragen

1. Likelihood-Budget Decompositie: Het tonen dat de globale Nested Sampling-beperking kan worden opgesplitst in per-blok budgetten die in constante tijd ($O(1)$) kunnen worden bijgewerkt via caching. Dit is een generalisatie van bestaande ideeën uit Gibbs-sampling, maar specifiek toegepast op de beperkte voorwaartse verdeling van Nested Sampling.
2. Efficiënte Kernel: De constructie van een "Constrained Slice-within-Gibbs" kernel die de complexiteit per vervanging verlaagt van $O(J^2)$ naar $O(J)$, en de totale algoritme-complexiteit van $O(J^3)$ naar $O(J^2)$ (onder standaard aannames).
3. Validatie en Schaalbaarheid: Een volledig gevalideerde implementatie die schaalbaarheid toont tot duizenden dimensies ($d \sim 2500$) en nauwkeurige bewijskracht-schattingen levert, zelfs op postie-geometrieën waar gradiënt-gebaseerde methoden (zoals NUTS) falen of divergeren.

Resultaten

De auteurs hebben NS-SwiG getest op vier benchmarks:

• Hiërarchisch Gaussisch Model:
  • Toont een schaling van $O(J)$ in "full-likelihood equivalents" (in plaats van $O(J^2)$ voor standaard Nested Sampling).
  • Bij $J=1000$ groepen is NS-SwiG aanzienlijk sneller dan standaard methoden, terwijl de nauwkeurigheid van de bewijskracht ($\log Z$) behouden blijft.
• Neal's Funnel:
  • Een klassiek probleem waarbij hyperparameters en lokale parameters sterk gekoppeld zijn.
  • In de "centered" parameterisatie (waar NUTS faalt door divergenties) levert NS-SwiG nauwkeurige resultaten.
  • NS-SwiG presteert vergelijkbaar met joint-space sampling maar met de verwachte schaalverbetering.
• Contextual Effects Model (Radon data):
  • Een regressiemodel met 85 groepen.
  • NS-SwiG levert nauwkeurige posterieurs en bewijskracht, zelfs in de uitdagende centered parameterisatie waar SMC-HMC een vertekende schatting gaf.
• Stochastische Volatiliteit (S&P 500):
  • Test met Markov-gestructureerde latent variabelen ($J=2516$).
  • NS-SwiG levert consistente posterieurs en bewijskracht, terwijl het minder rekentijd en evaluaties vereist dan NUTS, zelfs bij hoge dimensies.

Over het algemeen presteert NS-SwiG robuust zonder dat handmatige herschaling (reparameterization) nodig is, een veelvoorkomende vereiste voor gradiënt-gebaseerde methoden bij complexe geometrieën.

Betekenis en Conclusie

De paper biedt een cruciale doorbraak voor Bayesiaanse inferentie in de fysica en astronomie, waar hiërarchische modellen en grote catalogi van objecten (zoals gravitatiegolfbronnen of sterrenstelsels) steeds vaker voorkomen.

• Directe Bewijskracht: In tegenstelling tot veel schaalbare MCMC-methoden, biedt NS-SwiG direct een schatting van de Bayesiaanse bewijskracht, wat essentieel is voor modelvergelijking.
• Schaalbaarheid: Het verlaagt de rekentijd voor grote hiërarchische modellen aanzienlijk, waardoor analyse mogelijk wordt op schalen die voorheen onbereikbaar waren voor Nested Sampling.
• Robuustheid: Het algoritme is minder afhankelijk van handmatige tuning of specifieke herschalingen van het model dan gradiënt-gebaseerde methoden, wat het ideaal maakt voor complexe, niet-lineaire fysieke modellen.

Kortom, NS-SwiG maakt Nested Sampling opnieuw een praktische en schaalbare keuze voor moderne, hoog-dimensionale hiërarchische inferentieproblemen.