Rethinking U-net Skip Connections for Biomedical Image Segmentation

TL;DR

This paper investigates the role of skip connections in U-net architectures for biomedical image segmentation, revealing that removing the topmost skip connection improves robustness to domain shifts and enhances performance across multiple datasets.

Contribution

It demonstrates that selectively removing the uppermost skip connection in U-net improves domain robustness and segmentation accuracy, challenging the conventional use of all skip connections.

Findings

Removing the top skip connection increases domain robustness.

Deeper layers are less susceptible to domain shifts.

Performance improves up to 13% across datasets when removing the top skip connection.

Abstract

The U-net architecture has significantly impacted deep learning-based segmentation of medical images. Through the integration of long-range skip connections, it facilitated the preservation of high-resolution features. Out-of-distribution data can, however, substantially impede the performance of neural networks. Previous works showed that the trained network layers differ in their susceptibility to this domain shift, e.g., shallow layers are more affected than deeper layers. In this work, we investigate the implications of this observation of layer sensitivity to domain shifts of U-net-style segmentation networks. By copying features of shallow layers to corresponding decoder blocks, these bear the risk of re-introducing domain-specific information. We used a synthetic dataset to model different levels of data distribution shifts and evaluated the impact on downstream segmentation performance. We quantified the inherent domain susceptibility of each network layer, using the Hellinger distance. These experiments confirmed the higher domain susceptibility of earlier network layers. When gradually removing skip connections, a decrease in domain susceptibility of deeper layers could be observed. For downstream segmentation performance, the original U-net outperformed the variant without any skip connections. The best performance, however, was achieved when removing the uppermost skip connection - not only in the presence of domain shifts but also for in-domain test data. We validated our results on three clinical datasets - two histopathology datasets and one magnetic resonance dataset - with performance increases of up to 10% in-domain and 13% cross-domain when removing the uppermost skip connection.