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

This paper proposes a novel heterogeneous multi-agent reinforcement learning framework that leverages a shared critic mechanism and multi-head attention to enable cooperative, stable, and scalable automated feature transformation for structured data, overcoming the limitations of dynamic feature expansion and insufficient agent communication found in prior approaches.

The Gist

Imagine you are a chef trying to create the world's best soup. You have a pantry full of raw ingredients (your data features like onions, carrots, and spices). Sometimes, just throwing these ingredients into a pot isn't enough. To make the soup truly delicious, you need to transform them: maybe you roast the carrots first, mix the onions with a specific spice, or blend them into a paste. This process of mixing and transforming ingredients to create something better is called Feature Transformation.

For a long time, computers have been bad at figuring out which ingredients to mix and how to mix them. They either tried every single combination (which takes forever) or relied on guesswork.

This paper introduces a new, smarter way to do this using a team of AI chefs called HAFT. Here is how it works, broken down into simple concepts:

1. The Problem: A Growing, Chaotic Kitchen

Imagine your kitchen is expanding every time you add a new ingredient.

• The Chaos: As the AI creates new "recipes" (new features), the number of ingredients in the pantry grows. It becomes harder and harder for a single chef to remember everything or know which new ingredient is actually good.
• The Silos: In older systems, the chefs worked alone. Chef A picked an onion, Chef B picked a spice, but they never talked to each other. They didn't know what the others were thinking, leading to weird combinations (like putting chocolate in a savory soup).

2. The Solution: A Team of Specialized Chefs (Heterogeneous Agents)

The authors created a team of three specialized AI chefs who work together in a line (a "cascading" team):

• Chef 1 (Head Agent): Looks at the whole pantry and picks the first ingredient to use.
• Chef 2 (Operation Agent): Decides how to treat that ingredient. Should we chop it? Roast it? Mix it with acid?
• Chef 3 (Tail Agent): Looks at the first ingredient and the chosen method, then picks a second ingredient to combine with it.

Why "Heterogeneous"?
Just like in a real kitchen, you don't ask the person chopping vegetables to also be the one managing the oven.

• The Feature Chefs (1 and 3) have a special superpower: they use Attention. Imagine they have a magical spotlight that can instantly focus on the most important ingredients, even if the pantry has 1,000 items or only 5. This helps them handle the "growing pantry" problem without getting overwhelmed.
• The Operation Chef (2) has a simpler job: pick a tool from a fixed list (knife, blender, oven). They use a standard, efficient method for this.

3. The Secret Sauce: The Shared Critic (The Head Chef)

In the old days, the chefs only whispered to the person standing right next to them. In this new system, there is a Head Chef (The Shared Critic) standing on a balcony looking at the entire kitchen.

• The Head Chef sees everything: what the pantry looks like, what Chef 1 picked, what Chef 2 decided, and what Chef 3 is about to do.
• The Head Chef gives feedback to the whole team: "Great job! That combination is working!" or "Stop! That mix tastes terrible."
• Because they all listen to the same Head Chef, they stop fighting and start cooperating. They learn to make decisions that help the whole team win, not just themselves.

4. Keeping the Kitchen Stable (State Encoding)

Every time the team makes a new recipe, the pantry changes. If the kitchen layout keeps shifting wildly, the chefs get dizzy and can't learn.

• The authors invented a special "kitchen map" (State Encoding). Instead of showing the chefs the messy, changing pantry, this map summarizes the pantry into a neat, fixed-size report (like a checklist of "how many onions," "how spicy," etc.).
• This keeps the chefs calm and focused, allowing them to learn faster and more reliably.

5. The Result: A Better Soup

The paper tested this team on 23 different "cooking competitions" (real-world data problems like predicting loan risks or diagnosing diseases).

• Better Taste: The HAFT team consistently made better "soups" (more accurate predictions) than other methods.
• Faster Cooking: They were much faster than older methods that tried to group ingredients clumsily.
• Explainable: Unlike some AI that just gives a magic answer, this team can tell you exactly how they made the soup (e.g., "We took the 'Debt' ingredient, divided it by 'Income', and then took the square root"). This is crucial for doctors and bankers who need to trust the AI.

Summary

Think of HAFT as a highly organized, talking team of AI chefs.

• They have specialized roles (some pick ingredients, some pick tools).
• They use magic spotlights to handle huge pantries.
• They listen to a Head Chef who sees the big picture so they don't make mistakes.
• They use a stable map so they don't get confused when the pantry changes.

The result is a system that automatically invents better ways to combine data, leading to smarter and more reliable AI predictions.

Technical Summary

1. Problem Statement

The paper addresses the challenge of Automated Feature Transformation (AFT), specifically for structured (tabular) data. While deep learning excels in unstructured data, it often struggles to capture complex, high-order feature interactions in structured data without explicit feature engineering.

• The Goal: To automatically generate an optimal feature set ($F^*$) by applying mathematical operations (e.g., addition, multiplication, trigonometric functions) to original features, thereby maximizing downstream task performance (classification or regression).
• Key Challenges Identified:
  1. Dynamic Feature Space Expansion: As the RL agent iteratively generates new features, the feature space grows dynamically. This makes it difficult for standard agents to identify relevant features for the next crossing operation due to variable input sizes and increased complexity.
  2. Insufficient Agent Cooperation: Existing multi-agent approaches often rely on local information or sequential hand-offs, leading to suboptimal policies because agents lack a global view of the feature space and fail to coordinate effectively.

2. Methodology: The HAFT Framework

The authors propose HAFT (Heterogeneous multi-Agent reinforcement learning framework for cooperative and scalable Feature Transformation). The framework utilizes a cascading multi-agent architecture with three heterogeneous agents and a shared central critic.

A. Heterogeneous Cascading Agents

The framework employs three agents that collaborate in a specific sequence at each iteration $t$:

1. Head Feature Agent: Selects the first candidate feature ($f_1$) from the current feature set.
2. Operation Agent: Selects a mathematical operation ($o$) from a fixed set (e.g., $\sin, \log, +, \times$).
3. Tail Feature Agent: Selects the second candidate feature ($f_2$) based on the previous selections and the current state.

Key Architectural Innovations:

• Multi-Head Attention for Feature Agents: To handle the dynamically expanding feature space, the feature agents utilize a Transformer-based encoder with multi-head self-attention. This allows the agents to process variable-length feature sets, capture complex global interactions between features, and assign attention weights to the most relevant features without requiring fixed input dimensions.
• MLP-Based Operation Agent: Since the set of mathematical operations is static, this agent uses a standard Multi-Layer Perceptron (MLP). It incorporates a dynamic action mask to prevent invalid operations (e.g., preventing $\sqrt{x}$ if $x$ contains negative values).

B. Shared Central Critic & Advantage Decomposition

To solve the cooperation challenge, HAFT introduces a Shared Central Critic:

• Global View: The critic observes the entire feature space and the decisions of all agents, providing a unified baseline for value estimation.
• State Encoding: To stabilize learning against the shifting feature space, the critic input is encoded via a two-branch structure:
  1. Statistical Branch: Computes descriptive statistics (mean, variance, quartiles) of the feature set to create a fixed-size summary.
  2. Attention Branch: Uses multi-head attention to capture feature-feature interactions, followed by permutation-invariant pooling.
• Sequential Policy Update (HAPPO): The agents are updated sequentially (Head $\to$ Operation $\to$ Tail) using a customized Heterogeneous Actor-Propximal Policy Optimization (HAPPO) algorithm. This includes advantage decomposition, where the advantage of subsequent agents is adjusted based on the updated policies of preceding agents to ensure consistent credit assignment.

C. Reward Function

The agents are driven by a composite reward function that includes:

• Downstream Performance: The primary metric (e.g., F1-score, $R^2$) of the downstream ML model.
• Information Redundancy/Relevance: Penalizes high mutual information between generated features (redundancy) and encourages high mutual information between features and the target label (relevance).
• Entropy Regularization: Encourages exploration.

3. Key Contributions

1. Novel Framework (HAFT): A heterogeneous multi-agent RL framework specifically designed for scalable and cooperative feature transformation.
2. Scalable Agent Design: The use of Multi-Head Attention in feature agents allows the system to adapt to the continuously growing feature space, solving the scalability issue inherent in previous RL-based AFT methods.
3. Enhanced Cooperation: The introduction of a Shared Critic with a specialized two-branch state encoder significantly improves inter-agent communication and credit assignment, leading to more coherent transformation policies.
4. Stabilization Techniques: The proposed state encoding and sequential advantage decomposition mitigate the instability caused by dynamic state changes in the feature space.

4. Experimental Results

The authors evaluated HAFT on 23 real-world datasets (14 classification, 9 regression) from UCI, LibSVM, Kaggle, and OpenML.

• Overall Performance: HAFT consistently outperformed 8 state-of-the-art baselines (including RDG, ERG, LDA, AFAT, NFS, TTG, GRFG, and DIFER) across most datasets. It achieved the best F1 scores or 1-RAE in the majority of cases.
• Ablation Studies:
  • Shared Critic: Removing the shared critic (replacing it with individual critics) or removing advantage decomposition resulted in significant performance drops, confirming the necessity of global coordination.
  • State Encoding: A variant without the interaction-based attention branch in the critic state ($HAFT-u$) performed worse, proving that capturing feature interactions is crucial for the critic's evaluation.
• Robustness: HAFT maintained superior performance across different downstream models (Random Forest, XGBoost, SVM, KNN, Ridge), demonstrating its ability to generate generalizable feature spaces.
• Scalability & Efficiency: HAFT demonstrated significantly lower time complexity per iteration compared to the closest competitor (GRFG). While GRFG relies on time-consuming clustering, HAFT's attention mechanism allows for faster, more efficient feature selection.
• Interpretability: Case studies showed that HAFT generates traceable features (e.g., deriving [exudate6]-[exudate5] from raw inputs) that improve model performance by up to 13.9% and provide clear medical or domain-specific insights.

5. Significance

This paper makes a significant contribution to Automated Feature Engineering (AutoFE) and Multi-Agent Reinforcement Learning (MARL):

• Bridging the Gap: It effectively bridges the gap between the need for interpretable, structured feature engineering and the power of deep learning/RL.
• Solving the "Curse of Dimensionality" in AFT: By using attention mechanisms, it solves the critical problem of how to manage an ever-expanding feature space in an iterative RL setting, a limitation that plagued previous methods.
• Cooperative Learning: It demonstrates that heterogeneous agents with shared global information can learn complex, sequential decision-making policies more effectively than independent or locally-coordinated agents.
• Practical Impact: The method is not only theoretically sound but also practically efficient, offering a scalable solution for improving predictive models in finance, healthcare, and other domains relying on tabular data.