# Causally-Guided Diffusion for Stable Feature Selection

arXiv: https://arxiv.org/abs/2603.20930 · HTML: https://mlsanity.com/paper/2603.20930/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 16 of its 39 references that are in the map.
- First posted on arXiv 21 Mar 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 16 of 16)

- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Invariant Causal Prediction for Nonlinear Models](https://mlsanity.com/paper/1706.08576/) (2017)
- [Domain Adaptation by Using Causal Inference to Predict Invariant Conditional Distributions](https://mlsanity.com/paper/1707.06422/) (2017)
- [Invariance, Causality and Robustness](https://mlsanity.com/paper/1812.08233/) (2018)
- [Solving Inverse Problems with Score-Based Generative Priors learned from Noisy Data](https://mlsanity.com/paper/2305.01166/) (2023)
- [Feature Selection as Deep Sequential Generative Learning](https://mlsanity.com/paper/2403.03838/) (2024)
- [A Survey on Data-Centric AI: Tabular Learning from Reinforcement Learning and Generative AI Perspective](https://mlsanity.com/paper/2502.08828/) (2025)
- [Rethinking Spatio-Temporal Anomaly Detection: A Vision for Causality-Driven Cybersecurity](https://mlsanity.com/paper/2507.08177/) (2025)
- [Incremental Causal Graph Learning for Online Cyberattack Detection in Cyber-Physical Infrastructures](https://mlsanity.com/paper/2507.14387/) (2025)
- [DELTA: Variational Disentangled Learning for Privacy-Preserving Data Reprogramming](https://mlsanity.com/paper/2509.00693/) (2025)
- [(Non-) asymptotic properties of Stochastic Gradient Langevin Dynamics](https://mlsanity.com/paper/1501.00438/) (2015)
- [Causally-Guided Automated Feature Engineering with Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2602.16435/) (2026)

## Built on by (most cited first, 1 of 1)

- [Flat-Consensus Diffusion for Robust Data Reshaping under Noisy Evaluator](https://mlsanity.com/paper/2609.32696/) (2026)

## Abstract (arXiv)

Feature selection is fundamental to robust data-centric AI, but most existing methods optimize predictive performance under a single data distribution. This often selects spurious features that fail under distribution shifts. Motivated by principles from causal invariance, we study feature selection from a stability perspective and introduce Causally-Guided Diffusion for Stable Feature Selection (CGDFS). In CGDFS, we formalized feature selection as approximate posterior inference over feature subsets, whose posterior mass favors low prediction error and low cross-environment variance. Our framework combines three key insights: First, we formulate feature selection as stability-aware posterior sampling. Here, causal invariance serves as a soft inductive bias rather than explicit causal discovery. Second, we train a diffusion model as a learned prior over plausible continuous selection masks, combined with a stability-aware likelihood that rewards invariance across environments. This diffusion prior captures structural dependencies among features and enables scalable exploration of the combinatorially large selection space. Third, we perform guided annealed Langevin sampling that combines the diffusion prior with the stability objective, which yields a tractable, uncertainty-aware posterior inference that avoids discrete optimization and produces robust feature selections. We evaluate CGDFS on open-source real-world datasets exhibiting distribution shifts. Across both classification and regression tasks, CGDFS consistently selects more stable and transferable feature subsets, which leads to improved out-of-distribution performance and greater selection robustness compared to sparsity-based, tree-based, and stability-selection baselines.