Adapting Off-the-Shelf Source Segmenter for Target Medical Image Segmentation

Introduction

Accurate tumor segmentation is a critical step for early tumor detection and intervention, and has been significantly improved with advanced deep neural networks (DNN) . A segmentation model trained in a source domain, however, usually cannot generalize well in a target domain, e.g., data acquired from a new scanner or different clinical center, in implementation. Besides, annotating data in the new target domain is costly and even infeasible . To address this, unsupervised domain adaptation (UDA) was proposed to transfer knowledge from a labeled source domain to unlabeled target domains .

The typical UDA solutions can be classified into three categories: statistic moment matching, feature/pixel-level adversarial learning , and self-training . These UDA methods assume that the source domain data are available and usually trained together with target data. The source data, however, are often inaccessible, due to data storage or privacy issues, for cross-clinical center implementation . Therefore, it is of great importance to apply an “off-the-shelf” source domain model, without access to the source data. For source-free classification UDA, Liang et al. proposed to enforce the diverse predictions, while the diversity of neighboring pixels is not suited for the segmentation purpose. In addition, the class prototype and variational inference methods are not scalable for pixel-wise classification based segmentation. More importantly, without distribution alignment, these methods relied on unreliable noisy pseudo labeling.

Recently, the source relaxed UDA was presented to pre-train an additional class ratio predictor in the source domain, by assuming that the class ratio, i.e., pixel proportion in segmentation, is invariant between source and target domains. At the adaptation stage, the class ratio was used as the only transferable knowledge. However, that work has two limitations. First, the class ratio can be different between the two domains, due to label shift . For example, a disease incident rate could vary between different countries, and tumor size could vary between different subtypes and populations. Second, the pre-trained class ratio predictor used in is not typical for medical image segmentation, thereby requiring an additional training step using the data in the source domain.

In this work, to address the aforementioned limitations, we propose a practical UDA framework aimed at the source-free UDA for segmentation, without an additional network trained in the source domain or the unrealistic assumption of class ratio consistency between source and target domains. More specifically, our framework hinges on the batch-wise normalization statistics, which are easy to access and compute. Batch Normalization (BN) has been a default setting in the most of modern DNNs, e.g., ResNet and U-Net , for faster and more stable training. Notably, the BN statistics of the source domain are stored in the model itself. The low-order batch statistics, e.g., mean and variance, are domain-specific, due to the discrepancy of input data. To gradually adapt the low-order batch statistics from the source domain to the target domain, we develop a momentum-based progression scheme, where the momentum follows an exponential decay w.r.t. the adaptation iteration. For the domain shareable high-order batch statistics, e.g., scaling and shifting parameters, a high-order batch statistics consistent loss is applied to explicitly enforce the discrepancy minimization. The transferability of each channel is adaptively measured first, from which to balance the contribution of each channel. Moreover, the proposed unsupervised self-entropy minimization can be simply added on top of our framework to boost the performance further.

Our contributions are summarized as follows:

$\bullet$ To our knowledge, this is the first source relaxed or source free UDA framework for segmentation. We do not need an additional source domain network, or the unrealistic assumption of the class ratio consistency . Our method only relies on an “off-the-shelf” pre-trained segmentation model with BN in the source domain.

$\bullet$ The domain-specific and shareable batch-wise statistics are explored via the low-order statistics progression with an exponential momentum decay scheme and transferability adaptive high-order statistics consistency loss, respectively.

$\bullet$ Comprehensive evaluations on both cross-subtype (i.e., HGG to LGG) and cross-modality (i.e., T2 to T1/T1ce/FLAIR) UDA tasks using the BraTS 2018 database demonstrate the validity of our proposed framework and its superiority to conventional source-relaxed/source-based UDA methods.

Methodology

We assume that a segmentation model with BN is pre-trained with source domain data, and the batch statistics are inherently stored in the model itself. At the adaptation stage, we fine-tune the model based on the batch-wise statistics and the self-entropy (SE) of target data prediction. The overview of the different setups of conventional UDA and our “off-the-shelf (OS)” UDA is shown in Fig. 1. Below, we briefly revisit the BN in Subsec. 2.1 first and then introduce our OSUDA in Subsec. 2.2. The added unsupervised SE minimization and the overall training protocol are detailed in Subsec. 2.3.

In testing, the input is usually a single sample rather than a batch with $B$ samples. Therefore, BN stores the exponentially weighted average of the batch statistics at the training stage and used it in testing. Specifically, the mean and variance over the training are tracked progressively, i.e.,

where $\eta\in$ is a momentum parameter. After $K$ training iterations, $\overline{\mu}_{l,c}^{K}$, $\{\overline{\sigma}^{2}\}^{K}_{l,c}$, $\gamma_{l,c}^{K}$, and $\beta_{l,c}^{K}$ are stored and used for testing normalization .

2 Adaptive source-relaxed batch-wise statistics adaptation

Early attempts of BN for UDA simply added BN in the target domain, without the interaction with the source domain . Recent studies indicated that the low-order batch statistics, i.e., mean $\mu_{l,c}$ and variance $\{\sigma^{2}\}_{l,c}$, are domain-specific, because of the divergence of cross-domain representation distributions. Therefore, brute-forcing the same mean and variance across domains can lead to a loss of expressiveness . In contrast, after the low-order batch statistics discrepancy is partially reduced, with domain-specific mean and variance normalization, the high-order batch statistics, i.e., scaling and shifting parameters $\gamma_{l,c}$ and $\beta_{l,c}$, are shareable across domains .

However, all of the aforementioned methods require the source data at the adaptation stage. To address this, in this work, we propose to mitigate the domain shift via the adaptive low-order batch statistics progression with momentum, and explicitly enforce the consistency of the high-order statistics in a source-relaxed manner.

Low-order statistics progression with an exponential momentum decay scheme. In order to gradually learn the target domain-specific mean and variance, we propose an exponential low-order batch statistics decay scheme. We initialize the mean and variance in the target domain with the tracked $\overline{\mu}_{l,c}^{K}$ and $\{\overline{\sigma}^{2}\}^{K}_{l,c}$ in the source domain, which is similar to applying a model with BN in testing . Then, we progressively update the mean and variance in the $t$-th adaptation iteration in the target domain as

where $\eta^{t}=\eta^{0}\text{exp}(-t)$ is a target adaptation momentum parameter with an exponential decay w.r.t. the iteration $t$. ${\mu}_{l,c}^{t}$ and $\{{\sigma}^{2}\}^{t}_{l,c}$ are the mean and variance of the current target batch. Therefore, the weight of $\overline{\mu}_{l,c}^{K}$ and $\{\overline{\sigma}^{2}\}^{K}_{l,c}$ are smoothly decreased along with the target domain adaptation, while ${\mu}_{l,c}^{t}$ and $\{{\sigma}^{2}\}^{t}_{l,c}$ gradually represent the batch-wise low-order statistics of the target data.

Transferability adaptive high-order statistics consistency. For the high-order batch statistics, i.e., the learned scaling and shifting parameters, we explicitly encourage its consistency between the two domains with the following high-order batch statistics (HBS) loss:

where $\gamma_{l,c}^{K}$ and $\beta_{l,c}^{K}$ are the learned scaling and shifting parameters in the last iteration of pre-training in the source domain. $\gamma_{l,c}^{t}$ and $\beta_{l,c}^{t}$ are the learned scaling and shifting parameters in the $t$-th adaptation iteration. $\alpha_{l,c}$ is an adaptive parameter to balance between the channels.

We note that the domain divergence can be different among different layers and channels, and the channels with smaller divergence can be more transferable . Accordingly, we would expect that the channels with higher transferability contribute more to the adaptation. In order to quantify the domain discrepancy in each channel, a possible solution is to measure the difference between batch statistics. In the source-relaxed UDA setting, we define the channel-wise source-target distance in the $t$-th adaptation iteration as

Then, the transferability of each channel can be measured by $\alpha_{l,c}=\frac{L\times C\times(1+d_{l,c})^{-1}}{\sum_{l}\sum_{c}(1+d_{l,c})^{-1}}$. Therefore, the more transferable channels will be assigned with higher importance, i.e., with larger weight $(1+\alpha_{l,c})$ in $\mathcal{L}_{l,c}$.

3 Self-entropy minimization and overall training protocol

The training in the unlabeled target domain can also be guided by an unsupervised learning framework. The SE minimization is a widely used objective in modern DNNs to encourage the confident prediction, i.e., the maximum softmax value can be high . SE for pixel segmentation is calculated by the averaged entropy of the classifier’s softmax prediction given by

where $H_{0}$ and $W_{0}$ are the height and width of the input, and $\delta_{b,n}$ is the histogram distribution of the softmax output of the $n$-th pixel of the $b$-th image in a batch. Minimizing $\mathcal{L}_{SE}$ leads to the output close to a one-hot distribution.

At the source-domain pre-training stage, we follow the standard segmentation network training protocol. At the target domain adaptation stage, the overall training objective can be formulated as $\mathcal{L}=\mathcal{L}_{HBS}+\lambda\mathcal{L}_{SE}$, where $\lambda$ is used to balance between the BN statistics matching and SE minimization. We note that a trivial solution of SE minimization is that all unlabeled target data could have the same one-hot encoding . Thus, to stabilize the training, we linearly change the hyper-parameter $\lambda$ from 10 to 0 in training.

Experiments and Results

The BraTS2018 database is composed of a total of 285 subjects , including 210 high-grade gliomas (HGG, i.e., glioblastoma) subjects, and 75 low-grade gliomas (LGG) subjects. Each subject has T1-weighted (T1), T1-contrast enhanced (T1ce), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (FLAIR) Magnetic Resonance Imaging (MRI) volumes with voxel-wise labels for the enhancing tumor (EnhT), the peritumoral edema (ED), and the necrotic and non-enhancing tumor core (CoreT). Usually, we denote the sum of EnhT, ED, and CoreT as the whole tumor. In order to demonstrate the effectiveness and generality of our OSUDA, we follow two UDA evaluation protocols using the BraTS2018 database, including HGG to LGG UDA and cross-modality (i.e., T2 to T1/T1ce/FLAIR) UDA .

For evaluation, we adopted the widely used Dice similarity coefficient and Hausdorff distance metrics as in . The Dice similarity coefficient (the higher, the better) measures the overlapping part between our prediction results and the ground truth. The Hausdorff distance (the lower, the better) is defined between two sets of points in the metric space.

HGG and LGG have different size and position distributions for tumor regions . Following the standard protocol, we used the HGG training set (source domain) to pre-train the segmentation model and adapted it with the LGG training set (target domain) . The evaluation was implemented in the LGG testing set. We adopted the same 2D U-Net backbone in , sliced 3D volumes into 2D axial slices with the size of 128$\times$128, and concatenated all four MRI modalities to get a 4-channel input.

The quantitative evaluation results are shown in Table 1. Since the pixel proportion of each class is different between HGG and LGG domains, the class ratio-based CRUDA only achieved marginal improvements with its unsupervised learning objective. We note that the Dice score of the core tumor was worse than the pre-trained source-only model, which can be the case of negative transfer . Our proposed OSUDA achieved the state-of-the-art performance for source-relaxed UDA segmentation, approaching the performance of SEAT with the source data, which can be seen as an “upper-bound.”

We used OSUDA-AC and OSUDA-SE to indicate the OSUDA without the adaptive channel-wise weighting and self-entropy minimization, respectively. The better performance of OSUDA over OSUDA-AC and OSUDA-SE demonstrates the effectiveness of adaptive channel-wise weighting and self-entropy minimization. The illustration of the segmentation results is given in Fig. 2. We can see that the predictions of our proposed OSUDA are better than the no adaptation model. In addition, CRUDA had a tendency to predict a larger area for the tumor; and the tumor core is often predicted for the slices without the core.

2 Cross-modality T2 to T1/T1ce/FLAIR UDA

Because of large appearance discrepancies between different MRI modalities, we further applied our framework to the cross-modality UDA task. Since clinical annotation of the whole tumor is typically performed on T2-weighted MRI, the typical cross-modality UDA setting is to use T2-weighted MRI as the labeled source domain, and T1/T1ce/FLAIR MRI as the unlabeled target domains . We followed the UDA training (80% subjects) and testing (20% subjects) split as in , and adopted the same single-channel input backbone. We note that the data were used in an unpaired manner .

The quantitative evaluation results are provided in Table 2. Our proposed OSUDA outperformed CRUDA consistently. In addition, in CRUDA, the additional class ratio prediction model was required to be trained with the source data, which is prohibitive in many real-world cases. Furthermore, our OSUDA outperformed several UDA methods trained with the source data, e.g., CycleGAN and SIFA , for the two metrics. The visual segmentation results of three target modalities are shown in Fig. 3, showing the superior performance of our framework, compared with the comparison methods.

Discussion and Conclusion

This work presented a practical UDA framework for the tumor segmentation task in the absence of the source domain data, only relying on the “off-the-shelf” pre-trained segmentation model with BN in the source domain. We proposed a low-order statistics progression with an exponential momentum decay scheme to gradually learn the target domain-specific mean and variance. The domain shareable high-order statistics consistency is enforced with our HBS loss, which is adaptively weighted based on the channel-wise transferability. The performance was further boosted with the unsupervised learning objective via self-entropy minimization. Our experimental results on the cross-subtype and cross-modality UDA tasks demonstrated that the proposed framework outperformed the comparison methods, and was robust to the class ratio shift.

Acknowledgements

This work is partially supported by NIH R01DC018511, R01DE027989, and P41EB022544.

References