Causally-Guided Diffusion for Stable Feature Selection

1. Il Problema: Selezione delle Caratteristiche e Shift di Distribuzione

La selezione delle caratteristiche (feature selection) è fondamentale per costruire sistemi di intelligenza artificiale robusti e interpretabili. Tuttavia, la maggior parte dei metodi esistenti (come LASSO, Mutual Information, o metodi basati su alberi decisionali) ottimizza le prestazioni predittive sotto una singola distribuzione di dati.

• Limitazione principale: Questi metodi tendono a selezionare caratteristiche spurie che hanno una forte correlazione con il target solo all'interno della distribuzione di addestramento specifica.
• Conseguenza: Quando il modello viene distribuito in ambienti reali dinamici con shift di distribuzione (cambiamenti nelle condizioni dei dati), le prestazioni crollano perché le correlazioni spurie non si generalizzano.
• Sfida: La selezione delle caratteristiche è un problema combinatorio (lo spazio delle possibili sottoinsiemi è $2^p$) e i criteri di stabilità basati sulla causalità inducono paesaggi di ottimizzazione non convessi, rendendo difficile trovare ottimi globali con metodi greedy o basati su gradienti tradizionali.

2. Metodologia: CGDFS

Il paper propone CGDFS (Causally-Guided Diffusion for Stable Feature Selection), un quadro concettuale che riformula la selezione delle caratteristiche come un problema di inferenza approssimata della distribuzione a posteriori sui sottoinsiemi di caratteristiche, piuttosto che come una ricerca di un singolo sottoinsieme ottimale.

Il framework si basa su tre componenti chiave:

A. Rappresentazione e Prior Diffusivo

• Relassamento Continuo: Invece di ottimizzare direttamente su vettori binari discreti $\{0, 1\}^p$, CGDFS utilizza una maschera di selezione continua $z \in [0, 1]^p$.
• Prior Appreso: Viene addestrato un modello di diffusione (diffusion model) per apprendere una distribuzione a priori sui sottoinsiemi plausibili di caratteristiche.
  • Il modello viene addestrato su un "pool" di maschere candidate generate da vari metodi euristici (selezione casuale, LASSO, Random Forest, ecc.).
  • Questo prior cattura le dipendenze strutturali (complementarità e ridondanza) tra le caratteristiche, guidando l'esplorazione dello spazio combinatorio verso regioni significative senza bisogno di una scoperta causale esplicita.

B. Likelihood Basata sulla Stabilità

• Viene definita una funzione di verosimiglianza (likelihood) che valuta i sottoinsiemi di caratteristiche non solo in base all'accuratezza predittiva, ma anche alla stabilità attraverso diversi ambienti.
• L'obiettivo di stabilità $U(s)$ combina la perdita media tra gli ambienti e la varianza della perdita:
  $$U(s) = \bar{\ell}(s) + \lambda_{var} \cdot \text{Var}_e(\ell_e(s))$$
  Dove $\bar{\ell}$ è la perdita media e $\text{Var}_e$ è la varianza tra gli ambienti.
• Questo approccio incoraggia la selezione di caratteristiche il cui comportamento predittivo rimane invariato (causalmente invariante) al variare della distribuzione dei dati.

C. Campionamento Guidato e Inferenza

• Langevin Dynamics: Per campionare dalla distribuzione a posteriori $p(s|D) \propto p(D|s)p(s)$, il metodo utilizza la dinamica di Langevin riscaldata (annealed Langevin dynamics).
• Il gradiente per l'aggiornamento della maschera è la somma di due termini:
  1. Il gradiente del prior (dal modello di diffusione), che garantisce la plausibilità strutturale.
  2. Il gradiente della likelihood (dall'obiettivo di stabilità), che garantisce prestazioni robuste.
• Guida Causale (Soft): L'Invariant Causal Prediction (ICP) non viene usato come vincolo rigido, ma come segnale di guida "soft" per biasare la costruzione del pool di maschere, accelerando l'esplorazione senza compromettere la validità dell'inferenza se le assunzioni causali non sono perfette.
• Discretizzazione: I campioni continui vengono aggregati calcolando le frequenze di inclusione a posteriori per ogni caratteristica. Le caratteristiche finali sono quelle con le frequenze più alte, fornendo una selezione stabile e interpretabile.

3. Contributi Chiave

1. Formulazione Probabilistica: Trasformazione della selezione delle caratteristiche sotto shift di distribuzione in un problema di inferenza approssimata della distribuzione a posteriori, gestendo l'incertezza intrinseca invece di cercare un singolo ottimo deterministico.
2. Prior basato su Diffusione: Introduzione di un prior appreso tramite modelli di diffusione per esplorare efficientemente lo spazio combinatorio delle sottoinsiemi, catturando dipendenze complesse tra le feature.
3. Robustezza e Stabilità: Dimostrazione che l'uso di un obiettivo di stabilità multi-ambiente combinato con il campionamento guidato porta a selezioni di feature più robuste rispetto all'ottimizzazione puntuale.
4. Integrazione Causale Soft: Uso dell'invarianza causale come guida soft all'interno di un framework probabilistico, evitando i limiti delle assunzioni causali rigide.

4. Risultati Sperimentali

Gli autori hanno valutato CGDFS su 12 dataset reali (classificazione e regressione) con shift di distribuzione controllati, confrontandolo con baselines come LASSO, Mutual Information, Random Forest, Selezione di Stabilità, e metodi basati su causalità (ICP).

• Prestazioni Predittive (Out-of-Distribution): CGDFS ha ottenuto sistematicamente le migliori prestazioni (F1-macro più alto per la classificazione, MSE più basso per la regressione) su tutti i dataset testati, superando sia i metodi basati sulla correlazione che quelli basati sull'ottimizzazione diretta della stabilità.
• Stabilità: CGDFS ha mostrato una varianza delle prestazioni tra gli ambienti significativamente inferiore rispetto alle baselines, indicando una maggiore robustezza agli shift di distribuzione.
• Ablation Study:
  • Rimuovere il prior di diffusione degrada drasticamente le prestazioni, confermando che il prior strutturato è essenziale per guidare l'esplorazione.
  • Sostituire il campionamento con l'ottimizzazione diretta (Grad-Stab) porta a risultati meno stabili e più sensibili all'inizializzazione.
  • L'uso della guida causale (ICP) migliora la stabilità ma non è essenziale; il framework funziona bene anche senza di essa, ma ne beneficia quando integrata come segnale soft.
• Interpretabilità: Il metodo fornisce frequenze di inclusione che quantificano l'incertezza, offrendo una selezione più affidabile e riproducibile rispetto ai metodi deterministici.

5. Significato e Impatto

Il lavoro CGDFS rappresenta un passo avanti significativo nell'intersezione tra apprendimento causale, modelli generativi (diffusione) e selezione delle caratteristiche.

• Superamento dei limiti attuali: Risolve il problema della "brittleness" (fragilità) dei modelli di feature selection tradizionali quando affrontano dati non i.i.d. (non indipendenti e identicamente distribuiti).
• Nuovo Paradigma: Sposta il focus dalla ricerca di un singolo sottoinsieme "migliore" all'esplorazione di una distribuzione di sottoinsiemi plausibili, gestendo esplicitamente l'incertezza.
• Applicabilità: Offre un approccio scalabile e pratico per applicazioni ad alto rischio (sanità, finanza, equità algoritmica) dove la stabilità e la generalizzazione sono critiche quanto l'accuratezza.

In sintesi, CGDFS dimostra che l'inferenza bayesiana approssimata, guidata da modelli di diffusione e principi di invarianza causale, è una strategia superiore per selezionare caratteristiche robuste in ambienti dinamici.