Optimized Feature Generation for Tabular Data via LLMs with Decision Tree Reasoning

We start by framing column feature generation as an optimization problem for the rule generator, where the rule defines a mapping from the original set of features to a new feature. In this regard, our objective is to generate an informative one-dimensional column feature, $\mathcal{X}^{\prime}$, by optimizing the rule $r:\mathcal{X}\to\mathcal{X}^{\prime}$, i.e., a novel column feature that enhances the performance of the prediction model when trained with the new feature, $f:\mathcal{X}\oplus\mathcal{X}^{\prime}\rightarrow\mathcal{Y}$. Formally, based on the generated rule w.r.t to the training dataset $r:=g(\mathcal{D}_{\mathtt{train}})$, our optimization problem is as follows:

where $\mathcal{D}\oplus r:=\{\mathbf{x}_{i}\oplus r(\mathbf{x}_{i}),\mathbf{y}_{i}\}_{i=1}^{N}$, $\oplus$ denotes concatenation of the column features, $N$ is the number of samples in the dataset $\mathcal{D}\subseteq\mathcal{X}\times\mathcal{Y}$, and $g$ is the rule generator (i.e., LLM $\mathcal{M}$ in our case) applied to the training dataset $\mathcal{D}_{\mathtt{train}}$. $\mathcal{L}_{f}$ is an objective function evaluated using the prediction model $f$, such as the mean absolute error measured with a decision tree regressor for a regression task. In short, we optimize a rule to achieve the best validation score measured with $\mathcal{D}_{\mathtt{val}}\oplus r$ on $f^{*}$.

However, such a bi-level optimization is often non-differentiable (or extremely compute-heavy) as it requires computing gradients through the optimization process of $f^{*}$, making it challenging to learn the rule. Moreover, the generated rule may also require non-differentiable operations such as logical conjunction between categorical features (e.g., ‘$\mathtt{Is\ a\ Smoker}$ = $\mathtt{Has\ Fever}\ \wedge\ \mathtt{Has\ Difficulty\ Breathing}$’ in Figure 1). One way to tackle this issue is to use black-box optimizations, such as evolutionary strategies [39] or reinforcement learning [40]. However, these approaches still face several limitations. For instance, they require a manually pre-defined search space for the operations to be used (and it becomes computationally heavy as the search space grows) and may overlook the effective experimental design of feedback for improvement.

To tackle this issue, we propose using LLMs as optimizers to iteratively improve the rule by prompting them with a trajectory of feedback, i.e., the history of previously generated rules along with their corresponding validation scores and reasoning information. This optimization via prompting method [25] has emerged as an efficient and effective tool for black-box optimization. Furthermore, LLMs offer several advantages, e.g., leveraging the semantics of column and feature values for better optimization and the flexibility to operate without a restricted pre-defined search space for rules.

We now describe the core algorithm. First, we prompt the LLM to propose the name of a new column to generate based on the task description, e.g., ‘$\mathtt{Smoking}$ $\mathtt{Status}$’ in Figure 1. We then compute the validation score of the prediction model and extract reasoning information from the decision tree fitted with the training dataset to initialize the optimization trajectory, which is used as feedback for the LLM. Here, this novel decision tree reasoning effectively conveys knowledge of past experiments, which further facilitates the optimization process. Finally, we continue updating the trajectory as the optimization progresses, allowing the LLM to iteratively improve the rule.

Column name generation. OCTree first generates the column name $c_{\mathtt{new}}$ of a novel feature. OCTree achieves this by prompting the LLM $\mathcal{M}$ with a prompt $p_{\mathtt{col}}$ to recommend a new column name that would be useful for predicting the target (see Appendix A.1): $c_{\mathtt{new}}=\mathcal{M}(p_{\mathtt{col}}(\mathcal{C},c_{\mathtt{target}}))$. With their language understanding capabilities, LLMs generate reasonable column names. For example, the LLM might propose using trading volume as a new feature for predicting stock prices.

Initialize optimization and extract reasoning. OCTree then generates the initial rule $r_{0}$ for deriving a novel column feature from the original column names $\mathcal{C}$. OCTree achieves this by prompting the LLM $\mathcal{M}$ with $p_{\mathtt{init}}$ to propose a rule to predict $c_{\mathtt{new}}$ with $c_{1},\dots,c_{M}$ (see Appendix A.2): $r_{0}=\mathcal{M}(p_{\mathtt{init}}(\mathcal{C},c_{\mathtt{new}}))$. Then, we obtain the initial score $s_{0}$ evaluated with $f^{*}$: $s_{0}=\mathcal{L}_{f^{*}}(\mathcal{D}_{\mathtt{val}}\oplus r_{0})$. Also, we extract the decision tree reasoning $d_{0}$ with CART [41] fitted on training dataset:

CART is a binary tree that recursively splits data based on criteria, such as Gini impurity, to predict outcomes. We use CART because (i) tree-based algorithms, which typically use ensembles of simple decision trees like CART, still outperform deep learning on tabular prediction tasks, and (ii) it is more interpretable than tree ensembles and can be easily expressed in natural language. For example, as illustrated in Figure 1, CART can expressed in natural language using the if-else syntax. Intuitively, the decision tree reasoning extracted by CART provides valuable insights learned from the entire training dataset. The reasoning explicitly shows the columns that are considered more significant (as nodes in the tree) and the corresponding values (as thresholds of the nodes) used for prediction.

Optimization with decision tree reasoning. To optimize the rule, we describe the task in natural language and provide the trajectory $\mathcal{T}_{t}=\{(s_{i},d_{i},r_{i})\}_{i=0}^{t}$, which includes a history of previously proposed rules with the corresponding scores and reasonings. Specifically, we generate a new rule $r_{t+1}$ with the LLM $\mathcal{M}$ using the prompt $p_{\mathtt{gen}}$ (see Appendix A.3). The prompt $p_{\mathtt{gen}}$ asks $\mathcal{M}$ to generate a new rule that is not in $\mathcal{T}_{t}$ that would give a better score than the scores in $\mathcal{T}_{t}$: $r_{t+1}=\mathcal{M}(p_{\mathtt{gen}}(\mathcal{T}_{t},\mathcal{C},c_{\mathtt{target}}))$. Optimization trajectories are presented in ascending order of score since LLM is more likely to generate entries similar to those that appear later in the list [42, 25]. Finally, we append the score $s_{t+1}=\mathcal{L}_{f^{*}}(\mathcal{D}_{\mathtt{val}}\oplus r_{t+1})$, decision tree reasoning $d_{t+1}=\mathtt{CART}(\mathcal{D}_{\mathtt{train}}\oplus r_{t+1})$, and rule $r_{t+1}$ to the optimization trajectory $\mathcal{T}_{t}:\mathcal{T}_{t+1}=\mathcal{T}_{t}\cup\{(s_{t+1},d_{t+1},r_{t+1})\}$. We then optimize rule $r$ with a fixed number of iterations and select the rule with the highest validation score.

Generating multiple features. The optimization steps can be repeated in order to generate multiple useful features. For example, after generating additional column ‘Smoking Status,’ ‘Physical Activity Level’ can be generated based on the original features and ‘Smoking Status.’ Formally, we first generate a new column $\mathcal{X}^{\prime}=r_{\mathtt{opt}}(\mathcal{X})$, where $r_{\mathtt{opt}}$ is the optimized rule for generating a column feature $\mathcal{X}^{\prime}$. Then, one can generate a new input space $\mathcal{X}^{\mathtt{new}}=\mathcal{X}\oplus\mathcal{X}^{\prime}$. Using $\mathcal{D}^{\mathtt{new}}\subseteq\mathcal{X}^{\mathtt{new}}\times\mathcal{Y}$, OCTree iteratively performs the optimization process to generate new column features again. We sequentially introduce multiple new column features until the validation score stops improving.

Results on datasets with language descriptions. We first describe the details of the experiments when clear language descriptions are provided (e.g., column names and categorical features). In these cases, the LLM generates a logical rule in natural language. Since the logical rule is easily converted to Python code, we prompt the LLM to convert it. In this case, we use gpt-4o with a temperature of 0.0 (since the Llama 2 model often fails to do this; see Appendix A.4 for used prompt).

As shown in Table 1, our framework consistently improves various types of baseline models. For instance, when we generate column ‘$\mathtt{Trading}$ $\mathtt{Volume}$’ for the Tesla Stock dataset using Llama 2 during training of the XGBoost, the relative error decreases by 15.9%. Additionally, our framework is compatible with any kind of LLM. In particular, while Llama 2, the open-source model, improves the baseline models with a high margin, one can even improve with the advanced model. For example, using GPT-4o, the most recent LLM by OpenAI, our method reduces the relative error by 17.1% on the Tesla Stock dataset for the XGBoost. Remark that we select recently released datasets (e.g., Enefit is a recent Kaggle competition dataset) to curate datasets that LLMs would not have seen.

Results on datasets without language descriptions. In practice, language descriptions that clearly describe the given prediction task are not always available for all datasets. For example, feature names and values are often changed to arbitrary meaningless symbols in many financial and medical datasets to protect confidentiality [46]. Even though the language description does not exist, our framework can be easily extended to this kind of case by using arithmetic rules as feature generators. This is because the superiority of our framework comes from the optimization capability of LLMs [24, 25], using decision tree reasoning as explicit feedback, rather than utilizing language description.

Specifically, to ensure that the datasets do not contain language descriptions, we use ordinal encoding for categorical features and then normalized all features with a min-max scaler to convert the original numeric values to different values. In addition, we use generic indicators for the column names. For example, if the number of columns is $M=5$, we use $\mathcal{C}=\{\text{`}\mathtt{x1}\text{'},\text{`}\mathtt{x2}\text{'},\dots,\text{`}\mathtt{x5}\text{'}\}$ for the column names. For the initial rule, we simply use the product of the two most important features measured by XGBoost, e.g., $x_{6}=x_{1}\times\ x_{5}$. The generated rule is, therefore, in the form of an arithmetic formula.

As shown in Table 2, our framework improves the baseline models even without language descriptions. Specifically, our framework reduces errors by an average of 5.0% over the XGBoost classifier. We hypothesize that our method successfully finds good feature generation rules since LLM automatically generates informative new column features in a black-box manner from the large search space of arithmetic rules since our method does not need to manually specify the search space of the rules.

We experiment with several open LLMs to evaluate their performance when employed as the rule generator in our framework. As shown in Table 3, while all of these models result in improvements over the baseline on the datasets without language descriptions, we find our own model (i.e., Llama 2 7B fine-tuned on UltraChat) to be particularly effective. This demonstrates that our framework can be implemented effectively even with open models at a moderate scale, provided that they have strong chat capabilities. We suspect that this is due to the enhanced ability of such models to understand and generate contextually relevant and coherent rules, which leads to improved optimization outcomes. Furthermore, Code Llama also demonstrates competitive performance; it is trained on code datasets, which include a variety of arithmetic rules, therefore beneficial for suggesting rules for datasets lacking language descriptions. This suggests that incorporating code datasets alongside dialogue datasets may further enhance the performance of the rule generator LLM.

Incorporating previous automatic feature engineering methods. As OCTree is orthogonal to existing automatic feature engineering methods, our method naturally integrates with them. The simplest approach is to first generate features with OCTree and subsequently employ these existing methods to further enhance the feature set. In Table 4, we first compare the performance of using OCTree alone. Here, we compare using XGBoost and MLP to show that our framework outperforms the others for both tree-based and deep learning models. As can be seen from the table, our method outperforms state-of-the-art automatic feature engineering methods (i.e., AutoFeat [23] and OpenFE [17]). We exclude CAAFE [19] from our comparison since this method relies on a language-based context of the dataset, making it inapplicable to the datasets in Table 2. In contrast, other methods, including ours, are context-agnostic and can also be applied without any description (but our method benefits from clear language descriptions as shown in Table 1 if it is available). Furthermore, we highlight that using our method in combination with OpenFE (i.e., OCTree^† in Table 4) further improves the performance, resulting in a 7.9% reduction in relative error for XGBoost.

Ablation study of the proposed components. Our framework consists of two main components: (i) generating new additional column features (Gen. Feat.), and (ii) providing explicit decision tree reasonings as feedback (DT reasoning) during the optimization process. As shown in Table 5, both components are crucial to our framework. First of all, the rules for introducing new column features are successfully optimized even without using explicit decision trees for feedback (i.e., just providing a score as feedback to the LLM), therefore resulting in performance improvements. Secondly, one can get even better performance by providing the decision tree as feedback to the LLM. We hypothesize that providing a decision tree is actually the same as providing learned information from the entire training samples. Also, decision trees include information about important columns and their threshold values used for predictions in their nodes. In addition, deep learning models often lack interpretability (i.e., non-trivial to represent them in prompts), whereas decision trees are used in prompts for LLMs because they are easily converted to natural language using if-else syntax.

Transferring generated rules to other prediction models. While we optimize feature generation rules to improve the performance of a certain single model, it is able to use these generated features in other models to improve their performance. For example, it is highly practical to generate features by evaluating with XGBoost, which has a relatively short evaluation time compared to large deep neural networks, and then use the generated features to improve large models. To verify such transferability of our framework, we first optimize the column generation rules by evaluating with XGBoost, and then train MLP and HyperFast along with the generated features. As shown in Table 6, the generated features are useful for improving the performance of MLP and HyperFast, and show high practicality when features need to be generated in a given amount of time.

Examination on LLM for column generation. During the initialization, LLM recommends a new column feature that is not considered in the original dataset. Here, we conduct an experiment to verify whether LLM actually introduces the most beneficial feature that can improve prediction performance. We perform an analysis from two perspectives: (i) whether LLM can distinguish column features that are more relevant to the target task when given candidates for new column names, and (ii) whether it is actually beneficial to use the real values of the column features if they are obtainable.

Our motivation is that LLM can understand the relationship between the target task and the column feature so that LLM can introduce the new columns needed for the task. To verify this, we first verify that LLM can distinguish column features that are more important for the task. For the experiment, we first corrupt the dataset by removing two existing features. We then prompt LLM to rank these two features (by providing as candidates) in order of most informative to the target task. Finally, we compare the performance of adding each feature to the corrupted dataset. As shown in Table 7, LLM indeed distinguishes more informative column features. Here, in the Disease dataset, we remove ‘$\mathtt{Cholesterol}$ $\mathtt{Level}$’ and ‘$\mathtt{Cough}$’ from the original dataset. We then ask LLM what attribute is more important in predicting whether a patient has a disease. Of the two, LLM answers ‘$\mathtt{Cough}$,’ which shows more performance improvement than using ‘$\mathtt{Cholesterol}$ $\mathtt{Level}$’ as an additional feature. Therefore, generated column features from LLM are beneficial for the prediction.

Using this ability of LLM, we propose to generate a new column through LLM, which practitioners actually need for their target task, even without candidates. For example, on the Clinical Trial dataset, we verify that LLM indeed generates useful column features, which are beneficial for the target task (i.e., predicting the patient’s mortality). During the column name generation step, LLM introduces ‘$\mathtt{Age}$’ as an additional column name. First, we get the actual ages of patients for US National Library of Medicine. We then measure the performance gain by using the actual values obtained from the real world. As shown in Figure 3, imputing the generated column with real values shows an improvement, validating that OCTree actually recommends columns that are useful for the target task. Therefore, we suggest practitioners utilize OCTree (i) to identify additional column features that should be collected or, if this is not possible, (ii) optimize a column feature generation rule.