Can Race-sensitive Biomedical Embeddings Improve Healthcare Predictive Models?

Abstract

This reproducibility study presents an algorithm to weigh in race distribution data of clinical research study samples when training biomedical embeddings. We extracted 12,864 PubMed abstracts published between January 1^st, 2000 and January 1^st, 2022 and weighed them based on the race distribution data extracted from their corresponding clinical trials registered on ClinicalTrials.gov. We trained Word2vec and BERT embeddings and evaluated their performance on predicting length of hospital stay (LHS) and intensive care unit (ICU) readmission using MIMIC-IV electronic health record data. We observed that models trained using race-sensitive embeddings do not consistently outperform the neutral embeddings ones when used for LHS prediction (with similar Mean Absolute Error 1.975 vs. 2.008) or ICU readmission prediction (with similar accuracy 74.61% vs. 75.17% and the same AUC 0.775), respectively. We conclude that demographic sensitive embeddings do not necessarily significantly improve the accuracy of health predictive models as previously reported in the literature.

Introduction

Inadequate diversity or population representativeness of study samples of clinical trials is a known issue that compromises the generalizability of clinical trials research. FDA has commented that certain populations, such as women and racial and ethnic minorities, are underrepresented in clinical trials [1]. Chanstain et al. also identified the problem of racial disproportionality in COVID-19 clinical trials, which directly affects racial health disparities [2]. Underrepresented populations are one of the main reasons that trials’ outcomes are biased and compromise the generalizability of the results to the entire patient population. Feldman et al. [3] showed the gender bias against female participants in clinical studies did not improve over years after analyzing clinical studies from published articles from PubMed and records from Aggregate Analysis of ClinicalTrials.gov (AACT). The age bias was also frequently observed in certain disease areas, for instance, Schoenmaker et al. [4] showed that dementia trial patients are older than real-world dementia patients. Somerson et al. [5] demonstrated a lack of racial diversity in a set of 158 orthopedic clinical trials, in particular, African American and Hispanic populations were underrepresented. Consequently, the underlying recruitment biases can mislead clinical investigators to problematic conclusions about the interventions’ effects or efficacy in practice within different racial groups. The poor representation of certain groups of population in clinical trials, could potentially threaten real-world patient safety by causing unforeseeable post-marketing adverse drug effects [4, 6].

Word embeddings have been prevalent in Natural Language Processing (NLP) applications since they are able to capture useful semantic properties and linguistic relationships between words in vector format. Medical term embeddings, as a variant of word embeddings, involves learning vector representations of medical terms using NLP models. Such biomedical embeddings have been commonly leveraged as (part of) input features to train downstream machine learning models for different healthcare related inference tasks, such as drug-drug interaction extraction [7, 8], mortality prediction [9, 10], and drug adverse event detection [11-13], to name a few. On one hand, it usually requires large scale biomedical literature corpora to train domain-specific embeddings and assess their feasibility on various Biomedical Natural Language Processing (BioNLP) tasks. On the other hand, the quality of embeddings is directly related to the quality of the corpora they are trained from. For instance, if a collection of clinical notes containing primarily white patients was used to train embeddings, the semantics of racial bias will be statistically modeled into the trained embeddings. Downstream applications implemented from such embeddings will inevitably inherit this population bias and yield questionable conclusions in healthcare with deficit knowledge about various population subgroups. Therefore, data bias should be addressed in training embeddings for BioNLP tasks.

Existing biomedical embeddings rarely report the potential bias they inherited from the data due to the lack of knowledge or analysis of the bias. Fortunately, in the clinical trial literature, the enrollment metadata of clinical trials is publicly available on ClinicalTrials.gov, and one can leverage this information to retrospectively assess the impact of biases presented. A clinical trial’s demographic characteristics (e.g., gender, race, age groups, and ethnicity) are normally reported along with their findings. Extending a study [14] that examined gender-specific embeddings, we proposed an algorithm to train concept embeddings with the awareness of demographic characteristics of clinical trial studies. The idea is to weigh a study’s contribution to the embeddings by considering its demographic subgroups, such as the ratio between different races and age groups, to obtain different demographic-sensitive embeddings. To implement this method, we mapped clinical trials’ demographics from ClinicalTrials.gov and with their corresponding abstracts published on PubMed. Leveraging Black and White enrollment data in training the embeddings, we evaluated our race-sensitive embeddings’ efficacy on training machine learning models to predict patients’ lengths of hospital stay and probability of intensive care unit (ICU) readmission using Medical Information Mart for Intensive Care IV (MIMIC-IV) [15] data.

Related work

Developing meaningful representations of words, also known as word embeddings, has been the main drivers of advances and innovation in the field of NLP since its inception. Word2Vec [16] is an unsupervised learning technique to generate word embeddings, by learning continuous multi-dimensional vector representation for each word in the training corpus. The obtained vectors capture word senses and similarity in different contexts. Previous methods on training word embeddings were based on (1) averaged surrounding context words, such as the continuous bag-of-words (cbow) model in Word2Vec, (2) weighted context words, such as the skip-gram model in Word2Vec, (3) global co-occurrence statistics, such as GloVe [17], and (4) word n-grams, such as FastText [18]. Generating high quality embeddings for biomedical concepts is more challenging since biomedical entities or mentions have a high degree of ambiguity, high number of synonyms, and can span multiple words. Existing concept embeddings in the biomedical domain primarily focus on varying the training corpora (e.g., Electronic Health Records (EHR)[19, 20], medical claims [21, 22], biomedical corpora [23], or mixed [24, 25]) or method (e.g., Word2Vec [20, 22, 24, 26-28], GloVe [21], FastText [29], BERT [23], or adding temporal [27, 29] or task-related knowledge [26, 30, 31] to train concept embeddings). There were a handful of studies that evaluated the performance of the concept embeddings in terms of semantic similarity [32] or medicine fields [33]. However, only a few studies considered the impact of gender bias on the generated embeddings, and attempted to detect and neutralize unwanted stereotypical gender associations in semantics of the training corpora [34-37]. None of these works considered the impact of race bias.

Methods

Datasets

Data from ClinicalTrials.gov and PubMed were leveraged to train the race-sensitive embeddings. ClinicalTrials.gov is a publicly available repository for metadata on clinical trials. PubMed is a publicly available repository of 33 million citations in biomedical literature. We used all completed interventional trials which were submitted to ClinicalTrials.gov before 01/01/2022. Each trial is associated with a unique NCT ID, which we use to extract race enrollment data from the AACT database (https://aact.ctti-clinicaltrials.org/), a relational database that aggregates information (including result data elements) about every study registered in ClinicalTrials.gov. We then linked trials’ NCT IDs to PubMed articles (identified with PMIDs) using the latter’s API. The abstract of each identified PubMed article was extracted to form the corpus for training embeddings.

MIMIC-IV [15] datasets were used to formulate two prediction tasks, intensive care unit readmission and length of hospital stay, to evaluate the race-sensitive embeddings. MIMIC-IV is a free publicly available repository of de-identified data of patients admitted to an ICU or the emergency department of Beth Israel Deaconess Medical Center (BIDMC) between 2008 - 2019. It contains a variety of information entailing the demographics, medication, diagnosis, procedures, and hospital stays of the patients. In this work, MIMIC-IV version 1.0 was used.

Neutral and Race-sensitive embeddings

The corpus to train our language embeddings are PubMed abstracts that are linked to clinical trials with race enrollment data. Figure 1 illustrates this process. To train race-neutral baseline embeddings, each abstract was included only once in the corpus regardless of demographic participation in the trial. To form a comparison, the race-sensitive embeddings were created by upsampling abstracts with greater proportions of Black participants. For instance, the clinical trial NCT01224678 was linked to a PubMed article (PMID 33849913). There were 238 White patients and 35 Black or African American patients enrolled in this trial, yielding a ratio of 0.15. This ratio is fed into a heuristic bootstrap policy (explained later) to determine the number of repetitions of this abstract, denoted as n. If n=4, it means that this abstract will be reused 4 times in our training corpus.

Figure 1.

Training process of race-sensitive and natural baseline embeddings.

We experimented with several heuristics and calculated word-level embeddings similarity scores as optimizations to determine each abstract importance, and based on performance we chose the heuristic below:

$$bootstrappolicy(x)=\{\begin{array}{cc}0,& x=0\\ 1,& 0<x\le 0.1\\ 4,& 0.1<x\le 0.2\\ 8,& 0.2<x\le 0.4\\ 10,& 0.4<x\le 0.8\\ 20,& x\ge 0.8\end{array}$$

where x is the ratio of race in enrollment of a clinical trial. In our case, x is the percentage of Black participants among Black and White participants enrolled in a clinical trial. For instance, if such ratio is 0.15, the abstract for this trial will be reused 4 times in our corpus for training embeddings. By doing this, studies with a higher ratio of Black enrollment contribute more towards the final embeddings, making them more race sensitive.

We permuted the training corpora (race-sensitive and neutral) and embedding models (FastText and BERT) to train four types of embeddings: Fr: FastText Race-sensitive embeddings; Fn: FastText Neutral embeddings; Br: BERT Race-sensitive embeddings; Bn: BERT Neutral embeddings. For FastText-based embeddings, the vector dimension was set to 100, the window size was 10, minimum word count was 1 to cover some rare concepts, and training epochs was 20. For BERT-based embeddings, the pre-training masked ratio was 15%, learning rate 1e-4, with epochs set to 2. The embeddings were computed as the mean of the last two layers of the BERT model, with the vector dimension equal to 768 and maximum tokens set to 128.

Preprocessing MIMIC-IV data

Figure 2(a) describes the steps of feature engineering and encoding of MIMIC-IV patient records. For our evaluation tasks, we limited features to those available at a patient’s admission, without looking into a patient’s discharge information (such as last care unit or discharge location). Hence, we used seven categorical features: first care unit, admission type, admission location, insurance, marital status, ethnicity, gender; four numeric features: days since last admission, number of previous admissions, number of previous procedures, age; and four textual features: procedures, drug, diagnoses, previous diagnoses. We mapped ICD9/10 codes, which were used to encode procedures and diagnoses, to their string format using the provided code name mappings. Categorical features were encoded using OneHot encoder, numeric features were normalized by removing the mean and scaling to unit variance, and textual features were encoded using neutral and race-sensitive embeddings, respectively. All encoded vectors were concatenated to one vector, as the representation of one patient’s record. The encoded dataset was split to training and testing datasets in a ratio of 85:15 based on patient unique ids so that patient data used for training will not be used for testing.

Figure 2.

(a) Feature engineering and encoding of MIMIC-IV patient records and split to training and testing datasets. (b)Training and test process of predicting length of hospital stay and ICU readmission with MIMIC-IV dataset.

Predicting length of hospital stay (LHS)

To evaluate the quality of generated embeddings, we applied them on the task of predicting LHS (Figure 2(b)). In this task, the goal was to predict a patient’s length of stay in the hospital based on the patient’s diagnoses from the previous admissions, diagnosis from the current admission, and demographic features. Estimating a patient’s length of stay is important in hospital planning around the allocation of rooms and resources as well as predicting need for different levels of care and recovery. For a patient’s procedures (or diagnoses) in one visit, we mapped each ICD9/10 code to its code name in string, and then encoded it using trained embeddings. After encoding all procedures (or diagnoses) of one visit into vectors, we calculated the mean of all vectors to represent this visit’s procedures (or diagnoses). This encoded vector was concatenated to the aforementioned features. The combined feature vector was fed into eight regression models for predicting LHS. For each of the eight regression models, we trained one predictive model without embeddings as baseline, one using neutral embeddings, and another one using race-sensitive embeddings, to form a comparison. The performance of each model was tested on the same testing dataset, with the metrics of Mean Absolute Errors (MAE) to compute the differences between predicted LHS (in unit of days) with the reported hospital stay in MIMIC-IV.

Predicting ICU readmission

We also evaluated the generated embeddings on the task of classifying unplanned readmission of a patient to ICU (Figure 2(b)). Identifying high-risk patients likely to suffer from readmission before release is important for determining the medical quality of a healthcare system, and identifying predictors of ICU readmissions and hospital mortality. For this task, the goal was to classify whether a patient will be readmitted into the ICU within 30 days after discharge, based on the patient’s records (the same as in predicting LHS task). All features were concatenated into one combined vector that was fed into nine commonly used classification models. To form a comparison, for each model, we trained one classifier without embeddings as baseline, one using neutral embeddings, and another one using race-sensitive embeddings. The performance of each trained classifier was evaluated on the same testing dataset, with the metrics of accuracy and AUC.

Each type of estimator (classifier or regressor) was implemented in Scikit-learn with parameters specified in Table 1. For estimators with randomization, we repeatedly trained each estimator with five random_seeds and averaged the performance. For each estimator, we employed the two-tailed Mann-Whitney U test [38] to determine whether the performance difference was statistically significant (α=0.05) in terms of embeddings types (neutral or race-sensitive). We chose Mann-Whitney U test because it allows two groups to be compared without assuming that values are normally distributed.

Table 1.

Estimators and their parameters in Scikit-learn implementation.

Task	Estimator	Parameters
Length of Hospital Stay	LinearRegression	n_jobs=8
	ElasticNet	random_state=random_seed
	GradientBoosting Regressor	n_estimators=1000, max_depth=10, min_samples_split=8, loss=ls, learning_rate=0.001, max_features=sqrt, random_state=random_seed
	LGBM Regressor	random_state=random_seed, n_jobs=8
	LinearSVR	random_state=random_seed, tol=1e-5
	MLP Regressor	random_state=random_seed, max_iter=500, early_stopping=True
	PassiveAggressive Regressor	random_state=random_seed
	XGBRegressor	random_state=random_seed
Intensive Care Unit Readmission	KNeighbors Classifier	n_neighbors=100
	AdaBoost Classifier	random_state=random_seed
	DecisionTree Classifier	random_state=random_seed
	Linear SVC	C=0.1, dual=False, max_iter=1000, random_state=random_seed
	Logistic Regression	C=10, max_iter=1000, random_state=random_seed
	MLP Classifier	alpha=1, max_iter=1000, early_stopping=True, random_state=random_seed
	RandomForest Classifier	random_state=random_seed
	Ridge Classifier	alpha=1.0, solver=sparse_cg, random_state=random_seed
	SGD Classifier	loss=log_loss, alpha=1e-4, n_iter_no_change=1000, early_stopping=True, random_state=random_seed

Results

We obtained 18,417 completed interventional clinical trials with race enrollment reported in results before 01/01/2022. After linking each trial’s NCT IDs to PubMed articles, we downloaded 12,795 abstracts (linked to 8,126 NCT IDs) from PubMed. The race enrollment information for each abstract is available from its corresponding clinical trial. Since the race categories used for result analysis on ClinicalTrials.gov were self-reported with varying levels of granularity, normalization was needed to distinguish and aggregate different race terms into coarser groups such as White, Black/African American, Asian, Hispanic, etc. For instance, “Caucasian” and “European” were categorized into “White” while “South African” was categorized into “Black.” After race normalization, 8,615 abstracts (linked to 5,421 NCT IDs) with enrollment data on “White” or “Black” race were maintained as the training corpus. This corpus was fed to two language models FastText and BERT to train neutral and race-sensitive embeddings, separately.

Table 2 shows the frequency of clinical trials and corresponding PubMed abstracts based on their ratio of Black/African American population enrollment versus White race. We illustrated thresholds the same as the repetition policy we employed for training race-sensitive embeddings. Among the total 5,421 clinical trials included in our training corpus, 2,912(53.7%) studies had a black/white enrollment ratio less than 0.1, corresponding to 4,951 (57.5%) PubMed abstracts. 640 (11.8%) studies had a black/white enrollment ratio over 0.8, corresponding to 1,014 (11.8%) PubMed abstracts.

Table 2.

Count and frequency of clinical trials and PubMed abstracts in terms of race enrollment.

Black/African American race enrollment percentage (x)	No. of Clinical Trials	Percentage of total trials (5,421)	No. of PubMed abstracts	Percentage of PubMedabstracts (8,615)
x < 0.1	2912	53.7%	4951	57.5%
0.1 < x <= 0.2	711	13.1%	1109	12.9%
0.2 < x <= 0.4	695	12.8%	1055	12.2%
0.4 < x <= 0.8	412	7.6%	598	6.94%
x > 0.8	640	11.8%	1014	11.8%

Figure 3(a) shows the ethnicity distribution of patient count in MIMIC-IV data. Among the total 256,878 patients, over 62% are White, about 12% are Black, 6.4% Asian, 6.3% Unknown and 1.2% Unable to obtain. Figure 3(b) shows the number of patient records for the ICU readmission classification task from MIMIC-IV dataset. After splitting patients into the 85:15 ratio, there are 10,977 patient records with ICU readmission and 39,678 records without ICU readmission in the training dataset, while in the testing dataset, 1,963 records with ICU readmission and 6,965 without ICU readmission. Figure 3(c) shows the histogram of patients’ length of hospital stay in both training and testing datasets. Around half of patients stay less than 4 days in the hospital.

Figure 3.

MIMIC-IV data statistics. (a) Ethnicity distribution of patient count; (b) Number of patient records for ICU readmission; (c) Histogram of patients’ length of hospital stay in both training and testing datasets.

We reported the classification results of utilizing different embeddings in predicting length of hospital stay in Table 3. Prediction models without any embeddings, as baseline, performed consistently worse than the ones with embeddings. Limiting to the FastText embedding only, the best performing model is LGBMRegressor, with an average MAE score of 1.732 and MSE score 10.029. The best performing single model is the MLP Regressor leveraging BERT race-sensitive embeddings, with an average MAE score of 1.697 and MSE score 9.964. Averaging the eight models’ performance based on embedding types, the model trained with FastText neutral embeddings outperformed models with race-sensitive embeddings with a lower averaged MAE (2.186 vs 2.235) and MSE (14.619 vs 14.838) for predicting hospital stay. Conversely, models with BERT race-sensitive embeddings achieved a lower averaged MAE (1.975 vs 2.008) and MSE (12.225 vs 12.529), compared to its corresponding neutral ones. Only the GradientBoosting Regressor’s p-values for the Mann-Whitney U test are significant.

Table 3.

LHS prediction comparison between eight predictive models utilizing neutral, race-sensitive embeddings, or without embeddings. The best performance single model is bolded and underscored. The best averaged performance of eight models for four different embeddings are bolded. MAE: Mean Absolute Error; MSE: Mean Squared Error. Significant P-values are marked with an asterisk. “--” indicates that p-value calculation is not applicable. Mann-Whitney U test is performed on MAE.

	Without embeddings		FastText Embeddings					BERT Embeddings
	Baseline		Neutral		Race-sensitive		Mann-Whitney U test	Neutral		Race-sensitive		Mann-Whitney U test
	MAE	MSE	MAE	MSE	MAE	MSE	P-value	MAE	MSE	MAE	MSE	P-value
LinearRegression	2.459	18.693	2.184	13.386	2.205	13.418	--	2.037	11.557	2.002	11.169	--
ElasticNet	2.512	19.250	2.436	18.504	2.424	18.404	--	2.347	17.352	2.343	17.238	--
GradientBoosting Regressor	2.461	18.756	2.048	13.883	2.034	13.800	0.008*	2.008	12.875	1.995	12.700	0.008*
LGBMRegressor	2.458	18.788	1.772	10.507	1.732	10.029	--	1.816	11.086	1.797	10.921	--
LinearSVR	2.115	20.693	1.925	16.601	2.078	17.218	0.056	1.747	11.439	1.782	10.775	1.000
MLPRegressor	2.483	18.794	1.901	10.922	1.876	10.760	0.421	1.721	10.144	1.697	9.964	0.310
PassiveAggressive Regressor	3.671	30.648	3.251	21.584	3.556	23.654	0.841	2.481	13.354	2.308	12.686	0.421
XGBRegressor	2.486	19.323	1.975	11.565	1.978	11.422	--	1.906	12.430	1.881	12.350	--
Average	2.581	20.618	2.186	14.619	2.235	14.838	--	2.008	12.529	1.975	12.225	--

The results for ICU readmission classification are presented in Table 4. Classification models without any embeddings, as baseline, performed consistently worse than the ones with embeddings. For models using FastText embeddings, the best accuracy score is achieved by the RandomForest Classifier (81.47%) with race-sensitive embeddings, while the best AUC score (0.825) is achieved by LinearSVC with neutral embeddings. For models using BERT embeddings, the RandomForest Classifier model trained from neutral embeddings achieved the best averaged accuracy score (80.48%), while the best averaged AUC score (0.797) is achieved by the SGD classifier with race-sensitive embeddings. Among the nine types of classifiers, the BERT embedding-based models consistently underperform the FastText embedding-based models, except the SGD classifier. Averaging the nine models’ performance based on embedding types, the best averaged accuracy score is 75.17% with FastTest race-sensitive embedding, and the best averaged AUC score is 0.775, the same for FastText neutral or race-sensitive embeddings. None of the p-values for the Mann-Whitney U tests are significant.

Table 4.

ICU readmission classification comparison between nine classification models utilizing neutral, race-sensitive embeddings, or without embeddings. The best performance single model is bolded and underscored. The best averaged performance of nine models for four different embeddings are bolded. Acc: Accuracy; AUC: Area under curve. “--” indicates that p-value calculation is not applicable. Mann-Whitney U test is performed on accuracy.

Model type	Without embeddings		FastText Embeddings					BERT Embeddings
	Baseline		Neutral		Race-sensitive		Mann-Whitney U test	Neutral		Race-sensitive		Mann-Whitney U test
	Acc	AUC	Acc	AUC	Acc	AUC	P-value	Acc	AUC	Acc	AUC	P-value
KNeighbors Classifier	63.06%	0.630	74.00%	0.792	74.90%	0.796	--	67.60%	0.753	68.10%	0.758	--
AdaBoost Classifier	65.45%	0.644	76.15%	0.803	74.91%	0.805	--	72.67%	0.765	71.40%	0.758	--
DecisionTree Classifier	65.78%	0.507	74.24%	0.615	73.16%	0.592	0.079	72.56%	0.579	72.85%	0.582	0.095
Linear SVC	61.81%	0.634	79.18%	0.825	78.75%	0.823	--	74.35%	0.780	74.98%	0.785	--
Logistic Regression	62.28%	0.636	73.19%	0.782	72.95%	0.785	--	73.53%	0.774	74.43%	0.777	--
MLPClassifier	63.89%	0.646	75.43%	0.800	75.13%	0.802	0.690	73.47%	0.788	73.43%	0.792	1.000
RandomForest Classifier	68.88%	0.578	81.69%	0.808	81.74%	0.808	0.841	80.48%	0.764	80.44%	0.765	1.000
Ridge Classifier	60.57%	0.632	78.75%	0.812	78.46%	0.813	--	73.77%	0.772	74.25%	0.775	--
SGDClassifier	62.05%	0.636	58.82%	0.738	66.54%	0.749	0.421	72.33%	0.792	72.69%	0.797	0.841
Average	63.75%	0.616	74.61%	0.775	75.17%	0.775	--	73.42%	0.752	73.62%	0.754	--

We evaluated the same estimator’s performance on Black and White populations separately, to check if the race-sensitive embeddings could improve performance on race-specific test datasets. In Table 5, we reported two best performance single models from previous evaluations (i.e., MLP Regressor for LHS task, Linear SVC for ICU readmission task) on Black and White populations, respectively. For the LHS task, the MLP Regressor with BERT race-sensitive embeddings consistently outperform the neutral ones on both Black and White populations. In addition, the race-sensitive one achieved a lower MAE (1.594) on Black population compared to White population (1.634), yielding a 2.4% improvement, which is harder to achieve given the higher average LHS among Black patients. For the ICU readmission task, the LinearSVC model with FastText race-sensitive embeddings achieved a slightly better accuracy score (75.4% vs. 75.3%) and AUC score (0.817 vs 0.816) on Black population, compared to the neutral embedding one. For White population, the neutral embedding model outperforms the race-sensitive one in both accuracy (79.3% vs. 78.7%) and AUC (0.825 vs. 0.823) scores.

Table 5.

Comparison of two best estimators’ performance for LHS and ICU readmission tasks on Black and White populations. Acc: Accuracy; AUC: Area under curve. MAE: Mean Absolute Error; MSE: Mean Squared Error.

				BERT embeddings
Task	Estimator	Race	LHS (days)	Neutral		Race-sensitive
				MAE	MSE	MAE	MSE
Length of Hospital Stay	MLPRegressor	Black	3.22±4.77	1.711	11.565	1.594	10.256
		White	3.09±4.03	1.684	9.390	1.634	8.862
				FastText embeddings
Task	Estimator	Race	Support	Neutral		Race-sensitive
				Acc	AUC	Acc	AUC
ICU Readmission	LinearSVC	Black	947	75.3%	0.816	75.4%	0.817
		White	5996	79.3%	0.825	78.7%	0.823

Discussion

Biases presented in data will inevitably propagate to the machine learning models that are trained using such data. Medical literature contains varying demographic underrepresentation biases that implicitly affect downstream diagnostic and predictive healthcare applications. In this work, we experimented with a method that leverages clinical trial enrollment data to rectify underlying population imbalance in training language embeddings and evaluated whether race-sensitive embeddings can lead to performance improvement on two important clinical prediction tasks. Using Black/African American and White races as a study case, we observed that models with race-sensitive embeddings slightly outperformed models with neutral embeddings, where the improvement is not statistically significant. In the task of predicting length of hospital stay, the MLP Regressor leveraging BERT race-sensitive embeddings achieved the smallest average MAE score of 1.697 and MSE score 9.964. In the task of classifying patients’ unplanned ICU readmission, the classification accuracy scores for models with race-sensitive embeddings increased only 0.56% (or 0.2%) than neutral embeddings for FastText (or BERT). This negligible increment reflected the almost identical performance between the neutral models and race-sensitive embedding models. This observation is also confirmed by the Mann-Whitney U tests for estimators, where only the GradientBoosting Regressor’s p-value (=0.008) is significant. For other estimators, we cannot reject the hypothesis that the performances are the same for a model leveraging race-sensitive embeddings or neutral embeddings.

For the ICU readmission classification task, the results in Table 3 shows that BERT embedding-based models have lower errors than FastText embedding-based models for the predicting length of hospital stay task. However, from Table 4 we observed that the BERT embedding-based models consistently underperform the FastText embedding-based models, except the SGD classifier. This indicates that transformer-based language models are not always superior to the traditional neural network language model. Even though BERT embeddings normally have a longer vector size than FastText (768 vs. 100 in our case), this may introduce complicated word representation and diminish the effectiveness of embeddings for certain tasks. For example, if a task relies mostly on rigid numeric/categorical features rather than textual features, longer embeddings only introduce unpredictable noises to the results.

The results in Table 5 indicated that embeddings built from literature on race-biased clinical trials can lead to race-specific gaps in performance. For both LHS and ICU readmission tasks, the estimators consistently performed better on White population than on Black population. It is probably due to those estimators inherited population bias in the training data, as White population consists of over 62% of all patient data. Our attempt to rectify this bias using race-sensitive embeddings achieved a limited improvement for the LHS task on Black population, where the race-sensitive embeddings-based model performed better than the neutral embedding one in terms of MAE and MSE. The same trend is also observed for the ICU readmission task on Black population.

Limitations and Future work

There are several possible explanations for why the inclusion of the race-sensitive embeddings did not improve the performance of the model as suggested by previous literature. Self-reported race categories introduced data quality issues as described in the Results section. We simply chose the race enrollment ratios as the bootstrap policy in the hope of proliferating studies with more balanced population enrollment. The heuristic of enrollment ratios that were used for corpus construction was not rigorously evaluated. It remains unclear whether the embeddings generated from this corpus optimally capture the semantics for characterizing different race groups. It will require domain knowledge to design the evaluations that unveil the quality and differences between neutral and race-sensitive embeddings.

The sequence of diagnoses and procedures in a patient’s record, manifested as clinical events, carries temporal information, which plays an essential role in understanding the state of the patients and could enhance clinical information applications. In this work, this temporal dimension, such as the order of diagnoses or procedures, was not sufficiently maintained in the conversion to embeddings. In conversion of BERT-based embeddings, a list of diagnosis (or procedures) was input as one whole string, and the mean of BERT model’s last two layers’ output was considered as an encoded vector. For FastText, we simply took the sum of each token’s embeddings as the encoded vector. Future work should explore different textual encoding schemas to maintain as much temporal information as possible, with a potential of further enhancing the downstream applications’ performance.

We did not finetune each estimator’s (classifier or regressor) hyperparameters to optimize their performance for each different embedding model separately. Though this may lead to unfair comparisons between different types of estimators, it was sufficient to compare one type of estimator among different embeddings. In future work, exhaustive optimization for an estimator can be achieved by manually specifying the hyperparameter space. More advanced neural network based (such as Convolutional Neural Networks and LSTM) classification or regression models can also be implemented to discover the best performance single estimator.

Conclusion

We proposed a method for leveraging metadata from clinical trials and abstracts of clinical literature to build race-sensitive word2vec and BERT-based biomedical embeddings, and evaluated its efficacy for clinical prediction tasks with de-identified EHR data. Our results showed that the performance differences between race-sensitive embeddings and neutral embeddings are not statistically significant in predicting length of hospital stay and classifying ICU readmission based on MIMIC-IV datasets. This work demonstrated that language embeddings trained from corpus that profiling PubMed articles based on clinical trial metadata, specifically race enrollment data, do not necessarily significantly improve the accuracy of health predictive models as previously reported in the literature.