Retrieval-Augmented Feature Generation for Domain-Specific Classification

Feature Engineering. Feature engineering is an essential part of machine learning, including selecting, modifying or creating new features from raw data to improve the task performance [18]. The target of this process is to optimize the feature representation space. In this process, we tailor the best and most interpretable feature set for specific machine learning tasks for better model accuracy and generalization. Among feature engineering, feature generation is the approach that generates new features from existing ones in a dataset. This is usually through various mathematical or logical transformation operations [5, 20]. The goal of feature generation is to create complex latent feature spaces [28, 17]. However, existing feature generation methods often lack transparency and require a large amount of manual operations.

Retrieval Augmented Generation (RAG) [13] is a technique that integrates the information retrieval capabilities with the generative language models to enhance the performance of the final output. RAG effectively assists LLMs with tasks requiring extensive and specific domain knowledge [26, 9]. Given a domain-specific task, the LLMs can access a large external library, which might be a set of documents or knowledge related to the task [7]. Then, the LLMs generate a query according to the task information and use it for searching based on the similarity between the query and the candidate documents. This approach helps the LLMs to produce responses that are more accurate, domain-relevant, and reduce hallucinations [25]. However, existing RAG for feature generation lacks a reasoning process [22, 26, 13]. To generate interpretable features, the LLM needs to analyze the relevance and rationality of the retrieved information and the generated features. The reasoning process ensures the effectiveness of the new features by validating their effectiveness and explainability in the domain knowledge.

The set of documents $\mathcal{R}=\{\mathcal{R}_{k}\},k=1,2,\ldots$ consists of documents $\mathcal{R}_{k}$, which contains the information of potential features $\{g_{k}\},k=1,2,\ldots$ that can be accessed through specific calculations, combinations or judgments with existing features in $\mathcal{F}_{t-1}$:

• •

  Calculation involves arithmetic transformations, statistical measures, or algorithmic functions applied to derive $g_{t}$. In our case, we define a calculator for features as a function $\mathit{o}_{c}:\mathbb{R}^{n}\rightarrow\mathbb{R}$ applied to a vector of existing features, resulting in a scalar or vector output as the new feature.

• •

  Combination involves integrating two or more features into a single new feature through methods such as concatenation, averaging, or more complex fusion techniques. In our case, we define a calculator for combination as a function $\mathit{o}_{f}:\mathbb{R}^{n}\times\mathbb{R}^{m}\rightarrow\mathbb{R}^{p}$, where $n$, $m$, and $p$ are dimensions of the input features and the resulting feature, respectively.

• •

  Judgment involves logical or rule-based decision-making processes that integrate domain knowledge or empirical rules to form a new feature. In our case, we define a calculator for judgment as a decision function $\mathit{o}_{d}:\mathcal{X}\rightarrow\{0,1\}$ where $\mathcal{X}$ is the input space composed of one or more feature vectors, and the output is a new categorical feature based on predefined criteria.

In this stage, our objective is to generate a new feature and align it with the existing features in the data table. Now we have introduced external knowledge about features through RAG. This knowledge includes the deeper structural details and relationships between features, which are not directly reflected in the original dataset. Here we adopt the LLM to analyze and extract feature structures and relationships, and further generate new features to expand the feature space.

From the $k$ documents with potential features $\{g_{k}\}$, the LLM first assesses these potential features and then ranks them based on their potential to enhance the performance of downstream tasks according to LLM’s evaluation. Here, to help the model’s evaluation, we employ a “Chain of Thought” [24] strategy to encourage the LLM to think and formulate logical hypotheses and reasoning for the possible outcomes of integrating each potential feature into the data table step-by-step.

The LLM then selects the most promising document $\mathcal{R}_{t}\in\{\mathcal{R}_{k}\}$ from these candidates, and generate a new feature $g_{t}$ according to $\mathcal{R}_{t}$. We integrate $g_{t}$ into the data table:

Here, the $g_{t}$ is generated by the corresponding calculation, combination, or judgment method $o_{t}$ mentioned in document $\mathcal{R}_{t}$ with existing features and is of the same length as $f_{i}$.

Finally, we feed $\mathcal{D}_{t}$ into the downstream task and obtain the performance metric $\mathcal{P}_{t}$, which is then fed back to the LLM for performance evaluation.

In this stage, our objective is to test the effectiveness of the newly generated feature and decide whether to keep it. We decide whether to keep the feature be the improvement of $\mathcal{P}_{t}-\mathcal{P}_{t-1}$, where a higher value of $\mathcal{P}$ indicates better performance:

When the adding of new feature $g_{t}$ brings an improvement in the downstream task performance of the data table, we formally adopt $g_{t}$ as a new feature within the table. The LLM updates the dataset text information $\mathcal{C}=\mathcal{C}_{t}=p_{\phi}(\mathcal{C}_{t-1}|\mathcal{R}_{t})$ to incorporate the information related to $g_{t}$. The feature generation process expands the dimensionality of the feature space. As the feature space expands, the model can understand and interpret the data more comprehensively, further improving the performance of downstream tasks and enhancing generalization capabilities.

After that, we move forward to generate new queries to search for more potential features with the updated data table $\mathcal{D}_{t}$ and the data description $\mathcal{C}_{t}$. This iteration process continues until a predefined maximum number of iteration $T$ is reached, or the optimal feature set is found when achieving peak performance in the downstream task.

Datasets. We evaluate the RAFG method on four real-world datasets, including Parkinson’s Disease Classification (PDC) [16], Animal Information Dataset (AID) [3], Global Country Information Dataset 2023 (GCI) [1] and Diabetes Health Indicators Dataset (DIA) [21]. The detailed information is shown in Table 1.

Overall Performance. In Table 3 we show the overall results of RAFG across four domains-specific datasets.

(1) Compared with baseline methods, RAFG consistently achieves superior performance. For instance, in Random Forest, RAFG shows a significant improvement in accuracy across all datasets and has a $19.2\%$ increase on the AID dataset than raw data, and a $3.8\%$ increase than CAAFE, which is the second-highest accuracy baseline. The RAFG method shows particular strength in boosting model accuracy. The performance demonstrates RAFG’s enriching of feature space with domain-specific knowledge enhances the model’s understanding and boosts overall performance.

(2) Across different models and metrics, RAFG consistently outperforms other methods in all metrics. Notably, in GCI dataset under DT, RAFG achieves a precision of $66.9\%$, surpassing the highest precision of $60.7\%$. Moreover, for F1 scores, RAFG raises the score in the GCI under DT from $55.6\%$ to $68.1\%$. These results demonstrate the RAFG’s capability of reducing misclassification effectively.

(3) The generation of new features through RAFG significantly enhances model performance. As in the results, the addition of features, e.g., from 756 to 770 in the PDC dataset, leads to substantial increases in performance metrics, including accuracy and F1. This result suggests that the new features not only numerically expand the feature set, but also enhance the information richness of the dataset. In this way the models can capture more complex patterns and relationships within the data.

Searching relevant documents from external knowledge library effectively assists RAFG in acquiring useful information for feature extraction. In Figure 3, we compare the performance of RAFG with and without the use of an external database ("LLM Only" in the figure) across different datasets under the DT. The results demonstrate that incorporating external databases consistently enhances the accuracy of RAFG across all datasets.

In Table 4, we demonstrate the feature information and data description before and after RAFG, and in Table 5 we showcase the detailed information of newly generated features in the order of generation along with the functions and meanings. We first compare the changes in accuracy for downstream tasks and information gain when generating new features. Figure 5 shows that each new feature improves the metrics for downstream tasks and adds extra information to the dataset. The RAFG increases the dimensionality of the dataset’s feature space by incorporating external information. Thus, the model in a downstream task can better capture the complex relationships between features to improve task performance.