Bayesian Invariance Modeling of Multi-Environment Data

Problematica

Il documento affronta la sfida della selezione di variabili invarianti in dati multi-ambiente. In molti contesti statistici e di machine learning, i dati vengono raccolti attraverso ambienti distinti (ad esempio, diversi ospedali, condizioni sperimentali o periodi temporali). Un approccio standard alla selezione delle variabili identifica i predittori che spiegano un esito sotto la specifica distribuzione dei dati osservati. Tuttavia, ciò spesso confonde le relazioni causali stabili con associazioni specifiche dell'ambiente che non generalizzano a nuovi contesti.

L'obiettivo è identificare un sottoinsieme di caratteristiche la cui relazione condizionale con l'esito rimanga stabile (invariante) attraverso tutti gli ambienti, anche se le distribuzioni marginali delle caratteristiche stesse variano. Sebbene metodi precedenti abbiano affrontato questo problema tramite test di ipotesi (ad esempio, Invariant Causal Prediction, ICP) o ottimizzazione regolarizzata, tali approcci spesso faticano con la scalabilità computazionale in dimensioni elevate o mancano di un quadro probabilistico unificato per la quantificazione dell'incertezza.

Metodologia: Bayesian Invariant Prediction (BIP)

Gli autori propongono la Bayesian Invariant Prediction (BIP), un modello probabilistico che inquadra la selezione di caratteristiche invarianti come un problema di inferenza bayesiana.

1. Struttura del Modello e Assunzioni
BIP si basa su due assunzioni fondamentali riguardanti il processo generativo dei dati:

• Eterogeneità delle Caratteristiche: Le distribuzioni delle caratteristiche variano tra gli ambienti.
• Invarianza dell'Esito: L'esito $y$ dipende da un sottoinsieme specifico di caratteristiche (l'insieme invariante) nello stesso modo in tutti gli ambienti.

Il modello tratta l'insieme di caratteristiche invarianti come una variabile latente, codificata da un vettore selettore binario $z \in \{0, 1\}^p$. La distribuzione congiunta del selettore latente e dei dati osservati $D$ si fattorizza come:
$$p(z, D) = p(z) \prod_{e=1}^E \prod_{i=1}^{n_e} p_e(x_{z, ei}) \, g(y_{ei} | x_{z, ei}) \, p_e(x_{-z, ei} | x_{z, ei}, y_{ei})$$
Fondamentalmente, la distribuzione condizionale dell'esito date le caratteristiche selezionate, $g(y|x_z)$, è aggregata (pooled) attraverso tutti gli ambienti e condivisa, mentre le distribuzioni delle caratteristiche $p_e$ sono specifiche per ogni ambiente.

2. Formulazione della Posteriore
Il modello deriva una distribuzione a posteriori semplificata per il selettore invariante $z$:
$$p(z | D) \propto p(z) \prod_{e=1}^E \prod_{i=1}^{n_e} \frac{g(y_{ei} | x_{z, ei})}{p_e(y_{ei} | x_{z, ei})}$$
Questa espressione rivela che la probabilità a posteriori di un insieme di caratteristiche è proporzionale al prior moltiplicato per il prodotto dei rapporti di verosimiglianza tra la condizionale aggregata e le condizionali locali specifiche dell'ambiente. Se un insieme $z$ è realmente invariante, la condizionale aggregata coincide con le condizionali locali, massimizzando il rapporto. Se $z$ non è invariante, il disallineamento riduce la probabilità a posteriori.

3. Algoritmi di Inferenza

• Inferenza Esatta (BIP): Per contesti a bassa dimensionalità, gli autori enumerano tutti i $2^p$ insiemi di caratteristiche candidati per calcolare la posteriore esatta. Ciò comporta la stima delle condizionali locali e aggregate (ad esempio, tramite massima verosimiglianza in una famiglia Gaussiana lineare) e il calcolo dei rapporti di verosimiglianza.
• Inferenza Variazionale (BIP-VI): Per gestire dati ad alta dimensionalità dove l'enumerazione esatta è intrattabile, gli autori sviluppano BIP-VI. Questo algoritmo approssima la posteriore utilizzando una famiglia variazionale mean-field $q_\phi(z) = \prod \text{Bernoulli}(z_j | \text{sigmoid}(\phi_j))$. Ottimizza l'Evidence Lower Bound (ELBO) tramite gradiente stocastico e stimatori di gradiente non distorti (U2G) per variabili discrete. A differenza dei metodi esatti, BIP-VI non richiede una ricerca esaustiva o lo pre-screening delle caratteristiche.

Contributi Teorici

Il documento stabilisce garanzie teoriche rigorose per il framework BIP:

• Consistenza della Posteriore: Sotto le assunzioni di invarianza univoca e stima coerente delle condizionali, la posteriore di BIP si concentra sul vero selettore di caratteristiche invarianti $z^*$ quando il numero di ambienti ($E$) e la dimensione del campione per ambiente ($n$) tendono all'infinito.
• Tassi di Contrazione: Gli autori dimostrano che la posteriore si contrae esponenzialmente rispetto a $n$ ed $E$. Il tasso dipende da due fattori:
  1. Conoscenza a Priori: Prior informativi che penalizzano i set non invarianti accelerano la contrazione.
  2. Eterogeneità degli Ambienti: Una maggiore eterogeneità tra gli ambienti (misurata dalla discrepanza tra le condizionali locali e quelle aggregate) porta a una contrazione più rapida della posteriore.
• Campioni Finiti e Relax: La teoria si estende a campioni finiti e assunzioni rilassate (ad esempio, insiemi invarianti non univoci o misspecification del modello). In questi casi, la posteriore si concentra sui selettori "più approssimativamente invarianti" coerenti con il modello e il prior.

Risultati Sperimentali

Gli autori valutano BIP e BIP-VI su dataset sintetici e reali, confrontandoli con ICP, Hidden-ICP ed EILLS.

1. Studi Sintetici

• Bassa Dimensionalità ($p=10$): BIP e BIP-VI recuperano il vero insieme invariante con alta precisione, superando o eguagliando i metodi esistenti all'aumentare del numero di ambienti. La massa della posteriore si concentra sul vero insieme come previsto dalla teoria.
• Alta Dimensionalità ($p=450$): I metodi esistenti (ICP, EILLS) richiedono uno pre-screening (ad esempio, Lasso) per ridurre lo spazio delle caratteristiche, il che spesso porta a prestazioni scarse se lo screening manca le vere caratteristiche invarianti. Al contrario, BIP-VI opera sull'intero spazio delle caratteristiche senza pre-screening. Esso raggiunge tassi di scoperta esatta e scalabilità significativamente superiori rispetto alle versioni con screening di altri metodi.

2. Applicazione nel Mondo Reale: Perturbazione Genica
Il metodo viene applicato a un dataset di perturbazione genica in lievito ($p=6.169$ geni) per identificare i regolatori di un gene target.

• Prestazioni: BIP-VI raggiunge la precisione più elevata tra i metodi di predizione invariante, superando ICP-screen ed EILLS-screen. Sebbene il suo recall sia moderato (a causa della natura conservativa della posteriore e del numero limitato di ambienti), identifica un numero maggiore di predizioni corrette rispetto a ICP-screen (371 contro ~67 in un confronto specifico).
• Validazione: Le predizioni fatte da BIP-VI si allineano strettamente con i principali risultati della letteratura precedente (Meinshausen et al., 2016) e sono validate rispetto a esperimenti di knockout tenuti fuori dal modello (held-out).

Significato e Rivendicazioni

Il documento sostiene che BIP offre una prospettiva probabilistica sulla predizione invariante che colma il divario tra test di ipotesi e ottimizzazione. Il suo significato primario risiede in:

1. Quantificazione dell'Incertezza: A differenza dei metodi basati su stime puntuali, BIP fornisce una distribuzione a posteriori completa sugli insiemi invarianti, permettendo una quantificazione rigorosa dell'incertezza.
2. Scalabilità: BIP-VI supera la complessità esponenziale della predizione invariante esatta, consentendo l'analisi di dati ad alta dimensionalità senza dipendere da passaggi di pre-screening potenzialmente perdenti.
3. Approfondimento Teorico: Il lavoro fornisce prove formali di consistenza e tassi di contrazione, collegando esplicitamente la velocità di apprendimento all'eterogeneità degli ambienti.

Gli autori segnalano alcune limitazioni, tra cui la dipendenza da stime fisse per i modelli condizionali (che può introdurre bias in piccoli campioni) e l'uso dell'inferenza variazionale mean-field, che potrebbe avere difficoltà con posteri multi-modali in contesti con pochissimi ambienti. Si suggerisce come lavoro futuro l'integrazione della stima del modello condizionale direttamente nel framework probabilistico e l'esplorazione di famiglie variazionali più espressive.