Detecting Dataset Bias in Medical AI: A Generalized and Modality-Agnostic Auditing Framework

We first consider shortcut risks in vision-based tasks. The high dimensionality of image data creates many opportunities for shortcuts to exist without directly impacting the machine learning task. For instance, hospital-specific tokens placed in the chest X-ray field of view (e.g., [50]) may impact only a small number of pixels in the image but may constitute a shortcut when associated primarily with a specific disease condition in the training dataset. In cases like these, DNNs trained on the data may exploit salient features like tokens/watermarks (or other statistical regularities not task-related) to achieve low training error. We focus here on skin lesion classification where the construction of large-scale datasets may not be able to adequately balance across patient characteristics, clinical sites, and dermascopic imaging sensors/settings and where such features of the dataset and collection process may manifest as shortcuts.

The ISIC 2019 skin lesion dataset [43, 6, 7] was analyzed for bias with respect to the included metadata and patient attributes. The dataset consists of 25,331 training images and attributes age, race, sex, anatomical location, and skin color on the Fitzpatrick scale. Image metadata includes height, width, and year of collection. While the original task labels included seven diagnostic categories, we reduce the classification task to malignant vs. benign conditions (e.g., [51, 29, 31]). The data auditing procedure (see Sec. 4) was applied to the entire training dataset relative to the binary malignancy classification task and excluding any images with missing attribute values.

The main auditing results are found in Figure 2 where the height, width, and year attributes exhibit the highest combination of utility and detectability, indicating these attributes are more likely sources of bias within the dataset. While these particular attributes may seem unrelated to the task labels, this image metadata can act as a proxy for the camera type and/or clinical site. Furthermore, while all images are resized and cropped to a fixed resolution prior to running the auditing procedure, the high detectability scores indicate that some of this information is still retained in the images themselves, either directly or via proxy.

We next consider the text domain where the dimensionality of the data remains high and the variability of natural language create unique opportunities for shortcuts to exist. Similarly to the image domain, DNNs are often used to learn compact representations of natural language text for various tasks and may exploit shortcuts in the data that are not relevant to the clinical task yet allow them to achieve low training error. We apply the G-AUDIT procedure to the electronic medical record dataset introduced by [19] for the purpose of characterizing stigmatizing language usage by physicians. The dataset contains 5,201 annotated instances across 3 tasks, with each task focusing on a different thematic group of stigmatizing language — Credibility & Obstinacy, Compliance, and Descriptions of Appearance/Demeanor. Models are provided a window of text centered around a keyword or phrase which has been identified by domain experts as being a potential indicator of unconscious bias towards a patient. They are asked to characterize the implication of the input instance (e.g., “the patient claims to brush their teeth 2x daily”; “unable to track down insurance claims”). Each instance is associated with auxiliary attributes which indicate a patient’s race and gender, and the clinical setting from which the statement was drawn (e.g., OB-Gyn, Surgery).

The stigmatizing language dataset serves as an interesting case study for several reasons. First, the authors of [19] included an analysis which examined how well each attribute could be predicted based on the last embedding layer of the primary stigmatizing language task models. Their results provide a direct reference for our method. Second, as is often true in practice, the stigmatizing language dataset has an ambiguous causal structure. This allows us to evaluate the robustness of our method to potential errors in misspecification of the direction of dependency between attributes and labels. Finally, there are documented disparities in the prevalence of stigmatizing language between demographic groups which we expect to show up directly within our utility measure (e.g., Black patients are more likely than white patients to experience discrimination). To gather an unbiased estimate of the prevalence of stigmatizing language in the population, models should not make use of demographics as a predictive shortcut.

As shown in Figure 3, clinical specialty had a higher utility than both patient race and sex for the compliance and appearance / demeanor tasks. Prior work shows that downstream models for these tasks likely did not encode sex and race characteristics beyond what could be explained by a reliance on clinical specialty alone [19]. This is consistent with our observation that clinical specialty had a higher utility than race and sex, implying a stronger association with the task label. Importantly, this does not preclude performance disparities based on these sensitive attributes—clinical specialties in the JHM dataset include OB-GYN with an all-female patient population and Pediatrics with an approximately 95% Black patient population [19]. Instead, our results suggest that downstream models are more likely to exploit shortcuts related to the identification of different clinical domains than to directly encode race or sex to improve performance.

In terms of detectability, we find differences between conditioned and unconditioned measures to be fairly small. Interestingly, within the Credibility and Obstinacy task, while all attributes had relatively low utility, the detectability of sex was higher than that of any other EHR task and attribute evaluated. This presents a potential explanation for the findings of [19], which identified sex within the Credibility and Obstinacy task as the only demographic attribute recoverable above baselines levels across all three tasks.

Lastly, we consider the dataset auditing problem in the context of tabular data. While not as feature rich or high dimensional as images and text, tabular data presents its own unique challenges from an auditing perspective. While the features learned by task models in the image and text domains often lack interpretability (e.g., DNN embeddings), tabular data provides a direct mapping between attributes and their measured values. In these cases, attributes used as inputs to a task model can still act as shortcuts when they exhibit undesirable associations with the task labels that could arise due to sample bias or incorrect usage of the attribute itself.

For instance, if clinical sites differ in their test-ordering protocols, associations between disease conditions and clinical site may be reflected in the test orders/results provided to machine learning task models [42]. Furthermore, the gap in task performance between DNNs and more traditional ML models (e.g., SVMs, ensemble methods, etc.) is much smaller in the tabular case and this provides an opportunity to measure both detectability and task model performance over a wider class of machine learning models.

Here, we evaluated our auditing methodology on tabular data extracted from the publicly available Medical Information Mart for Intensive Care (MIMIC) III Dataset [23]. Specifically, we extracted a dataset for predicting mortality in Intensive Care Unit (ICU) patients using features involved in the Simplified Acute Physiology II (SAPS II) score [25, 39], a score used to measure disease severity in patients after their first 24 hours in the ICU. These features include patient age as well as summary statistics of heart rate, systolic blood pressure, temperature, labwork indicators/results, ventilator usage, and Glasgow Coma Scale (GCS) during the first 24 hours of the patient’s ICU stay.

Our task is to predict patient mortality based on the tabularized data of a given patient. The final processed dataset consists of 34,386 patient records and 40 features, detailed in Appendix B. Features are one-hot encoded and include medication details, missing variables (e.g., GCS or lab test results), and patient demographics such as race and insurance coverage. We select an FT-Transformer [18] with default hyperparameters for use as a task model in all experiments. A benefit of this selection is the ability to apply the SPLIT approach [17, 15] to empirically test how well our detectability measure corresponds with other baseline approaches for measuring detectability [17, 15].

As shown in Figure 5, we find that, across model classes, G-AUDIT-based detectability correlates strongly with SPLIT-measured recoverability, with Spearman’s $\rho$ values of .92, .92, .76, .84, .75, and .86 for decision tree, random forest, logistic regression, FT-Transformer, XGBoost, and naive Bayes task models respectively (all $p<.0001$). Importantly, both obvious shortcuts such as ‘temperature missing’ -—which is easily detectable because it is explicitly represented in the input data as a negative placeholder value for temperature—- as well as less obvious instances like dopamine, norepinephrine, vasopressin, ventilator and IV usage are highly detectable. We do note some variability in performance that does not always seem correlated with model strength as measured in Figure 4. This suggests using a small ensemble of models of varying complexity may provide a more holistic view of detectability.

By running our synthetic calibration process (Sec. 4.4), we identify that some of the attributes in our set could cause worst case performance drops as great as  0.2 AUC for an average task model assuming full detectability (see Figure 6).

The core objective of our data auditing procedure is to identify and quantify the degree to which each attribute of the data represents a potential learning shortcut.

For our purposes, datasets consist of samples ($X$), associated metadata or attributes ($A$), and task labels ($Y$). We establish the existence (not direction) of a relationship between $A\leftrightarrow Y$ through a measure we call utility. The utility of an attribute refers to our ability to infer the value of the task label $Y$ simply by observing $A$. In the extreme case, if $A$ is perfectly correlated with $Y$, then machine learning models simply need to detect $A$ in order to correctly solve the task.

However, a large value for an attribute’s utility is not sufficient to consider it as a shortcut. For this, we require to know the attribute’s detectability. Since machine learning models will typically only take $X$ as input, detectability measures the extent to which the value of $A$ can be inferred from $X$.

Attributes with high-utility and high-detectability represent the greatest risk for biasing downstream task models. However, while we use utility and detectability as a proxy for risk, features that are causally relevant to the task and have high utility and detectability are useful, rather than representing possible shortcuts. We must rely on domain expertise to determine whether high-risk attributes are reasonable features or unanticipated dataset flaws.

We use information theory and principles of causal inference to measure the utility and detectability of dataset attributes. Utility is measured as the mutual information, $MI(A;Y)=H(Y)-H(Y|A)$ after adjusting for chance as per [44, 45] and represents the reduction in uncertainty about the value of $Y$ by observing $A$ after adjusting for the entropy of the underlying distributions and chance based on the number of categories. We rely on the faithfulness assumption which implies that if there is a relationship between $A,Y$, then $MI(A;Y)>0$.

For detectability, the objective is to determine the ability to infer the value $A$ from the data and we accomplish this by first training a surrogate model $f:A\rightarrow X$ to predict $A$ from $X$. We next leverage clinician and domain expertise to establish the likely directions of dependency between data and labels. The primary role of this step is to understand whether the relationship is $Y\rightarrow X$ or $X\rightarrow Y$. Our ultimate goal is to identify whether there is a relationship $A\leftrightarrow X$ and if so measure its strength. However, we need to control for the potential of information leaking through $Y$ and this process relies on the assumed/known relationship between $Y,X$.

In anti-causal scenarios (i.e., $Y\rightarrow X$), we know it could be the case that $A\rightarrow Y\rightarrow X$ and as a result information could be leaked about the attribute into $X$ not by a shortcut, but by solving the task itself. To control for this, we condition on $Y$ during the training process of our surrogate attribute prediction model $f:X\rightarrow\hat{A}$. Specifically, we assume either $Y$ is already discrete or define $Y^{D}$ as a sufficiently granular discretization of $Y$. Then, we partition $X,Y^{D}$ into disjoint training subsets $S$ s.t. for given subset $S_{i}$, all $y^{D_{j}}\in S_{i}$ have the same value. For each subset $S_{i}\in S$, we train separate surrogates $f_{i}(X)=\hat{A}$. By training separate surrogates in this manner, we ensure that differences in task labels $Y$ that could be used by $f_{i}$ to predict a given attribute are minimized. Based on test set predictions obtained using cross-validation with all $f_{i}$, we obtain a full set of predictions ($\hat{A}$) for the entire dataset and we can similarly measure $MI(A;\hat{A})$ to understand how recoverable $A$ is from $X$ while reducing the impact of task relevant information.

In the causal case (i.e., $X\rightarrow Y$), we are able to directly estimate the relationship without conditioning. In fact, we cannot condition on $Y$ since $Y$ represents a collider (i.e., $Y$ may be dependent on both $X$ and $Y$). In that case, measuring the strength of relationship between $X\leftrightarrow A$ given $Y$ could give falsely inflated results since conditioning on a collider $Y$ creates an otherwise non-existent association between its parents ($X,A$). For example, if two conditions both increase mortality rate through different pathways, then given that a patient has died, knowledge of either attribute provides information about the other (e.g., given $A$ is the presence or absence of trauma wounds and $X$ contains information relating to disease, a patient who died but did not have trauma wounds is more likely to have had disease: $A\leftrightarrow X|Y$ even though $A\perp X$ ). As a result, we do not condition for cases where we expect $X\rightarrow Y$ and instead directly estimate $MI(A;\hat{A})$.

A key aspect of the data auditing procedure is that every sample in the dataset contributes to the calculation of attribute utility and detectability. For determining utility, only the labels and metadata are required and can be estimated without any model training. However, detectability requires attribute prediction, so we ensure that every sample in the dataset contributes to the detectability estimate using a cross-validation procedure that ensures unbiased predictions of $\hat{A}$.

We first partition the full dataset into $K$ disjoint folds. For each fold, we hold out the data of that fold for testing and use the data from the $K-1$ remaining folds to train or finetune a sufficiently expressive machine learning model to predict $\hat{A}$ from $X$. Given that trained model, we predict $\hat{A}$ on the held-out fold’s data. After repeating the train/test procedure for each fold, and, as necessary due to the dependency between attribute and label, each value in $Y^{D}$, we aggregate all predictions together to compute detectability measures.

While utility and detectability measures provide a quantitative measure of bias, how to interpret them in the context of potential model performance risks is less clear. The magnitudes of these measures are not easily comparable across datasets so we cannot rely on thresholds or conventions to directly assess risk. We could instead look to the relative rankings of attributes with respect to utility/detectability to understand attribute risk, but cannot directly translate values to drops in task-relevant metrics like AUC.

To address this concern, we provide a supplementary method for generating an upper bound on performance risk with which attribute utility can be interpreted. To construct the bound, we first create a synthetic attribute which we can insert with 100% detectability (e.g., a visible watermark in a fixed image position, a single token added to a text input, an added column to tabular data). We then vary the utility of this synthetic attribute and train a task model for each case. We evaluate the task models on synthetic and counterfactual data distributions and measure the resulting performance drop between the two distributions. The counterfactual data distribution is constructed in such a way as to create the worst case scenario where the synthetic attribute is anti-correlated with the true label and any shortcut exploited during training will yield worst-case behavior at test time.

Formally, let $X$ be the feature data, $Y$ be the true binary task labels (of length $N_{Y}$), and $A$ be the synthetic artifact values. For simplicity, we assume $A$ is also binary and initialize the values of $A$ to have the same values as $Y$ in the dataset (i.e., the normalized utility would be $1.0$ when $A=Y$). To vary the utility, we randomly select an index set ${\mathcal{I}}$ for $N<\frac{N_{Y}}{2}$ rows of $A$ and flip the values in those rows (i.e., if $i\in{\mathcal{I}}$ then $A[i]=\neg Y[i]$). This preserves the existence of the relationship between $Y,A$ but reduces its strength. For each $i\in{\mathcal{I}}$, we also insert a fully detectable artifact into $X$. We run the same training and testing procedure as described in Section 4.5 to get the baseline task performance numbers for the synthetic distribution.

To construct the counterfactual distribution, we create an additional test set as follows. For each image in the test set, we insert the synthetic artifact $A_{C}$ to be anti-correlated with the label such that $Y=0\Rightarrow A_{C}=1$ and vice versa. In the worst case, if the task model exploited the synthetic artifact shortcut during training, then at test time, it will be more likely to predict $\hat{Y}=1$ when $A_{C}=1$ resulting in more errors and lower AUC. As the utility of $A$ increases in the training set, the risk of performance degradation in this context also increases.

We view this as an approximate upper bound on risk because of the fact that we have controlled for the detectability of the synthetic attribute and thus can understand the performance risks associated with each value of utility. To assess risk for attributes in the original dataset, we can look at the worst-case AUC drop relative to the measured utility for the original attributes.