When Raw Data Prevails: Are Large Language Model Embeddings Effective in Numerical Data Representation for Medical Machine Learning Applications? (2024)

Abstract: The introduction of Large Language Models (LLMs) has advanced data representation and analysis, bringing significant progress in their use for medical questions and answering. Despite these advancements, integrating tabular data, especially numerical data pivotal in clinical contexts, into LLM paradigms has not been thoroughly explored. In this study, we examine the effectiveness of vector representations from last hidden states of LLMs for medical diagnostics and prognostics using electronic health record (EHR) data. We compare the performance of these embeddings with that of raw numerical EHR data when used as feature inputs to traditional machine learning (ML) algorithms that excel at tabular data learning, such as eXtreme Gradient Boosting. We focus on instruction-tuned LLMs in a zero-shot setting to represent abnormal physiological data and evaluating their utilities as feature extractors to enhance ML classifiers for predicting diagnoses, length of stay, and mortality. Furthermore, we examine prompt engineering techniques on zero-shot and few-shot LLM embeddings to measure their impact comprehensively. Although findings suggest the raw data features still prevail in medical ML tasks, zero-shot LLM embeddings demonstrate competitive results, suggesting a promising avenue for future research in medical applications.

## introduction numerical data plays a pivotal role across various domains. for instance, much of the data used for analytics from electronic health records (ehrs) are numerical values in tabular formats, documenting patient demographics (e.g., age), vital signs, laboratory tests, and nurse assessments. utilizing recent work shows large language models (llms)' vast potential on text generation over structured data input, including chain-of-thought (cot) reasoning over tabular data @xcite , classification on diseases @xcite . llms have also exhibited exceptional promise in medical nlp tasks, evident in their stellar performance in the united states medical licensing examination (medqa) @xcite . however, the use of embedding representations, particularly for medical diagnostics and outcome predictions using standard ehr numerical data, remains largely unexplored. in these areas, raw data inputs have traditionally dominated feature representation for ml algorithms before the era of llms. this is exemplified by their use in critical applications such as mortality prediction and early sepsis warnings @xcite , and patient infection @xcite . the potential of llmderived features as a viable alternative to raw data features in ml applications is still unclear. this study aims to address the knowledge gap in applying llm embeddings for ehr numerical data representation in ml algorithms. while llms have demonstrated remarkable performance in textbased medical tasks, few studies explore their use in numerical-based tasks, especially in ehr predictive modeling, such as mortality prediction and clinical deterioration detection. these tasks, which directly impact clinical workflow, offer a critical testbed for evaluating llm capabilities and setting appropriate expectations for their use in practical applications. furthermore, existing work on llms shows potential in handling numerical reasoning tasks, suggesting that they can effectively represent numerical data @xcite . leveraging this potential, our study probes whether llm embeddings can enhance performance in domain-specific tasks. moreover, using llms for tabular data representation provides an opportunity to unify structured and unstructured data within ehrs, seamlessly integrating information across modalities such as clinical notes and embedded tables @xcite . our work uniquely examines the impact of different formats and embedding methods on llm last layers and ml classifiers, focusing on open-source, zero-shot llms suitable for single-gpu systems to address the resource limitations prevalent in many hospitals and academic research settings. to test whether llms possess knowledge of clinical reference ranges for numerical values, we conducted a preliminary study to by directly prompting them with questions about twenty-four critical clinical features. these questions focused on the numerical ranges, units of measurement, and contextual reasoning behind these values for each feature, such as vital signs and lab test results. the experiments were done on mistral-7b-instruct @xcite and llama2-13b-chat @xcite , two of the most popular open-source llms. the findings, depicted in figure 1 , suggest that llms, despite being trained on general-domain data, demonstrate a strong understanding of the expected physiological ranges and correct units of measurement for these features. building on these initial findings, our study then explores the application of llm embeddings in clinical prediction tasks. we utilize three clinical prediction tasks derived from two independent ehrs and four ml classifier input settings. we investigate the impact of table-to-text conversion formats, embedding extraction methods, prompt engineering, and few-shot techniques, along with early results from parameter-efficient fine-tuning, on the quality of llm embeddings. our main contributions are the following: results show that, despite external evidence indicating that llms possess extensive knowledge of medical facts, extracting usable representations of this knowledge for downstream tasks will require significant additional methodological progress. ## related work recent studies highlight llms in tabular data analysis @xcite introduces tablellm, which converts tables to text using a manual template. @xcite studies cot reasoning over tables. @xcite examines the abilities of llms on numerical data understanding. @xcite , closest to our work, explores zero-shot llm for structured longitudinal ehr data and finds that gpt-4 can outperform xgb on clinical prediction tasks. our study, however, uniquely focuses on open-box llm embeddings for enhancing ml algorithms. raw ehr data are commonly used in medical ml applications, as found by a survey on medical ml research @xcite . they noted that labs and vital signs as frequent data types for patient representation learning. @xcite introduces an xgb algorithm predicting clinical deteriorations using ehr features like demographics and lab values. @xcite used 104 clinical ehr features across various ml algorithms to establish baselines for clinical tasks such as mortality predictions. our work uses the same dataset and tasks as @xcite to compare llm embeddings against traditional ml classifier outcomes on the same raw data feature baseline. shwartz-ziv and armon (2022) compares xg-boost models with deep learning models across eleven tabular datasets, covering tasks from gesture phase segmentation to higgs boson monte carlo simulation. their findings suggest that xg-boost outperformed deep learning models on these tasks. similarly, this paper reaches the same conclusion in a different context, where we compare xgboost using raw data features for medical machine learning applications and extend the comparison to evaluate raw data features against llm embeddings. notably, we find that zero-shot llm embeddings, when used with xgboost, have the potential to perform as well as raw data features. ## methods and experiment setup figure 2 illustrates the study overview and experiment setup. we began with a patient's tabular data input, represented using the pandas dataframe data structure (raw data). this raw data was converted to text using four distinct conversion methods, detailed in §4.1, and llm encoded the converted text, with the last hidden states extracted to generate embedding features ( §4.2). these embeddings were subsequently used to train various ml classifiers on two datasets for binary prediction tasks. we started with zero-shot, off-the-shelf llms for experiments ( §4.3). we then investigated the impact of prompt-engineering techniques and fewshot learning configurations on the embeddings and subsequent predictions ( §4.4). an initial investigation was also conducted to assess the effects of parameter-efficient fine-tuning on llm embeddings for ml tasks, focusing on two of the models ( §4.5). as baselines, we included traditional ml classifiers trained directly on raw tabular data inputs. to benchmark the effectiveness of llm embeddings, we used randomly initialized embeddings of the same size as the llm-generated embeddings. ## discussion to understand the discrepancy between the two data representations, we examined the training effectiveness of raw data features and llm embeddings by controlling the training set size. figure 6 compares the performance of the raw data xgb baseline model with the mistral and meditron embeddings across different training set sizes for two tasks in the mimic dataset. the raw data xgb baseline model shows a significant increase in au-roc scores with larger training sets, achieving high performance. in contrast, both the mistral and meditron embeddings paired with xgb models exhibit much smaller improvements, consistently performing lower than the raw data xgb baseline. this highlights the greater effectiveness of xgb when learning from raw data features compared to llm embeddings for these prediction tasks. our findings suggest that raw data features provide more informative input for ml models compared to llm-generated embeddings. while llm embeddings capture complex representations, they may not be as tailored for binary medical predic-tion tasks. additionally, computing efficiency is an important consideration, as llms require significantly more gpu memory than raw data features. for instance, training xgb with raw data features used less than 1000 megabytes of gpu memory, whereas encoding input to llm embeddings required between 1500 megabytes (in zero-shot settings) and over 8000 megabytes per sample, with even higher usage for longer inputs, such as those from mimic-extract. however, zero-shot llm embeddings achieve comparable performance in certain scenarios, highlighting their potential for rapid deployment without extensive training. a promising direction is distilling these embeddings into a smaller space while retaining their extensive knowledge @xcite recently proposes llm2vec, a method to train decoder-only llms as text encoders with unsupervised training, which merits further investigation. ## conclusion we present the first analysis of llm embeddings for numerical ehr data features in medical ml applications, showing the opportunity and challenges of using llm embeddings as a substitute of raw data features. we hope to encourage future research on improving llm embeddings, particularly for imbalanced label prediction, and advancing health predictions with multi-modal data, while addressing interpretability and bias.