[2602.16435] Causally-Guided Automated Feature Engineering with Multi-Agent Reinforcement Learning
Summary
The paper presents CAFE, a novel framework for automated feature engineering that combines causal discovery with multi-agent reinforcement learning to enhance feature robustness and efficiency.
Why It Matters
This research addresses the limitations of existing automated feature engineering methods that often produce fragile features. By integrating causal structures, CAFE improves the reliability of feature construction, making it a significant advancement for AI systems that rely on robust data representations.
Key Takeaways
- CAFE reformulates automated feature engineering as a causally-guided decision process.
- The framework achieves up to 7% improvement over traditional AFE methods across various benchmarks.
- CAFE effectively reduces performance drops under covariate shifts by approximately 4x compared to non-causal methods.
- Utilizing causal structures as soft priors enhances feature stability and compactness.
- The multi-agent deep Q-learning architecture allows for efficient feature selection and transformation.
Computer Science > Artificial Intelligence arXiv:2602.16435 (cs) [Submitted on 18 Feb 2026] Title:Causally-Guided Automated Feature Engineering with Multi-Agent Reinforcement Learning Authors:Arun Vignesh Malarkkan, Wangyang Ying, Yanjie Fu View a PDF of the paper titled Causally-Guided Automated Feature Engineering with Multi-Agent Reinforcement Learning, by Arun Vignesh Malarkkan and 2 other authors View PDF HTML (experimental) 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...