Dual-Agent Reinforcement Learning for Automated Feature Generation

TL;DR

This work tackles automated feature generation for tabular data by addressing redundancy, state representation, and discrete/continuous feature interactions. It introduces DARL, a dual-agent reinforcement learning framework that pairs a feature-generation agent with a feature-discrimination agent, augmented by a Transformer-based state embedding and operation differentiations for feature types. The agents optimize downstream-task performance via mutual-information-informed rewards and Deep Q-Network training, iterating over multiple rounds to build a compact, high-quality feature set. Empirical results on 21 datasets show DARL consistently surpasses baselines, with robust performance across various downstream models and improved efficiency, highlighting practical impact for scalable, data-driven feature engineering.

Abstract

Feature generation involves creating new features from raw data to capture complex relationships among the original features, improving model robustness and machine learning performance. Current methods using reinforcement learning for feature generation have made feature exploration more flexible and efficient. However, several challenges remain: first, during feature expansion, a large number of redundant features are generated. When removing them, current methods only retain the best features each round, neglecting those that perform poorly initially but could improve later. Second, the state representation used by current methods fails to fully capture complex feature relationships. Third, there are significant differences between discrete and continuous features in tabular data, requiring different operations for each type. To address these challenges, we propose a novel dual-agent reinforcement learning method for feature generation. Two agents are designed: the first generates new features, and the second determines whether they should be preserved. A self-attention mechanism enhances state representation, and diverse operations distinguish interactions between discrete and continuous features. The experimental results on multiple datasets demonstrate that the proposed method is effective. The code is available at https://github.com/extess0/DARL.

Figures (4)

• Figure 1: After data preprocessing, the dataset is transformed using a feature generation algorithm, resulting in a new dataset that significantly improved the F1-score in downstream machine learning tasks compared to the original dataset.
• Figure 2: Overview of DARL. The dataset is transformed into feature embedding representations through a self-attention mechanism. Subsequently, a feature generation agent produces a sequence of operators, and a feature discrimination agent generates a discriminator sequence. These two sequences are combined with the original feature set to generate a new feature set. The updated feature set is then input into downstream tasks for evaluation, and the results are fed back to the two agents. This process iterates until the best feature set is discovered or the maximum number of iterations is reached.
• Figure 3: Time comparison of different algorithms.
• Figure 4: The convergence of the algorithm.