Agentic Feature Augmentation: Unifying Selection and Generation with Teaming, Planning, and Memories

TL;DR

This paper introduces a unified agentic framework for feature engineering that combines selection and generation through a multi-agent system with reinforcement learning, improving AI model performance by balancing redundancy reduction and meaningful feature addition.

Contribution

It proposes the MAGS system, integrating selection, generation, and strategic coordination with reinforcement learning for the first time in feature engineering.

Findings

Achieves superior task performance across multiple benchmarks.

Effectively balances feature redundancy reduction and meaningful feature augmentation.

Demonstrates the benefits of agentic teaming and planning in feature engineering.

Abstract

As a widely-used and practical tool, feature engineering transforms raw data into discriminative features to advance AI model performance. However, existing methods usually apply feature selection and generation separately, failing to strive a balance between reducing redundancy and adding meaningful dimensions. To fill this gap, we propose an agentic feature augmentation concept, where the unification of feature generation and selection is modeled as agentic teaming and planning. Specifically, we develop a Multi-Agent System with Long and Short-Term Memory (MAGS), comprising a selector agent to eliminate redundant features, a generator agent to produce informative new dimensions, and a router agent that strategically coordinates their actions. We leverage in-context learning with short-term memory for immediate feedback refinement and long-term memory for globally optimal guidance. Additionally, we employ offline Proximal Policy Optimization (PPO) reinforcement fine-tuning to train the router agent for effective decision-making to navigate a vast discrete feature space. Extensive experiments demonstrate that this unified agentic framework consistently achieves superior task performance by intelligently orchestrating feature selection and generation.