Domain Generalization on Medical Imaging Classification using Episodic Training with Task Augmentation

TL;DR

This paper introduces a novel episodic training approach with task augmentation for domain generalization in medical imaging, improving model robustness across unseen domains with limited source data.

Contribution

It proposes a new meta-learning scheme combining episodic training and task augmentation to enhance generalization in medical imaging classification tasks.

Findings

Improved generalization on unseen medical imaging domains.

Effective mitigation of overfitting with limited source domains.

Validated on histopathological and CT image datasets.

Abstract

Medical imaging datasets usually exhibit domain shift due to the variations of scanner vendors, imaging protocols, etc. This raises the concern about the generalization capacity of machine learning models. Domain generalization (DG), which aims to learn a model from multiple source domains such that it can be directly generalized to unseen test domains, seems particularly promising to medical imaging community. To address DG, recent model-agnostic meta-learning (MAML) has been introduced, which transfers the knowledge from previous training tasks to facilitate the learning of novel testing tasks. However, in clinical practice, there are usually only a few annotated source domains available, which decreases the capacity of training task generation and thus increases the risk of overfitting to training tasks in the paradigm. In this paper, we propose a novel DG scheme of episodic training with task augmentation on medical imaging classification. Based on meta-learning, we develop the paradigm of episodic training to construct the knowledge transfer from episodic training-task simulation to the real testing task of DG. Motivated by the limited number of source domains in real-world medical deployment, we consider the unique task-level overfitting and we propose task augmentation to enhance the variety during training task generation to alleviate it. With the established learning framework, we further exploit a novel meta-objective to regularize the deep embedding of training domains. To validate the effectiveness of the proposed method, we perform experiments on histopathological images and abdominal CT images.