안정적인 특징 선택을 위한 인과 관계 기반 확산 모델

Causally-Guided Diffusion for Stable Feature Selection

특징 선택은 견고한 데이터 중심 AI의 기본이며, 대부분의 기존 방법은 단일 데이터 분포 하에서 예측 성능을 최적화합니다. 이는 종종 분포 변화 시 실패하는 가짜 특징을 선택하게 만듭니다. 인과적 불변성의 원리에 기반하여, 우리는 특징 선택을 안정성 관점에서 연구하고, 안정적인 특징 선택을 위한 인과 관계 기반 확산 모델(CGDFS)을 소개합니다. CGDFS에서, 우리는 특징 선택을 특징 집합에 대한 근사 사후 추론으로 공식화하며, 이 사후 확률은 낮은 예측 오류와 낮은 교차 환경 분산을 선호합니다. 우리의 프레임워크는 세 가지 주요 아이디어를 결합합니다. 첫째, 우리는 특징 선택을 안정성을 고려한 사후 샘플링으로 공식화합니다. 여기서, 인과적 불변성은 명시적인 인과 관계 발견보다는 부드러운 귀납적 편향으로 작용합니다. 둘째, 우리는 학습된 prior를 가진 확산 모델을 훈련하여, 안정성을 고려한 likelihood와 결합하여 사용합니다. 이 확산 prior는 특징 간의 구조적 의존성을 포착하고, 조합적으로 매우 큰 선택 공간을 확장 가능하게 탐색할 수 있도록 합니다. 셋째, 우리는 확산 prior와 안정성 목표를 결합한 가이드된 어닐링 Langevin 샘플링을 수행하여, 추상화된 불확실성을 고려한 사후 추론을 가능하게 하며, 이 방식은 이산 최적화를 피하고 견고한 특징 선택을 제공합니다. 우리는 CGDFS를 분포 변화를 보이는 공개된 실제 데이터 세트에서 평가했습니다. 분류 및 회귀 작업 모두에서 CGDFS는 일관되게 더 안정적이고 전이 가능한 특징 집합을 선택하며, 이는 sparsity 기반, 트리 기반 및 안정성 선택 기준과 비교하여 더 나은 out-of-distribution 성능과 더 큰 선택의 견고성을 제공합니다.

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.