LLM -Forest: Ensemble Learning of LLM s with Graph-Augmented Prompts for Data Imputation (2025)

Abstract: Missing data imputation is a critical challenge in various domains, such as healthcare and finance, where data completeness is vital for accurate analysis. Large language models (LLMs), trained on vast corpora, have shown strong potential in data generation, making them a promising tool for data imputation. However, challenges persist in designing effective prompts for a finetuning-free process and in mitigating biases and uncertainty in LLM outputs. To address these issues, we propose a novel framework, LLM-Forest, which introduces a “forest” of few-shot learning LLM “trees” with their outputs aggregated via confidence-based weighted voting based on LLM self-assessment, inspired by the ensemble learning (Random Forest). This framework is established on a new concept of bipartite information graphs to identify high-quality relevant neighboring entries with both feature and value granularity. Extensive experiments on 9 real-world datasets demonstrate the effectiveness and efficiency of LLM-Forest. The implementation is available at https://github.com/Xinrui17/LLM-Forest

## introduction handling missing data is a fundamental challenge across various domains, including datasets such as clinical reports, financial records, product characteristics, and survey data in a tabular format @xcite @xcite . however, these datasets frequently suffer from missing values, which can severely impact the reliability of critical decisionmaking processes that depend on complete and accurate data. to address this issue, feature imputation @xcite @xcite , the process of estimating missing values based on † corresponding authors. common approaches for feature imputation can be broadly categorized into two groups. (1) statistical methods: techniques such as mean or mode imputation and mice (little and rubin, 2019) are simple and widely used due to their ease of implementation. while these methods are straightforward, they can introduce bias. (2) deep learning-based methods: models like gain @xcite , diffimpute @xcite , and grape @xcite aim to model complex data distributions and dependencies leveraging advanced neural network architectures to capture non-linear relationships within the data. however, they typically require large volumes of data and can be challenging to train, particularly in scenarios like healthcare where data are often limited, imbalanced, or subject to strict privacy regulations. recently, large language models (llms) @xcite @xcite have demonstrated strong potential in handling complex tasks by using external knowledge and contextual learning, offering a promising solution for enhancing performance on tabular data tasks @xcite @xcite . anam @xcite began exploring the adoption of @xcite as a data imputer by formulating text questions and prompting chatgpt to respond with imputed values for each missing cell. similarly, claim @xcite introduced an llm-based approach to generate text-specific descriptions of missing data. however, both methods require finetuning processes, which can be time-consuming and resource-intensive, limiting their scalability across different datasets. in contrast, in-context learning @xcite @xcite has excelled in tasks such as tabular classification (e.g., cancergpt @xcite ) and reasoning (e.g., chain-of-table @xcite ). yet the use of llms' broad knowledge and reasoning ability for data imputation, without relying on task-specific fine-tuning, remains largely unexplored. in this paper, we explore fine-tuning-free approaches to harness the power of pre-trained llms for data imputation. we aim to overcome two primary challenges: (1) extracting high-quality information from observed data to align with the preferences of the target entries to effectively guide llms, and (2) reduce potential output bias and uncertainty in current llms. to address these challenges, we propose a novel framework, llm-forest. it consists of multiple trees utilizing few-shot in-context learning with llms, inspired by ensemble learning (random forest), reducing the chance of bias and uncertainty compared to relying on a single llm. each tree in the forest represents an llm with a particular few-shot prompt learning process. to construct high-quality yet diverse prompts, we design a graph-based retrieval algorithm to explore highly relevant neighboring entries in preparation for constructing llm-forest. this graph algorithm is based on information-theoretic bipartite graphs to quantify relevance at both the featurelevel and value-level granularity. meanwhile, it reduces the computational complexity in the graph construction process and allows for the subsequent merging process to be easily implemented in parallel. this makes it scalable for very large datasets. we then carefully design prompts to integrate the retrieved neighbors and inject domain-level guidance, maximizing the utility of pre-trained llms without any fine-tuning. finally, we aggregate the outputs of the llm trees by weighting each tree's prediction according to its self-assessed confidence level, ensuring more robust and reliable results. in summary, our main contributions are : the results demonstrate the effectiveness and efficiency of llm-forest in imputing missing data. our framework outperforms traditional imputation methods, highlighting its potential for practical applications in data analysis. since the construction of each llm tree is independent, we can compute the predictions of llm trees in parallel. moreover, the graph retrieval component is efficient, enabling the framework to handle large-scale datasets by retrieving a small number of high-quality data objects for each llm tree. this ensures that the inference time of llm-forest does not increase significantly in large-scale datasets. ## notations and problem definitions let [n] = {1, 2, . . . , n} and define @xmath0. suppose the matrix @xmath1. the matrix @xmath2, a 2 , . . . , a n ] ⊤ , where each entry a @xmath3, @xmath4, is a d-dimensional vector. for example, in a diabetes dataset, a ⊤ i could represent the feature set of the patient i. let a ij , @xmath5, denote the j-th feature value of the i-th entry a ⊤ i . for convenience, we use {c 1 , c 2 , . . . , c d } to represent the d-dimensional features (columns) of x, where c @xmath6, @xmath7, @xmath8, denotes the ith entry of the j-th feature. let r j represent the set of distinct values in feature c j , which can be viewed as the collection of attribute values for all patients in the dataset. additionally, define a matrix @xmath9, where a cell has the value 1 if the corresponding entry in x is missing, and 0 otherwise. the objective is to predict the missing values in x based on the available information in the dataset. ## proposed method: llm-forest in this section, we present our proposed framework, llm-forest, which consists of four key components, as illustrated in figure 1 . the first three steps form the information-based graph retrieval module: (1) table to bipartite information graphs: table input bipartite graph construction p0 p0: patient 0 *: classified level p3 0 high bp p1 1 high bp p2 ? p0 p3 low* bmi p1 high* bmi p2 p0 p3 1 veggie p2 0 veggie p1 p2 low* gen hlth p0 high* gen hlth p3 ? p1 … 𝐺 1 𝐺 2 𝐺 3 𝐺 4 merge p0 p3 p1 p2 cumulative weighted p0 p3 p1 p2 … … random walk merge bipartite graph random walk llm tree-to-forest merge [mask 1] [mask 2] p0 p1 (1) neighbors (2) instruction (3) meta strategies (4) correlations/distribution (5) feature description prompt [mask 1] is 0 [mask 2] is 3 [mask 1] is 1 [mask 2] is 1 [mask 1] is 1 [mask 2] is 3 output 1 output 2 output 3 confidence 1 confidence 2 confidence 3 weighted voting [mask 1] is 0 [mask 2] is 3 … … … … final output imputed table (1) (2) (3) (4) p1 p2 (1) neighbors prompt p1 p2 (1) neighbors prompt patient id 0 1 2 3 … high bp bmi veggie genhlth … 0 33 1 4 … 1 32 0 1 … [𝐌𝐀𝐒𝐊 𝟏] 60 1 2 … 0 25 1 [mask 2] … … … … … … diabetes no yes yes no … features index label 1 high bp high* bmi 0 high bp low* bmi 1 veggie high* genhlth 0 veggie low* genhlth imputed table llm-forest patient id 0 1 2 3 … high bp bmi veggie genhlth … 0 33 1 4 … 1 32 0 1 … 0 60 1 2 … 0 25 1 3 … … … … … … diabetes no yes yes no … downstream tasks llm-forest e.g., table classification llm-tree llm-tree llm-tree (2) instruction (3) meta strategies (4) correlations/distribution (5) feature description (2) instruction (3) meta strategies (4) correlations/distribution (5) feature description cumulative weighted cumulative weighted this step converts tabular data into bipartite information graphs connecting dimensional features and their corresponding entries, where edges are weighted based on self-information to quantify the informativeness of feature values, enriching the representation of feature-entry relationships. (2) hierarchical bipartite graph merging: the bipartite graphs are then hierarchically merged, aggregating shared information across different features. (3) graph retrieval via random walk: random walks are performed on the merged bipartite graphs to identify correlated and diverse neighbor nodes for each entry (node), preparing for the construction of the llm-tree. finally, (4) construct llm-tree to forest: based on the identified neighbors, a tailored few-shot prompt is constructed for each llm, forming a forest with confidence-weighted voting. we detailed each procedure below. ## experiments in this section, we evaluate llm-forest across various real-world datasets under different settings. ## conclusion in this paper, we propose a novel framework llm-forest designed to enhance the data imputation via large language models. the framework combines multiple llm trees employing diverse few-shot learning, forming an llm forest to achieve robust and accurate imputations. specifically, to extract high-quality contextual information for llm trees, we developed a graph-based retrieval algorithm. to mitigate llm bias and uncertainty, we introduce a confidence-weighted voting mechanism that aggregates outputs from all llm trees based on their self-assessed confidence levels. comprehensive experiments on 9 realworld datasets in various domains demonstrate that llm-forest effectively improves imputation accuracy while maintaining scalability and efficiency. the results highlight the potential of combining ensemble learning of llms with graph-based retrieval methods to tackle the complex problem of data imputation.