[2603.20930] Causally-Guided Diffusion for Stable Feature Selection

Computer Science > Machine Learning arXiv:2603.20930 (cs) [Submitted on 21 Mar 2026] Title:Causally-Guided Diffusion for Stable Feature Selection Authors:Arun Vignesh Malarkkan, Xinyuan Wang, Kunpeng Liu, Denghui Zhang, Yanjie Fu View a PDF of the paper titled Causally-Guided Diffusion for Stable Feature Selection, by Arun Vignesh Malarkkan and 4 other authors View PDF HTML (experimental) Abstract: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...