Tokenised Flow Matching for Hierarchical Simulation Based Inference

1. Il Problema: Inefficienza Computazionale nell'Inferenza Basata su Simulazione Gerarchica

L'Inferenza Basata su Simulazione (SBI) permette la stima dei parametri bayesiani quando la funzione di verosimiglianza è intrattabile, ma è disponibile un simulatore. Tuttavia, in contesti gerarchici (dove i dati provengono da molteplici siti o entità, come paesi in epidemiologia o pazienti in medicina), l'efficienza diventa un collo di bottiglia critico.

• Contesto: I modelli gerarchici condividono parametri globali ($\theta_g$) tra tutti i siti, mentre ogni sito $s$ ha i propri parametri locali ($\eta_s$) e osservazioni ($y_s$).
• Sfida: I metodi SBI tradizionali richiedono di simulare tutti i siti per ogni campione di addestramento. Se ci sono $n_s$ siti, il costo computazionale cresce linearmente con $n_s$, rendendo l'inferenza proibitiva per modelli complessi o dataset su larga scala.
• Limiti degli approcci esistenti: Le tecniche di "fattorizzazione della posterior" (Posterior Factorisation) riducono il numero di simulazioni ma introducono errori di propagazione o richiedono ancora simulazioni multiple per campione. Altre approcci non gerarchici non scalano bene con il numero di siti.

2. Metodologia Proposta: TFMPE

Gli autori propongono TFMPE (Tokenised Flow Matching for Posterior Estimation), un metodo che combina tre idee chiave per migliorare l'efficienza dei campioni e gestire osservazioni complesse.

A. Campionamento con Fattorizzazione della Verosimiglianza (Likelihood Factorisation - LF)

Invece di fattorizzare direttamente la distribuzione a posteriori (come fanno i metodi precedenti), il metodo propone di fattorizzare la verosimiglianza.

1. Fase 1 (Surrogato per sito): Viene addestrato un "surrogato neurale per sito" ($q_{\phi_l}$) che approssima la distribuzione delle osservazioni di un singolo sito dato i parametri globali e locali ($p(y_s | \theta_g, \eta_s)$). Questo richiede una sola simulazione per campione di addestramento.
2. Fase 2 (Stimatore della Posterior): Una volta appreso il surrogato, vengono generate osservazioni sintetiche multi-sito ($\hat{y}$) campionando dal surrogato. Queste vengono usate per addestrare uno stimatore della posterior ($q_{\phi_p}$) che mappa le osservazioni sintetiche ai parametri globali e locali.

• Vantaggio: Riduce drasticamente il numero di chiamate al simulatore costoso (da $n_s$ a 1 per campione di addestramento iniziale).

B. Tokenizzazione per Osservazioni Funzionali

Per gestire dati eterogenei (es. serie temporali irregolari, dati spaziali, variabili di dimensioni diverse), il metodo utilizza uno schema di tokenizzazione:

• Parametri e osservazioni sono mappati in una sequenza di "token".
• Ogni token include: identificatore della variabile, posizione intra-variabile, identificatore del gruppo (sito), e input funzionale (es. tempo o coordinate spaziali).
• Questo permette di gestire osservazioni a valore funzionale (funzioni) e dati irregolari senza doverli aggregare in statistiche riassuntive fisse, preservando l'informazione strutturale.

C. Flow Matching per la Stima della Posterior

Al posto delle reti generative tradizionali o dei flussi normalizzanti discreti, TFMPE utilizza il Flow Matching (adattamento di Lipman et al., 2022).

• Viene appreso un campo vettoriale che trasporta i campioni da una distribuzione base a quella target (la posterior).
• Architettura: Viene utilizzato un Transformer Encoder-only che processa la sequenza di token tramite self-attention. Questo permette di modellare le dipendenze complesse tra parametri globali, locali e osservazioni in modo scalabile.
• Stabilità: Il Flow Matching offre dinamiche di addestramento più stabili rispetto allo score matching e percorsi probabilistici più semplici.

3. Contributi Chiave

1. Strategia LF Sampling: Una nuova strategia di addestramento che riduce il costo delle simulazioni a un singolo sito per campione, mitigando il collo di bottiglia computazionale nei modelli gerarchici.
2. TFMPE: Un approccio unificato che integra tokenizzazione, Flow Matching e fattorizzazione della verosimiglianza, supportando nativamente osservazioni a valore funzionale e strutture gerarchiche complesse.
3. Benchmark Gerarchico: Introduzione del primo benchmark specifico per l'SBI gerarchico, adattato dal benchmark SBI esistente, per valutare sistematicamente l'efficienza dei campioni e la scalabilità.

4. Risultati Sperimentali

Il metodo è stato valutato su un nuovo benchmark, un modello stagionale SEIR (epidemiologia) e un modello di fluidodinamica computazionale (CFD) per l'emodinamica.

• Benchmark SBI:

  • TFMPE con LF sampling ha mostrato la migliore efficienza dei campioni (sample efficiency) su quasi tutti i task, specialmente quando il numero di siti ($n_s$) è elevato rispetto al budget di simulazione.
  • Confrontato con metodi non gerarchici (NPE, SNPE) e metodi di fattorizzazione della posterior (PF), TFMPE ha mantenuto posteriori ben calibrati con un costo computazionale inferiore.
  • In alcuni task complessi (es. "Two Moons"), la fattorizzazione della posterior (PF) ha funzionato meglio a causa della difficoltà di approssimare la verosimiglianza non lineare, ma TFMPE ha dimostrato superiorità nei task dove il modello di osservazione per sito è più semplice da approssimare rispetto alla mappatura globale.

• Modelli Reali:

  • Epidemiologia (SEIR): TFMPE ha gestito con successo dati osservativi irregolari su 100 siti, producendo posteriori calibrati dove i metodi MCMC (NUTS) fallivano o richiedevano tempi di calcolo proibitivi.
  • Emodinamica (CFD): In un modello di flusso sanguigno 1D, TFMPE ha ridotto il costo computazionale di un fattore ~1850x per paziente rispetto alla simulazione diretta, permettendo l'inferenza su coorti di pazienti con una velocità di addestramento 2x superiore rispetto ai metodi standard, pur mantenendo l'accuratezza della posterior.

5. Significato e Implicazioni

Il lavoro di TFMPE rappresenta un passo avanti significativo per l'inferenza scientifica su larga scala:

• Scalabilità: Rende fattibile l'inferenza bayesiana su modelli gerarchici con centinaia o migliaia di siti, un regime precedentemente inaccessibile a causa dei costi di simulazione.
• Flessibilità dei Dati: La tokenizzazione permette di incorporare dati reali "sporchi" o irregolari (serie temporali con frequenze diverse, dati mancanti) direttamente nel modello senza pre-elaborazione aggressiva.
• Efficienza: Dimostra che separare l'apprendimento del surrogato per sito dalla stima della posterior globale può bilanciare efficacemente errore di approssimazione e costo computazionale.
• Applicabilità: Il metodo è particolarmente rilevante per campi come l'epidemiologia globale, la medicina personalizzata e la fisica computazionale, dove i simulatori sono costosi e i dati provengono da molte fonti diverse.

In sintesi, TFMPE risolve il problema del "collo di bottiglia delle simulazioni" nei modelli gerarchici, offrendo un framework scalabile, stabile e capace di gestire la complessità dei dati reali.