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

Arun Vignesh Malarkkan; Wangyang Ying; Yanjie Fu

cs.AIFeb 18, 2026

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

Authors: Arun Vignesh Malarkkan, Wangyang Ying, Yanjie Fu

Organizations: School of Computing, Arizona State University

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.

Figures & tables

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.

Dataset	Source	C/R	Samples	Features	RDG	ERG	LDA	AFT	NFS	TTG	GRFG	ELLM-FT	CAFE
Amazon Employee	Kaggle	C	32769	9	0.744	0.740	0.920	0.943	0.935	0.806	0.946	0.946	0.947
SVMGuide3	LibSVM	C	1243	21	0.703	0.747	0.683	0.829	0.831	0.766	0.831	0.836	0.846
German Credit	UCIrvine	C	1001	24	0.695	0.661	0.627	0.751	0.765	0.731	0.772	0.775	0.793
Messidor_features	UCIrvine	C	1150	19	0.673	0.635	0.580	0.678	0.746	0.726	0.757	0.757	0.774
Ionosphere	UCIrvine	C	351	34	0.919	0.926	0.730	0.827	0.949	0.938	0.960	0.963	0.976
Wine Quality Red	UCIrvine	C	999	12	0.599	0.611	0.600	0.658	0.666	0.647	0.686	0.685	0.707

Table 1: Overall performance comparison. ‘C’ for classification and ‘R’ for regression. Bold indicates the best result.

Panel of Figure 2 — Figure 2: Comparison of different CAFE variants in terms of F1 or 1-RAE.

Panel of Figure 4 — Figure 4: SHAP Explanation Stability. Lower variation values indicate more stable model explanations under noise.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Algorithm	SHD $\downarrow$	F1-Score $\uparrow$	Runtime (s)	CAFE Performance
PC Algorithm	$8.3\pm 2.1$	$0.64\pm 0.08$	$12.4\pm 3.2$	$0.742\pm 0.034$
GES	$6.7\pm 1.8$	$0.71\pm 0.07$	$45.7\pm 8.9$	$0.756\pm 0.028$
LiNGAM	$7.2\pm 2.0$	$0.68\pm 0.09$	$8.9\pm 2.1$	$0.748\pm 0.031$
NOTEARS-Lasso	$4.9\pm 1.3$	$0.78\pm 0.06$	$23.1\pm 5.4$	$0.773\pm 0.025$

Appendix

Table 2: Causal discovery algorithm comparison on synthetic data (10 datasets,

d=20

n=1000

Dataset	CAFE-Linear	CAFE-MLP	MLP Gain	Sample Efficiency
Quadratic	0.768	0.794	+2.6%	-7% episodes
Exponential	0.724	0.761	+3.7%	-12% episodes
Mixed	0.747	0.779	+3.2%	-11% episodes

Appendix

Table 3: CAFE with linear vs. non-linear causal discovery on synthetic SCMs.

Graph Quality	SHD Range	CAFE Performance	vs. Oracle	vs. Random
Excellent	$[0.0,2.0]$	$0.798\pm 0.021$	$-0.8\%$	$+12.4\%$
Good	$(2.0,5.0]$	$0.773\pm 0.025$	$-3.9\%$	$+8.8\%$
Fair	$(5.0,8.0]$	$0.751\pm 0.031$	$-6.7\%$	$+5.7\%$
Poor	$(8.0,\infty)$	$0.728\pm 0.038$	$-9.5\%$	$+2.4\%$

Appendix

Table 4: CAFE performance degradation vs. graph quality (SHD quartiles).

Dataset	CAFE (single)	CAFE (ensemble)	Improvement
Wine Red	0.707	0.714	+0.7%
German Credit	0.793	0.804	+1.0%
OpenML 586	0.810	0.819	+0.9%
Ionosphere	0.976	0.978	+0.2%

Appendix

Table 5: Ensemble DAG results.

Comparison	Wilcoxon $p$	Effect Size ( $r$ )	Interpretation
CAFE vs GRFG	$<0.001$	0.68	Large effect
CAFE vs ELLM-FT	$0.003$	0.51	Medium effect
CAFE vs AutoFeat	$<0.001$	0.62	Large effect
CAFE vs NFS	$<0.001$	0.59	Large effect
CAFE vs Original	$<0.001$	0.84	Large effect

Appendix

Table 6: Statistical significance analysis with multiple testing correction.

Shift Type	Intensity	CAFE	GRFG	ELLM-FT	AutoFeat
Covariate	Low	$-4.1\pm 1.8$	$-12.3\pm 3.2$	$-11.7\pm 2.9$	$-15.2\pm 4.1$
	Medium	$-12.8\pm 3.4$	$-31.8\pm 5.7$	$-29.4\pm 5.1$	$-34.7\pm 6.8$
	High	$-22.1\pm 4.9$	$-58.7\pm 8.2$	$-54.3\pm 7.6$	$-61.2\pm 9.1$
Label	Low	$-7.3\pm 2.1$	$-15.2\pm 3.8$	$-14.8\pm 3.5$	$-18.4\pm 4.7$
	Medium	$-18.2\pm 4.2$	$-34.1\pm 6.3$	$-32.6\pm 5.9$	$-37.8\pm 7.2$
	High	$-31.4\pm 6.1$	$-52.8\pm 7.9$	$-49.7\pm 7.3$	$-55.3\pm 8.6$

Appendix

Table 7: Robustness under mechanism-preserving distribution shifts (performance degradation %). Lower magnitude indicates stronger robustness. Results are averaged across datasets and seeds. These shifts modify feature or label distributions while preserving underlying causal mechanisms.

Method	ID Performance	OOD Performance	Drop (%)
GRFG	$0.734\pm 0.038$	$0.528\pm 0.067$	$-28.1$
ELLM-FT	$0.741\pm 0.035$	$0.503\pm 0.072$	$-32.1$
CAFE (ours)	$\mathbf{0.773\pm 0.032}$	$\mathbf{0.718\pm 0.044}$	$\mathbf{-7.1}$

Appendix

Table 8: In-distribution (ID) vs. OOD performance and relative degradation (mean

\pm

std).

Method	SHAP Stab. $\downarrow$	Feat. Stab. $\uparrow$	Causal Cons. $\uparrow$	Compactness $\uparrow$	Overall
GRFG	$0.089\pm 0.012$	$0.42\pm 0.08$	$0.31\pm 0.06$	$0.23\pm 0.05$	$2.8/5$
ELLM-FT	$0.078\pm 0.011$	$0.38\pm 0.07$	$0.28\pm 0.05$	$0.31\pm 0.06$	$3.0/5$
AutoFeat	$0.094\pm 0.013$	$0.51\pm 0.09$	$0.25\pm 0.04$	$0.19\pm 0.04$	$2.7/5$
CAFE	$0.043\pm 0.008$	$0.68\pm 0.11$	$0.73\pm 0.09$	$0.47\pm 0.08$	$4.2/5$

Appendix

Table 9: Comprehensive interpretability metrics comparison.

Figure 5: Convergence Comparison: CAFE vs GRFG.

Parameter	Range Tested	Optimal	Performance Range	Sensitivity	Selection Rule
$\lambda$ (NOTEARS)	[0.001, 0.1]	0.03	[0.734, 0.773]	Medium	3-fold CV
$\alpha$ (causal bonus)	[0.0, 1.0]	0.5	[0.742, 0.773]	Low	Fixed at 0.5
Episodes $E$	[10, 50]	30	[0.751, 0.775]	Low	Early stopping
Steps per episode $S$	[5, 25]	15	[0.743, 0.771]	Low	Fixed
$k_{g}$ (group size)	[5, 20]	10	[0.758, 0.773]	Very Low	$\min(\sqrt{d},15)$
Learning rate	[1e-4, 1e-2]	1e-3	[0.745, 0.773]	Medium	Fixed

Appendix

Table 10: Hyperparameter sensitivity analysis (mean performance across datasets).

Challenge	CAFE	Best Method	Best Score	Gap	Root Cause
High-dim, low-n ( $d>n$ )	$0.612\pm 0.045$	GRFG	$0.635\pm 0.041$	$-3.6\%$	Unreliable CD
Sparse graphs	$0.687\pm 0.031$	ELLM-FT	$0.694\pm 0.028$	$-1.0\%$	Weak causal signal
Non-linearities	$0.723\pm 0.027$	AutoFeat	$0.734\pm 0.025$	$-1.5\%$	Linear CD assumption
Hidden confounders	$0.641\pm 0.044$	GRFG	$0.668\pm 0.039$	$-4.0\%$	Sufficiency violation
Temporal dynamics	$0.658\pm 0.038$	ELLM-FT	$0.671\pm 0.035$	$-1.9\%$	Static DAG

Appendix

Table 11: Performance under challenging conditions where CAFE shows limitations.

Method	XGBoost	Random Forest	Linear	SVM (RBF)	Average
Original	$0.682\pm 0.047$	$0.671\pm 0.052$	$0.598\pm 0.078$	$0.648\pm 0.055$	$0.650$
GRFG	$0.734\pm 0.038$	$0.718\pm 0.043$	$0.627\pm 0.071$	$0.691\pm 0.041$	$0.693$
ELLM-FT	$0.741\pm 0.035$	$0.725\pm 0.041$	$0.634\pm 0.068$	$0.698\pm 0.038$	$0.700$
AutoFeat	$0.728\pm 0.039$	$0.712\pm 0.044$	$0.618\pm 0.074$	$0.681\pm 0.042$	$0.685$
CAFE	0.773 ± 0.032	0.756 ± 0.037	0.662 ± 0.063	0.724 ± 0.035	0.729

Appendix

Table 12: Performance across different downstream models (mean F1/1-RAE over 5-fold CV).

Method	Episodes to Converge	Final Performance	Efficiency Ratio
GRFG	$28.4\pm 4.7$	$0.734\pm 0.038$	$1.00$
CAFE	18.3 $\pm$ 3.1	0.773 $\pm$ 0.032	1.55

Appendix

Table 13: Convergence characteristics (episodes to reach 95% of final performance).

Automated Feature Engineering

Automated Feature Engineering creates, transforms, and selects predictive inputs with less manual design across data types and applications.

See all 54 papers

FAMOSE: A ReAct Approach to Automated Feature Discovery
Feb 19, 2026Keith Burghardt, Jienan Liu, Sadman Sakib +2Automated Feature Engineering Large Language Model-Based Feature Engineering
Amazon.com, Inc., 410 Terry Ave N, Seattle, WA 98109, United States.
Actionable CBFI: Integrating Structural Decomposition and Causal Counterfactual Recourse for Tabular Machine Learning
Aug 28, 2026Sejong OhFeature Attribution Counterfactual Explanations
Dankook University, Republic of Korea
Adaptive Graph-of-Islands Evolution for Automatic Feature Engineering with LLMs
Jul 25, 2026Sha Li, Naren RamakrishnanEvolutionary Optimization Large Language Model-Guided Optimization
Department of Computer Science Virginia Tech Blacksburg, VA 24060 · Department of Computer Science Virginia Tech Alexandria, VA 22305

Feb 19, 2026J/K move · Enter open · S save

FAMOSE: A ReAct Approach to Automated Feature Discovery

Keith Burghardt, Jienan Liu, Sadman Sakib, Yuning Hao+1

Amazon.com, Inc., 410 Terry Ave N, Seattle, WA 98109, United States.

Automated Feature Engineering Large Language Model-Based Feature Engineering Tabular ML

Open paper AskPDF arXiv

Feature engineering remains a critical yet challenging bottleneck in machine learning, particularly for tabular data, as identifying optimal features from an exponentially large feature space traditionally demands substantial domain expertise. To address this challenge, we introduce FAMOSE (Feature AugMentation and Optimal Selection agEnt), a novel framework that leverages the ReAct paradigm to autonomously explore, generate, and refine features while integrating feature selection and evaluation tools within an agent architecture. To our knowledge, FAMOSE represents the first application of an agentic ReAct framework to automated feature engineering, especially for both regression and classification tasks. Extensive experiments demonstrate that FAMOSE is at or near the state-of-the-art on classification tasks (especially tasks with more than 10K instances, where ROC-AUC increases 0.23% on average), and achieves the state-of-the-art for regression tasks by reducing RMSE by 2.0% on average, while remaining more robust to errors than other algorithms. We hypothesize that FAMOSE's strong performance is because ReAct allows the LLM context window to record (via iterative feature discovery and evaluation steps) what features did or did not work. This is similar to a few-shot prompt and guides the LLM to invent better, more innovative features. Our work offers evidence that AI agents are remarkably effective in solving problems that require highly inventive solutions, such as feature engineering.

Abstract

Automated Feature Engineering

Explore similar work

FAMOSE: A ReAct Approach to Automated Feature Discovery

Actionable CBFI: Integrating Structural Decomposition and Causal Counterfactual Recourse for Tabular Machine Learning

Adaptive Graph-of-Islands Evolution for Automatic Feature Engineering with LLMs