LEKA: LLM-Enhanced Knowledge Augmentation

1 Introduction

Humans are good at identifying relevant sources of knowledge. This is an ability that is rooted in our capability for analogical reasoning and knowledge management. In contrast, artificial intelligence models do not inherently possess this intuition: they require explicit instructions and systematic training to identify and utilize relevant information. This gap presents a challenge in enriching domain knowledge and enhancing data augmentation.

∗Corresponding author.

Knowledge augmentation is vital for improving model performance, especially in domains with limited information or complex data structures [Tang et al., 2020]. Knowledge transfer is a crucial method within knowledge augmentation for improving learning performance by transferring knowledge from external information sources [Khodaee et al., 2024]. In data-limited scenarios, effective knowledge sourcing can bridge domain information gaps and enhance model robustness using relevant external information. In these scenarios, knowledge augmentation, especially knowledge transfer, can reduce reliance on extensive target domain data by strategically selecting sources matching target needs.

Despite its potential, traditional knowledge augmentation methods often involve manual intervention [Ringwald and Stiefelhagen, 2021; Nam et al., 2024; Zhang et al., 2024d], where human experts select source domains based on their subjective interpretation of the target domain's requirements. Using tabular learning as an example highlights several challenges: (1) structural and format differences across domains hinder data alignment and integration; (2) discrepancies in tasks and content between the selected source and target domains can diminish the success of knowledge augmentation; and (3) extensive data preprocessing is often required to properly match the chosen source domain dataset. Additionally, relying on human expertise for source domain selection can lead to biases and inefficiencies, as these choices are typically made based on prior knowledge rather than a rigorous, datadriven analysis.

A natural idea is to construct an automated search for datasets with relevant knowledge in a database that has a similar structure to the target data, such as utilizing RetrievalAugmented Generation (RAG) [Gao et al., 2023; Lewis et al., 2020]. RAG combines retrieval systems with generative language models, and it enhances the model's capabilities by providing access to an external library [Guo et al., 2017; Zhang et al., 2025b; Shi et al., 2018]. For dataset retrieval, some existing works [Fleischer et al., 2024; Siriwardhana et al., 2023] use the retrieved documents as contextual information for generation and significantly enrich the model's knowledge base.

However, this direct method is often impractical due to the high costs of embedding entire domains or datasets [Seemakhupt et al., 2024; Jin et al., 2024]. Additionally, retrieved knowledge may not align effectively with the intended application [Edge et al., 2024; Li and Ramakrishnan, 2025]. Such misalignments typically result from variations in data distribution, feature spaces, or contextual differences between source and target data. Furthermore, retrieved data quality, like noise or incomplete features, can also hinder its effectiveness for augmenting data via knowledge augmentation. Thus, in such situations, identifying the optimal knowledge source providing the most relevant, high-quality data tailored to specific target data needs via LLM capabilities remains a substantial challenge.

Our Targets. We aim to address three main challenges in retrieving source domain data: (1) how to effectively extract key information from the target domain with limited computational cost; (2) how to design an efficient, automatic source domain retrieval for knowledge augmentation; and (3) how to automatically harmonize the retrieved source data with the target domain to optimize transfer efficiency and improve learning outcomes.

Our Approach. To address the challenges above, we design LLM-Enhanced Knowledge Augmentation (LEKA), a novel and automated data retrieval and augmentation method by knowledge augmentation. Specifically, (1) we utilize an LLM to extract the key textual information of the target domain; (2) we deploy dataset-RAG in an external database to efficiently extract knowledge relevant to the target domain dataset. The RAG libraries can contain large amounts of continuously updated datasets, ensuring that the data retrieved from these libraries is more precise and timely; (3) then the LLM automatically harmonizes the retrieved source data with the target domain in feature space and marginal probability measures to enhance downstream task learning performance. The data harmonization reduces the structural and semantic differences between the source and target data. This data harmonization is crucial for reducing domain shift, which optimizes transfer efficiency and improves the overall learning outcomes of the model.

As shown in Figure 1, consider a target domain focused on predicting a rare cancer, so direct learning is challenging due to insufficient sample information. To augment the data with external information would typically require human experts with knowledge of both cancer physiological indicators and machine learning to retrieve datasets. In contrast, our method can automate the retrieval of a common cancer dataset as the knowledge source with extensive feature information. This retrieved dataset possesses a more complete and transparent feature space with a structure similar to the target domain. The LLM then adjusts the retrieved data based on the characteristics of the target domain. In this way, we achieve a highly harmonized source dataset for knowledge augmentation.

In summary, our contribution includes:

1. We introduce a novel automated data augmentation method, LEKA, which utilizes an automated retrieval method for external data and harmonizes it with target data for knowledge augmentation.
2. We develop a novel paradigm by incorporating an LLM into the data harmonization process to optimize data
3. space and structure to enhance downstream machine learning performance.
4. We conduct a series of experiments to validate the effectiveness and robustness of our LEKA method across different tasks. Experimental results demonstrate that our method has clear advantages over existing methods.

2 RelatedWork

2.1 Knowledge Transfer

Knowledge transfer is a knowledge augmentation method that improves learning on new tasks by transferring knowledge from a related task [Alyafeai et al., 2020; Wang et al., 2022; Wang et al., 2025]. Knowledge transfer allows cross-domain knowledge transformation despite data distribution or feature space differences. Specifically, knowledge transfer adapts models developed for one task to perform better on different but related tasks, as it adjusts the feature mappings and decision boundaries to suit new knowledge [Han et al., 2021; Yordanov et al., 2021]. Finding suitable pre-training knowledge is crucial for knowledge transfer because it significantly enhances the model's effectiveness by providing a relevant starting point. While many methods exist to adjust existing knowledge within the same or across different domains to improve transfer learning outcomes, the process of retrieving certain knowledge still heavily relies on manual effort.

2.2 Data Harmonizing with LLMs

Data harmonization applying LLMs offers significant benefits by leveraging their natural language capabilities to standardize and integrate diverse datasets and further enhance model performance through improved data consistency [Cao et al., 2009; Zhang et al., 2024a; Liu et al., 2021; Durante et al., 2024]. However, this approach has several challenges, including the high computational costs of training LLMs and potential biases in training data, which can adversely affect the harmonization process [Feng et al., 2021; Xie et al., 2024; Zhang et al., 2025a]. Furthermore, the risk of overfitting remains a concern, as models may become overly specialized in training data nuances, reducing their effectiveness on new datasets.

2.3 Retrieval Augmented Generation for Knowledge Augmentation

Retrieval Augmented Generation (RAG) [Lewis et al., 2020] for knowledge augmentation is a method that enhances the capabilities of generative language models by integrating information retrieval with model generation [Hu and Lu, 2024]. This technique aids LLMs in tasks that demand deep and specific domain knowledge [Zhang et al., 2024b; Huang and Huang, 2024]. RAG for knowledge augmentation provides access to expansive external libraries (collections of documents or domain-relevant knowledge), making it especially effective for transferring knowledge across different domains [Siriwardhana et al., 2023]. RAG for knowledge augmentation embeds extensive databases directly into the generative process so that specific, domain-related information is both accessible and effectively utilized. In this way, it can significantly enhance the performance for complex tasks in knowledge augmentation scenarios.

3 Preliminary

3.1 Definition

Definition 1. (Domain) A domain D is an ordered pair consisting of a feature space X and a marginal probability measure P defined on this feature space. In other words, D = (X, P), where X = {x|xi ∈ X, i = 1, . . . , n} is an instance set. And P is a probability measure that describes the probability of occurrence of feature vectors x ∈ X. This probability measure P makes (X , B(X ), P ) a probability space, where B(X) is the Borel σ-algebra on X .

Source Domain. The source domain DS consists of instances paired with labels y, represented as DS = {(x, y)|xi ∈ XS, yi ∈ YS, i = 1, . . . , nS }. Here, XS is the feature space and YS is the label space for the source domain. This domain provides labeled data used to train models in preparation for transfer learning tasks.

Target Domain. The target domain DT typically contains a mix of unlabeled instances and a smaller set of labeled instances, denoted as DT = {x ∈ XT} ∪ {(x, y)|(x, y) ∈ XT ×YT}. Here, XT is the feature space and YT is the label space for the target domain. We aim to evaluate and fine-tune the transfer learning models with data in the target domain.

Definition 2. (Task) A task T consists of a label space Y and a decision function f, formally noted as T = (Y, f), where Y is a metric space that contains all possible labels, and f is a mapping from the feature space X to a set of conditional probability measures on the label space Y.

Source Task. The learning task of the source task TS is typically represented as learning a target function fS : XS → YS, where YS is the label space of the source task.

Target Task. The learning task of the target task TT is typically represented as learning a target function fT : XT → YT , where YT is the label space of the target task.

Definition 3. (Knowledge Transfer) Knowledge transfer is an augmentation method that adopts observations from source domains and tasks, denoted as {(DSi , TSi )|i = 1, . . . ,mS }. Here mS ∈ N+ represents the number of source domains and tasks. Similar observations from target domains and tasks are denoted as {(DTj , TTj )|j = 1, . . . , mT }, where mT ∈ N+. The goal of knowledge transfer is to utilize the knowledge embedded in the source domains DSi to enhance the performance of the learned decision functions fTj across the target domains DTj , for j = 1, . . . ,mT .

3.2 ProblemFormulation

We formulate the task to enhance the performance of knowledge transfer through automated data retrieval and harmonization using an LLM. Concretely, we improve the performance of the learned decision function fTj by reconstructing and refining the source domain DS. In this way, we better utilize the knowledge implied in DS as we mitigate domain shifts and facilitate a more effective transfer of learned models. Concretely, our optimization objective is to retrieve and reconstruct a source domain D∗S :

where P is the performance indicator of fTj and ˆ DS is a reconstructed source domain aligned with DT .

4 Methodology

In this section, we introduce our novel knowledge augmentation method named LLM-Enhanced Knowledge Augmentation (LEKA), designed to dynamically and automatically retrieve and refine source data to transfer knowledge and further enhance target data learning. Specifically, our study focuses on tabular datasets as the concrete form of data. The LEKA leverages an LLM to extract essential information from the target dataset, including data structure, feature names, and descriptions. Within this framework, the LLM uses its capability to analyze and synthesize data to optimize the selection and refinement of source datasets. It extracts and summarizes keywords for retrieval tailored to the target domain and task. After retrieval, the LLM refines and harmonizes these datasets in feature space and marginal probability measures with the target dataset. Figure 2 illustrates the overview of the LEKA framework, which comprises three stages: (1) Dataset Retrieval; (2) Data Harmonization; and (3) Knowledge Transfer and Evaluation.

4.1 Dataset Retrieval

In this phase, our primary objective is to identify and retrieve a source dataset DS similar to the target dataset DT for effective pre-training in knowledge transfer scenarios. To achieve this, we utilize an LLM A to analyze and retrieve data that matches the structural properties and purposes of DT . A straightforward method is to embed the target dataset DT by A to capture its essential features. However, this approach can be computationally intensive and prone to errors. Inspired by the strategies outlined in [Zhang et al., 2024c], we turn to leverage the textual information, including the data descriptions CT and feature names lT of DT . These textual elements provide insights into the structure and purpose of DT , and the LLM can precisely focus on the semantic content of DT .

First, the LLM A constructs a query Q based on this textual information:

and then embeds the query:

where embed(·) is the embedding process that transforms the input into a vector representation.

We then retrieve from a library L of datasets for the top-k most relevant datasets {Dk} to query Q. We evaluate the relevance of all datasets in L to the query Q by the cosine similarity of the embeddings and select the top-k datasets with the highest similarity:

where dk = A(embed(Dk)) is the embedding of a potential dataset Dk's textual information, and ∥ · ∥ ∈ R is the norm function.

In this phase, we map textual information to a highdimensional vector space. Essentially, we approximate the probability distributions of the datasets in the library to retrieve k source datasets most similar in structure and purpose to the target dataset. We leverage the geometric properties of vectors in high-dimensional space to assess and quantify dataset similarities. This prepares us for further dataset processing for transfer learning.

4.2 Data Harmonization

Now that we have k potential source domain datasets, we turn to align these datasets with the target dataset DT with an LLM. This alignment process adjusts the feature space and marginal probability measures of the source dataset Dk to closely match those of the target dataset DT . This process involves transformations of features and adjustments of distributions. Here we take Dk as an example.

Feature Space Transformation. We denote Xk and XT as the feature spaces of the source dataset Dk and the target dataset DT , respectively. Our LLM constructs a mapping function f : Xk → XT for alignment. This function transforms the features in Xk to minimize the distance between the transformed source features and the target features. Here, we aim to minimize the distance d(f(Xk), XT ), where d is a kernel distance:

Here, κ(·, ·) is a kernel function, xki and xkj are the i-th and j-th feature vectors in the source dataset Dk. We adopt a Gaussian kernel κ(x, y) = exp(-γ∥x - y∥2 ), with γ being a positive bandwidth parameter. This kernel function can effectively measure the similarity between points in highdimensional spaces and captures both linear and non-linear relationships. Thus, it can handle the complexities inherent in high-dimensional data. Compared with Euclidean distances, kernel distances can capture the geometric structure of the data manifold. It provides a more robust and informative similarity measure for the source-target alignment.

Marginal Probability Measures Harmonization. Our LLM aligns the marginal probability measures Pk and PT of the source and target datasets. The LLM analyzes and refines textual information such as labels, feature names, and classification probability distributions. To measure the differences between Dk and DT , we adopt the Wasserstein distance:

where Γ(Pk,PT ) represents the set of all joint distributions with marginals Pk and PT on X × X . The Wasserstein distance minimizes the transportation cost while preserving the geometric structure of the data distributions. Specifically, the Wasserstein distance calculates the minimum "geographical" cost required to move data from one distribution to another, where "geographical" cost refers to the cost of moving data from one point to another in the feature space.

Moreover, the Wasserstein distance is advantageous when dealing with distributions whose support sets (i.e., the effective range or set of the distributions) do not fully overlap. In real-world data applications, it is common for the source and target datasets to originate from different distributions. Their data points do not align perfectly and cover the same areas. In these cases, traditional distance metrics like the Euclidean distance poorly reflect the actual differences between the two datasets, as they merely measure differences in position but ignore the overall structure of the data distributions.

With this source-target alignment process, the LLM improves the efficiency of knowledge transfer from the source domain to the target domain. Thus, it enhances the overall performance of our transfer learning models. Further, we optimize data handling and boost the adaptability and accuracy of transfer learning for adoption in various applications.

4.3 Knowledge Transfer and Evaluation

After aligning the source dataset D′k with the target dataset DT , and adjusting the marginal probability measures P ′k to PT , we proceed to integrate these aligned datasets into the transfer learning model. At this stage, our goal is to transfer knowledge from D′k to DT by minimizing a defined loss function while adapting the model to the target domain.

To achieve this, the transfer learning process focuses on updating the decision function fTj . Specifically, we aim to minimize the expected loss over the target dataset while incorporating the knowledge transferred from the source dataset:

where L is the loss function, x represents the features, y represents the labels in DT , and θ denotes the parameters of fTj that are being optimized. We calculate the expectation by the probability distribution PT , which has been closely aligned with P ′k to ensure consistency and maximize the efficacy of the knowledge transfer.

During update, the optimization of model parameters θ leverages both the target dataset DT and the aligned source dataset D′k , incorporating domain-specific characteristics to reduce domain shift. We define the optimization process as:

This formula aims to fine-tune the decision function fTj by minimizing the weighted sum of expected losses across the datasets. The loss function L(fTj (x; θ), y) evaluates prediction accuracy, guiding the adjustment of parameters. The expectations E(x,y)∈D′k and E(x,y)∈DT represent the mean losses over the source and target datasets, respectively. The weighting factor α adjusts the relative influence of each dataset and enables flexible adaptation between leveraging established knowledge from the source and integrating new data from the target. In this way, we apply the transferred knowledge to enhance model performance on the target task.

Following the optimization of the model parameters, the backpropagation process then updates θ by minimizing the total loss. This loss is a weighted sum calculated from the losses on D′k and DT . This involves calculating the gradient of the loss function with respect to θ and updating θ using gradient descent methods:

where η is the learning rate. The gradients are computed based on both datasets, which allows the model to learn from both the aligned source data and the target data.

After the parameter optimization and backpropagation, we evaluate the effectiveness of the transfer learning process; we adopt a performance metric ϕ on the target dataset:

where ϕ measures the performance of the optimized model f ∗ Tj on DT .

In this way, we systematically improve the model based on empirical performance metrics.

5 Experiments

In this section, we present four experiments to demonstrate the effectiveness and impacts of the LEKA. First, we compare the LEKA against several baseline methods on four downstream tasks. Second, we present the correlations between several target domains and their retrieved source domains. Finally, we discuss the reason for performance improvement.

5.1 Experiment Settings

Datasets and Domains. We evaluate our method on four datasets of medical and economic domains: (1) Breast Cancer Wisconsin (Diagnostic) (BCW) [Wolberg et al., 1995], (2) Heart Disease (HD) [Janosi et al., 1989], (3) Vehicle Insurance Data (VID) [Bhatt, 2019], and (4) Telco Customer Churn (TCC) [BlastChar, 2018]. We show the detailed information about the features of the datasets in Table 1.

Metrics and Models. We evaluate the model performance by the following metrics: Overall Accuracy (Acc) measures the proportion of true results (both true positives and true negatives) in the total dataset. Precision (Prec) reflects the ratio of true positive predictions to all positive predictions for each class. Recall (Rec), also known as sensitivity, reflects the ratio of true positive predictions to all actual positives for each class. F-Measure (F1) is the harmonic mean of precision and recall, calculated here as the macro-average. We apply the LEKA across a range of models: 1) Tabnet (TN) [Arik and Pfister, 2021]; 2) TabTransformer (TT) [Huang et al., 2020]; 3) Random Forest (RF) [Rigatti, 2017]; 4) Gradient Boosting Decision Trees (GBDT) [Lin et al., 2023] ; 5) XGBoost (XB) [Chen and Guestrin, 2016]. We compare the performance in these tasks both with and without our method.

Baseline Models. We compare the LEKA with six baseline methods, including: 1) Raw Data: using vanilla data; 2) TIFG [Zhang et al., 2024c]: feature generation with LLM; 3) KPDDS [Huang et al., 2024]: data synthesis with LLM using key point examples; 4) GReaT [Borisov et al., 2022]: data synthesis with LLM simulating data subset distributions; 5) FTT [Levin et al., 2022]: A TabTransformer pretrained on the source domain and then fine-tuned on the target domain; 6) TTab [Wang and Sun, 2022]: A transferable tabular Transformer capable of learning frommultiple tabular datasets.

Implementation Details. In our setup for data synthesis and model training, we utilize GPT-4o [OpenAI, 2024] as the query generator, combined with the Exa API [Exa, 2024] to fetch web pages containing datasets from Kaggle [Kaggle, 2024] and the UCI Machine Learning Repository [University of California, Irvine, 2024] that may be suitable for knowledge transfer. We extract dataset descriptions from web pages, and GPT-4o assesses their potential for knowledge transfer. Additionally, GPT-4o serves as a generator for executable code, performing up to five code generations. For our models, we configure TN, TT, and FTT with a batch size of 512 for the VID and TCC datasets and a batch size of 32 for the BCW dataset, a maximum of 100 epochs, and employ early stopping with a patience of 20. The learning rate is set at the default 0.02 for pytorch tabnet. For the RF and GBDT models, the number of trees is set to 100, with GBDT also configured with a learning rate of 0.1 and a max depth of 3. TTab is set with a maximum of 50 epochs, a learning rate of 1 × 10-3 , and a weight decay of 1 × 10-4 .

5.2 Experiment Results

Overall Performance. The results are shown in Table 3, comparing our LEKA across four datasets. In summary:

(1) Compared with baseline models, LEKA significantly improves accuracy, surpassing baseline methods like TIFG and GReaT, with enhancements up to 4.4% in models like GBDT. LEKA consistently achieves top results in precision and recall, indicating its precision in correctly identifying relevant cases while minimizing false positives. The F1 scores under LEKA are notably high, reflecting its effective balance between precision and recall across various models and datasets. These results validate the effectiveness of LEKA 's retrieval and harmonization strategies and its robustness in diverse application scenarios. Overall, LEKA's strategic approach to knowledge transfer is particularly advantageous in complex data environments, showcasing its adaptability and efficiency compared to traditional methods.

(2) The LEKA outperforms all baseline methods in almost all metrics and datasets. Specifically, LEKA shows enhancements in accuracy by 2 - 5% over other methods. It demonstrates notable improvements of up to 4.4% in GBDT for the VID dataset and 5.4% in RF for the BCW dataset. For improvement of overall accuracy, we demonstrate the LEKA's effective retrieval and harmonization of feature space and marginal probability measures with the target datasets. LEKA demonstrates exceptional precision, improving the TT metric for the TCC dataset by over 10% compared to baselines, effectively reducing false positives. Meanwhile, the LEKA consistently outperforms baselines in reducing misclassification rates with the highest F1 scores. These results prove that LEKA's retrieval and data harmonization reduce misclassification by forming a deeper and clearer understanding of the potential relationship between features.

Comparison with transfer learning methods. We then compare our LEKA method with transfer learning methods to demonstrate its effectiveness in enhancing model performance in complex domains. In these scenarios, the transfer learning methods show their reliance on manual data selection and alignment processes. The results demonstrate that LEKA outperforms transfer learning methods, FFT, and TTab across various metrics and datasets. For example, in the BCW dataset, LEKA improves accuracy by 3.5% and recall by 3.3% compared to the competing methods. This superior performance is consistent across other datasets like VID, HD, and TCC, with notable enhancements in precision and recall, emphasizing LEKA's effective data harmonization capabilities. This adaptability and efficiency in handling diverse datasets underscore LEKA's robustness and advantages.

6 Conclusion

In this paper, we introduce LLM-Enhanced Knowledge Augmentation (LEKA), a novel retrieval and harmonization framework that dynamically refines source data for effective knowledge transfer. This structure significantly enhances data augmentation by leveraging advanced LLM capabilities to automatically align and optimize data retrieved from diverse external libraries. Extensive experiments across various tasks demonstrate superiority to existing methods, especially in improving model adaptability and accuracy in complex data environments. By adopting LEKA on data-scarce domains, we achieve substantial improvements in learning performance and domain-specific task accuracy. For future work, we plan to extend this framework to include more varied data types and conduct a thorough empirical analysis to understand its underlying mechanisms and impacts.

7 Limitations

We acknowledge the following limitations: (1) the current work has only been tested on tabular data, and more complex test scenarios have not yet been involved; (2) despite its automation, LEKA can be computationally intensive. This could limit its applicability in resource-constrained environments or require substantial computational resources to maintain operational efficiency; (3) adopting the LEKA in tasks with unique requirements may have inherent limitations.

Acknowledgements

Dr. Kunpeng Liu is supported by the National Science Foundation (NSF) via the grant numbers 2426339 and 2348485. Dr. Yanjie Fu is supported by the National Science Foundation (NSF) via the grant numbers 2426340, 2416727, 2421864, 2421865, 2421803, and National Academy of Engineering Grainger Foundation Frontiers of Engineering Grants.