Diffusion posterior sampling for simulation-based inference in tall data settings

1. Il Problema: Inferenza su Modelli Non Lineari e Dati "Tall"

Il lavoro affronta la sfida di identificare i parametri $\theta$ di modelli non lineari complessi che spiegano dati osservati $x$. In molti campi scientifici (es. neuroscienze computazionali), questi modelli si basano su simulatori costosi dove la verosimiglianza $p(x|\theta)$ è intrattabile, rendendo metodi classici come MCMC inapplicabili o proibitivi.

La Simulation-Based Inference (SBI) risolve questo problema addestrando modelli generativi profondi per approssimare la distribuzione a posteriori $p(\theta|x)$ utilizzando dati simulati. Tuttavia, l'articolo si concentra specificamente sul setting "Tall Data", dove sono disponibili multiple osservazioni indipendenti e identicamente distribuite (i.i.d.) $x^*_{1:n} = (x^*_1, \dots, x^*_n)$.
In questo contesto, la posteriori a dati "tall" è definita come:
$$p(\theta | x^*_{1:n}) \propto \lambda(\theta)^{1-n} \prod_{j=1}^n p(\theta | x^*_j)$$
L'obiettivo è ottenere una posteriori più precisa e identificabile man mano che $n$ aumenta.

Limiti delle soluzioni esistenti:

• NPE/NLE/NRE: Richiedono spesso dataset di addestramento aumentati (costosi in termini di simulazioni) o non scalano bene con dimensioni variabili di $n$.
• F-NPSE (Factorized Neural Posterior Score Estimation): Un approccio recente basato su score (gradiente del log-probabilità) che combina i punteggi di singole osservazioni. Sebbene efficiente e flessibile, F-NPSE soffre di un collo di bottiglia critico: poiché la posteriori composta non corrisponde a un processo di diffusione noto, il campionamento richiede dinamiche di Langevin (MCMC). Questo rende il campionamento lento, instabile e sensibile all'iperparametrizzazione.

2. Metodologia Proposta

Gli autori propongono un nuovo algoritmo che mantiene i vantaggi della composizione degli score (amortization e flessibilità) ma elimina la necessità delle dinamiche di Langevin. La soluzione si basa sull'approssimazione esplicita del processo di diffusione della posteriori a dati "tall".

Fondamenti Teorici

1. Composizione degli Score: Partendo dalla formula della posteriori fattorizzata, gli autori derivano una formula esatta per lo score della posteriori diffusa a dati "tall" ($\nabla_{\theta_t} \log p_t(\theta_t | x^*_{1:n})$). Questa formula include un termine di correzione ($\nabla \log L_\lambda$) che manca nelle approssimazioni precedenti come F-NPSE.
2. Approssimazione di Secondo Ordine: Il termine di correzione è intrattabile. Gli autori lo approssimano utilizzando il framework di Tweedie, assumendo che i kernel di diffusione inversa (backward kernels) per la prior e per ogni singola posteriori siano distribuzioni Gaussiane.
   • Questo porta a una formula chiusa per lo score approssimato della posteriori tall, che è una combinazione lineare pesata degli score individuali e della prior.
   • La formula risultante (Eq. 16 nel paper) dipende dalle matrici di covarianza dei kernel inversi.

Algoritmi Proposti

Per rendere l'algoritmo pratico, vengono proposti due approcci per gestire le matrici di covarianza $\Sigma_{t,j}$:

1. JAC (Jacobian Approximation):

   • Calcola la covarianza approssimata prendendo la Jacobiana dello score appreso ($\nabla_{\theta_t} s_\phi$).
   • Vantaggi: Teoricamente molto preciso in assenza di rumore.
   • Svantaggi: Computazionalmente costoso ($O(m^2)$ per la matrice Jacobiana) e numericamente instabile, specialmente quando lo score appreso è rumoroso.

2. GAUSS (Gaussian Approximation):

   • Assume che le matrici di covarianza siano costanti e stimate empiricamente.
   • Procedura: Si esegue una breve inferenza (circa 100 passi DDIM) per ogni osservazione singola per stimare la covarianza empirica delle campioni risultanti. Questa covarianza viene poi usata come costante nell'algoritmo di campionamento finale.
   • Vantaggi: Estremamente stabile, veloce e robusto al rumore. Elimina il termine di residuo $F$ nell'equazione dello score.

Entrambi gli algoritmi permettono di utilizzare campionatori deterministici (come DDIM) invece delle costose dinamiche di Langevin, garantendo una generazione di campioni più rapida e stabile.

3. Risultati Sperimentali

Gli autori hanno valutato i metodi su tre livelli di complessità, confrontandoli con F-NPSE (implementato con Langevin) e altri baselines.

• Modelli Toy (Gaussiani e GMM):

  • Su modelli Gaussiani (dove la soluzione è analitica), GAUSS è quasi perfetto e 2.5 volte più veloce di Langevin.
  • Su modelli GMM (non Gaussiani), GAUSS mantiene una robustezza superiore al rumore, mentre JAC diventa instabile se lo score non è perfetto.
  • Speed-up: GAUSS è significativamente più veloce (fino a 2.5x) rispetto a Langevin, mantenendo un'accuratezza superiore (misurata con Sliced Wasserstein distance).

• Benchmark SBI (SLCP, SIR, Lotka-Volterra):

  • In scenari dove lo score deve essere appreso dai dati (NPSE), l'accumulo di errori di approssimazione con $n$ crescente è un problema critico.
  • GAUSS supera costantemente Langevin e JAC in termini di accuratezza e stabilità, specialmente per grandi $n$ (fino a 30 osservazioni).
  • Langevin tende a divergere o degradare rapidamente all'aumentare di $n$ e del rumore.

• Caso Reale: Modello Jansen-Rit (Neuroscienze):

  • Applicazione su un modello complesso di massa neurale (JRNMM) per l'inversione di segnali EEG.
  • Il problema presenta un'indeterminazione (accoppiamento tra parametri) che richiede dati "tall" per essere risolta.
  • GAUSS è l'unico metodo che passa il test di validazione locale (ℓ-C2ST) in entrambi i casi (3D e 4D), dimostrando di poter concentrare la posteriori attorno ai parametri veri e rivelare le dipendenze parametriche nascoste.

4. Contributi Chiave

1. Nuovo Algoritmo di Campionamento: Introduzione di un metodo che approssima esplicitamente il processo di diffusione inverso per la posteriori a dati "tall", eliminando la necessità delle dinamiche di Langevin.
2. Formulazione Teorica: Derivazione di una formula esatta per lo score della posteriori diffusa fattorizzata, includendo un termine di correzione spesso ignorato, e la sua successiva approssimazione tramite kernel Gaussiani.
3. Stabilità ed Efficienza: Dimostrazione che l'approccio GAUSS offre il miglior compromesso tra velocità, stabilità numerica e accuratezza, superando i limiti di F-NPSE.
4. Scalabilità: Validazione su problemi reali complessi, mostrando come l'approccio sia scalabile a scenari del mondo reale con modelli non lineari e costosi.

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo nell'ambito della Simulation-Based Inference (SBI) per dati "tall". Risolve il principale collo di bottiglia computazionale e di stabilità dei metodi basati su score (come F-NPSE) permettendo l'uso di campionatori deterministici veloci (DDIM).

La capacità di aggregare efficientemente informazioni da multiple osservazioni indipendenti senza dover riaddestrare modelli o eseguire costose simulazioni MCMC apre nuove possibilità per:

• L'inferenza in tempo reale in applicazioni scientifiche critiche.
• La risoluzione di problemi inversi mal posti (ill-posed) grazie alla maggiore identificabilità dei parametri fornita dai dati "tall".
• L'applicazione di metodi di inferenza basati su deep learning a problemi di grandi dimensioni in neuroscienze, fisica e ingegneria.

In sintesi, il paper trasforma l'inferenza su dati "tall" da un problema computazionalmente oneroso e instabile a una procedura efficiente, stabile e scalabile, rendendo l'approccio basato su score-based diffusion una soluzione praticabile per problemi scientifici complessi.