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

arXiv: https://arxiv.org/abs/2602.16435 · HTML: https://mlsanity.com/paper/2602.16435/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 35 references that are in the map.
- First posted on arXiv 18 Feb 2026.

## Builds on (most cited first, 14 of 14)

- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Feature Engineering for Predictive Modeling using Reinforcement Learning](https://mlsanity.com/paper/1709.07150/) (2017)
- [The autofeat Python Library for Automated Feature Engineering and Selection](https://mlsanity.com/paper/1901.07329/) (2019)
- [LLM-FE: Automated Feature Engineering for Tabular Data with LLMs as Evolutionary Optimizers](https://mlsanity.com/paper/2503.14434/) (2025)
- [Group-wise Reinforcement Feature Generation for Optimal and Explainable Representation Space Reconstruction](https://mlsanity.com/paper/2205.14526/) (2022)
- [Evolutionary Large Language Model for Automated Feature Transformation](https://mlsanity.com/paper/2405.16203/) (2024)
- [Dynamic and Adaptive Feature Generation with LLM](https://mlsanity.com/paper/2406.03505/) (2024)
- [A Unified View of Causal and Non-causal Feature Selection](https://mlsanity.com/paper/1802.05844/) (2018)
- [A Survey on Data-Centric AI: Tabular Learning from Reinforcement Learning and Generative AI Perspective](https://mlsanity.com/paper/2502.08828/) (2025)
- [Incremental Causal Graph Learning for Online Cyberattack Detection in Cyber-Physical Infrastructures](https://mlsanity.com/paper/2507.14387/) (2025)
- [DELTA: Variational Disentangled Learning for Privacy-Preserving Data Reprogramming](https://mlsanity.com/paper/2509.00693/) (2025)
- [Causal Feature Selection for Responsible Machine Learning](https://mlsanity.com/paper/2402.02696/) (2024)
- [Automation of Feature Engineering for IoT Analytics](https://mlsanity.com/paper/1707.04067/) (2017)

## Built on by (most cited first, 3 of 3)

- [Causally-Guided Diffusion for Stable Feature Selection](https://mlsanity.com/paper/2603.20930/) (2026)
- [SymboLLM-FE: LLM-Accelerated Symbolic Regression for Automated Feature Engineering on Tabular Data](https://mlsanity.com/paper/2608.28408/) (2026)
- [Vision: Data-Centric Anchoring for Robust and Interpretable Agentic AI](https://mlsanity.com/paper/2609.08216/) (2026)

## Abstract (arXiv)

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.