Interpretable epistemic uncertainty decomposition in sequential generative models via polynomial chaos surrogates

Riepilogo Tecnico: Decomposizione Interpretabile dell'Incertezza Epistemica nei Modelli Generativi Sequenziali tramite Surrogati di Caos Polinomiale

Enunciato del Problema

I modelli generativi sequenziali, come le Reti di Flusso Generativo (GFlowNets), sono sempre più centrali nella scoperta scientifica guidata dall'IA, nell'ottimizzazione combinatoria e nella progettazione biologica. Questi modelli operano spesso in condizioni di incertezza epistemica, in particolare quando le funzioni di ricompensa che guidano la loro generazione sono esse stesse incerte (ad esempio, surrogati di reti neurali addestrati su dati sperimentali limitati). Sebbene i metodi esistenti come gli ensemble profondi e le reti neurali bayesiane possano quantificare la magnitudine dell'incertezza di previsione, non riescono a fornire una decomposizione interpretabile di tale incertezza. Di conseguenza, i praticanti non possono determinare quali componenti specifiche dell'incertezza della ricompensa in ingresso guidano quali decisioni generative, né possono distinguere tra segnali chimici robusti e artefatti statistici derivanti da bias del set di addestramento.

Metodologia

Gli autori propongono un framework di modellazione surrogata che propaga l'incertezza epistemica attraverso le GFlowNets utilizzando Espansioni di Caos Polinomiale (PCE). L'approccio procede in tre fasi:

1. Addestramento dell'Ensemble e Parametrizzazione: Un modesto ensemble di GFlowNets viene addestrato in condizioni di ricompensa variate. Lo spazio delle funzioni di ricompensa viene parametrizzato in uno spazio latente a bassa dimensione $\mu \in \mathbb{R}^d$ utilizzando l'Analisi delle Componenti Principali (PCA) o espansioni analitiche (ad esempio, Karhunen–Loève). La PCA è preferita perché le sue componenti ortogonali approssimano input gaussiani indipendenti, prerequisito per una valida decomposizione analitica di Sobol.
2. Adattamento del Surrogato PCE: Per ogni azione $k$ al passo di traiettoria $t$, viene adattato un PCE per mappare i parametri di ricompensa a bassa dimensione $\mu$ sulle probabilità di azione trasformate tramite log-rapporto (utilizzando una trasformazione Additive Log-Ratio rispetto a un'azione di riferimento). I coefficienti sono stimati tramite regressione ridge regolarizzata.
3. Decomposizione Analitica: I coefficienti PCE adattati permettono il calcolo analitico degli indici di sensibilità di Sobol. A differenza dei metodi Monte Carlo che richiedono un'integrazione costosa, questi indici sono derivati direttamente dai coefficienti polinomiali, fornendo una decomposizione esatta basata sulla varianza di quanto ogni componente dell'incertezza della ricompensa contribuisca alla varianza nella distribuzione della politica.

Il framework include garanzie formali di convergenza per il surrogato e i suoi indici di Sobol. Quattro delle cinque proprietà teoriche chiave (convergenza PCE, convergenza degli indici di Sobol, preservazione del simplesso e limiti di propagazione dell'incertezza) sono state formalmente verificate nell'assistente di prova Lean 4.

Contributi Chiave

1. Framework Generale: Un metodo computazionalmente efficiente per propagare e decomporre l'incertezza epistemica nei modelli generativi sequenziali condizionati utilizzando surrogati PCE adattati a piccoli ensemble.
2. Verifica Formale: Garanzie teoriche di convergenza per il surrogato e i suoi indici di sensibilità, con quattro dei cinque teoremi formalmente verificati in Lean 4.
3. Decomposizione Interpretabile della Sensibilità: Il primo metodo analitico per decomporre l'incertezza della politica nei modelli generativi, rivelando esattamente quali componenti della ricompensa guidano specifiche decisioni generative. Questa capacità non è disponibile dagli ensemble profondi standard, dalle reti neurali bayesiane o dal dropout Monte Carlo.
4. Validazione Empirica: Dimostrazioni su sette compiti di complessità crescente, inclusi tre applicazioni reali e quattro compiti di validazione sintetici.

Risultati e Applicazioni

Il framework è stato valutato su tre principali domini reali:

• Ottimizzazione delle Reazioni Chimiche (Buchwald–Hartwig): Sul dataset Doyle–Dreher, il framework ha analizzato la selezione di catalizzatore, base, alogenuro arilico e additivo. Ha rivelato che la selezione del catalizzatore è robusta (basso indice di Sobol, $D \approx 71$), indicando un segnale chimico genuino, mentre la selezione dell'additivo è fragile (alto indice di Sobol, $D \approx 179$, 2,5 volte superiore), indicando suscettibilità al bias del set di addestramento. Questo affronta direttamente il "dibattito Chuang–Keiser" riguardante gli artefatti statistici nei predittori di resa.
• Scoperta Causale Bayesiana (Segnalazione Proteica Sachs): Applicato alla rete Sachs, l'analisi ha mostrato che, sebbene tutti gli spigoli presentassero un'alta varianza a causa di dati limitati (200 osservazioni), l'incertezza si è decomposta in regimi distinti. Gli spigoli della via MAPK erano guidati da una componente principale della variazione di covarianza, mentre gli spigoli hub PKA/PKC erano guidati da un'altra. Questo fornisce una mappa mirata per esperimenti di perturbazione.
• Progettazione Molecolare Basata su Frammenti: In un compito di progettazione di farmaci, il framework ha ribaltato l'assunzione convenzionale secondo cui gli scaffold sono robusti e le decorazioni sono fragili. I risultati hanno mostrato che la posizione del linker era il componente più fragile ($D \approx 28,2$), mentre le posizioni dello scaffold e delle decorazioni erano più robuste. Ciò suggerisce che gli sforzi sperimentali dovrebbero dare priorità alla diversità dei linker.

Nei compiti sintetici (grid-worlds, regressione simbolica e proxy di ragionamento LLM), il surrogato PCE ha raggiunto una copertura di calibrazione del 95% tra 0,97 e 1,00 per i passi dominanti. Computazionalmente, il surrogato ha valutato 10.000 campioni di politica in millisecondi, offrendo un aumento di velocità di $10^3$-$10^4\times$ rispetto al riaddestramento esaustivo dell'ensemble.

Significato e Affermazioni

Il documento afferma che questo framework offre una capacità unica: decomposizione interpretabile dell'incertezza. Mentre altri metodi possono affermare che una previsione è incerta, questo approccio spiega perché attribuendo la varianza a specifiche componenti della ricompensa. Ciò permette agli scienziati di:

• Distinguere tra segnali robusti e artefatti statistici nei dati di addestramento.
• Identificare quali parti di uno spazio di progettazione (ad esempio, linker rispetto agli scaffold) richiedono ulteriore validazione sperimentale.
• Allocare le risorse sperimentali in modo efficiente mirando alle specifiche fonti di incertezza epistemica.

Gli autori sottolineano che il metodo è agnostico rispetto all'architettura della GFlowNet e all'obiettivo di addestramento. Evidenziano limitazioni, come la dipendenza da una struttura a basso rango nel panorama della ricompensa affinché la PCA sia efficace e la necessità di riadattare il surrogato per diverse traiettorie (sebbene non sia necessario riaddestrare l'ensemble). Il lavoro si posiziona come un ponte tra la quantificazione dell'incertezza nell'ingegneria computazionale e l'IA generativa moderna, fornendo una via scalabile per comprendere la "scatola nera" del processo decisionale sequenziale in condizioni di incertezza.