Anatomy of Domain Shift Impact on U-Net Layers in MRI Segmentation

TL;DR

This paper introduces SpotTUnet, a CNN architecture that automatically determines which layers to fine-tune for MRI segmentation under domain shift, achieving high accuracy with minimal annotated data and providing insights into layer-wise domain shift effects.

Contribution

We propose SpotTUnet, a novel CNN that learns to select layers for fine-tuning in domain adaptation, enhancing performance with limited data and offering interpretability of domain shift impacts.

Findings

SpotTUnet matches top fine-tuning methods in accuracy.

It effectively operates with scarce annotated data.

Provides layer-wise visualization of domain shift impact.

Abstract

Domain Adaptation (DA) methods are widely used in medical image segmentation tasks to tackle the problem of differently distributed train (source) and test (target) data. We consider the supervised DA task with a limited number of annotated samples from the target domain. It corresponds to one of the most relevant clinical setups: building a sufficiently accurate model on the minimum possible amount of annotated data. Existing methods mostly fine-tune specific layers of the pretrained Convolutional Neural Network (CNN). However, there is no consensus on which layers are better to fine-tune, e.g. the first layers for images with low-level domain shift or the deeper layers for images with high-level domain shift. To this end, we propose SpotTUnet - a CNN architecture that automatically chooses the layers which should be optimally fine-tuned. More specifically, on the target domain, our method additionally learns the policy that indicates whether a specific layer should be fine-tuned or reused from the pretrained network. We show that our method performs at the same level as the best of the nonflexible fine-tuning methods even under the extreme scarcity of annotated data. Secondly, we show that SpotTUnet policy provides a layer-wise visualization of the domain shift impact on the network, which could be further used to develop robust domain generalization methods. In order to extensively evaluate SpotTUnet performance, we use a publicly available dataset of brain MR images (CC359), characterized by explicit domain shift. We release a reproducible experimental pipeline.