The Devil is in the Statistics: Mitigating and Exploiting Statistics Difference for Generalizable Semi-supervised Medical Image Segmentation

TL;DR

The paper tackles semi-supervised domain generalization for medical image segmentation under cross-institution domain shifts that introduce varying feature statistics and degrade pseudo-labels. It introduces statistics-individual branches (SIBs) to normalize features per domain and a statistics-aggregated branch (SAB) to learn domain-invariant representations, coupled with multi-level consistency regularization that includes histogram matching and a novel random BN perturbation (RandBN) at the feature level. The approach demonstrated state-of-the-art performance across three medical imaging benchmarks (Prostate, Fundus, M&Ms), with ablations confirming the effectiveness of SIAB, histogram matching, and RandBN in improving robustness and segmentation accuracy. The method maintains practical efficiency by discarding SIBs at inference while leveraging perturbations to expand the training distribution, and code is publicly available.

Abstract

Despite the recent success of domain generalization in medical image segmentation, voxel-wise annotation for all source domains remains a huge burden. Semi-supervised domain generalization has been proposed very recently to combat this challenge by leveraging limited labeled data along with abundant unlabeled data collected from multiple medical institutions, depending on precisely harnessing unlabeled data while improving generalization simultaneously. In this work, we observe that domain shifts between medical institutions cause disparate feature statistics, which significantly deteriorates pseudo-label quality due to an unexpected normalization process. Nevertheless, this phenomenon could be exploited to facilitate unseen domain generalization. Therefore, we propose 1) multiple statistics-individual branches to mitigate the impact of domain shifts for reliable pseudo-labels and 2) one statistics-aggregated branch for domain-invariant feature learning. Furthermore, to simulate unseen domains with statistics difference, we approach this from two aspects, i.e., a perturbation with histogram matching at image level and a random batch normalization selection strategy at feature level, producing diverse statistics to expand the training distribution. Evaluation results on three medical image datasets demonstrate the effectiveness of our method compared with recent SOTA methods. The code is available at https://github.com/qiumuyang/SIAB.

Figures (7)

• Figure 1: (a) Images sampled from 4 domains of the Fundus dataset. (b) Differences in feature statistics captured by batch normalization layers on the Fundus dataset. (c) Comparison of pseudo-label quality (measured by dice similarity coefficient against ground truth) between separate and mixed domain feature normalization.
• Figure 2: Overview of our method. The left figure illustrates the structure of the statistics-individual and aggregated branches. The right figure illustrates the consistency regularization framework extended with diverse perturbations. $u_w, u_s, u_h$ denote the weak, strong, and histogram matching-augmented unlabeled images respectively, and $p_w, p_s, p_h$ are the corresponding predictions. $p_r$ is the prediction of $u_w$ through random forward. Arrows symbolize the forward processes based on the proposed branches during prediction. Dashed lines signify that one branch is randomly selected each time. $\bigoplus$ denotes element-wise addition.
• Figure 3: Illustration of histogram matching. The top row represents the images from the Prostate dataset. The bottom row visualizes their grayscale distribution.
• Figure 4: T-SNE visualization of image features after applying different perturbations. Different colors represent different domains. Best viewed in color.
• Figure 5: Visualization results from the Prostate dataset.