Distributionally Robust Transfer Learning with Structurally Missing Covariates, with Application to Cross-National Cardiac Arrest Prediction

Technical Summary: Distributionally Robust Transfer Learning with Structurally Missing Covariates

Problem Statement
The paper addresses a critical challenge in deploying clinical prediction models across diverse healthcare systems: the scenario where a model trained in a high-resource "source" setting relies on a rich set of covariates $(X, A)$, but the "target" deployment setting lacks specific, highly predictive covariates $A$ entirely due to structural differences in data collection infrastructure. Unlike standard missing data where individual observations are incomplete, this "structural missingness" represents a population-level absence of variables. Furthermore, the target domain often lacks labeled outcome data, precluding supervised adaptation or retraining.

Existing approaches fail in this setting. Imputation methods require untestable assumptions about the distribution of $A$ in the target domain. Covariate shift corrections (e.g., importance weighting) only address differences in the observed variables $X$ and cannot recover information lost from the absence of $A$. Standard distributionally robust optimization (DRO) typically constructs uncertainty sets over the joint distribution of all variables, which is overly conservative and fails to distinguish between stable observed covariates and structurally missing ones.

Methodology: The DRUM Framework
The authors propose DRUM (Distributionally Robust Unsupervised transfer learning with structurally Missing covariates), a framework designed to transfer prediction models to unlabeled target populations where $A$ is absent. The core philosophy is to optimize for worst-case predictive performance over an uncertainty set of the unknown target conditional distribution $A|X$, rather than attempting to impute $A$ or assuming $A|X$ remains stable.

The methodology proceeds in three stages:

1. Problem Formulation:
   The goal is to learn a predictor $m(X)$ using only the shared covariates $X$. The authors formulate this as a minimax problem:
   $$\min_{m(\cdot)} \max_{P_{A|X} \in \mathcal{C}(\delta)} \mathbb{E}_{X \sim Q_X, A \sim P_{A|X}, Y \sim P^S_{Y|X,A}} \left[ Y^2 - (Y - m(X))^2 \right]$$
   Here, $\mathcal{C}(\delta)$ is an uncertainty set of conditional distributions for $A|X$ in the target domain, controlled by a robustness parameter $\delta$. The objective maximizes the worst-case reduction in squared error relative to a null predictor.

2. Generator-Based Parameterization and DRO:
   Direct optimization over the infinite-dimensional space of conditional distributions is intractable. DRUM parameterizes the conditional distribution $A|X$ using a neural network generator $g_\phi(X, \epsilon)$, where $\epsilon$ is latent noise.

   • Uncertainty Set: The uncertainty set is defined using Energy Distance (via the Engression framework). The set $\mathcal{C}_g(\delta)$ contains generators whose energy score deviates from the source conditional generator $g^S$ by at most $\delta$. This allows the model to explore conditional distributions that depart from the source without requiring density estimation.
   • Optimization: The problem is solved via a primal-dual approach. A "worst-case" generator is trained to minimize the expected squared prediction error on the target $X$ (marginalizing over the generated $A$), subject to the energy distance constraint.

3. Bias Correction (Neyman Orthogonality):
   The authors identify that the "plug-in" estimator (averaging the source outcome model $\hat{f}(X, A)$ over the generated $A$) inherits first-order biases from the estimation of $\hat{f}$. To mitigate this, they derive a bias-corrected estimator based on Neyman-orthogonal estimation and cross-fitting.

   • They construct a pseudo-outcome $F_i$ that combines source residuals (reweighted by a density ratio $\hat{\omega}$ to match the generator-induced distribution) and target plug-in predictions.
   • Regressing this pseudo-outcome on $X$ yields an estimator that is robust to first-order errors in the nuisance functions ($\hat{f}$ and the density ratio $\hat{\omega}$), achieving parametric convergence rates under product-of-rates conditions.

Key Contributions

1. A New DRO Framework for Structural Missingness: DRUM isolates uncertainty specifically to the conditional distribution $A|X$ rather than the joint distribution, using a generative model to parameterize the uncertainty set. This avoids the intractability of infinite-dimensional optimization.
2. Bias-Corrected Estimation: The paper derives a debiased estimator that reduces sensitivity to nuisance function estimation errors, significantly improving finite-sample performance in simulations and real-world applications.
3. Application to Cross-National OHCA Prediction: The framework is applied to predicting neurological outcomes after out-of-hospital cardiac arrest (OHCA), transferring models from the US-ROC registry (rich prehospital data) to Asian PAROS registries (where prehospital variables like epinephrine dose and blood pH are structurally missing).

Results

• Simulations: In settings with linear and nonlinear conditional relationships and varying dimensions of missing covariates, DRUM (particularly the debiased variants) substantially outperformed standard baselines (ERM, DRO, importance weighting) and pseudo-label imputation methods. It achieved lower mean and worst-case normalized Mean Squared Error (MSE) under varying degrees of distributional shift.
• Real-World Application (OHCA):
  • Calibration: When applied to Singapore, Japan, and Korea, DRUM produced significantly better-calibrated predictions than baselines. Standard baselines systematically over-predicted favorable outcomes due to the prevalence shift and missing covariates, whereas DRUM tracked closer to observed event rates.
  • Performance: DRUM-Debiased variants achieved the lowest Brier scores and Expected Calibration Error (ECE) across all sites.
  • Clinical Utility: At fixed probability thresholds relevant to clinical triage, DRUM-Debiased demonstrated higher precision and specificity, indicating that its calibrated probabilities lead to more reliable positive classifications compared to baselines which suffered from high false-positive rates due to miscalibration.

Significance and Claims
The paper claims that DRUM provides a principled alternative to the current dichotomy in clinical AI deployment: either restricting models to sites with identical data infrastructure or discarding predictive information by ignoring missing covariates. By explicitly accounting for uncertainty in the unobserved conditional distribution $A|X$ without relying on untestable assumptions, DRUM enables the deployment of high-resource models in low-resource settings.

The authors emphasize that the method's primary advantage lies in calibration and risk stratification (Brier score, ECE, fixed-threshold precision) rather than global discrimination (AUROC), which is consistent with the worst-case objective penalizing squared error. The framework is presented as a general solution for any transfer learning problem where predictive covariates are available during training but absent or restricted at deployment, with potential extensions to algorithmic fairness and pharmaceutical stability prediction.

The paper remains modest regarding limitations, noting that the current generator parameterization assumes continuous covariates, the framework currently handles a single source population, and the derivation focuses on squared-error and cross-entropy losses. Future directions include handling discrete/mixed-type variables, federated learning with multiple sources, and structured outcomes like time-to-event data.