Causally-Guided Diffusion for Stable Feature Selection

Probleemstelling

Bestaande methoden voor kenmerkselectie (feature selection) zijn fundamenteel gericht op het optimaliseren van voorspellende prestaties binnen één enkele data-distributie. Dit leidt vaak tot het selecteren van spurious features (schijnbare correlaties) die wel sterk correleren met het doel binnen de trainingsdata, maar falen wanneer het model wordt ingezet in nieuwe omgevingen met distributieveranderingen (distribution shifts).

De huidige uitdagingen zijn:

1. Brittleness: Methodes die alleen vertrouwen op in-distribution prestaties, bieden geen robuustheidsgaranties bij verschuivingen in de data.
2. Combinatorische complexiteit: Het vinden van de beste subset van kenmerken is een combinatorisch probleem (exponentieel groot zoekruimte), wat standaard continue optimalisatietechnieken moeilijk toepasbaar maakt.
3. Niet-convexe optimalisatie: Doelstellingen gebaseerd op stabiliteit leiden vaak tot niet-convexe landschappen waar gradiënt-based methoden vastlopen in lokale minima.
4. Onzekerheid: Bestaande methoden kiezen vaak voor één deterministische oplossing, terwijl er onder distributieveranderingen meerdere plausibele subsets kunnen bestaan die vergelijkbare stabiliteit bieden.

Methodologie: CGDFS

Het auteurs stellen Causally-Guided Diffusion for Stable Feature Selection (CGDFS) voor. Dit is een raamwerk dat kenmerkselectie formuleert als benaderende posterior-inferentie over subsets van kenmerken, in plaats van het zoeken naar één enkel optimum. Het raamwerk combineert drie kerncomponenten:

1. Representatie en Diffusion Prior (3.2)

• Continue Maskers: In plaats van te werken met discrete binaire vectoren ($z \in \{0,1\}^p$), worden subsets gemodelleerd als continue selectiemaskers ($s \in [0,1]^p$). Dit maakt differentiatie en gradiënt-based inferentie mogelijk.
• Geleerde Prior: Een diffusiemodel (score-based diffusion) wordt getraind om een prior te leren over plausibele kenmerkmaskers. Dit model wordt getraind op een "mask pool" die bestaat uit kandidaat-subsets gegenereerd door diverse heuristieken (zoals LASSO, Mutual Information, Random Forests).
• Doel: De prior vangt structurele afhankelijkheden tussen kenmerken (zoals complementariteit en redundantie) en stuurt de inferentie naar zinvolle gebieden in de combinatorische zoekruimte zonder expliciete causale grafen te hoeven ontdekken.

2. Stabiliteitsgeoriënteerde Likelihood (3.3)

• Multi-omgeving Evaluatie: Het model wordt geëvalueerd over $E$ verschillende omgevingen (datasets met covariaat-verschuiving).
• Doelfunctie: Voor een gegeven masker $s$ wordt een stabiliteits-objectief $U(s)$ gedefinieerd dat de gemiddelde voorspellende fout ($\bar{\ell}$) en de variantie van de fout over de omgevingen ($\text{Var}_e$) combineert:
  $$U(s) = \bar{\ell}(s) + \lambda_{var} \cdot \text{Var}_e(\ell_e(s))$$
• Likelihood: De likelihood wordt impliciet gedefinieerd als $p(D|s) \propto \exp(-U(s))$. Dit straalt causaliteit uit door subsets te belonen die consistent presteren over verschillende omgevingen, zonder expliciete causale interventie-data te vereisen.

3. Geleide Posterior Sampling (3.4 & 3.5)

• Annealed Langevin Dynamics: Om de posterior $p(s|D) \propto p(D|s)p(s)$ te benaderen, wordt er gebruikgemaakt van geleide Langevin-dynamiek. Dit combineert de gradiënten van de geleerde diffusion prior met de gradiënten van de stabiliteits-likelihood.
• Causale Gids (ICP): Om de inferentie verder te sturen, worden Invariant Causal Prediction (ICP) scores gebruikt als "zachte" (soft) gids. Kenmerken met hoge invariantie-scores krijgen een hogere kans om in de mask-pool te worden opgenomen tijdens het trainen van de prior. Dit is geen harde constraint, maar een inductieve bias die de zoekruimte efficiënter maakt.
• Aggregatie: De methode genereert meerdere steekproeven (chains). De definitieve selectie wordt bepaald door de posterior inclusion frequencies te berekenen (hoe vaak een kenmerk in de top-$k$ voorkomt over alle steekproeven). Dit levert een robuuste, interpreteerbare selectie op die onzekerheid kwantificeert.

Kernbijdragen

1. Probabilistische Formulering: Het formuleren van kenmerkselectie onder distributieverandering als posterior-inferentie, waarbij stabiliteit en prestatie tegelijkertijd worden geoptimaliseerd.
2. Diffusion-Based Prior: De introductie van een diffusiemodel als een geleerde prior over continue selectiemaskers, wat schaalbare exploratie van de combinatorische ruimte mogelijk maakt zonder discrete zoekopdrachten.
3. Zachte Causale Gids: Het integreren van causale invariantie (via ICP) als zachte bias in de prior, in plaats van als harde constraint, waardoor het systeem robuust blijft tegen imperfecte causale aannames.
4. Empirische Validatie: Uitgebreide experimenten op synthetische en real-world tabulaire datasets (classificatie en regressie) die aantonen dat CGDFS superieure out-of-distribution prestaties levert.

Resultaten

De auteurs hebben CGDFS getest op 12 datasets (o.a. UCI en OpenML) met gecontroleerde distributieveranderingen. De resultaten tonen aan:

• Superieure Prestaties: CGDFS behaalde consistent de hoogste F1-macro scores (classificatie) en de laagste MSE (regressie) vergeleken met baselines zoals LASSO, Mutual Information, Random Forest, Stability Selection en Invariant Causal Prediction (ICP).
• Stabiliteit: CGDFS vertoonde de laagste variantie in prestaties over verschillende test-omgevingen, wat aantoont dat de geselecteerde kenmerken robuust zijn tegen distributieveranderingen.
• Ablatie Studies:
  • Zonder de diffusion prior daalde de prestatie aanzienlijk, wat aantoont dat de prior essentieel is voor het navigeren in de zoekruimte.
  • Vervanging van posterior sampling door directe gradiënt-optimalisatie (Grad-Stab) leidde tot minder stabiele resultaten en hogere variantie.
  • ICP als zachte gids verbeterde de stabiliteit, maar was niet strikt noodzakelijk; het werkte het beste in combinatie met het posterior-ramek.
• Onzekerheidskwantificatie: CGDFS levert een verdeling van kenmerken op, waarbij de "posterior inclusion frequency" een interpreteerbare maatstaf biedt voor de betrouwbaarheid van een kenmerk.

Significantie

Dit werk is significant omdat het een brug slaat tussen causal learning, diffusion models en feature selection.

• Het lost het probleem op van "brittle" kenmerkselectie door stabiliteit over omgevingen als primair criterium te nemen.
• Het vermijdt de valkuilen van deterministische optimalisatie door onzekerheid expliciet te modelleren via posterior sampling.
• Het biedt een schaalbare oplossing voor het combinatorische probleem van kenmerkselectie in hoge dimensies, zonder de complexiteit van volledige causale grafiek-discovery.
• De methode is direct toepasbaar in kritieke domeinen zoals gezondheidszorg en financiën, waar modellen robuust moeten zijn tegen veranderingen in de realiteit.

Kortom, CGDFS bewijst dat het benaderen van kenmerkselectie als een probabilistisch inferentieprobleem, geleid door zowel data-gedreven patronen (diffusion) als stabiliteitsprincipes, leidt tot aanzienlijk robuustere en generaliseerbaardere machine learning systemen.