Causality-Encoded Diffusion Models for Interventional Sampling and Edge Inference

1. Il Problema

I modelli di diffusione (Diffusion Models) sono diventati strumenti statistici potenti per la stima flessibile di distribuzioni complesse, eccellendo nella generazione di dati e nell'apprendimento non parametrico. Tuttavia, presentano due limitazioni fondamentali nel contesto dell'inferenza causale:

1. Agnosticismo Causale: I modelli di diffusione standard apprendono la legge congiunta dei dati senza codificare le asimmetrie direzionali necessarie per l'interpretazione causale. Di conseguenza, non supportano nativamente query interventionali (es. $do(X=x)$) né l'analisi di strutture causali.
2. Mancanza di Inferenza Statistica Rigorosa: Esistono metodi per l'inferenza sui bordi dei grafi aciclici diretti (DAG), ma spesso si basano su assunzioni restrittive (es. linearità, rumore additivo) o faticano a scalare in spazi ad alta dimensione senza perdere accuratezza.

L'obiettivo del lavoro è colmare questo divario sviluppando un framework che integri la struttura causale nota (un DAG) direttamente nel processo di generazione dei dati, permettendo sia il campionamento interazionale che l'inferenza statistica sui bordi del grafo.

2. Metodologia: CEDM e CEDMI

Gli autori propongono due componenti principali: CEDM (Causality-Encoded Diffusion Models) per la generazione e CEDMI (CEDM-based Inference) per il test statistico.

A. CEDM: Modelli di Diffusione Codificati per la Causalità

Il framework CEDM modifica il processo di diffusione standard per rispettare la fattorizzazione di un DAG dato $G = (V, E)$.

• Processo Inverso Condizionale: Invece di apprendere un unico score congiunto $\nabla \log p_t(x)$, il modello apprende una serie di score condizionali locali: $\nabla_{x'_j} \log p_t(x'_j | x_{pa(j)})$, dove $pa(j)$ sono i genitori del nodo $j$ nel DAG.
• Campionamento Ordinato Topologicamente: Durante la fase di generazione (reverse diffusion), i nodi vengono generati in ordine topologico. Per ogni nodo $j$, il processo stocastico inverso evolve condizionatamente solo ai valori già generati dei suoi genitori.
• Campionamento Interventionale ($do$-sampling): Una volta addestrato, il modello può simulare distribuzioni interventionali. Per eseguire un'intervento $do(X_S = x^*_S)$, i nodi in $S$ vengono "congelati" ai valori fissati $x^*_S$ e il processo di diffusione inversa viene eseguito solo sui nodi discendenti, propagando l'effetto attraverso il grafo secondo la fattorizzazione causale.
• Apprendimento: Gli score condizionali sono approssimati tramite reti neurali (ReLU) addestrate minimizzando un rischio di score matching denoising, adattato alla struttura condizionale del DAG.

B. CEDMI: Inferenza sui Bordi

Sfruttando la capacità di campionamento interazionale, gli autori sviluppano un test di ipotesi per verificare la presenza di bordi diretti specifici.

• Ipotesi: Testare se un arco candidato $(k, j)$ appartiene al vero DAG. L'ipotesi nulla $H_0$ assume che l'arco sia assente.
• Procedura di Resampling:
  1. Si addestra un CEDM sul grafo "nulla" (il grafo di lavoro con l'arco candidato $(k, j)$ rimosso).
  2. Si generano campioni nulli simulando $X_j$ condizionatamente ai genitori rimasti (escluso $k$) e ai valori osservati di $k$ e degli altri genitori, utilizzando l'intervento $do$.
  3. Si calcola una statistica di dipendenza condizionale (un coefficiente multivariato basato su MCODEC) tra $X_j$ e $X_k$ dati gli altri genitori.
  4. Si confronta la statistica osservata con la distribuzione nulla ottenuta dai campioni simulati per calcolare un p-value.
• Coefficiente di Dipendenza (MCODEC): Viene introdotto un nuovo coefficiente di dipendenza condizionale multivariata, estensione del coefficiente di Azadkia-Chatterjee, che è senza parametri (tuning-free) e basato sui ranghi, garantendo stabilità nel calcolo Monte Carlo.

3. Contributi Chiave

1. Framework Generativo Causale: Introduzione di CEDM, un framework di generazione basato su score che incorpora esplicitamente la struttura DAG, permettendo il campionamento sia osservazionale che interazionale.
2. Garanzie Statistiche di Convergenza:
   • Viene dimostrato che l'errore di recupero della distribuzione (in distanza di variazione totale) è governato dalla dimensione locale massima (dimensione di un nodo più i suoi genitori, $q = \max_j (d_j + |pa(j)|)$) e non dalla dimensione ambientale totale $d$.
   • Questo fornisce una riduzione dimensionale strutturale: per grafi sparsi, la complessità statistica è significativamente inferiore rispetto ai metodi generativi standard.
   • Viene provato il controllo asintotico del tasso di errore di Tipo I per il test CEDMI.
3. Validazione Empirica:
   • Recupero Distribuzione: CEDM supera i baseline (diffusione agnostica e VACA) nel recupero di distribuzioni interventionali, specialmente in grafi sparsi.
   • Potenza e Calibrazione: CEDMI mostra un controllo dell'errore di Tipo I vicino al livello nominale e una potenza superiore rispetto a metodi esistenti (come CCIT, NNLSCIT, CDCIT) in scenari non lineari.
   • Applicazione Reale: Applicazione ai dati di citometria a flusso (rete di segnalazione proteica di Sachs), dove il metodo ha identificato interazioni biologiche plausibili (es. Erk $\to$ Akt) che erano assenti o controverse nelle reti di riferimento.

4. Risultati Principali

• Simulazioni: In scenari sintetici (catene, hub, grafi casuali e rete di Sachs), CEDM ha mostrato un errore di recupero (MMD quadrato) inferiore rispetto ai metodi concorrenti. La superiorità è più marcata quando la struttura del grafo è sparsa, confermando la teoria sulla dipendenza dalla dimensione locale.
• Inferenza: Il test CEDMI mantiene tassi di rifiuto corretti sotto l'ipotesi nulla (calibrazione) e raggiunge potenze vicine a 1 per effetti moderati con dimensioni campionarie ragionevoli ($n \ge 500$), superando metodi basati su kernel o classificatori che spesso soffrono di errori di calibrazione o bassa potenza in contesti non lineari.
• Dati Reali: Nell'analisi della rete di segnalazione delle cellule T, CEDMI ha fornito evidenze statistiche contro alcuni bordi del "consensus network" e a favore di un'interazione Erk-Akt, supportata da letteratura biologica ma non presente nel grafo originale, dimostrando l'utilità pratica per la validazione di link causali disputati.

5. Significato e Impatto

Questo lavoro rappresenta un passo significativo verso l'integrazione tra l'apprendimento profondo generativo (in particolare i modelli di diffusione) e l'inferenza causale formale.

• Riduzione della Complessità: Dimostra che l'incorporazione di vincoli strutturali (DAG) non è solo un modo per imporre vincoli, ma riduce attivamente la complessità statistica del problema di apprendimento, permettendo di scalare a dimensioni più elevate.
• Inferenza Non Parametrica: Offre un metodo potente per testare l'indipendenza condizionale e la causalità in setting non lineari e non parametrici, superando le limitazioni dei modelli lineari o gaussiani tradizionali.
• Strumento per la Scienza: Fornisce uno strumento pratico per gli scienziati per validare o falsificare ipotesi causali in reti complesse (come quelle biologiche o finanziarie) utilizzando dati osservazionali, sfruttando la potenza dei moderni modelli generativi per simulare scenari controfattuali.

In sintesi, il paper propone un ponte teorico e pratico tra la generazione di dati ad alta fedeltà e l'inferenza causale rigorosa, trasformando i modelli di diffusione da semplici generatori in simulatori causali capaci di rispondere a domande "what-if" e di testare la struttura stessa delle relazioni causali.