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

1. Il Problema

Il Nested Sampling (NS) è diventato uno strumento fondamentale per l'inferenza bayesiana nelle scienze fisiche, poiché permette di stimare direttamente l'evidenza bayesiana ($Z$), essenziale per il confronto tra modelli e il test di ipotesi. Tuttavia, l'applicazione del Nested Sampling a modelli ad alta dimensionalità, in particolare ai modelli bayesiani gerarchici, presenta sfide computazionali significative.

Nei modelli gerarchici, si hanno tipicamente $J$ gruppi di dati, ciascuno con parametri locali $\theta_j$ e parametri iperparametri condivisi $\psi$. La verosimiglianza (likelihood) si fattorizza come prodotto di contributi locali: $L = \prod L_j$.
Il collo di bottiglia principale risiede nel campionamento dalla prior vincolata (la distribuzione a priori troncata a una soglia di verosimiglianza $\ell^*$).

• Approccio standard: Per ogni proposta di un parametro locale, è necessario verificare il vincolo globale $\sum \ell_j > \ell^*$. In un'implementazione ingenua, questo richiede il ricalcolo della somma totale dei $J$ termini a ogni passo, con un costo di $O(J)$. Poiché sono necessari $O(J)$ aggiornamenti per un ciclo completo (sweep), il costo totale per sostituzione di un punto diventa $O(J^2)$.
• Conseguenza: La complessità complessiva dell'algoritmo scala come $O(J^3)$ (o peggio), rendendo il metodo proibitivo per cataloghi con centinaia o migliaia di gruppi (es. sorgenti di onde gravitazionali, galassie).
• Alternativa Gradient-Based: I metodi basati su gradienti (come HMC/NUTS) scalano meglio ma faticano con vincoli di verosimiglianza complessi e geometrie posteriori patologiche (es. "funnel"), e non forniscono direttamente l'evidenza senza tecniche aggiuntive costose.

2. Metodologia: NS-SwiG

L'autore propone NS-SwiG (Nested Sampling with Slice-within-Gibbs), un algoritmo che combina l'anello esterno del Nested Sampling con un kernel interno di Slice-within-Gibbs che sfrutta la struttura di fattorizzazione del modello.

A. Decomposizione del Budget di Verosimiglianza

L'idea centrale è decomporre il vincolo globale in "budget" locali.
Dato lo stato corrente con log-verosimiglianza totale $S = \sum \ell_j$ e soglia $\ell^*$, se si propone una modifica solo per il gruppo $k$ (da $\theta_k$ a $\theta'_k$), il vincolo globale $\sum \ell_j > \ell^*$ può essere riscritto come:
$$\ell'_k > \ell^* - (S - \ell_k) \equiv B_k$$
Dove $B_k$ è il budget locale per il gruppo $k$.

• Vantaggio: Una volta calcolato $S$ e memorizzato il vettore dei contributi $\{\ell_j\}$, il calcolo del budget $B_k$ richiede tempo costante $O(1)$.
• Aggiornamento: Se la proposta è accettata, si aggiorna $S$ e $\ell_k$ in $O(1)$.
• Risultato: Il costo di un intero sweep (aggiornamento di tutti i $J$ blocchi) scende da $O(J^2)$ a $O(J)$.

B. Kernel Slice-within-Gibbs Vincolato

L'anello interno di campionamento utilizza un approccio a blocchi alternati:

1. Aggiornamento degli Iperparametri ($\psi$): Si campiona $\psi$ dalla sua condizione completa vincolata. Questo richiede la valutazione di tutte le $J$ likelihood (costo $O(J)$), ma viene fatto solo una volta per ciclo.
2. Sweep Interno sui Parametri Locali ($\theta_j$): Per ogni gruppo $j$, si campiona $\theta_j$ usando il Slice Sampling vincolato al budget locale $B_j$. Poiché ogni controllo richiede solo un calcolo di likelihood locale, il costo è $O(1)$ per blocco.

L'algoritmo è implementato in JAX (framework BlackJAX), permettendo la vettorizzazione e il parallelismo su più particelle (live points).

3. Contributi Chiave

1. Decomposizione del Budget: Dimostrazione che il vincolo globale del Nested Sampling può essere trasformato in vincoli locali aggiornabili in tempo costante tramite caching, estendibile anche a variabili latenti con struttura Markoviana.
2. Kernel Constrained Slice-within-Gibbs: Costruzione di un kernel MCMC specifico per la prior vincolata che riduce la complessità per sweep da $O(J^2)$ a $O(J)$.
3. Riduzione della Complessità: Sotto assunzioni standard, la complessità complessiva dell'algoritmo scende da cubica ($O(J^3)$) a quadratica ($O(J^2)$), rendendo fattibile l'analisi di modelli con migliaia di dimensioni.
4. Validazione su Benchmark: Implementazione completa e validazione su quattro benchmark, dimostrando scalabilità fino a $d \approx 2500$ dimensioni.

4. Risultati Sperimentali

L'articolo valuta NS-SwiG su quattro scenari rispetto a metodi basati su gradienti (NUTS, SMC-HMC) e al Nested Sampling standard (NSS):

• Modello Gaussiano Gerarchico:
  • NS-SwiG conferma la scalabilità $O(J)$ in termini di valutazioni della likelihood completa.
  • Rispetto a NSS standard, NS-SwiG è ordini di grandezza più veloce per $J$ grandi (es. $J=1000$).
  • Le stime dell'evidenza rimangono accurate e coerenti con i valori analitici.
• Distribuzione "Funnel" (Neal):
  • Un caso classico di geometria patologica dove i metodi basati su gradienti (NUTS) falliscono o richiedono reparametrizzazione.
  • NS-SwiG cattura correttamente la geometria del funnel anche nella parametrizzazione "centrata" (dove NUTS produce divergenze), fornendo stime di evidenza accurate senza bisogno di trasformazioni non centrate.
• Modelli ad Effetti Contestuali (Radon):
  • Su un dataset reale di 85 contee, NS-SwiG fornisce stime di posteriori e evidenze accurate sia in parametrizzazione centrata che non centrata.
  • SMC-HMC mostra segni di degenerazione delle particelle e stime di evidenza distorte nel caso centrato.
• Volatilità Stocastica (S&P 500):
  • Test su serie temporali con struttura Markoviana (che rompe l'indipendenza condizionale semplice).
  • L'estensione di NS-SwiG gestisce la struttura Markoviana mantenendo il costo di controllo $O(1)$ tramite budget locali per i vicini.
  • NS-SwiG supera NUTS in termini di efficienza computazionale e fornisce stime di evidenza, anche su serie lunghe ($J=2516$).

5. Significato e Implicazioni

• Scalabilità per l'Astrofisica: Il metodo è particolarmente rilevante per l'astrofisica e la cosmologia, dove i cataloghi di oggetti (onde gravitazionali, galassie) stanno crescendo esponenzialmente. NS-SwiG permette di scalare l'inferenza gerarchica mantenendo la capacità di calcolare l'evidenza bayesiana, cruciale per il confronto dei modelli.
• Robustezza Geometrica: A differenza dei metodi basati su gradienti, NS-SwiG non richiede una parametrizzazione specifica (come quella non centrata) per funzionare bene su geometrie complesse o patologiche.
• Alternativa ai Metodi Gradient-Free: Offre una soluzione scalabile per problemi dove i gradienti sono difficili da calcolare o dove la struttura del modello favorisce approcci basati su blocchi.
• Limiti: L'algoritmo richiede una struttura di fattorizzazione nella likelihood. Sebbene comune, non è universale. Inoltre, il numero di particelle necessarie per mantenere la precisione dell'evidenza cresce con la dimensionalità, il che può diventare un collo di bottiglia di memoria oltre certe soglie ($d \gtrsim 10^3$).

In sintesi, NS-SwiG rappresenta un avanzamento significativo nell'efficienza computazionale del Nested Sampling, trasformandolo da un metodo limitato a problemi a bassa dimensionalità o con geometrie semplici, a uno strumento praticabile per l'inferenza gerarchica su larga scala in presenza di vincoli complessi.