[Paper Review] Domain Generalization for Medical Imaging Classification with Linear-Dependency Regularization

Haoliang Li, Yufei Wang+4 more|arXiv (Cornell University)|Sep 27, 2020

Domain Adaptation and Few-Shot Learning45 references81 citations

TL;DR

Introduces Linear-Dependency Domain Generalization (LDDG) for medical imaging by enforcing latent-space rank constraints and distribution alignment to improve cross-domain generalization without target-domain data.

ABSTRACT

Recently, we have witnessed great progress in the field of medical imaging classification by adopting deep neural networks. However, the recent advanced models still require accessing sufficiently large and representative datasets for training, which is often unfeasible in clinically realistic environments. When trained on limited datasets, the deep neural network is lack of generalization capability, as the trained deep neural network on data within a certain distribution (e.g. the data captured by a certain device vendor or patient population) may not be able to generalize to the data with another distribution. In this paper, we introduce a simple but effective approach to improve the generalization capability of deep neural networks in the field of medical imaging classification. Motivated by the observation that the domain variability of the medical images is to some extent compact, we propose to learn a representative feature space through variational encoding with a novel linear-dependency regularization term to capture the shareable information among medical data collected from different domains. As a result, the trained neural network is expected to equip with better generalization capability to the "unseen" medical data. Experimental results on two challenging medical imaging classification tasks indicate that our method can achieve better cross-domain generalization capability compared with state-of-the-art baselines.

Motivation & Objective

• Motivate and address domain shift in medical imaging when target-domain data are unavailable.
• Learn a representative latent feature space that shares information across domains.
• Regularize latent features to have a fixed rank equal to the number of categories.
• Align latent feature distributions to a Gaussian prior to reduce overfitting to source domains.

Proposed method

• Model uses a feature extractor Q_theta, a variational encoder F_omega, and a linear classifier T_phi.
• Impose a rank-based regularization on latent features so rank(Z)=C (number of classes) to capture category information.
• Apply KL-divergence based distribution alignment to map latent features from multiple source domains to a Gaussian prior N(0,1).
• Train with a composite loss: L_obj = sum L_c(y_hat, y) + lambda1 L_rank + lambda2 KL(q(Z|X) || N(0,1)).
• Provide theoretical upper bounds on the target-domain risk under Assumptions 1-2 and prove related theorems.
• Experiment on skin lesion classification and spinal cord gray matter segmentation to demonstrate cross-domain generalization gains.

Experimental results

Research questions

• RQ1 Can a latent-space rank constraint capture shared information across multiple medical imaging domains to improve generalization to unseen domains?
• RQ2 Does aligning latent features to a predefined prior via KL divergence reduce overfitting to source domains and improve target-domain performance?
• RQ3 How does the proposed LDDG method compare to state-of-the-art domain generalization baselines on skin lesion classification and spinal cord segmentation?
• RQ4 What is the impact of varying the latent-space rank on generalization performance?

Key findings

• LDDG outperforms baselines (DeepAll, MASF, MLDG, CCSA) on cross-domain skin lesion classification across multiple target domains.
• LDDG achieves the best performance among compared methods on averaged results over target domains in Table 1.
• Ablation shows both Rank and KL components improve generalization, with rank(C) (equal to number of classes) yielding peak accuracy in many settings.
• Rank regularization with nuclear norm underperforms compared to the proposed rank(C) constraint, highlighting the benefit of incorporating category information.
• In ablations, combining rank and KL yields better performance than either component alone.