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

2025-11-26

AI Summary

To address reinforcement learning instability caused by dynamic feature-space expansion and insufficient agent collaboration in structured-data feature transformation, this paper proposes HMA-FT, a heterogeneous multi-agent reinforcement learning framework. HMA-FT decouples feature selection and cross-feature generation via a heterogeneous agent specialization mechanism; enhances inter-agent coordination through a shared critic and multi-head attention; and improves policy stability under dynamic feature spaces via state encoding. Experiments demonstrate that HMA-FT consistently outperforms existing automated feature transformation methods on classification and regression tasks, achieving significant gains in effectiveness, inference efficiency, out-of-distribution robustness, and decision interpretability. The framework establishes a novel paradigm for adaptive feature engineering on high-dimensional structured data.

Technology Category

• Machine Learning: Feature Construction/Reformulation
• Multiagent Systems: Multiagent Learning
• Humans and AI: Human-in-the-loop Machine Learning

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Automated feature transformation methods are inefficient due to heuristic or exhaustive search approaches
• Reinforcement learning approaches suffer from dynamic feature space expansion causing instability
• Insufficient agent cooperation leads to suboptimal feature crossing and degraded performance

Innovation

• Multi-agent reinforcement learning with attention for feature transformation
• Shared critic mechanism enhances agent communication and cooperation
• State encoding stabilizes learning in dynamic feature spaces