Privacy-Preserving Constrained Domain Generalization via Gradient Alignment

TL;DR

This paper introduces a privacy-preserving domain generalization method for medical imaging that uses gradient alignment to improve model generalization across different hospitals without compromising patient privacy.

Contribution

The paper proposes a novel gradient alignment loss for federated learning, enhancing cross-domain generalization in medical imaging while preserving privacy.

Findings

Outperforms state-of-the-art federated learning methods in cross-domain tasks

Improves model generalization to unseen medical image domains

Connects gradient alignment with MMD for distribution matching

Abstract

Deep neural networks (DNN) have demonstrated unprecedented success for medical imaging applications. However, due to the issue of limited dataset availability and the strict legal and ethical requirements for patient privacy protection, the broad applications of medical imaging classification driven by DNN with large-scale training data have been largely hindered. For example, when training the DNN from one domain (e.g., with data only from one hospital), the generalization capability to another domain (e.g., data from another hospital) could be largely lacking. In this paper, we aim to tackle this problem by developing the privacy-preserving constrained domain generalization method, aiming to improve the generalization capability under the privacy-preserving condition. In particular, We propose to improve the information aggregation process on the centralized server-side with a novel gradient alignment loss, expecting that the trained model can be better generalized to the "unseen" but related medical images. The rationale and effectiveness of our proposed method can be explained by connecting our proposed method with the Maximum Mean Discrepancy (MMD) which has been widely adopted as the distribution distance measurement. Experimental results on two challenging medical imaging classification tasks indicate that our method can achieve better cross-domain generalization capability compared to the state-of-the-art federated learning methods.