Riassunto Tecnico: TabPFN-CFM

Definizione del Problema

Determinare la struttura causale e prevedere gli esiti causali da dati osservazionali rimane una sfida significativa nel machine learning e nella statistica. I metodi esistenti affrontano tipicamente questi compiti in modo isolato: alcuni si concentrano sulle previsioni strutturali (apprendimento del grafo causale), mentre altri sulla previsione degli esiti (stima l'effetto degli interventi). Inoltre, molti approcci faticano a gestire variabili di confondimento non osservate, il che può portare a grafi non identificabili e stime degli esiti distorte. Analogamente, gli attuali modelli spesso mancano della capacità di supportare l'intero spettro delle query causali definite dalla Gerarchia Causale di Pearl: osservativa, interventista e controfattuale.

Metodologia

Gli autori introducono TabPFN-CFM, un modello di fondazione causale progettato per prevedere congiuntamente la struttura causale e gli esiti da dati osservazionali. L'approccio si basa sul framework Bayesian Prior-Over-Functions (PFN).

Framework Core

Il modello assume che il sistema sottostante segua un Modello Causale Strutturale (SCM) $\psi = \{U, V, F\}$, dove $U$ rappresenta le variabili non osservate, $V$ le variabili osservate e $F$ le equazioni strutturali. Il modello è addestrato per approssimare la distribuzione a posteriori $P(\psi | D_{obs}, G_{est})$, dove $D_{obs}$ è il dataset osservazionale e $G_{est}$ è la conoscenza a priori opzionale del grafo.

Generazione dei Dati e Addestramento

• Prior Sintetici: Il modello è addestrato su dataset sintetici generati da una distribuzione a priori su SCM. Questo prior include strutture di grafi casuali (DAG), nodi mancanti, diverse distribuzioni di rumore ed equazioni strutturali definite da reti neurali con pesi casuali e non linearità.
• Formulazione del Task: Ogni campione di addestramento consiste in un set di fitting ($D_{fit}$), una query di predizione ($D_{pred}$) e un campione controfattuale ($D_{causal}$). Il modello è addestrato per gestire simultaneamente tre tipi di query:
  1. Query Osservativa: Prevedere $P(y^* | x^*, T=t^*)$.
  2. Query Interventista: Prevedere $P(y^* | x^*, do(T=1-t^*))$.
  3. Query Controfattuale: Prevedere $P(y^* | x^*, do(T=1-t^*), y=y^*)$.
• Gestione del Grafo: Il modello utilizza Grafi Misti Diretti Aciclici (ADMG) per rappresentare il confondimento non osservato tramite archi bidiretti. Ciò consente al modello di prevedere strutture di confondimento invece di assumere che tutte le variabili siano osservate.

Architettura

TabPFN-CFM estende l'architettura Do-PFN con un esplicito condizionamento al grafo:

• Encoder: Codifica il set di fitting, la query di predizione e la struttura prior dell'ADMG (matrice di adiacenza $A$, matrice bidirezionale $C$ e matrice ancestrale $R$). La codifica del grafo fornisce l'accesso diretto alle relazioni di genitore, figlio, antenato e discendente.
• Transformer Backbone: Utilizza meccanismi di attenzione row-wise e column-wise. Fondamentalmente, il modulo di predizione presta attenzione alle embedding del grafo, permettendo alle predizioni di incorporare il prior del grafo.
• Output:
  • Predizione dell'Esito: Una testa produce i logit su bucket di esito discretizzati.
  • Predizione della Struttura: Un decoder genera predizioni elemento per elemento per la matrice di adiacenza diretta ($\hat{A}$), la matrice di correlazione bidirezionale ($\hat{C}$) e la matrice ancestrale ($\hat{R}$).
• Efficienza di Addestramento: Gli autori hanno implementato diverse ottimizzazioni architettoniche e di addestramento (es. ottimizzatore Muon, attivazione ReLU2, pre-layer normalization) che hanno migliorato la velocità di addestramento di circa 3-4 volte rispetto alle configurazioni baseline.

Funzione di Perdita

Il modello è ottimizzato utilizzando una funzione di perdita congiunta:
$$L = L_{pred} + \lambda \cdot L_{graph}$$
Dove $L_{pred}$ è la perdita di cross-entropy per la predizione dell'esito, e $L_{graph}$ è la perdita di binary cross-entropy per le predizioni strutturali (matrici di adiacenza, ancestrale e di confondimento). Le prove teoriche nell'appendice dimostrano che includere l'input del grafo come condizione può solo migliorare l'accuratezza a posteriori (o lasciarla invariata) e che minimizzare la log-probabilità è equivalente a minimizzare la divergenza KL rispetto alla vera distribuzione a posteriori.

Contributi Chiave

1. Modello di Fondazione Causale Unificato: TabPFN-CFM è il primo modello a prevedere congiuntamente sia la struttura causale che gli esiti attraverso tutti i tre livelli della Gerarchia Causale di Pearl all'interno di un singolo framework di modello di fondazione.
2. Gestione del Confondimento Non Osservato: Utilizzando gli ADMG, il modello predice esplicitamente le variabili di confondimento non osservate, affrontando un limite dei metodi esistenti di apprendimento della struttura causale tramite deep learning.
3. Inferenza Condizionata al Grafo: Il modello sfrutta le strutture di grafo note quando disponibili per migliorare l'accuratezza della predizione, rimanendo comunque robusto quando il grafo è sconosciuto.
4. Miglioramenti dell'Efficienza: La procedura di addestramento raffinata e le modifiche architettoniche risultano in una convergenza dell'addestramento quasi 4 volte più veloce rispetto alle precedenti baseline.

Risultati Sperimentali

Il modello è stato valutato su dataset sintetici (sia in-distribuzione che out-of-distribution) e dataset reali (Amazon Sales e Law School Admissions).

• Predizione dell'Esito:
  • Sui SEM Out-of-Distribution (OOD) sintetici, TabPFN-CFM ha ottenuto un Mean Squared Error (MSE) inferiore rispetto a meta-learner stabiliti (S-, T-, X-, DR-learners) e Do-PFN in contesti interventisti e controfattuali.
  • Nel dataset Law School Admissions, il modello ha migliorato significativamente l'accuratezza nei contesti controfattuali utilizzando l'informazione controfattuale, superando le baseline che non possono effettuare tali predizioni.
• Predizione della Struttura:
  • TabPFN-CFM ha superato forti baseline (FCI, GES, LiNGAM, PC, AVICI) nella predizione delle matrici di adiacenza, ancestrale e di confondimento.
  • Nello specifico, sui dati sintetici OOD, il modello ha raggiunto un AUROC di 0.889 per l'adiacenza e di 0.828 per le matrici di confondimento, significativamente superiore al secondo miglior performer (AVICI con rispettivamente 0.623 e 0.563).
• Condizionamento al Grafo: Fornire la vera struttura del grafo come input ha migliorato le prestazioni nei contesti interventisti e controfattuali, ma ha avuto un impatto trascurabile sulle predizioni osservative, in linea con la teoria causale.

Significato

Il paper sostiene che TabPFN-CFM rappresenti un passo avanti nell'unificare i compiti di inferenza causale. Addestrando un singolo modello su un diversificato prior di sistemi causali, l'approccio dimostra che i modelli di fondazione possono generalizzare a problemi causali del mondo reale senza overfitting a specifiche distribuzioni a priori. La capacità di gestire il confondimento non osservato e di supportare query controfattuali suggerisce una via verso un ragionamento causale più robusto nel machine learning. Gli autori osservano che, sebbene il lavoro attuale si concentri su strutture acicliche, estensioni future potrebbero affrontare contesti temporali e causali ciclici.