MLA-BIN: Model-level Attention and Batch-instance Style Normalization for Domain Generalization of Federated Learning on Medical Image Segmentation

TL;DR

This paper introduces MLA-BIN, a novel approach combining model-level attention and style normalization to enhance domain generalization in federated learning for medical image segmentation, addressing unseen domain performance issues.

Contribution

The paper proposes MLA-BIN, integrating model-level attention and batch-instance style normalization to improve unseen domain generalization in federated learning for medical imaging.

Findings

Outperforms state-of-the-art methods in medical image segmentation tasks.

Effectively generalizes to unseen domains in federated learning settings.

Reduces the impact of inter-domain style differences on model performance.

Abstract

The privacy protection mechanism of federated learning (FL) offers an effective solution for cross-center medical collaboration and data sharing. In multi-site medical image segmentation, each medical site serves as a client of FL, and its data naturally forms a domain. FL supplies the possibility to improve the performance of seen domains model. However, there is a problem of domain generalization (DG) in the actual de-ployment, that is, the performance of the model trained by FL in unseen domains will decrease. Hence, MLA-BIN is proposed to solve the DG of FL in this study. Specifically, the model-level attention module (MLA) and batch-instance style normalization (BIN) block were designed. The MLA represents the unseen domain as a linear combination of seen domain models. The atten-tion mechanism is introduced for the weighting coefficient to obtain the optimal coefficient ac-cording to the similarity of inter-domain data features. MLA enables the global model to gen-eralize to unseen domain. In the BIN block, batch normalization (BN) and instance normalization (IN) are combined to perform the shallow layers of the segmentation network for style normali-zation, solving the influence of inter-domain image style differences on DG. The extensive experimental results of two medical image seg-mentation tasks demonstrate that the proposed MLA-BIN outperforms state-of-the-art methods.