Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Arun Vignesh Malarkkan, Wangyang Ying, Yanjie Fu

TL;DR

CAFE is introduced, a framework that reformulates AFE as a causally-guided sequential decision process, bridging causal discovery with reinforcement learning-driven feature construction, and underscores that causal structure, used as a soft inductive prior rather than a rigid constraint, can substantially improve the robustness and efficiency of automated feature engineering.

Abstract

Automated feature engineering (AFE) enables AI systems to autonomously construct high-utility representations from raw tabular data. However, existing AFE methods rely on statistical heuristics, yielding brittle features that fail under distribution shift. We introduce CAFE, a framework that reformulates AFE as a causally-guided sequential decision process, bridging causal discovery with reinforcement learning-driven feature construction. Phase I learns a sparse directed acyclic graph over features and the target to obtain soft causal priors, grouping features as direct, indirect, or other based on their causal influence with respect to the target. Phase II uses a cascading multi-agent deep Q-learning architecture to select causal groups and transformation operators, with hierarchical reward shaping and causal group-level exploration strategies that favor causally plausible transformations while controlling feature complexity. Across 15 public benchmarks (classification with macro-F1; regression with inverse relative absolute error), CAFE achieves up to 7% improvement over strong AFE baselines, reduces episodes-to-convergence, and delivers competitive time-to-target. Under controlled covariate shifts, CAFE reduces performance drop by ~4x relative to a non-causal multi-agent baseline, and produces more compact feature sets with more stable post-hoc attributions. These findings underscore that causal structure, used as a soft inductive prior rather than a rigid constraint, can substantially improve the robustness and efficiency of automated feature engineering.

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

TL;DR

CAFE is introduced, a framework that reformulates AFE as a causally-guided sequential decision process, bridging causal discovery with reinforcement learning-driven feature construction, and underscores that causal structure, used as a soft inductive prior rather than a rigid constraint, can substantially improve the robustness and efficiency of automated feature engineering.

Abstract

Automated feature engineering (AFE) enables AI systems to autonomously construct high-utility representations from raw tabular data. However, existing AFE methods rely on statistical heuristics, yielding brittle features that fail under distribution shift. We introduce CAFE, a framework that reformulates AFE as a causally-guided sequential decision process, bridging causal discovery with reinforcement learning-driven feature construction. Phase I learns a sparse directed acyclic graph over features and the target to obtain soft causal priors, grouping features as direct, indirect, or other based on their causal influence with respect to the target. Phase II uses a cascading multi-agent deep Q-learning architecture to select causal groups and transformation operators, with hierarchical reward shaping and causal group-level exploration strategies that favor causally plausible transformations while controlling feature complexity. Across 15 public benchmarks (classification with macro-F1; regression with inverse relative absolute error), CAFE achieves up to 7% improvement over strong AFE baselines, reduces episodes-to-convergence, and delivers competitive time-to-target. Under controlled covariate shifts, CAFE reduces performance drop by ~4x relative to a non-causal multi-agent baseline, and produces more compact feature sets with more stable post-hoc attributions. These findings underscore that causal structure, used as a soft inductive prior rather than a rigid constraint, can substantially improve the robustness and efficiency of automated feature engineering.
Paper Structure (68 sections, 4 theorems, 26 equations, 5 figures, 13 tables, 2 algorithms)

This paper contains 68 sections, 4 theorems, 26 equations, 5 figures, 13 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Problem Statement
  3. Methodology
  4. Theoretical Foundation: Why Causal Structure Guides Robust Feature Engineering
  5. Phase I: Causal Graph Discovery as a Soft Inductive Prior for Feature Grouping
  6. Phase II: Multi-Agent Reinforcement Learning with Causal Guidance
  7. Experimental Study
  8. Experimental Protocol
  9. Experimental Results
  10. Overall Performance Analysis (RQ1)
  11. Ablation Study and Component Analysis (RQ2)
  12. Robustness Analysis (RQ3)
  13. Interpretability Analysis (RQ4)
  14. Related Works
  15. Conclusion
  16. ...and 53 more sections

Key Result

Proposition 1

Let $S\subseteq\mathcal{F}$ contain a subset $S^\star$ that suffices for predicting $Y$ across environments (e.g., $S^\star \supseteq \mathrm{PA}_Y$ and the shift is mechanism-preserving). If a transformation $\phi$ preserves the information in $S^\star$ (e.g., is injective on $S^\star$ or is a suff If $S$ excludes any such invariant set, no general invariance guarantee is available.

Figures (5)

  • Figure 1: CAFE Framework Overview. Phase I learns a causal graph and derives soft causal priors that group features by their relation to the target: direct causes, indirect causes, and non-causal features. Phase II employs cascaded multi-agent reinforcement learning with causally shaped rewards and adaptive exploration strategies, balancing causal coherence with statistical discovery.
  • Figure 2: Comparison of different CAFE variants in terms of F1 or 1-RAE.
  • Figure 3: Robustness Study of CAFE.
  • Figure 4: SHAP Explanation Stability. Lower variation values indicate more stable model explanations under noise.
  • Figure 5: Convergence Comparison: CAFE vs GRFG.

Theorems & Definitions (6)

  • Proposition 1: Informal invariance for causally sufficient summaries
  • Definition 1: Causal Ancestry
  • Proposition 2: Causal Feature Invariance
  • proof : Proof Sketch
  • Corollary 1: Causal Hierarchy for Feature Engineering
  • Proposition 3: Multi-Agent Factorization Benefits