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

🕒 Published (v1): 2025-11-26 21:45 UTC · Source: Arxiv · Venue: KDD 2026 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

HAFT is a heterogeneous multi-agent reinforcement learning framework for automated feature transformation on structured/tabular data. Three specialized RL agents (two attention-based feature selectors, one MLP-based operation selector) cooperate via a shared central critic with advantage decomposition to iteratively cross features, outperforming 8 baselines on 23 datasets.

Problem

Existing RL-based automated feature transformation methods suffer from two key limitations: (1) the feature space dynamically expands each iteration, making it hard for fixed-architecture RL agents to select relevant features at scale; (2) agents coordinate only through partial local information (e.g., prior agent's selection), leading to insufficient cooperation and suboptimal feature crossing policies.

Method

HAFT uses three cascading heterogeneous agents in each iteration: a head feature agent (Transformer encoder with multi-head self-attention over the variable-length feature set → softmax distribution for selection), an MLP-based operation agent (fixed input size; dynamic validity mask prevents illegal operations like sqrt on negatives), and a tail feature agent (same attention architecture as head, conditioned on prior selections). The selected feature pair and operation generate one new feature per step.

A shared central critic receives a two-branch fixed-size state: a distributional-statistics branch (7 per-feature statistics → normalized R^{1×49}) and an attention-based interaction branch (per-feature token embeddings → permutation-invariant pooling → R^{1×49}). Agents are updated sequentially via a customized HAPPO algorithm with multi-agent advantage decomposition: each subsequent agent's advantage is re-weighted by the policy ratio of already-updated agents, ensuring monotonic improvement guarantees. The reward signal augments downstream task performance with mutual-information-based redundancy and relevance terms plus an entropy bonus.

Key Contributions

• Heterogeneous cascading agent design: attention-based feature agents for dynamic variable-length input; MLP operation agent for fixed-action-space selection
• Shared central critic with two-branch fixed-size state encoding to stabilize training under dynamic feature expansion
• Sequential HAPPO-based advantage decomposition for principled credit assignment across interdependent agents
• Mutual-information-augmented reward shaping (redundancy penalty + relevance reward + entropy term)

Results

• HAFT achieves best or second-best performance on 21 of 23 datasets (F1 for classification, 1-RAE for regression) against 8 baselines: RDG, ERG, LDA, AFAT, NFS, TTG, GRFG, DIFER
• Representative gains over best baseline: SpectF F1 0.876 vs. 0.957 (ERG/Ionosphere ties); German Credit 0.740 vs. 0.667 (RDG); Messidor 0.730 vs. 0.692 (ERG); Openml_586 0.707 vs. 0.646 (GRFG)
• Case study (Messidor_features): HAFT-transformed feature space improves F1 from 0.641 → 0.730 (+13.9%); all top-10 features by mutual information were HAFT-generated
• Time complexity per iteration substantially lower than GRFG (the most similar baseline), which spends the majority of its time on clustering; HAFT's attention mechanism replaces clustering with no performance penalty
• Ablation: HAFT > HAFT-a (no advantage decomposition) ≈ HAFT-c (separate critics) across regression and classification tasks; HAFT > HAFT-u (statistics-only critic state) in most settings

Limitations

• Evaluation restricted to tabular/structured data; generalization to unstructured modalities is not addressed
• Fixed operation set of 16 operations; domain-specific operations require manual extension
• Maximum iteration cap of 100 episodes with 25 steps may not be sufficient for very high-dimensional datasets (e.g., AP-omentum-ovary with 10,936 features), where scalability claims are less thoroughly validated
• The short-term vs. long-term gain trade-off experiment covers only one dataset (Airfoil) with 5 steps, limiting the generalizability of this analysis
• Hyperparameter sensitivity (entropy coefficient range 0.01–0.3, three learning rate choices) is reported but not systematically ablated

Relevance to Agentic AI / LLM Agents

HAFT is a concrete instantiation of multi-agent RL where heterogeneous agents with different architectures and roles must cooperate on a shared, evolving state—a core challenge in agentic systems. The shared critic with advantage decomposition addresses the non-stationarity problem that arises when multiple agents act sequentially on a shared environment, directly relevant to MARL foundations underlying multi-agent LLM orchestration. The attention-based architecture for handling dynamically growing action/state spaces mirrors challenges in tool-augmented or memory-augmented LLM agents where context grows over time. The framework demonstrates that principled credit assignment and role specialization (heterogeneous agents) yield more stable cooperative behavior than homogeneous or locally-coordinated approaches.