Source-linked AI summary
Guided Uncertainty-Aware Robust Domain Transfer
Xin Xiong, Zijian Guo, Tianxi Cai
TL;DR
Clinical prediction models can degrade across time and institutions, while limited labels and regulatory constraints make retraining difficult. GUARD combines target-anchored robust transfer, semi-supervised estimation, and uncertainty-aware source weighting; simulations and a rheumatoid arthritis EHR application show accuracy holding up as deployment moves away from the training target, with more reliable intervals than competing methods.
Problem
Clinical models face temporal and institutional distribution shifts, while scarce labeled outcomes and regulatory constraints make full retraining difficult.
Source: §1 · p. 2Method
GUARD updates a deployed model through target-anchored robust transfer, cross-fitted estimation using unlabeled target covariates, and uncertainty-aware source aggregation.
Source: §1 · p. 4Results
Simulations and the rheumatoid arthritis EHR application show accuracy holding up as deployment moves away from the training target, with more reliable intervals than competing methods.
Source: §7 · p. 28Takeaways & Limitations
GUARD holds the approved rule fixed, updates a low-dimensional head using newly labeled observations, and reports intervals for the updated quantity.
Source: §7 · p. 30The framework assumes labels are missing completely at random; handling outcome-dependent labeling requires further research.
Source: §7 · p. 28
Abstract
from arXiv · showhide
Artificial intelligence models deployed as clinical decision support tools often suffer substantial performance degradation over time and across institutions due to covariate shift, concept drift, and cross-system heterogeneity. Fully retraining complex models is frequently infeasible, particularly in EHR settings where labeled outcome data are scarce and regulatory constraints limit model modification. We propose GUARD (Guided and Uncertainty-Aware Robustness to Domain shift), a unified statistical framework for updating an existing deployed model using limited labeled target data and multiple heterogeneous source populations while guarding against future distributional shifts. GUARD formulates robust multi-source transfer learning as a semi-supervised, adversarial optimization problem anchored at the current model, and employs double machine learning with cross-fitting to obtain debiased, efficient estimates that leverage abundant unlabeled target covariates. It then constructs an uncertainty-aware, guided group DRO estimator that combines source and target information while accounting for sampling variability in the source models. Our framework jointly provides principled robustness to domain shift, efficient semi-supervised estimation, and valid statistical inference for multi-source transfer of clinical prediction models. We demonstrate the utility of GUARD via extensive simulation experiments and a real world application to predicting future disease activity in rheumatoid arthritis from EHR data spanning more than a decade. GUARD recalibrates a pretrained model with only a small number of labeled records per year and remains accurate over multi-year horizons where target-only and standard transfer estimators degrade sharply.
1 Introduction
Clinical models degrade across time and institutions, while scarce labels and regulatory constraints make full retraining difficult. GUARD addresses this by combining target-guided robust transfer from heterogeneous sources with semi-supervised estimation and uncertainty-aware inference.
1 Introduction
- EHR models face covariate and concept drift over time and institutional heterogeneity, challenging long-term reliability and transportability.Changes in populations, health systems, coding, and clinical practice can shift features and feature-outcome relationships.Source: §1 · p. 2
- Scarce labeled outcomes and regulatory limits make full retraining of complex clinical models difficult, motivating updates that use limited target labels and heterogeneous source data.The paper also emphasizes robustness to future shifts during updating.Source: §1 · p. 2
- Existing transfer methods typically select individually target-proximate sources and do not protect against continued drift beyond the observed target.A mixture of heterogeneous sources may approximate the target even when no individual source does.Source: §1 · pp. 2–3
- Group DRO guards against plausible future populations but can sacrifice accuracy on a specific target, while prior guided approaches do not account for unlabeled target covariates or source-estimation uncertainty.These approaches use target data to guide or restrict robust prediction, but generally treat source risks or predictors as fixed.Source: §1 · p. 3
- GUARD combines semi-supervised multi-source adversarial transfer anchored to the deployed model, cross-fitted double machine learning, uncertainty-aware source aggregation, and resampling-based inference.It uses limited labeled target data, abundant unlabeled target covariates, and partially labeled source populations.Source: §1 · p. 4
2 GUARD Population Framework
GUARD updates a deployed predictor in a shared target-scale feature space and robustly selects mixtures of target and source outcome mechanisms. Its uncertainty calibration favors mixtures that are target-compatible and precisely estimated, including mixtures of sources that may not transfer individually.
2.1 Notations and data setting
- GUARD works with a local update to the deployed predictor in an augmented feature space, using a lower-dimensional linear update to recalibrate the model with limited target labels.The feature map can include the deployed score, selected covariates, and engineered or learned representations.Source: §2.1 · pp. 4–5
- Source outcome mechanisms are projected onto the target update scale, allowing comparisons even when both covariate and conditional outcome distributions differ across populations.The target coefficient is the natural benchmark for post-deployment updating.Source: §2.1 · pp. 4–5, §2.2 · p. 5
2.2 Target-anchored source-mixture robustness
- GUARD anchors robustness to the current target update and uses a coefficient-library uncertainty matrix to favor source mixtures that are both target-compatible and reliably estimated.The matrix reflects finite-sample uncertainty from source-specific sample sizes and noise levels.Source: §2.2 · pp. 5–6
2.3 GUARD as uncertainty-calibrated perturbed-control DRO
- GUARD’s perturbed-control DRO adjusts target-guided excess loss for mixture uncertainty, discounting poorly estimated source directions relative to similarly compatible, more precise mixtures.The outer optimization considers future mixtures; the inner adjustment is calibrated to estimation uncertainty.Source: §2.3 · pp. 6–7
2.4 Transfer Learning Interpretation
- The transfer interpretation selects the least-uncertain mixture whose induced update remains close to the target, enabling joint borrowing even when no individual source is close.The penalized form tunes the tradeoff between target proximity and mixture uncertainty.Source: §2.4 · pp. 7–8, Figure 1 · p. 8
3 GUARD Estimation and Inference
GUARD estimates target-projected source coefficients and their joint uncertainty using semi-supervised, cross-fitted double machine learning, then solves a robust optimization problem. For inference, it screens Bayesian-bootstrap perturbations and unions fixed-weight confidence intervals across retained resamples.
3 GUARD Estimation and Inference
- GUARD estimation has three components: estimating the target-projected coefficient library, its uncertainty matrix, and the robust update; inference resamples this procedure and forms a union confidence interval over stable draws.The weight optimizer is a constrained quadratic program and can be nonsmooth in its estimated inputs.Source: §3 · p. 8
3.1.1 Estimating the target-projected coefficient library B
- Unlabeled target covariates estimate target geometry and evaluate source mechanisms on the target scale, while cross-fitted orthogonal moments reduce first-order nuisance-estimation bias.Flexible outcome-regression and density-ratio estimators can be used without introducing first-order bias in the target-projected source coefficient.Source: §3.1.1 · pp. 8–9
3.1.2 Estimation of V
- The coefficient uncertainty matrix captures shared target-side covariance across source projections as well as labeled-sample contributions, which mainly affect its diagonal.The shared unlabeled target sample makes the matrix generally non-diagonal.Source: §3.1.2 · p. 10
3.1.3 Estimation of GUARD
- GUARD combines the estimated coefficient library and uncertainty matrix by optimizing weights that balance target compatibility against uncertainty-aware source borrowing.The estimation procedure uses cross-fitting and returns the robust coefficient together with its weights and estimated inputs.Source: §3.1.3 · p. 11
3.2 Perturbation inference for GUARD
- A screened Bayesian bootstrap addresses nonsmoothness from the constrained weight optimization by resampling the estimation-and-optimization procedure and forming a union confidence interval.The procedure perturbs both the coefficient library and uncertainty inputs, then retains stable resamples.Source: §3.2 · p. 12
3.2.1 Perturbation weights
- Perturbation weights reweight empirical averages in the cross-fitted estimating equations while keeping the original nuisance estimates and sample splits fixed.Weights are generated independently within each independent data pool.Source: §3.2.1 · p. 12
3.2.2 Screening resamples and union confidence interval
- Screening removes resamples with unusually large objective perturbations; inference then treats each retained candidate weight as fixed and unions its confidence interval.The screening targets perturbations of the GUARD objective that drive changes in the constrained optimizer.Source: §3.2.2 · pp. 13–14
4 Theoretical Properties
GUARD’s theoretical results characterize transferability through source mixtures, bound estimation error, and justify screened-union confidence intervals under stated assumptions.
4.1 Model Assumptions
- The theory assumes independent labeled and unlabeled pools, moment and overlap conditions, and cross-fitted nuisance accuracy as sample sizes grow.Additional assumptions control approximation remainders and gradient regularity.Source: §4.1 · p. 14
4.2 Population transferability under source mixing
- A source is transferable when it can receive meaningful weight in some target-compatible mixture, even if it is not individually close to the target.The separation condition prevents non-transferable sources from improving target compatibility through cancellation.Source: §4.2 · p. 16
- GUARD approximates the clean-source oracle without explicitly identifying non-transferable sources; its error depends on the compatibility radius, participation threshold, and covariance geometry.The oracle is the variance-minimizing target-compatible mixture using only the target and transferable sources.Source: §4.2 · pp. 16–17
4.3 Estimation error bound
- The estimation analysis bounds fixed-penalty GUARD error and links calibration of the compatibility constraint to gains from transferable sources.The bound accounts for estimated-input error and higher-order nuisance error.Source: §4.3 · p. 17
4.4 Inference validity for screened union CI
- Screening retains bootstrap perturbations compatible with sampling variation, while a union of retained fixed-weight intervals protects coverage when the optimizer changes nonsmoothly.With asymptotic probability at least 1 − α1, the screened set contains a resample whose optimizer is close to the population GUARD optimizer.Source: §4.4 · p. 18
- The coverage and interval-length theorems formalize validity and length properties for the screened union confidence interval under the stated conditions.The supplied passages name these results but do not provide their full bounds.Source: §4.4 · p. 18
5 Simulation studies
The simulations compare GUARD with transfer, pooling, and robust benchmarks across temporal drift and heterogeneous-source settings. GUARD generally retains strong future and current-target accuracy, with favorable interval quality across the considered settings.
5.1 Simulation settings, benchmarks and evaluation metrics
- The simulations include temporal future generalization, heterogeneous-source current-target transfer, and clean or corrupted source libraries with mixture and precision differences.Design I generates four future-population types; Designs II and III test source mixing and robustness to corrupted sources.Source: §5.1 · p. 19
- Future generalization is evaluated by worst-case regret over generated future coefficients, while current-target estimation uses coefficient error and inference uses coverage and interval length or interval score.For future inference, the Winkler interval score jointly reflects interval length and coverage.Source: §5.1 · p. 20
5.2.1 Estimation and Prediction Performance
- 0.017 versus 0.036 for source–target mixing, 0.038 for maximin, 0.11 for transGLM, and 0.16 for target only: GUARD had the lowest current-target MSE and retained the strongest regret performance under several future drifts.Under random drift, GUARD’s regret rose from 0.017 to 0.062 as δ increased from 0 to 0.3; under perturbed-hull drift, it rose from 0.022 to 0.046 as εv increased from 1 to 6.Source: Figure 2 · p. 21, §5.2.1 · p. 21
- 0.0084 versus 0.15 for target only in Design I, and 0.0099 versus 0.028 for source–target mixing and 0.068 for target only in Design II: GUARD had lower MSE under exact source mixing.In Design III, GUARD’s MSE changed from 0.0007 to 0.0013 in (i) and from 0.0054 to 0.0065 in (ii) as κ increased from 0 to 2.Source: Figure 3 · p. 22, §5.2.1 · p. 22
5.2.2 Inference Validity and Interval Quality
- GUARD had the lowest interval score in most future-drift settings, with a 0.226 average interval width and 0.993 coverage at the current target.For very large drift, wider methods became competitive in some settings.Source: Figure 4 · p. 23, §5.2.2 · p. 23
- In current-target inference, GUARD coverage stayed near 95% in Designs I–II as α increased, with the shortest intervals among methods that achieved valid coverage.In Design I at α = 0, mean interval length was 0.21 for GUARD, versus 0.33 for source–target mixing, 0.26 for maximin, and 0.40 for target only.Source: Figure 5 · p. 24, §5.2.2 · p. 24
6 Robust Disease Activity Prediction in Rheumatoid Arthritis
This application uses a pretrained EHR representation and limited target-year labels to predict future rheumatoid arthritis disease activity under substantial temporal change. GUARD retained low excess risk across longer horizons and produced narrower uncertainty intervals than several alternatives.
6.1 Study Cohort, pre-trained feature representation, and evaluation design
- The study used 671,003 post-2009 patient-windows, 100 labeled target-year records per training year, and a pretrained 22-dimensional RA EHR representation evaluated through 2020.The pretrained transformer used data through 2009; the 22-dimensional PCA representation retained about 95% of variation.Source: §6.1 · p. 25
6.2.1 Prediction performance
- For 2012 training, GUARD’s excess risk was 0.0134 in 2012 and 0.0146 in 2020, versus target-only’s 0.0699 and 0.5232.Source–target mixing had slightly lower excess risk in 2012, but its advantage disappeared by 2014.Source: §6.2.1 · p. 26
- GUARD’s excess risk rose 16% (0.014 to 0.016) by horizon 6 versus 630% for target-only, while retaining higher R2 and AUC except GraphPath’s best AUC at horizon 6.Figure 6 summarizes excess risk, Pearson correlation, and AUC across horizons pooled over six target years.Source: §6.2.1 · p. 26, Figure 6 · p. 27
6.2.2 Uncertainty Quantification
- For 2012, GUARD’s average confidence-interval length was 0.101 versus 0.234 for maximin, 1.304 for source–target mixing, and 2.306 for target-only, detecting 11 coefficients different from zero.Eight detected coordinates were embeddings; selected embedding directions were driven mainly by longitudinal EHR code patterns, accounting for 80.4% of attribution.Source: §6.2.2 · p. 27
- GUARD mean-response intervals widened 2.53-fold from typical to extreme covariate profiles, versus 2.99-fold for maximin and target-only and 3.33-fold for source–target mixing.GUARD interval lengths ranged from 0.076 to 0.191; target-only ranged from 0.740 to 2.216.Source: §6.2.2 · pp. 27–28
- Prediction intervals had approximately nominal coverage across covariate-extremity deciles, while GUARD’s lengths rose only from about 1.50 to 1.52.Target-only intervals reached 2.61 at decile 10.Source: §6.2.2 · p. 28
7 Discussion
GUARD anchors robust updating to the current target population and calibrates source borrowing to estimation uncertainty. The discussion describes its scope, extensions, and open limitations in settings involving covariate shift, labeling, representation choice, and recursive updates.
7 Discussion
- Simulations and the rheumatoid arthritis EHR application show accuracy holding up as deployment moves from the training target, with more reliable intervals than competing methods.GUARD anchors the update to the current target population while using the uncertainty matrix to direct source borrowing toward estimable sources.Source: §7 · p. 28
- Labels are assumed missing completely at random; missingness at random given X can be handled with propensity weighting, while formal treatment of outcome-dependent labeling remains future work.The application mitigates labeled–unlabeled covariate imbalance by matching each labeled window to twenty unlabeled windows.Source: §7 · p. 28
- Transferability depends on feature map A, which must be rich enough to correct drift yet small enough to estimate from a few hundred target labels.Building a theoretically justified representation step into GUARD remains future work.Source: §7 · p. 28
- When the target sample is not larger than every source sample, partitioning it can make V diagonal by construction, but reduces target-side estimation precision.Each target block then carries roughly n0/L observations.Source: §7 · pp. 28, 30
- The squared-loss mixture argument does not extend directly to cross-entropy, and holding target covariates fixed leaves GUARD robust to concept drift but not future covariate shift.Extending to nonlinear links would require a local quadratic expansion or direct treatment of conditional group DRO.Source: §7 · p. 30
- Under sparsity and regularity conditions, the high-dimensional extension uses subvector or low-dimensional inference; target-coefficient inference additionally requires conditions and λ to control transfer bias.The extension requires a restricted eigenvalue condition and continues to treat L as fixed.Source: §7 · p. 30
- GUARD holds the approved rule fixed and updates a low-dimensional head in a prespecified feature space using few new labels, with intervals for the updated quantity.Source: §7 · p. 30
- Recursive yearly updating breaks the cross-population independence assumption, leaving error accumulation across repeated cycles an open problem.Source: §7 · p. 31