Two-Stage Feature Generation with Transformer and Reinforcement Learning

1 Introduction

Feature generation is a critical aspect of the construction of high-performance machine learning models. High-quality features can enhance the robustness, generalization capability, and interpretability of models [Zheng and Casari, 2018; Dong and Liu, 2018]. For instance, as Figure 1 shows, combining weight and height features into a Body Mass Index (BMI) [Obese, 1998] feature can significantly improve a model's ability to predict health conditions. As the scale and complexity of data increase, traditional feature engineering methods often struggle to uncover the complex patterns and underlying relationships within the data. Consequently, automatic feature generation techniques, especially those based on deep learning, have become a focal point of research [He et al., 2021]. Recently, reinforcement learning has also provided new ideas for feature generation [Liu et al., 2019; Wang et al., 2022]. Despite their ability to extract complex features, deep learning methods still face challenges in adapting to different tasks and achieving efficient optimization [LeCun et al., 2015]. This study introduces a two-stage feature generation framework that integrates pre-training with reinforcement learning-based fine-tuning to achieve more efficient and precise feature generation.

∗Corresponding author

The novelty of this framework lies in its staged training method, which enables stepwise optimization of feature generation through pre-training and reinforcement learning fine-tuning. In the first stage of pre-training, we employ an Encoder-Decoder architecture [Sutskever, 2014], optimizing the model using cross-entropy loss or mean squared error loss. The encoder extracts latent information from the raw dataset, while the decoder reconstructs this information into logically coherent sequences of feature combinations. For example, initial features might appear as [V 1, EOS, V 2, EOS, V 3, STOP], and during decoding, more intricate feature combinations like [+V 1,+V 2, EOS,-V 3,×V 1, STOP] are generated. The goal of this stage is to build a stable initial model for feature generation.

Upon completing the first stage of pre-training, the second stage incorporates the Proximal Policy Optimization (PPO) algorithm [Schulman et al., 2017] from reinforcement learning to fine-tune the model. This phase dynamically adjusts the model weights through a reward mechanism tied to specific downstream tasks, ensuring that the generated features better align with the requirements of these tasks. Specifically, the reward value of the generated feature combinations will be calculated. Based on the reward signal, the model optimizes its generation strategy, progressively learning highquality feature representations. The PPO algorithm excels at balancing the magnitude of updates, preventing issues such as gradient explosion during the training process.

This two-stage training strategy offers several advantages. In the first stage, the Encoder-Decoder structure generates features step by step and computes the loss at each step, ensuring the stability and consistency of the generated features. In the second stage, the introduction of the PPO algorithm endows the feature generation process with dynamic optimization capabilities, enhancing task adaptability.

In summary, this paper introduces a novel feature generation framework, which leverages an encoder-decoder architecture integrated with Proximal Policy Optimization (PPO). The proposed framework enhances feature generation by dynamically adjusting the optimization strategy and improving adaptability across different tasks. Using a two-stage training process, we address the challenges of feature redundancy and inefficient exploration in traditional methods. Our main contributions include:

• Encoder-Decoder Architecture for Efficient Feature Generation: We propose an encoder-decoder architecture that effectively models the feature generation process, allowing for the extraction of meaningful latent representations from raw data. This ensures high-quality feature combinations for downstream tasks.
• Two-Stage Training with PPO Fine-Tuning: We introduce a two-stage method where pre-training initializes the model, and reinforcement learning fine-tuning optimizes feature generation based on specific task requirements. This allows the framework to dynamically adjust its strategy and enhance task performance.
• Enhanced Adaptability and Efficiency: By incorporating reinforcement learning and task-specific feedback, our framework can efficiently generate features tailored to diverse datasets and models, offering better performance and computational efficiency compared to existing methods.

2 RelatedWork

Automated feature generation plays a critical role in enhancing the performance of machine learning models [Domingos, 2012]. Common methods in this field can generally be divided into expansion-reduction methods and searchbased methods, both aiming to generate meaningful and highquality features for predictive modeling.

Expansion-reduction methods generate a large pool of candidate features through various transformations and then reduce redundancy by selecting the most relevant subset. For example, ExploreKit [Katz et al., 2016] follows a three-step workflow: candidate feature generation, ranking, and evaluation. It applies transformation functions to the entire dataset and evaluates the performance of the predictive model to select the most valuable features. Similarly, FEADIS [Dor and Reich, 2012] generates new features by randomly combining original features with mathematical functions, while Autofeat [Horn et al., 2020] expands the feature space using nonlinear transformations before selecting a small subset of features for inclusion. OpenFE [Zhang et al., 2023] is an efficient expansion-reduction method that combines FeatureBoost and a two-stage pruning algorithm, enabling it to rapidly and accurately identify useful new features on large-scale datasets.

Despite their contributions, these methods often face challenges related to the computational burden and redundancy caused by the exponential growth of candidate features. As the number of features increases, selecting truly relevant ones becomes increasingly difficult, This prompts researchers to develop alternative approaches that enhance representation completeness [Wang et al., 2021].

Search-based methods adopt a more targeted method by using search algorithms to explore potential feature transformations and combinations. These methods have shown potential in overcoming the limitations of expansion-reduction techniques. For instance, GRFG [Wang et al., 2022] utilizes group reinforcement learning to select operations for groups of features and performs operations between them, effectively managing feature dependencies. TransGraph [Khurana et al., 2018] uses a transformation graph and Q-learning algorithm to generate higher-order features. DIFER [Zhu et al., 2022] introduces a feature optimizer within an encoder-predictordecoder structure that maps features to a continuous vector space, optimizes embeddings along the gradient direction induced by the predictor, and recovers better features from the optimized embeddings via the decoder.

In addition, some studies on multi-label feature selection have proposed relevance metrics and feature decomposition strategies to ensure task-specific feature quality while reducing redundancy [Zhang et al., 2021; Zhang and Gao, 2021; Li et al., 2024].

In summary, while expansion-reduction methods can create a vast array of potential features, search-based methods offer a more strategic exploration of the feature space [Komer et al., 2014], potentially leading to more effective and efficient feature engineering.

3 The ProposedMethod

3.1 Problem Settings

Given a dataset D = {F, y}, where F represents the original feature set and y represents the target label set. The feature set F = {V1, V2, . . . , VN} consists of N discrete features.

Proximal Policy Optimization Pre-Trained Model

We aim to find the optimal feature set F ∗ that maximizes the performance metric of downstream tasks, defined as:

where A represents a downstream machine learning model (e.g., random forests, SVM, or neural networks), and M denotes the evaluation metric, such as accuracy, F1-score, or mean squared error.

For continuous features, mathematical operations include "absolute value", "square", "inverse", "logarithm", "square root", "cube", "addition", "subtraction", "multiplication", and "division". For discrete features, mathematical operations include "cross" [Luo et al., 2019] and "addition".

3.2 Overall Framework

The framework combines reinforcement learning, an encoder-decoder architecture, and a feature transformation mechanism, aiming to generate optimized features to address challenges such as avoiding redundant feature generation and optimizing feature transformation sequences. The workflow is based on an encoder-decoder architecture as the main structure, a pre-training phase for initialization, and a reinforcement learning phase for fine-tuning.

At the core of the framework is an encoder-decoder architecture that can efficiently explore the feature space and generate new features. We select the Transformer architecture [Vaswani et al., 2017] as the encoder-decoder model due to its superior capability in handling sequential problems [Raffel et al., 2020]. Its self-attention mechanism effectively captures long-range dependencies and complex patterns within the data [Bahdanau, 2014]. The encoder processes the input dataset, mapping it into a latent representation that captures the structure and interdependencies of the original features. This latent representation serves as the foundation from which the decoder generates sequences of feature transformations.

The decoder is responsible for generating feature transformation sequences, which are produced step-by-step through an autoregressive mechanism [Graves, 2013a]. The output of each step depends on the result of the previous step, ensuring that the feature transformation process proceeds coherently and effectively captures the dependencies among features. To avoid generating invalid features, operations are combined with features to form a unified token. For example, operations like +V 1, -V 2, etc., are treated as tokens. The sequence output by the decoder is similar to [+V 1, +V 2, EOS, -V 3, ×V 1, STOP], where EOS represents the end-of-feature token and STOP represents the endof-sequence token. In this way, we generate two new features: V 1 + V 2 and -V 3 × V 1, thereby expanding the feature set and improving the data representation for downstream machine learning tasks.

The pre-training phase focuses on initializing the encoderdecoder model. The encoder extracts latent representations of the dataset, and the decoder generates feature transformation sequences based on these representations. These transformations generate an augmented feature set, which is optimized by evaluating its performance on a pre-trained MLP model. The optimization objective is to minimize the cross-entropy loss or mean squared error loss, thereby enhancing the quality of feature generation and its adaptability to tasks.

To encourage exploration during training, a sampling mechanism with temperature scaling is adopted [Guo et al., 2017]. This mechanism adjusts the sampling probabilities of certain transformations, enabling the model to explore a wide range of feature combinations. The pre-trained encoderdecoder provides a solid foundation for the reinforcement learning phase, accelerating convergence and enhancing performance.

After pre-training, the model is fine-tuned using the Proximal Policy Optimization (PPO) algorithm. This phase focuses on iteratively improving the feature generation process to maximize cumulative rewards. The reward function is based on improvements in the performance of the downstream task, such as increased classification accuracy or F1 score. In each iteration of the PPO phase, similar to the pretraining stage, the encoder generates a latent representation of the dataset, and the decoder proposes candidate feature transformations based on this representation. These transformations are evaluated using external performance metrics, and a reward is calculated for the generated sequence. The PPO algorithm adjusts the parameters of the encoder-decoder to maximize the expected reward, balancing exploration and exploitation.

The PPO fine-tuning enables the framework to dynamically adapt to the data, generating task-optimized feature sets. This method not only ensures that the generated features enhance performance but also maintains their interpretability and relevance to the problem domain.

3.3 Pre-Training Phase for Encoder-Decoder Model

The goal of the pre-training phase is to enable the encoderdecoder model to learn an initial strategy. A new feature set is generated by the transformation sequence, which is then processed by the pre-trained MLP to compute predictions and calculate the prediction loss based on the true labels. By minimizing the prediction loss between the generated transformation sequences and the true labels, the model learns how to enhance the representations of features. The encoder-decoder is implemented based on the Transformer architecture.

Input: The input for the pre-training phase includes a dataset Dtrain = {F, y}, where F represents the original feature set, and y is the target variable. Additionally, a validation dataset Dval is used to evaluate the model's performance and generalization ability during training.

Output: The output is a pre-trained encoder-decoder model capable of generating feature transformation sequences for creating new features.

The encoder maps the input features F into a latent space Z, where the latent representation z ∈ Z captures the structural relationships among the features. This ensures that the decoder has a comprehensive understanding of the data, allowing it to extract hierarchical feature representations. Formally, this process is expressed as:

where θe represents the parameters of the encoder. The latent space Z effectively captures the complex relationships between features, providing a compact yet informative representation for subsequent transformation.

The decoder then receives the latent representation z and predicts the transformation operation a, which is applied to the features. This process can be formalized as:

where θd represents the parameters of the decoder.

Furthermore, the decoder uses a temperature parameter T to adjust the probability distribution of the output sequence [Graves, 2013b].

where T is the temperature parameter, ai represents one of all possible actions. The temperature parameter controls the smoothness of the output probability distribution. When the temperature T is high, the generated features become more random, leading to a more exploratory method. Conversely, when T is low, the model tends to select features with higher probabilities, resulting in more deterministic outputs.

In this process, the decoder generates new feature transformation operations. These operations are applied to the original training dataset Dtrain, creating a new dataset Dnew train . The dataset Dnew train is then fed into a pre-trained MLP, which outputs the predicted probability distribution for each sample's class. Finally, the loss between these predicted probabilities and the true labels is calculated to evaluate the quality of the generated feature transformations. The parameters of the encoder-decoder model are adjusted through backpropagation to minimize the loss function. The framework inputs the data into the MLP to compute the loss after generating a complete feature (i.e., when the EOS token is generated), rather than waiting to generate the entire transformation sequence. This step-by-step evaluation provides timely feedback on the effectiveness of each generated feature, thereby promoting faster convergence during training and further improving the quality of the generated features.

For classification tasks, the cross-entropy loss is defined as:

where N is the number of training samples. yi is the true label of the i-th training sample. ˆ yi is the predicted probability for the i-th sample, obtained by the pre-trained MLP after applying the current generated features to the training dataset Dtrain.

For regression tasks, the mean squared error (MSE) loss is defined as:

where N is the number of training samples. ˆ yj is the predicted value of the j-th feature. yj is the true value of the target feature for the j-th sample.

3.4 Proximal Policy Optimization (PPO)

The Proximal Policy Optimization (PPO) fine-tuning phase is designed to refine the feature generation process by leveraging reinforcement learning. While the pre-training phase provides a strong initialization for the encoder-decoder architecture, the PPO phase ensures that the model continues to optimize feature transformations based on feedback from downstream tasks. This phase is critical for dynamically adapting to specific datasets and tasks. The final new dataset is generated at this stage.

The objective of the PPO fine-tuning phase is to iteratively improve the feature generation process by maximizing a reward function that reflects the performance improvement of downstream tasks. Unlike the pre-training phase, where training is performed after generating each feature, in the PPO phase, the model parameters are optimized based on feedback from downstream tasks after generating the complete transformation sequence. This method comprehensively evaluates the overall performance of the complete feature set and the synergistic relationships among features, ensuring that the optimization direction aligns more closely with the requirements of downstream tasks, thereby producing higher-quality feature sets with greater global coherence.

The feature generation process is modeled as a Markov Decision Process (MDP). The encoder-decoder serves as the policy network, with the state, action, and reward defined as follows:

State (s): The current state represents the feature set.

Action (a): An action corresponds to a transformation operation selected by the decoder. For example, +V1, -V2, ∗V2.

Reward (r): The reward is the performance improvement of the new dataset in the downstream task. The reward can be expressed as:

where Dnew is the dataset with the newly generated feature set.

Policy (π): The policy determines which action to apply, given the current state. PPO is used to iteratively optimize the policy to maximize cumulative rewards.

The PPO fine-tuning workflow begins with initializing the encoder-decoder model using the weights obtained during the pre-training phase. This ensures that the model starts finetuning with a well policy for feature generation.

At each iteration, the encoder generates latent representations of the current dataset, which are passed to the decoder. The decoder samples candidate transformation sequences based on the current policy πθ(a|s). The sequence is applied to the original dataset to generate a new dataset. The updated dataset is evaluated using a downstream machine learning model (e.g., a classifier or regressor). The model calculates the sequence's reward based on the performance improvement achieved by the new dataset.

Using the computed rewards, PPO adjusts the policy parameters θ to maximize the expected cumulative reward. The policy objective is defined as:

where rt(θ) = πθ(at|st) πθold (at|st) is the probability ratio, ˆ At is the advantage function, and ϵ is a clipping parameter to limit the magnitude of policy updates. To encourage the exploration Algorithm 1 Encoder-Decoder Training with PPO FineTuning

• 1: Input: Dataset D = {F, y}, where F = {V1, V2, . . . , VN} is the feature set and y is the target label.
• 2: Output: Optimized feature set F ∗ .
• 3: Pre-Training Phase:
• 4: Train the encoder-decoder model using a supervised loss function (cross-entropy or MSE).
• 5: For each feature:
• 6: Compute loss Lclassification or Lregression.
• 7: Backpropagate and update model parameters.
• 8: Store the best model with minimal validation loss as the initialization for PPO fine-tuning.
• 9: PPO Fine-Tuning Phase:
• 10: For each iteration t = 1 to Tmax:
• 11: Sample transformation operations a from the policy πθ (a|s).
• 12: Apply transformations to D to generate a new feature set Dnew .
• 13: Evaluate the performance of Dnew on downstream task using a performance metric (e.g., accuracy, F1-score).
• 14: Compute reward r based on the performance improvement:

• 15: Compute PPO objective:

16: End for PPO Iterations

of diverse transformations, an entropy regularization term is added to the PPO objective [Cai et al., 2023]. This term prevents the policy from converging prematurely to suboptimal solutions by promoting randomness in action selection. The process of sampling, evaluation, and policy updates is repeated for a predefined number of iterations or until convergence criteria are met.

The PPO fine-tuning phase offers several advantages. Stability is achieved through PPO's clipping mechanism, which ensures stable training by limiting large updates to the policy. The reward-driven method allows the model to dynamically adapt to specific datasets and tasks.

4 Experiments

4.1 Experimental Setup

DataDescription. We conducte experiments on 13 datasets from UCI [Public, 2024b], Kaggle [Howard, 2024], and OpenML [Public, 2024a], LibSVM [Lin, 2024], comprising 9 classification tasks and 4 regression tasks. Table 2 shows the statistics of the data.

Evaluation Metrics. For classification tasks, we use F1score, accuracy, and precision to evaluate the dataset. For regression tasks, we use 1-relative absolute error (1-RAE) and R2 to evaluate the dataset. The equation for 1-RAE is as fol-

lows:

y∗ i is the actual target value of the i-th observation, yi is the predicted target value of the i-th observation, and ym is the mean of all actual target values.

Baseline Methods. We compare our method with 5 widely used feature generation methods, as well as random generation and feature dimension reduction methods: (1) Base: using the original dataset without feature generation. (2) GRFG [Wang et al., 2022]: iteratively generates new features and reconstructs an interpretable feature space through groupgroup interactions. (3) DFS [Kanter and Veeramachaneni, 2015]: an expansion-reduction method that first expands and then selects feature, automatically generated features for the dataset. (4) DIFER [Zhu et al., 2022]: Performs automated feature engineering in a continuous vector space, introducing an encoder-predictor-decoder to optimize features. It optimizes embeddings along the gradient direction and recovers improved features from the optimized embeddings. (5) OpenFE [Zhang et al., 2023]: Offers expert-level automated feature generation by integrating novel enhancement methods and a two-stage pruning algorithm, effectively identifying significant features.

Hyperparameter Settings. We utilize the Adam optimizer [Kingma and Ba, 2015] to optimize the PPO algorithm, with a learning rate set to 1 × 10-4 . The clipping parameter ϵ for PPO is set to 0.2, and an entropy coefficient of 1 × 10-4 was applied to promote exploration. The number of PPO iterations is 10. The split ratios for the training set, validation set, and test set are 0.6:0.2:0.2. The encoder-decoder model incorporated 8 attention heads, with an embedding vector dimension of 128 and a model hidden layer dimension of 128. The maximum sequence length for feature transformations was set to 100. During pre-training, a temperature scaling parameter T ranging from 0.1 to 1 is used to balance exploration and exploitation.

4.2 Overall Comparison

In this experiment, we compare the performance of TSFG and baseline models for feature transformation. We evaluate classification tasks using F1-score, accuracy, and precision, and regression tasks using 1-RAE and R2 . Table 1 shows the comparison results. We can see that in most cases, TSFG performs the best. By integrating the pre-training and fine-tuning stages, TSFG can accurately capture the internal patterns of the features, thereby identifying the optimal feature space. In conclusion, this experiment demonstrates the effectiveness of TSFG in feature transformation.

4.3 Ablation Study

To evaluate the contribution of each component in our framework, we perform ablation experiments. This experiment aims to verify whether each component of our method indeed has a positive impact on the final results. Therefore, we have developed two variants: (1) No Pre-Training: The model is trained without the pre-training phase. This variant is referred to as "TSFG+". (2) No PPO Fine-Tuning: The model is evaluated without reinforcement learning fine-tuning. This variant is referred to as "TSFG#".

As shown in Figure 3, the results indicate that pre-training and PPO fine-tuning contribute significantly to the overall performance.

4.4 Runtime Comparison

As shown in Table 3, we compare the runtime efficiency of the proposed framework with baseline methods. Despite incorporating reinforcement learning, our method achieves competitive runtime efficiency due to the use of pretraining. Compared with GRFG and DIFER, both of which use deep learning methods, the proposed method strikes a balance between computational cost and performance improvement.

4.5 Comparison on Different DownstreamTasks

To evaluate the adaptability and effectiveness of the proposed framework, we compare its performance across different downstream machine learning models, including Random Forest, XGBoost, Support Vector Machines (SVM), and CatBoost (CAT). The experiment is conducted on the diabetes dataset.

As shown in Table 4, the proposed framework performs well across different downstreammodels. This highlights the framework's ability to generate high-quality features whose effectiveness is not dependent on the complexity of the task or the choice of downstream model. Furthermore, it demonstrates the framework's adaptability to various modeling requirements and its potential for broad applications in the field of feature generation.

4.6 Effectiveness of Generated Features

We evaluate the quality of the generated features by analyzing their importance in downstream models. Feature importance scores from tree-based models (e.g., Random Forest) are used to measure the contribution of generated features. Figure 4 shows that the generated features have high importance scores, often surpassing original features, indicating their effectiveness in capturing complex relationships in data.

5 Conclusion

In this paper, we propose a novel automated feature generation framework that integrates reinforcement learning with a Transformer-based encoder-decoder architecture. The framework addresses the challenges of feature redundancy, inefficient exploration of the feature space, and adaptability to diverse datasets and tasks. Our method dynamically generates high-quality features, enhancing the performance of downstream machine learning models. Extensive experiments on various datasets validate the effectiveness of the proposed framework. The framework provides a solid foundation for advancing automated feature generation.

Acknowledgements

This work was supported by the Science Foundation of Jilin Province of China under Grant YDZJ202501ZYTS286, and in part by Changchun Science and Technology Bureau Project under Grant 23YQ05.