# Two-Stage Feature Generation with Transformer and Reinforcement Learning

arXiv: https://arxiv.org/abs/2505.21978 · HTML: https://mlsanity.com/paper/2505.21978/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL with diffusion and flow models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 16 of its 36 references that are in the map.
- First posted on arXiv 28 May 2025.

## Placed in

- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [AutoML: A Survey of the State-of-the-Art](https://mlsanity.com/paper/1908.00709/) (2019)
- [Feature Engineering for Predictive Modeling using Reinforcement Learning](https://mlsanity.com/paper/1709.07150/) (2017)
- [The autofeat Python Library for Automated Feature Engineering and Selection](https://mlsanity.com/paper/1901.07329/) (2019)
- [OpenFE: Automated Feature Generation with Expert-level Performance](https://mlsanity.com/paper/2211.12507/) (2022)
- [AutoCross: Automatic Feature Crossing for Tabular Data in Real-World Applications](https://mlsanity.com/paper/1904.12857/) (2019)
- [Group-wise Reinforcement Feature Generation for Optimal and Explainable Representation Space Reconstruction](https://mlsanity.com/paper/2205.14526/) (2022)
- [Generative Partial Multi-View Clustering](https://mlsanity.com/paper/2003.13088/) (2020)
- [DIFER: Differentiable Automated Feature Engineering](https://mlsanity.com/paper/2010.08784/) (2020)
- [A Survey on Deep Reinforcement Learning for Data Processing and Analytics](https://mlsanity.com/paper/2108.04526/) (2021)

## Abstract (arXiv)

Feature generation is a critical step in machine learning, aiming to enhance model performance by capturing complex relationships within the data and generating meaningful new features. Traditional feature generation methods heavily rely on domain expertise and manual intervention, making the process labor-intensive and challenging to adapt to different scenarios. Although automated feature generation techniques address these issues to some extent, they often face challenges such as feature redundancy, inefficiency in feature space exploration, and limited adaptability to diverse datasets and tasks. To address these problems, we propose a Two-Stage Feature Generation (TSFG) framework, which integrates a Transformer-based encoder-decoder architecture with Proximal Policy Optimization (PPO). The encoder-decoder model in TSFG leverages the Transformer's self-attention mechanism to efficiently represent and transform features, capturing complex dependencies within the data. PPO further enhances TSFG by dynamically adjusting the feature generation strategy based on task-specific feedback, optimizing the process for improved performance and adaptability. TSFG dynamically generates high-quality feature sets, significantly improving the predictive performance of machine learning models. Experimental results demonstrate that TSFG outperforms existing state-of-the-art methods in terms of feature quality and adaptability.