Causally-Guided Automated Feature Engineering with Multi-Agent Reinforcement Learning

1. Problem Statement

Automated Feature Engineering (AFE) aims to autonomously construct high-utility representations from raw tabular data to improve machine learning performance. However, existing AFE methods face critical limitations:

• Brittleness under Distribution Shift: Most methods rely on statistical correlations (heuristics). When input distributions shift (e.g., in industrial or chemical processes), these correlations often break, leading to spurious features that fail to generalize.
• Combinatorial Explosion: The space of possible feature transformations grows exponentially with dimensionality, making greedy search computationally expensive and prone to local optima.
• Lack of Causal Guidance: Current approaches (including RL-based and LLM-based methods) generally ignore causal structures, treating feature generation as a sequence modeling or statistical optimization problem rather than a mechanism-preserving one.

Goal: To develop an AFE framework that constructs features preserving causal relations with the target, ensuring robustness under covariate shifts and mechanism-preserving changes, while maintaining computational tractability.

---

2. Methodology: The CAFE Framework

CAFE reformulates AFE as a causally-guided sequential decision process. It integrates causal discovery with a cascading multi-agent Deep Q-Network (DQN) architecture. The framework operates in two phases:

Phase I: Causal Graph Discovery (Soft Inductive Prior)

Instead of treating causal discovery as a rigid constraint, CAFE uses it to generate soft inductive priors to guide the search.

• Algorithm: Uses NOTEARS with Lasso regularization to learn a sparse Directed Acyclic Graph (DAG) over features and the target variable.
• Feature Grouping: Based on the learned DAG, features are categorized into three causal roles:
  1. Direct Causes ($C_{direct}$): Features with a direct edge to the target ($X \to Y$).
  2. Indirect Causes ($C_{indirect}$): Features connected via multi-hop paths ($X \rightsquigarrow Y$).
  3. Non-Causal/Other ($C_{other}$): Features with no causal path to the target.
• Screening: Within each group, features are screened (e.g., by Mutual Information) to manage dimensionality, creating a structured search space.

Phase II: Multi-Agent Reinforcement Learning

Phase II navigates the exponential transformation space using three specialized agents operating in a cascade:

1. Primary Group Agent ($\pi_1$): Selects a causal feature group (Direct, Indirect, or Other) to transform.
2. Operator Agent ($\pi_o$): Selects a transformation operator (e.g., arithmetic, log, trigonometric) from a predefined library.
3. Secondary Group Agent ($\pi_2$): Selects a partner group for binary operations (e.g., interaction terms).

Key Technical Components:

• State Representation: Agents receive hierarchical statistical summaries of the current feature set and group-specific information (e.g., mean, std, quartiles, mutual information rankings).
• Causally-Guided Exploration: The exploration strategy is a weighted mix of:
  • Causal-Hierarchical: Prioritizes transformations respecting the causal structure (e.g., Direct $\times$ Indirect).
  • Top-k Mutual Information: Hedges against causal discovery errors by selecting features with strong empirical associations.
  • Diversity Sampling: Prevents premature convergence.
• Reward Shaping: The reward function ($R_t$) balances three objectives:
  1. Performance ($R_{perf}$): Improvement in validation score.
  2. Causal Alignment ($\Psi_{causal}$): A bonus for using features from causal groups (weighted $w_{direct} > w_{indirect} > w_{other}$).
  3. Complexity & Diversity: Penalties for excessive feature count and rewards for diverse transformations.
• Robustness Guards: Includes numerical safeguards (e.g., protected division, log of absolute values) to prevent overflow and NaNs.

---

3. Key Contributions

1. Causally-Guided AFE Formulation: The first framework to explicitly model AFE as a sequential decision process guided by causal structure, moving beyond correlation-based heuristics.
2. Soft Inductive Bias: Introduces the concept of using causal graphs as soft priors rather than hard constraints. This maintains flexibility when causal discovery is imperfect while providing principled guidance.
3. Cascading Multi-Agent Architecture: Proposes a novel factorization of the AFE action space into three agents (Group, Operator, Partner), significantly reducing the effective action space and variance compared to joint optimization.
4. Causally-Shaped Rewards: Designs a reward function that explicitly incentivizes causal consistency alongside predictive utility, ensuring intervention stability.

---

4. Experimental Results

The authors evaluated CAFE on 15 public benchmark datasets (classification and regression) against 10 baselines, including statistical methods (AutoFeat), RL-based methods (GRFG, NFS), and LLM-based methods (ELLM-FT).

• Predictive Performance:

  • CAFE achieved superior performance on 13 out of 15 datasets.
  • Demonstrated up to 7% improvement in macro-F1 (classification) and inverse relative absolute error (regression) over strong baselines.
  • Showed consistent gains in high-dimensional regression and small-sample classification tasks.

• Robustness under Distribution Shift:

  • Under controlled covariate shifts, CAFE reduced performance degradation by approximately 4x compared to the non-causal multi-agent baseline (GRFG).
  • Example: Under high shift intensity, GRFG degraded by ~58%, while CAFE degraded by only ~22%.
  • This confirms that causal features preserve utility better than spurious correlations when mechanisms are preserved but distributions shift.

• Interpretability and Stability:

  • SHAP Stability: CAFE-generated features exhibited 58.6% lower variance in SHAP explanations under input noise compared to statistical baselines.
  • The features were more compact and semantically aligned with the underlying data-generating process.

• Efficiency:

  • Despite higher computational cost per episode (due to causal discovery), CAFE required 40–70% fewer episodes to converge to optimal performance, resulting in comparable or better total time-to-target.

---

5. Significance and Implications

• Safety-Critical Applications: CAFE addresses a fundamental gap in high-stakes domains (chemical engineering, healthcare, finance) where distribution shifts are common. By prioritizing causal mechanisms over correlations, it produces models that are more reliable under operational changes.
• Paradigm Shift: It establishes a new paradigm for AFE where causal discovery acts as a guide for reinforcement learning, rather than a separate post-hoc analysis step.
• Theoretical Insight: The work validates the hypothesis that features constructed from causal ancestors (direct/indirect) possess invariance properties that statistical correlations lack, providing a theoretical foundation for robust feature engineering.
• Limitations & Future Work: The approach currently assumes linear additive noise models and static causal graphs. Future work aims to integrate non-linear causal discovery (e.g., neural causal models) and handle temporal dynamics/feedback loops.

In summary, CAFE demonstrates that integrating causal reasoning into the automated feature engineering loop significantly enhances model robustness, interpretability, and efficiency, particularly in dynamic environments where data distributions are non-stationary.