UdonCare: Hierarchy Pruning for Unseen Domain Discovery in Predictive Healthcare

TL;DR

UdonCare tackles unseen-domain generalization in predictive healthcare by discovering latent patient domains from medical ontologies and decomposing domain-invariant information for robust predictions. It combines a hierarchy-guided domain discovery module with mutual forward learning to separate domain-specific from label-related features during training. The method is validated on MIMIC-III and MIMIC-IV across mortality, readmission, drug recommendation, and diagnosis tasks, outperforming eight baselines with notable AUPRC gains. The results demonstrate that integrating structured medical knowledge into domain generalization can improve generalization without excessive computational cost, suggesting a practical path for more robust clinical prediction models.

Abstract

Healthcare providers often divide patient populations into cohorts based on shared clinical factors, such as medical history, to deliver personalized healthcare services. This idea has also been adopted in clinical prediction models, where it presents a vital challenge: capturing both global and cohort-specific patterns while enabling model generalization to unseen domains. Addressing this challenge falls under the scope of domain generalization (DG). However, conventional DG approaches often struggle in clinical settings due to the absence of explicit domain labels and the inherent gap in medical knowledge. To address this, we propose UdonCare, a hierarchy-guided method that iteratively divides patients into latent domains and decomposes domain-invariant (label) information from patient data. Our method identifies patient domains by pruning medical ontologies (e.g. ICD-9-CM hierarchy). On two public datasets, MIMIC-III and MIMIC-IV, UdonCare shows superiority over eight baselines across four clinical prediction tasks with substantial domain gaps, highlighting the untapped potential of medical knowledge in guiding clinical domain generalization problems.

Figures (4)

• Figure 1: The Overall Framework of UdonCare. The forward structure adds a domain pathway for mutual learning, extending beyond the backbone pathway of conventional predictive models. During training, we first feed patient data $\mathbf{x}$ into the backbone pathway, which learns patient features $\mathbf{p}$ through $f_{\phi,k}(\cdot)$ and produces the output prediction $\hat{y}_p$. In parallel, we obtain $\hat{y}_h$ from invariant features $\mathbf{h}$ along the domain pathway by applying $\mathrm{DiscoveryAlgo}(\cdot)$, $g_{\theta,k}(\cdot)$, and $h(\cdot)$. Here we iteratively adapt latent domains in $\mathbf{M}$ and update parameters on both pathways by ground truths $y$.
• Figure 2: A simple illustration of hierarchy-guided domain discovery.
• Figure 3: Effectiveness of Domain Discovery. The left figure shows the effect of the number of iteration on AUPRC in MIMIC-III dataset, and the right one shows comparison among variants upon UdonCare.
• Figure 4: Performance by Training Size. We evaluate drug recommendation on MIMIC-IV, and values on the x-axis indicate % of the entire training data. The dotted lines divide two ranges: [1, 10] and [10, 100].