Heterogeneous Multi-Agent Reinforcement Learning with Attention for Cooperative and Scalable Feature Transformation

TL;DR

HAFT develops a heterogeneous multi-agent reinforcement learning framework for scalable automated feature transformation on tabular data. It employs three cascading agents (two feature selectors and one operation selector) coordinated by a shared critic, aided by multi-head attention and a two-branch state encoding to manage dynamic feature spaces. Ablation and robustness studies demonstrate the necessity of the shared critic, state encoding, and attention mechanism, while scalability analyses show faster performance than clustering-based rivals. The approach yields improved downstream performance, interpretable transformed features, and strong generalization across diverse datasets.

Abstract

Feature transformation enhances downstream task performance by generating informative features through mathematical feature crossing. Despite the advancements in deep learning, feature transformation remains essential for structured data, where deep models often struggle to capture complex feature interactions. Prior literature on automated feature transformation has achieved success but often relies on heuristics or exhaustive searches, leading to inefficient and time-consuming processes. Recent works employ reinforcement learning (RL) to enhance traditional approaches through a more effective trial-and-error way. However, two limitations remain: 1) Dynamic feature expansion during the transformation process, which causes instability and increases the learning complexity for RL agents; 2) Insufficient cooperation and communication between agents, which results in suboptimal feature crossing operations and degraded model performance. To address them, we propose a novel heterogeneous multi-agent RL framework to enable cooperative and scalable feature transformation. The framework comprises three heterogeneous agents, grouped into two types, each designed to select essential features and operations for feature crossing. To enhance communication among these agents, we implement a shared critic mechanism that facilitates information exchange during feature transformation. To handle the dynamically expanding feature space, we tailor multi-head attention-based feature agents to select suitable features for feature crossing. Additionally, we introduce a state encoding technique during the optimization process to stabilize and enhance the learning dynamics of the RL agents, resulting in more robust and reliable transformation policies. Finally, we conduct extensive experiments to validate the effectiveness, efficiency, robustness, and interpretability of our model.

Figures (8)

• Figure 1: Framework Overview. HAFT consists of three heterogeneous agents: two feature agents that select candidate features and one operation agent that determines a mathematical operation. A central critic leverages global information from the current feature space to coordinate the agents and enhance policy learning. This feature transformation process continues iteratively until an optimal feature set is identified or the iteration limit is reached.
• Figure 2: Two‑branch state encoder. A statistics branch expands the dynamic feature set and normalizes it to $\mathbb{R}^{1 \times 49}$; an attention branch embeds features, applies multi‑head attention, and pooling to obtain $z_{attn}$. Finally create input $s^{\text{critic}} = [S_{\text{stats}}, z_{\text{attn}}]$.
• Figure 3: Influence of Shared Critic and Advantage Decomposition
• Figure 4: Influence of Different State Encoding
• Figure 5: Time Complexity Comparison between HAFT and GRFG