Exploring trade-offs in equitable stroke risk prediction with parity-constrained and race-free models

Abstract

A recent analysis of common stroke risk prediction models showed that performance differs between Black and White subgroups, and that applying standard machine learning methods does not reduce these disparities. There have been calls in the clinical literature to correct such disparities by removing race as a predictor (i.e., race-free models). Alternatively, a variety of machine learning methods have been proposed to constrain differences in model predictions between racial groups. In this work, we compare these approaches for equitable stroke risk prediction. We begin by proposing a discrete-time, neural network-based time-to-event model that incorporates a parity constraint designed to make predictions more similar between groups. Using harmonized data from Framingham Offspring, MESA, and ARIC studies, we develop both parity-constrained and unconstrained stroke risk prediction models, then compare their performance with race-free models in a held-out test set and a secondary validation set (REGARDS). Our evaluation includes both intra-group and inter-group performance metrics for right-censored time to event outcomes. Results illustrate a fundamental trade-off in which parity-constrained models must sacrifice intra-group calibration to improve inter-group discrimination performance, while the race-free models strike a balance between the two. Consequently, the choice of model must depend on the potential benefits and harms associated with the intended clinical use. All models as well as code implementing our approach are available in a public repository. More broadly, these results provide a roadmap for development of equitable clinical risk prediction models and illustrate both merits and limitations of a race-free approach.

INTRODUCTION

Stroke is a leading cause of death and disability in the United States. Almost 800k Americans experience a stroke each year, most of which are first attacks.¹ Stroke risk is strongly associated with social determinants of health, including income and education,¹ and Black Americans both experience stroke at higher rates (>50%) and have worse post-stroke outcomes.^2,3

The American Heart Association/American Stroke Association recommend using stroke risk prediction models to guide screening and interventions.⁴ Whereas the Framingham Stroke Risk Profile was developed in a predominantly White population,⁵ data from recent prospective cohort studies allow stroke risk prediction models to be developed and evaluated in racially and ethnically diverse populations. Such studies include the atherosclerosis risk in communities (ARIC) study,⁶ the multi-ethnic study of atherosclerosis (MESA),⁷ and the reasons for geographic and racial differences in stroke (REGARDS) study.⁸

Recently, our team harmonized data from the ARIC, MESA, REGARDS, and Framingham Offspring⁹ studies, then investigated the performance of stroke risk prediction algorithms including in Black and White subpopulations.¹⁰ Our evaluation included the Framingham Stroke Risk Profile, the Revised Framingham Stroke Risk Profile based on self-reported measures from REGARDS,¹¹ and the pooled cohort equations, which predict risk of cardiovascular disease events including stroke.^12,13 Findings revealed disparities between groups in model discrimination performance, with all models performing better for White individuals than Black individuals. Notably, use of standard machine learning risk prediction techniques, including random survival forests¹⁴ and regularized Cox proportional hazards,¹⁵ did not substantially improve performance overall or lessen these disparities. This result aligns with other studies suggesting that machine learning does not substantially improve performance when applied to tabular clinical data, including in cardiovascular disease risk prediction¹⁶ and prediction of death after myocardial infarction (MI);¹⁷ as well as work showing that common machine learning models have notable deficits in fairness across racial/ethnic groups.¹⁸ More generally, there is a recognize need to develop fairness methods for a wider range of machine learning algorithms, and to apply such methods more broadly, including in healthcare.¹⁹

A commonly proposed approach to improve the fairness of machine learning models is to remove race and other sensitive attributes as predictors (i.e., race-free models),^20,21 recognizing that including race in clinical prediction models can harm disadvantaged groups.²² However, removing race can have limited impact due to associations between race and other available predictors,²³ therefore a variety of machine learning methods have proposed to explicitly constrain differences in model predictions between racial groups, ensuring that race does not impact predictions, either directly or indirectly.²⁴ Specifically, representation learning methods can be used to limit information about race (and potentially other sensitive attributes) among latent features used to make predictions, leading them to be more similar (in distribution) between racial groups.²⁵ However, while this approach can improve the model’s ability to discriminate between cases and controls both within and between groups, it also tends to produce less accurate risk predictions (i.e., poorer calibration), and can reduce discrimination performance across the population as a whole.^26,27 Indeed, there is a fundamental trade-off between calibration and discrimination performance such that achieving equal error rates in all groups necessarily implies imperfect calibration.^28–30 Depending on the benefits and harms associated with a particular use case, imperfect calibration can be justified and the corresponding model fair,³¹ and in general, defining fairness depends on the particular clinical context.³² In some cases, however, poorer calibration seems to be clear violation of fairness, for example when underpredicting risk in disadvantaged groups.³³

With these considerations in mind, our goal was to explore the use of representation learning methods to address fairness in stroke risk prediction and compare to a race-free approach. We hypothesized that this approach would allow us to select a trade-off between discrimination and calibration that is appropriate for a given use case. We further hypothesized a race-free approach would not yield meaningful improvements in fairness due to associations between race and other available predictors known to be valuable when predicting stroke risk (e.g., hypertension).

We proceed by first proposing a machine learning method for time-to-event prediction that is parityconstrained, meaning that it is designed to encourage predictions to be independent of an attribute of interest, in this case racial status (i.e., racial parity). Our approach couples representation learning with a discrete-time, neural network-based model previously shown to improve time-to-event prediction in other clinical settings.^34–36 This approach is particularly advantageous for clinical settings in which factors affecting event occurrence may be distinct from those affecting the corresponding event times.³⁶ Using harmonized data from Framingham Offspring, MESA, and ARIC studies, we develop parity-constrained stroke risk prediction models and compare their performance with race-free models both in a held-out test set and in a secondary validation set (REGARDS).

Our evaluation includes standard time to event metrics as well as newly developed metrics designed to quantify fairness in a time to event setting with censoring. Specifically, we use the xCI, a group-conditional modification of Harrell’s concordance index (CI).³⁷ The xCI extends the xAUC metric, which was recently proposed to evaluate the bias of binary prediction models,³⁸ to the survival analysis setting. We also include race-conditional modifications of several other discrimination metrics suitable for models predicting right-censored events. Although many competing fairness definitions and metrics have been proposed,²⁴ the xAUC and xCI are well suited to evaluation of clinical decision support. Like the AUC (area under the receiving operating characteristic curve) and CI, they are agnostic to the specific decision threshold. This allows them to be used throughout the model life cycle, including in situations where the decision threshold has not yet been determined or depends on clinical context. Furthermore, these metrics have an intuitive interpretation: they quantify how often the model correctly ranks individuals by risk when the higher-risk individual is in one subgroup, and the lower-risk individual is in a different subgroup.

Equipped with these metrics, we conducted a rigorous evaluation to compare the performance and fairness of our proposed approach with models not incorporating a parity constraint, including race-free models. The results of our evaluation demonstrate that our approach allows us to control the aforementioned trade-off between fairness in calibration and fairness in discrimination (i.e., equal error rates) without loss of performance across the population as a whole (i.e., test set, REGARDS) and confirm that a race-free approach does not result in parity in predictions between groups. However, consistent with results in related work,²⁷ application of the parity constraint led to poorer calibration, notably including underprediction of risk in the Black subpopulation. This suggests that parity-constrained models should be used with caution and may not be appropriate for stroke risk prediction specifically. By the same token, however, they are likely to be effective in promoting parity in settings where unconstrained models tend to underpredict risk in disadvantaged groups.

These results provide a roadmap for development and tuning of parity-constrained clinical risk prediction models, and illustrate limitations of a race-free approach. Our models are available in a public repository for use as clinical decision support, and source code needed to implement the proposed model and metrics are also available in a public repository.

METHODS

Data

All study procedures were approved by the Duke Heath Institutional Review Board.

Data used in this analysis were obtained from the Biologic Specimen and Data Repository Information Coordinating Center (BioLINCC) and harmonized across multiple cohorts using the same procedures used in previous work.¹⁰ Here we briefly summarize the exclusion and exclusion criteria, predictors, and preprocessing steps, including imputation of missing values. A comprehensive description of these steps has been previously reported.¹⁰

Included participants were of White or Black race, free of stroke or transient ischemic attack, and 45 years of age or older. Data from examinations 3 (years 1983–1987), 6 (years 1995–1998), and 8 (years 2005–2008) in Framingham Offspring; 1 (years 1987–1989) and 4 (years 1996–1999) in ARIC; and 1 (years 2000–2002) in MESA were combined to create our analysis sample. In studies with multiple examinations, some participants contributed multiple data points when they were free of stroke at the start of follow-up. Predictors include sociodemographics (age, sex, race), vital status (weight, height, body mass index, systolic and diastolic blood pressure), medical history (diabetes, hypertension, history of heart disease including prior myocardial infarction, atrial fibrillation, left ventricular hypertrophy), family history of stroke, labs (lipids, fasting blood glucose, creatinine), medications (statins, any cholesterol medications, insulin, aspirin), smoking and alcohol, general health status, fruits and vegetables consumption (quintiles), sodium intake (quintiles), physical activity, income and education. Categorial variables were imputed to their cohort-specific mode, and continuous variables were imputed to their cohort-specific median. The outcome for this analysis was ischemic or hemorrhagic stroke following the definition used in previous analysis.¹⁰ Time to the first stroke or censoring was measured in days since the baseline examinations. Participants with more than 12 years of follow-up were censored at 12 years.

Discrete-Time Neural Network Model

We use a discrete-time, neural-network-based time-to-event (i.e., survival analysis) model to predict the probability of stroke occurrence within each of the next i years, where i ranges from 1 to 12, as well as the probability that stroke will not occur before the 12-year time horizon. By predicting the probability of event occurrence in pre-defined time intervals – in this case, 1-year intervals – the discrete-time approach avoids strong assumptions made by other models, such as proportional hazards³⁹ or a particular parametric form⁴⁰ for the event density. Consequently, it has been shown to improve performance across a range of prediction settings.^34–36

Our neural network is comprised of a) an encoder, which maps input features to a set of latent features; and b) a decoder that predicts the distribution of stroke occurrence based on these latent features. As described in the next section, this structure contributes to fair prediction of stroke occurrence by providing latent features to which our parity constraint can be applied. The constraint encourages the latent features to have similar distributions in Black and White participants, implying that they do not contain information needed to distinguish between the two groups, and thus promoting fairer predictions. For the encoder, we use a single fully connected layer with 256 neurons and ReLU activation. For the decoder, we use a single hidden layer with 256 neurons and ReLU activation followed by a prediction head with softmax activation. The prediction head predicts the probability of event occurrence in each time interval as well as the probability of event occurrence beyond the 12-year time horizon.

Our prediction loss, including handling of right-censoring, which we assume is non-informative, derives from the discrete failure time model described by Kalbfleisch and Prentice.⁴¹ We aim to find parameters that minimize the negative log-likelihood and therefore maximize the likelihood of the observed data under this model. More complete descriptions with full mathematical details are provided in Sections 1.1 and 1.2 of the Supplemental Material.

Encouraging Racial Parity through Balanced Latent Representations

To encourage our model to make its predictions more similar (in distribution) between groups, we predict stroke risk from latent features that are derived from the original features (i.e., predictors) by our neural network, but constrained to contain limited information about racial status. This is achieved by incorporating an additional loss term based on the Maximum Mean Discrepancy (MMD),⁴² a statistical measure that can be used to quantify the difference between two (unknown) probability distributions from samples. Here we use the MMD to quantify the difference between the latent representations of Black and White individuals in our neural network. By minimizing the MMD together with the negative log-likelihood of our model, we aim to make predictions more similar between groups without substantially compromising overall prediction performance. The tradeoff between parity in error rates, as quantified by the xCI, and overall discrimination performance, is tuned as a hyperparameter of the model, as described in the next section. A more complete description with full mathematical details is provided in Section 1.3 of the Supplemental Material.

Model Development Procedure

Framingham Offspring, ARIC and MESA were divided at random into a training set (60%) used to optimize model parameters, a validation set used to tune hyperparameters (see below) and select final models (20%), and a test set used to estimate out-of-sample prediction performance (20%). Data from the REGARDS study was used as a secondary evaluation set.

In addition to the negative log-likelihood (See Section 1.2 of the Supplemental Material, our training objective includes the aforementioned parity constraint and an additional parameter shrinkage term (L2 regularizer) to mitigate overfitting. Hyperparameters tuned during model development therefore included the learning rate as well as weights associated with the MMD-based parity constraint (λ_MMD) and L2 (i.e., ridge) regularization (λ_L2) on all model parameters. Training took place in mini-batches of size 1000 using the Adam optimizer⁴³ for up to 100 epochs, with early stopping based on loss on the validation set. All model development was conducted using TensorFlow 2.4.

To explore the effects of removing race in addition to the effects of our parity constraint, we trained a first set of models that included race as a predictor, and a second set of models that did not. Three models were selected for final evaluation on the test set and REGARDS. The first, which we call the unconstrained model, is the model with highest CI on the validation set. Importantly, however, xCI within or between racial subgroups may be substantially lower. The second, which we call the race-free unconstrained model, is the model with highest CI on the validation set among all race-free models. The third, which we call the parity-constrained model, is the model with highest minimum xCI. In other words, this is the model for which the worst observed xCI performance (either within or between racial subgroups) is highest.

We also compare performance to race-aware and race-free L2-regularized Cox proportional hazards models, which are a common baseline in time-to-event settings. These models were implemented using the scikit-survival library.⁴⁴ These models were trained, tuned, and evaluated using the same data partition (i.e., training set, validation set, test set) used for the other models.

Source code and final models are available at https://github.com/engelhard-lab/equitable-stroke-risk-prediction.

Performance Metrics for Evaluating Fairness

To quantify fairness in discrimination performance, we use an extension of the Concordance Index (CI)³⁷ inspired by the xAUC concept proposed by Kallus et al.³⁸ In intuitive terms, this metric, the xCI, quantifies a model’s ability to maintain consistent discrimination performance when comparing individuals across different subgroups (e.g., racial subgroups) within the population. Specifically, it assesses whether the model correctly ranks individuals from one subgroup relative to individuals from another subgroup based on their observed event times.

The xCI is calculated by finding all pairs of individuals i and j from subgroups A and B, respectively, for which (1) individual i had an observed event at some time t_i; and (2) individual j had an event or was censored at a later time t_j > t_i. This metric quantifies the probability that a randomly selected individual from group A with an observed event will be (correctly) ranked higher than a randomly selected individual from group B for whom the event happened later or not at all. In stroke risk prediction, for example, the Black-White xCI quantifies the probability that a Black patient who had a stroke at time t would be correctly prioritized for intervention over a White patient who had not yet had a stroke by that time. For a fair model, we would expect that the probability that one patient is correctly prioritized over another does not depend on the group status of either patient. This condition is satisfied when the xCI is equal across all pairs of groups (i.e., Black-White, White-Black, Black-Black, White-White).

As shown by Uno et al.,⁴⁵ an unbiased estimate of the CI may be obtained by weighting each pair of individuals’ contribution to the CI calculation by the inverse probability of censoring. We therefore use an inverse probability of censoring weighted xCI (IPCW xCI), with weights obtained following Uno’s approach,⁴⁵ as our primary evaluation metric.

We also use cross-group extensions of several additional discrimination metrics: (a) the x-receiver operating characteristic (xROC) curve, (b) the x-precision-recall (xPR) curve, (c) the xAUC, and (d) the x-average precision (xAP) to the survival analysis setting by defining time-varying versions of each. For example, the xAUC_t is defined as the xAUC when defining cases as individuals with an observed event at time τ ≤ t and controls as individuals who had an event or were censored at a later time τ > t.⁴⁶

Full mathematical details for all metrics are presented in Sections 2.1 and 2.2 of the Supplemental Material. Source code to calculate all fairness metrics is available at https://github.com/engelhard-lab/fair-tte-performance.

Additional Performance Metrics

In addition to the previously described xCI, we use the following metrics to quantify performance overall and within racial subgroups:

Concordance Index (CI):

Harrell’s concordance index³⁷ quantifies the degree to which model-predicted stroke risk is consistent with the order of observed events.

Inverse Probability of Censoring Weighted Concordance Index (IPCW CI):

Incorporating inverse probability of censoring weights when estimating the CI, as well as when estimating the xCI, yields unbiased estimates by compensating for biases due to the censoring distribution.⁴⁵

Calibration in the large quantifies the overall agreement between the average predicted probabilities and the average observed event rates, with values close to zero indicating good agreement. The calibration slope quantifies the model’s ability to differentiate between different risk levels. A slope of 1 indicates perfect calibration, while values less than 1 indicate underprediction of risk, and values greater than 1 indicate overprediction of risk. Both calibration metrics were calculated by first dividing model predicted risk into deciles. In each decile, average model-predicted risk at time t was compared to the observed event rate at time t, as determined via Kaplan-Meier estimator. Calibration in the large and the calibration slope were then reported as the intercept and slope, respectively, of the linear regression of observed event rates on model-predicted risk.

AUC_t:

The area under the receiver operating characteristic (ROC) curve at time t (AUCt)⁴⁶ quantifies the model’s ability to distinguish between individuals with observed events before t versus individuals whose observed event time or censoring time is later than t, indicating that the event did not occur by time t. A value of 0.5 indicates that prediction is no better than random guessing, while a value of 1 represents perfect discrimination. Similar to the CI, The AUCt can be weighted to correct for effects of right censoring, but here we use the unweighted estimator.

AP_t:

The average precision at time t (APt) approximates the area under the precision recall curve at time t (PRt) by quantifying the model’s precision (i.e., positive predictive value) averaged across all recall (i.e., sensitivity) levels when distinguishing between individuals with observed events before t versus individuals whose observed event time or censoring time is later than t, indicating that the event did not occur by time t. at a specific time point (t). A value equal to the cumulative event rate at time t indicates that prediction is no better than random guessing, while a value of 1 represents perfect precision. Similar to the CI, The APt can be weighted to correct for effects of right censoring, but here we use the unweighted estimator.

The Brier Score at time t is the mean squared difference between the model-predicted cumulative event probability at time t and a binary variable indicating whether the event was observed by that time. The Brier score is affected by both discrimination and calibration performance,⁴⁷ with a score of 0 corresponding to perfect predictions.

Confidence intervals for all metrics were obtained by via bootstrapping. All performance were evaluated using model predictions at t = 10 years.

RESULTS

Description of Cohort and Dataset

Baseline characteristics and outcome rates for a similar harmonized cohort have been previously described.¹⁰ However, the current work is based on data obtained from BioLINCC and harmonized across multiple cohorts.

The development dataset included 8,217 samples from Framingham Offspring, 24,509 samples from ARIC, and 4,490 samples from MESA. Each sample reflects a single patient visit. In the development dataset, the mean age was 59.0 years (standard deviation, SD 8.3), 54.1% were female, and 20.9% were Black. In REGARDS, the mean age was 64.4 years (SD 9.4), 55.3% were female, and 40.8% were Black. The average stroke rate per 100 participant-years of follow-up was 0.342 (1267 events) over the 12-year period. Baseline characteristics and outcome rates are presented in the Supplemental Material (sTable 1).

Rates of missingness were low (<5%) for all predictors except fasting glucose, health status, family history of stroke, left ventricular hypertrophy, consumption of fruits and vegetables, creatinine, and sodium intake where the missingness rates varied between 8.8% and 24.4%. Missingness rates for all predictors are presented in the Supplemental Material (sTable 2).

Discrimination Performance and Fairness

Discrimination performance across the test set as a whole was similar between all models; Differences in discrimination performance were not statistically significant in both evaluation samples (test set, REGARDS) despite the large sample sizes. In the test set, IPCW CI ranged from 0.75 (0.72–0.77) (race-free Unconstrained) to 0.74 (0.71–0.77) (Parity-Constrained). In REGARDS, IPCW CI ranged from 0.69 (0.68–0.70) (race-free Unconstrained) to 0.68 (0.67–0.69) (Parity-Constrained). Discrimination performance was also similar between all models in Black and White subgroups in both the test set and REGARDS (see Table 1).

Table 1:

Model performance summary

	Cox Model	Unconstrained Model	Race-Free Cox Model	Race-Free Unconstrained Model	Parity-Constrained Model
	TEST	TEST	TEST	TEST	TEST
IPCW CI mean (95% CI)
All	0.743 (0.718–0.771)	0.741 (0.709–0.767)	0.743 (0.718–0.772)	0.745 (0.716–0.771)	0.739 (0.710–0.767)
Black	0.742 (0.698–0.784)	0.730 (0.682–0.770)	0.741 (0.696–0.785)	0.733 (0.683–0.775)	0.732 (0.686–0.775)
White	0.734 (0.709–0.768)	0.737 (0.694–0.763)	0.733 (0.707–0.768)	0.741 (0.700–0.772)	0.742 (0.705–0.771)
Brier Score mean (95% CI)
All	0.031 (0.026–0.035)	0.030 (0.027–0.034)	0.031 (0.026–0.035)	0.030 (0.027–0.034)	0.030 (0.027–0.034)
Black	0.044 (0.036–0.053)	0.044 (0.035–0.052)	0.044 (0.036–0.053)	0.044 (0.035–0.052)	0.044 (0.036–0.054)
White	0.027 (0.023–0.031)	0.027 (0.022–0.030)	0.027 (0.023–0.031)	0.026 (0.022–0.030)	0.027 (0.023–0.030)
Calibration Slope mean (95% CI)
All	0.791 (0.649–0.973)	0.803 (0.622–0.975)	0.829 (0.674–1.013)	0.862 (0.681–1.030)	0.965 (0.713–1.129)
Black	0.664 (0.476–0.934)	0.769 (0.491–1.027)	0.758 (0.534–1.095)	0.892 (0.573–1.186)	1.288 (0.819–1.748)
White	1.004 (0.793–1.236)	0.853 (0.609–1.072)	0.883 (0.705–1.091)	0.882 (0.650–1.106)	0.875 (0.659–1.117)
Calibration Intercept mean (95% CI)
All	0.006 (0.001–0.011)	0.006 (0.001–0.011)	0.005 (0.000–0.009)	0.003 (−0.001–0.008)	−0.002 (−0.007–0.004)
Black	0.010 (−0.002–0.025)	0.007 (−0.003–0.020)	0.009 (−0.005–0.022)	0.003 (−0.007–0.017)	−0.002 (−0.017–0.018)
White	0.001 (−0.004–0.007)	0.005 (0.001–0.010)	0.003 (−0.002–0.008)	0.002 (−0.003–0.008)	−0.003 (−0.011–0.003)
REGARDS	REGARDS	REGARDS	REGARDS	REGARDS	REGARDS
IPCW CI mean (95% CI)
All	0.689 (0.675–0.703)	0.687 (0.676–0.700)	0.690 (0.676–0.702)	0.690 (0.678–0.702)	0.682 (0.671–0.694)
Black	0.677 (0.657–0.698)	0.670 (0.648–0.694)	0.678 (0.658–0.699)	0.674 (0.653–0.700)	0.673 (0.656–0.696)
White	0.698 (0.673–0.715)	0.697 (0.679–0.714)	0.697 (0.672–0.713)	0.695 (0.676–0.712)	0.690 (0.668–0.710)
Brier Score mean (95% CI)
All	0.044 (0.042–0.046)	0.043 (0.041–0.045)	0.044 (0.042–0.046)	0.043 (0.041–0.045)	0.043 (0.041–0.045)
Black	0.047 (0.044–0.050)	0.045 (0.043–0.049)	0.047 (0.044–0.050)	0.045 (0.042–0.049)	0.045 (0.042–0.049)
White	0.041 (0.038–0.044)	0.041 (0.038–0.044)	0.042 (0.038–0.044)	0.041 (0.039–0.044)	0.041 (0.039–0.044)
Calibration Slope mean (95% CI)
All	0.477 (0.410–0.544)	0.686 (0.623–0.785)	0.491 (0.426–0.551)	0.757 (0.689–0.860)	0.838 (0.753–0.941)
Black	0.390 (0.302–0.474)	0.553 (0.457–0.688)	0.416 (0.322–0.509)	0.676 (0.549–0.809)	0.839 (0.687–1.011)
White	0.648 (0.533–0.748)	0.837 (0.737–0.970)	0.559 (0.456–0.659)	0.845 (0.737–0.973)	0.853 (0.731–0.974)
Calibration Intercept mean (95% CI)
All	0.024 (0.021–0.028)	0.019 (0.015–0.023)	0.023 (0.019–0.027)	0.014 (0.009–0.017)	0.012 (0.007–0.016)
Black	0.027 (0.021–0.034)	0.023 (0.016–0.028)	0.029 (0.023–0.036)	0.018 (0.011–0.024)	0.018 (0.011–0.025)
White	0.018 (0.014–0.023)	0.016 (0.011–0.020)	0.018 (0.014–0.024)	0.010 (0.005–0.014)	0.008 (0.003–0.013)

However, the predictions of the Parity-Contrained model were much more similar between Black and White individuals with an observed stroke, as well as between Black and White individuals without an observed event, compared to all other models (see Figure 1), suggesting that the parity constraint was effective. As a consequence, stark differences in IPCW xCI were observed between models (see Figure 2). In the test set, Black-White IPCW xCIs were 0.89 (0.86–0.91) and 0.87 (0.82–0.90) for the Cox and Unconstrained models, respectively, whereas White-Black IPCW xCIs were 0.53 (0.48–0.58) and 0.53 (0.48–0.58) for the Cox and Unconstrained models, respectively, yielding absolute differences of 0.33 and 0.34. These differences were reduced in the race-free models. Black-White IPCW xCIs were 0.86 (0.82–0.90) and 0.85 (0.81–0.90) for the race-free Cox and race-free Unconstrained models, respectively, whereas White-Black IPCW xCIs were 0.58 (0.54–0.62) and 0.57 (0.52–0.61) for the race-free Cox and race-free Unconstrained models, respectively, yielding absolute differences of 0.28 and 0.28. However, this difference was almost eliminated for the Parity-Constrained model: Black-White and White-Black IPCW xCIs were 0.75 (0.67–0.80) and 0.72 (0.69–0.75), yielding an absolute difference of 0.03.

Figure 1: Comparison of predicted risk between models:

The panels show the race-stratified distributions of model-predicted 10-year stroke risk in individuals with and without observed stroke events for the Cox model, Unconstrained model, race-free Cox model, race-free Unconstrained model, and Parity-Constrained model in the test set (top) and secondary validation set (REGARDS; bottom).

Figure 2: Comparison of within- and between-group discrimination performance between models.

The bars show the IPCW xCI, including 95% confidence interval obtained via bootstrapping, for the Cox model, Unconstrained model, race-free Cox model, race-free Unconstrained model, and Parity-Constrained model across all four group combinations (Black or White cases versus Black or White controls) in the test set (top) and secondary validation set (REGARDS; bottom). The dotted lines show the IPCW CI for each model across the population as a whole.

In REGARDS, Black-White IPCW xCIs were 0.79 (0.78–0.81) and 0.80 (0.78–0.82) for the Cox and Unconstrained models, respectively, whereas White-Black IPCW xCIs were 0.57 (0.55–0.59) and 0.55 (0.52–0.57) for the Cox and Unconstrained models, respectively, yielding absolute differences of 0.23 and 0.25. As before, these differences were reduced in the race-free models. Black-White IPCW xCIs were 0.74 (0.71–0.76) and 0.75 (0.72–0.76) for the race-free Cox and race-free Unconstrained models, respectively, whereas White-Black IPCW xCIs were 0.64 (0.61–0.65) and 0.62 (0.60–0.64) for the race-free Cox and race-free Unconstrained models, respectively, yielding absolute differences of 0.10 and 0.13. Again the difference was almost eliminated for the Parity-Constrained model: Black-White and White-Black IPCW xCIs were 0.66 (0.63–0.69) and 0.71 (0.69–0.72), yielding an absolute difference of 0.05.

The xROCt and xAUCt at t = 10 years show the same trend, with all four xROCt curves (Black-White, White-Black, Black-Black, White-White) being much more similar for the Parity-Constrained model compared to the other models on both the test set and REGARDS (see sFigure 1). In the test set, the largest absolute difference in xAUCt is 0.31 and 0.32 for the Cox and Unconstrained models, respectively. This difference is smaller for the race-free models, with largest absolute differences of 0.25 and 0.26 for race-free Cox and race-free Unconstrained, respectively. However, this difference is almost eliminated (0.01) for the Parity-Constrained model. In REGARDS, the largest absolute difference in xAUCt is 0.21 and 0.23 for the Cox and Unconstrained models, respectively, compared to 0.09 and 0.12 for the race-free Cox and race-free Unconstrained models, respectively. Again this difference is much smaller (0.05) for the Parity-Constrained model.

The xPR curves and distribution of model-predicted 10-year stroke risk for all models without conditioning on an observed event are shown in the Supplemental Material (see sFigures 2–3).

Calibration and Fairness

The Brier score was similar for all models both in the test set and in REGARDS: no statistically significant differences were observed in the full population, the Black subpopulation, or the White subpopulation (see Table 1).

However, large and statistically significant differences in calibration slope but not calibration intercept were observed in both evaluation sets (see Figure 3 and Table 1). In the test set, the Parity-Constrained model was most calibrated across the population as a whole (slope = 0.97; 95% CI = 0.71–1.13), whereas the Cox model was unconstrained calibrated in the White subpopulation (slope = 1.00; 95% CI = 0.79–1.24). In the Black subpopulation, the Parity-Constrained model underpredicted risk (slope = 1.29; 95% CI = 0.82–1.75), whereas the other models overestimated risk, with the race-free Unconstrained model having unconstrained calibration slope (slope=0.89; 95% CI = 0.57–1.19).

Figure 3: Comparison of within-group calibration between models.

The panels show calibration at 10 years for Cox, Unconstrained, race-free Cox, race-free Unconstrained, and Parity-Constrained models in the test set (top) and secondary validation set (REGARDS; bottom).

In REGARDS, the Parity-Constrained model was best calibrated both across the population as a whole (slope = 0.84; 95% CI = 0.75–0.94) and in Black (slope = 0.84; 95% CI = 0.69–1.01) and White (slope = 0.85; 95% CI = 0.73–0.97) subpopulations, in each case overestimating risk by a similar margin. In contrast, the Cox and Unconstrained models overestimated risk by a much larger margin across the population as a whole (Cox slope = 0.48; 95% CI = 0.41–0.54; Unconstrained slope = 0.69; 95% CI = 0.62–0.79), in the Black subpopulation (Cox slope = 0.39; 95% CI = 0.30–0.47, Unconstrained slope = 0.55; 95% CI = 0.46–0.69), and in the White subpopulation (Cox slope = 0.65; 95% CI = 0.53–0.75, Unconstrained slope = 0.84; 95% CI = 0.74–0.97). The race-free models overestimated risk by a greater margin than the Parity-Constrained model but were less extreme than their race-aware counterparts (see Table 1).

Model Selection and Generalization

Our model selection process further illustrates that the proposed parity constraint is effective in pushing toward fair discrimination performance within and across subgroups (see Figure 4). In this figure, fair discrimination performance is defined as equality between (a) the CI in the population as a whole, and (b) the minimum xCI within or across all subgroups. Note that the minimum xCI cannot exceed the population CI, because the latter is a weighted average of the xCIs within and across subgroups. Stark differences were observed in the distribution of the minimum xCI between models that included the parity constraint, illustrated in blue, and models that did not include this constraint, illustrated in red.

Figure 4: Model selection with the IPCW CI and minimum xCI.

The plots show the effect of our parity constraint (i.e., latent distribution matching) on within- and between-group discrimination performance, as quantified via the minimum xCI, and illustrate our model selection process. The parity-constrained model (red) with highest minimum xCI on the validation set (leftmost plot) was chosen as our Parity-Constrained model. The non-parity-constrained model (blue) with highest overall CI on the validation set was chosen as our Unconstrained model. The minimum xCI cannot exceed the CI, as the latter is a weighted average of all four xCI values (Black or White cases versus Black or White controls); fair discrimination performance is achieved when the two are equal.

DISCUSSION

Machine learning tends to yield only modest improvement in performance when applied to tabular data, with notable examples being cardiovascular disease risk prediction¹⁶ and prediction of death after MI.¹⁷ Consistent with these results, we previously found that machine learning did not substantially improve stroke risk prediction compared to traditional models, including Framingham stroke, REGARDS self-report, and the pooled cohort equations.¹⁰ However, we hypothesized that machine learning methods would provide greater control over the balance between potential benefits and harms associated with stroke risk prediction by allowing us to tune the trade-off between fair discrimination performance and fair calibration performance. In this work, we explored this hypothesis by developing a machine learning method for parity-constrained time-to-event prediction, then train parity-constrained stroke risk prediction models using previously harmonized data from Framingham Offspring, MESA, and ARIC studies. We then evaluate their performance both in a held-out test set and in a secondary validation set (REGARDS). Importantly, we compare our method to the much simpler approach commonly recommended in the clinical literature, that is, the removal of race from the prediction model.

To quantify fair discrimination performance, we use the xCI, an extension of the xAUC³⁸ metric recently proposed to quantify the fairness of binary classification models. The xCI_A,B between groups A and B – in this case, Black and White subpopulations, respectively – may be interpreted as the probability that an individual in group A with an observed event is (correctly) ranked higher than an individual from group B for whom the event happened at a later time or not at all. If the model we are evaluating is fair, we would hope that this probability would not depend on group status, so that xCI_A,B, xCI_B,A, xCI_A,A, and xCI_B,B are all similar. Put simply, the probability that two individuals are correctly ranked by the model should not depend on the group status of either individual. As our previous results¹⁰ and current results (see Figure 2) show, this was not the case for the traditional models.

We successfully addressed this shortcoming by developing a machine learning approach in which a set of learned, latent features – which are derived via neural network from available clinical features (i.e., predictors) – are constrained to contain limited information about a sensitive variable of interest. In the current application, the sensitive variable is Black versus White race. It is from these latent features that the model then predicts the outcome. Importantly, although the model is encouraged to remove information about group status from the latent features, this information can still be retained to the extent that it benefits prediction performance more than it compromises fairness; and how much fairness the model will sacrifice to improve prediction performance is a hyperparameter that can be tuned to suit the clinical application. This approach is relevant to any clinical setting in which biased predictions are a concern and can be applied to any binary or categorical attribute of interest, including sex, race, and ethnicity. Equipped with this approach, we selected the model with unconstrained worst-case xCI performance, which might be described intuitively as “leaving no group left behind”.

Results clearly demonstrate that our approach is effective in reducing differences in discrimination performance across groups. Compared with our baseline approach, a regularized Cox proportional hazards model, the Parity-Constrained model ranked individuals correctly (i.e., concordant with outcomes) at similar rates regardless of their group status (see Figure 2). Moreover, it achieves this result without loss of discrimination performance across the populations as a whole (i.e., not race-stratified) both in our test set and in REGARDS (see Table 1). In contrast, removing race as a predictor was not sufficient to achieve parity in discrimination performance across groups. Although differences in xCI were smaller when race was removed, they remained substantial (see Figure 2). More broadly, removing race tends to be effective in promoting prediction parity only for models that incorporate a very limited number of predictors, none of which is strongly correlated with race; whereas it is ineffective to the extent that associations between race and outcomes can be alternatively explained by other predictors. In this sense, race-free models represent a balance between the extremes of including race and prediction parity, respectively, because they remove only those associations with race that cannot be explained by other available predictors.

Our results also clearly illustrate the known, fundamental trade-off between fair discrimination performance and fair calibration performance.^28–30 In general, whether such parity is desirable depends on the particular clinical context, and imperfect calibration is not inherently unfair.³¹ In stroke risk prediction specifically, however, achieving fair discrimination performance by applying our parity constraint reduces model-predicted risk among Black patients, bringing the risk distribution in Black patients closer to the risk distribution in White patients (see sFigures 3–4). The result is that the Parity-Constrained models underpredict risk in the Black subpopulation, which is undesirable and likely to cause harm rather than reducing it.

By the same token, however, this result underscores that the approach is effective in promoting parity and equalizing discrimination performance within and between subgroups, which is desirable in the many other settings where unconstrained models tend to underpredict risk in disadvantaged groups. In general, however, designing equitable risk prediction algorithms requires careful attention to consequences associated with a specific clinical use case.³² It is therefore critical that the merits of the Parity-Constrained model (and our approach more broadly) be considered in light of the redistribution of resources, benefits, and potential harms likely to occur if it were used in a particular clinical setting.

Finally, we hope that these results – notably the profound impact of group status on the model’s ability to correctly rank individuals correctly – will draw further attention to an oversight that is common in the clinical literature when evaluating bias of clinical decision support tools. Often, fairness is evaluated by dividing the population into subgroups, then evaluating a range of performance metrics in each subgroup (e.g., Black individuals, White individuals). This approach is necessary but incomplete, because it does not compare predictions between groups. We might find, for example, that AUC and other ranking-based metrics are adequate in both Black and White subpopulations but fail to recognize that all White individuals are ranked higher than all Black individuals. When considering impacts on clinical decision-making, therefore, between-group comparison is critical.

Limitations

This work has important limitations. First and foremost, our analysis was limited to Black and White individuals, as other racial and ethnic groups were not adequately represented across all of our harmonized cohorts. Further work will be needed to develop and evaluate fair prediction models for a diverse range of races and ethnicities, and to extend the approach to multi-group fairness, in which there are multiple sensitive attributes that may be partially overlapping. Second, our evaluation relies on the accurate and complete recording of outcomes, and results would be affected by any disparities in the rate at which outcomes were recorded (i.e., differential observability).⁴⁸ Third, our approach sacrifices simplicity for fairness, and as a result may be (a) less interpretable and (b) more difficult to implement compared to traditional approaches. Whether this tradeoff is justified may depend on the specific clinical use case. If interpretability is paramount, alternative approaches such as Fair Cox⁴⁹ might be more appropriate, whereas our approach is better suited for settings in which the effects of predictors on the hazard rate are complex and potentially time-varying. Finally, this approach has not yet been implemented in clinical practice, so its clinical and real-world benefits including on fair clinical outcomes are currently unknown and remain to be explored in subsequent work.

Another key limitation of this work is the assumption of non-informative censoring, which is controversial due to the presence of the competing risk of non-stroke death. To address this concern, we present additional results in the Supplementary Material showing that accounting for this competing risk in our performance evaluation has limited effect on the resulting performance metrics and does not influence trends between racial groups. Importantly, however, accounting for competing risks and carefully calibrating the resulting models to a particular clinical setting would be critical in a real-world deployment.

Finally, it is important to recognize that identifying and mitigating sources of bias affecting clinical risk prediction models goes beyond the specific choice of algorithm. To ensure fairness, a much broader range of factors must be considered throughout the model development process, including adequate representation of minority groups in the training data.^50,51

Conclusion

These results show that the machine learning method we propose achieves fair within- and between-group discrimination performance without compromising discrimination performance across the population as a whole. However, they also illustrate a fundamental trade-off in which parity-constrained models must sacrifice intra-group calibration to improve inter-group discrimination performance. In the current setting, stroke risk prediction, applying a parity constraint therefore leads to underprediction of risk in the Black subpopulation, which is undesirable and likely to cause harm rather than reducing it. In settings where unconstrained models tend to underpredict risk in disadvantaged groups, however, this effect is desirable and likely to reduce harm. As expected, a race-free approach did not lead to parity between groups, but it may strike an appropriate balance between fair discrimination and fair calibration. Exploring the trade-off between competing notions of fairness is critical to mitigate harms associated with clinical risk prediction models, and the choice of model must depend on the intended clinical use. All models as well as code implementing our approach are available in a public repository. More broadly, these results provide a roadmap for development of equitable clinical risk prediction models and illustrate both merits and limitations of a race-free approach.

Highlights.

• Common stroke risk prediction models do not perform equally well in Black and White subgroups.

• There have been calls in the clinical literature to correct such disparities by removing race as a predictor. Alternatively, machine learning methods can be used to constrain differences in model predictions between racial groups.

• In this work, we compare these approaches for equitable stroke risk prediction in a large cohort previously harmonized by our team.

• Results illustrate a fundamental trade-off between intra-group calibration and inter-group discrimination performance, and show that in this case, race-free models strike a balance between the two.

• Consequently, the choice of model must depend on the potential benefits and harms associated with the intended clinical use.

• Our results provide a roadmap for development of equitable clinical risk prediction models and illustrate both merits and limitations of a race-free approach.