JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

Sintesi Tecnica: JADAI – Ammortamento Congiunto di Design Adattivo e Inferenza Bayesiana

Problema
Molti domini scientifici e ingegneristici comportano l'inferenza di parametri non osservati $\theta$ da osservazioni $x$ generate da simulatori complessi. Sebbene l'Inferenza Basata su Simulazione (SBI) abbia fatto avanzare la capacità di approssimare distribuzioni a posteriori $p(\theta|x)$ senza valutazioni esplicite della verosimiglianza, questi metodi assumono tipicamente un design sperimentale fisso. Al contrario, il Design Sperimentale Bayesiano (BED) e il Design Adattivo (BAD) si concentrano sulla selezione di variabili di design $\xi$ per massimizzare il guadagno di informazione (ad esempio, l'Expected Information Gain, EIG), ma spesso trattano l'inferenza come un passaggio separato e computazionalmente costoso eseguito dopo la raccolta dei dati. Gli approcci esistenti spesso de-accoppiano questi compiti: i metodi BAD spesso si affidano a un'inferenza non ammortizzata e lenta o a semplici famiglie parametriche, mentre i metodi SBI non ottimizzano l'acquisizione dei dati. Questa separazione porta a un guadagno di informazione subottimale e a un'inefficienza computazionale, particolarmente in contesti ad alta dimensionalità o multimodali.

Metodologia
Gli autori introducono JADAI (Jointly Amortizing Adaptive Design and Bayesian Inference), un framework che unifica il design adattivo e l'inferenza a posteriori in un unico sistema addestrabile end-to-end. L'innovazione principale è l'ammortamento congiunto di tre componenti:

1. Una Rete di Policy ($\pi_\phi$): Propone il design successivo $\xi_t$ basandosi sulla storia sperimentale corrente.
2. Una Rete di Storia/Riassunto ($\eta_\omega$): Codifica la sequenza di coppie design-osservazione passate $\{(\xi_k, x_k)\}_{k=1}^t$ in una statistica riassuntiva a dimensione fissa $h_t$.
3. Una Rete di Inferenza ($q_\psi$): Approssima la distribuzione a posteriori $p(\theta|h_t)$ ad ogni step della sequenza sperimentale.

Obiettivo di Ottimizzazione
JADAI minimizza una funzione di perdita generica che aggrega le riduzioni incrementali dell'errore a posteriori lungo le sequenze sperimentali. Invece di massimizzare direttamente l'Expected Information Gain (EIG) — che richiederebbe valutazioni di densità — il framework utilizza una somma telescopica delle perdite a posteriori.

• L'obiettivo di addestramento è derivato da un limite inferiore variazionale sul Total Expected Information Gain (TEIG).
• La perdita al passo $t$, $\ell_t$, è definita dalla funzione di perdita dell'estimatore a posteriori (ad esempio, diffusione o flow matching).
• La funzione di utilità $u_T$ è costruita come una somma di differenze tra perdite consecutive: $u_T = \sum (\ell^*_{t-1} - \ell_t)$, dove $\ell^*$ denota una versione "detached" (con stop-gradient) della perdita precedente.
• Questa formulazione assicura che la policy e le reti di riassunto siano addestrate per massimizzare la riduzione della perdita a posteriori ad ogni passo, massimizzando efficacemente il guadagno di informazione senza richiedere valutazioni esplicite delle densità della distribuzione a posteriori.

Scelte Architetturali Chiave

• Modelli di Diffusione: La rete di inferenza $q_\psi$ è istanziata utilizzando modelli di diffusione (o flow matching), consentendo al framework di approssimare distribuzioni a posteriori ad alta dimensionalità e multimodali ad ogni step sperimentale, una capacità spesso mancante nei precedenti metodi BAD limitati ai normalizing flows.
• Stabilità dell'Addestramento: Per evitare la backpropagation attraverso il tempo (BPTT) nidificata e garantire un addestramento stabile, il framework impiega:
  • Stati di Storia Detached: I design sono generati da uno stato di riassunto "detached" per prevenire percorsi di gradiente ciclici tra la rete di policy e la rete di storia.
  • Apprendimento Curriculum (Curriculum Learning): La lunghezza del rollout viene aumentata gradualmente durante l'addestramento, e viene utilizzata una strategia di "design prior mix-in" nelle fasi iniziali per esporre le reti a design diversificati prima di fare affidamento esclusivamente sulla policy appresa.

Risultati
JADAI è stato valutato su tre benchmark che rappresentano una crescente complessità nella struttura della policy e della distribuzione a posteriori:

1. Localizzazione della Sorgente (Location Finding, LF): Un compito che richiede l'inferenza delle posizioni delle sorgenti da misurazioni di intensità rumorose. JADAI ha ottenuto punteggi sPCE (sequential Prior Contrastive Estimation) superiori o competitivi rispetto ai baseline allo stato dell'arte (ad esempio, DAD, iDAD, RL-BOED), specialmente per orizzonti lunghi ($T > 10$). La policy appresa ha bilanciato efficacemente esplorazione ed esploitazione, e la distribuzione a posteriori basata sulla diffusione ha catturato correttamente la natura multimodale e scambiabile delle sorgenti.
2. Elasticità di Sostituzione Costante (Constant Elasticity of Substitution, CES): Un problema di design impegnativo che riguarda l'inferenza delle preferenze, dove i design informativi si trovano in un ristretto "punto ottimale" (sweet spot). JADAI ha superato tutti i baseline esistenti all'orizzonte standard $T=10$, dimostrando robustezza in scenari in cui i design casuali sono non informativi.
3. Scoperta di Immagini MNIST: Un compito ad alta dimensionalità in cui l'obiettivo è ricostruire un'immagine da misurazioni parziali. JADAI ha raggiunto punteggi state-of-the-art per l'indice di similarità strutturale (SSIM) e l'errore quadratico medio normalizzato (NRMSE), superando il baseline CoDiff. Il metodo ha identificato con successo le classi di cifre precocemente e ha raffinato i dettagli locali, mantenendo le prestazioni anche in presenza di rumore di misurazione additivo.

Significatività e Rivendicazioni
L'articolo sostiene che JADAI rappresenti un passo avanti significativo unificando design e inferenza, eliminando così la necessità di passaggi di inferenza separati e costosi durante il processo adattivo. I contributi chiave includono:

• Ammortamento Congiunto: È il primo framework che ammorta congiuntamente il design adattivo e l'inferenza a posteriori utilizzando un'utilità generale che aggrega le riduzioni incrementali dell'errore a posteriori, evitando la necessità di valutazioni di densità.
• Scalabilità a Posteriori Complesse: Sfruttando i modelli di diffusione, JADAI può gestire distribuzioni a posteriori ad alta dimensionalità e multimodali ad ogni step sperimentale, superando i limiti di espressività dei precedenti metodi basati su normalizing flows o miscele gaussiane.
• Prestazioni: Il framework raggiunge prestazioni superiori o competitive attraverso diversi benchmark, inclusa la ricostruzione di immagini ad alta dimensionalità, mantenendo l'efficienza computazionale durante l'impiego (i rollout vengono completati in millisecondi).

Gli autori osservano che, sebbene JALDA attualmente si basi su simulatori differenziabili, l'approccio è ben adatto al crescente panorama dei framework di simulazione differenziabile. Riconoscono che il campionamento a posteriori rimane un collo di bottiglia per l'impiego sub-secondario in casi ad alta dimensionalità, ma suggeriscono la distillazione post-hoc come una possibile direzione futura. Il lavoro non pretende di risolvere l'ottimizzazione di simulatori black-box senza gradienti, ma si pone come una soluzione robusta per contesti differenziabili dove le pipeline classiche sono computazionalmente impraticabili.