Technical Summary: Correcting Heterogeneous Diagnostic Bias Using Causal Hidden Markov Models

Problem Statement

Clinical prediction models developed using routine electronic health records (EHR) often suffer from heterogeneous diagnostic bias. In routine care, individuals identified as high-risk (due to clinical factors) or belonging to specific demographic groups (e.g., sex, ethnicity, socioeconomic status) are frequently tested more often than others. This results in differential diagnostic delay, where certain groups are underdiagnosed or diagnosed later than they would be under equal testing conditions.

Consequently, models trained on such data learn to predict the observed diagnosis rate rather than the true underlying disease risk. This leads to systematic underestimation of risk for underdiagnosed groups, causing calibration errors and potential unfairness in clinical decision-making. Standard approaches often fail to address this because the "true" disease status is unobserved for undiagnosed individuals, and the bias is driven by the interaction between patient attributes and the healthcare system's testing behavior.

Methodology

The authors propose a causal inference framework utilizing a Hidden Markov Model (HMM) to correct for this bias. The approach involves three main components:

1. Causal Modeling via Directed Acyclic Graphs (DAG)

The authors define a DAG where:

• $X_n$: Baseline covariates (risk factors).
• $A_n$: Observability attributes (e.g., sex, ethnicity, comorbidities) that influence both disease risk and the probability of testing.
• $\tilde{A}_n$: The presentation of $A_n$ specifically affecting the testing decision conditional on disease stage.
• $S_{n,t}$: Latent disease stage (0: No disease, 1: Early stage, 2: Late stage).
• $R_{n,t}$: Test results (Negative, Early Positive, Late Positive, or No Test).
• $D_{n,t}$: Diagnosis (occurs if a test is positive).

The goal is to estimate a counterfactual estimand: the probability of diagnosis $P(D_{n,T}^{\tilde{A}_n=0} = 1 | X_n, A_n)$, representing the risk if the individual were subject to the testing rate of a reference group (typically the group with the shortest diagnostic delays), while retaining their actual disease progression and risk factors.

2. Hidden Markov Model Formulation

The longitudinal disease process is modeled as an HMM with:

• Transition Kernel ($Q$): Defines the probability of moving between disease stages (0 $\to$ 1 $\to$ 2). It assumes a progressive disease where regression is impossible. The transition from "No disease" to "Early stage" is governed by a hazard function $h(X_n, A_n, t)$, while the transition from "Early" to "Late" is a constant probability $\beta_{S=1 \to 2}$.
• Emission Probabilities ($\Gamma$): Defines the probability of observing a test result $R_{n,t}$ given the latent stage $S_{n,t}$ and observability attribute $A_n$. Crucially, testing rates for early-stage disease ($\beta_{S=1, A_n}^R$) are allowed to vary by $A_n$, capturing the diagnostic heterogeneity. Late-stage testing rates are assumed independent of $A_n$ due to severe symptoms.
• Initial State ($\pi$): The probability distribution of the disease stage at baseline.

The model parameters are estimated via maximum likelihood using a forward algorithm and gradient-based optimization (L-BFGS in PyTorch), marginalizing over the unobserved latent states.

3. Counterfactual Outcome Imputation

Once the HMM is fitted, the method generates a counterfactual dataset to train a standard prediction model:

1. Reconstruction: The HMM estimates the probability of the latent disease trajectory for each individual given their observed testing history.
2. Intervention: The emission probabilities are modified to reflect the testing rate of the reference group ($\tilde{A}_n = 0$).
3. Imputation: Counterfactual diagnosis probabilities ($p_{n,T}^{cf}$) are calculated. To avoid inflating incidence, the authors apply a monotonicity assumption: if an individual was diagnosed in reality ($D=1$), they remain diagnosed in the counterfactual ($D^{cf}=1$). Only undiagnosed individuals ($D=0$) in the underdiagnosed group are re-imputed based on their counterfactual probability.
4. Recalibration: The counterfactual probabilities for re-imputed cases are rescaled to ensure the total counterfactual incidence matches the estimated counterfactual population incidence.
5. Training: A standard logistic regression model is trained on this re-imputed dataset to predict the counterfactual outcome.

Key Contributions

• Causal HMM Framework: The paper introduces a novel application of causal HMMs to explicitly model and correct for differential diagnostic delay in progressive diseases, moving beyond simple label correction.
• Identifiability and Robustness: Through simulations, the authors demonstrate that the HMM parameters are identifiable and that the approach remains robust even when assumptions (such as perfect test sensitivity or uniform progression rates) are violated.
• Real-World Application: The method is applied to a large-scale UK Biobank dataset to predict 5-year Chronic Kidney Disease (CKD) incidence, successfully correcting bias across multiple subgroups.
• Data-Centric Fairness: The approach aligns with "data-centric" fairness by modeling the data-recording mechanisms (testing behavior) that generate bias, rather than simply removing protected attributes from the model.

Results

Simulation Studies

In simulations with 50,000 patients, the proposed method significantly outperformed standard "Naïve" (trained on observed data) and "Blind" (trained without protected attributes) models:

• Calibration: The Observed:Expected (O:E) ratio for the underdiagnosed group improved from 1.34 (Naïve) to 1.02 (Proposed), indicating near-perfect calibration.
• Robustness: Even when assumptions were violated (e.g., imperfect test sensitivity or varying progression rates), the proposed method maintained better calibration than standard models, though parameter estimation bias increased.

Chronic Kidney Disease (CKD) Case Study

Using 123,460 individuals from the UK Biobank:

• Diagnostic Heterogeneity: The HMM identified diabetes as the strongest driver of observability (Odds Ratio: 10.36 for testing rates), followed by hypertension. Being female was associated with lower testing rates (OR: 0.69).
• Bias Correction:
  • The Naïve model severely underpredicted risk for patients without diabetes (O:E ratio: 1.55).
  • The Imputed model corrected this, achieving an O:E ratio of 1.01 for the non-diabetic group, while maintaining calibration for the diabetic group (O:E: 1.01).
• Clinical Utility: Decision curve analysis showed improved net benefit across both groups, particularly for the underdiagnosed population, without reducing performance for well-monitored groups.
• Risk Stratification: The corrected model increased estimated risk for patients without diabetes, reflecting their true underlying risk rather than their lower likelihood of being tested.

Significance and Claims

The authors claim that their approach offers a theoretically grounded and empirically validated method to reduce bias and unfairness in clinical prediction models without requiring external incidence data or fixed assumptions about disease prevalence.

By explicitly modeling the "diagnostic process" as a latent variable influenced by patient attributes, the method allows for the estimation of risk under a fair counterfactual scenario where testing rates are equalized. The authors emphasize that this leads to more accurate risk stratification, ensuring that underdiagnosed populations are not systematically overlooked in preventive care. They position this work as a step toward "data-centric" fairness, addressing the root causes of bias in routine clinical data rather than just mitigating the symptoms in the final model.

The paper concludes modestly, noting that while the method relies on specific assumptions (e.g., Markovian progression, constant late-stage testing rates), it provides a flexible framework that can be extended to more complex disease trajectories and multiple test types in future work.