Hyperbolic U-Net for Robust Medical Image Segmentation

Swasti Shreya Mishra, Max van Spengler, Erwin Berkhout, Pascal Mettes

Keywords: Medical Image Segmentation, Hyperbolic Neural Networks, Robust Medical Image Segmentation, U-Net, Robustness

TL;DR: We propose Hyperbolic U-Net, which brings U-Net's proven encoder-decoder structure entirely to the hyperbolic space, resulting in a network highly robust to out-of-the-box noise.

Abstract: The U-Net architecture is a leading network in medical image segmentation. Despite its strong segmentation performance, U-Net struggles when dealing with noise in image data, such as random interference and brightness variations. While a common occurrence, the presence of random noise leads to strong performance degradation in U-Net, hampering its clinical integration and robustness. In this work, we investigate the role of geometry in U-Net. All U-Net variations share the same geometric foundations, namely Euclidean geometry. Here, we propose Hyperbolic U-Net, which maintains U-Net's proven encoder-decoder structure while operating entirely in the Poincar\'e ball of hyperbolic space. We identify two main roadblocks for training a fully Hyperbolic U-Net and propose a solution for each: (i) fully hyperbolic literature has so far focused on encoders, limiting their applicability to segmentation. We introduce hyperbolic 2D transpose convolution and hyperbolic bilinear upsampling layers that make it possible to create decoders; and (ii) existing hyperbolic parameter initializations are not suitable for hyperbolic decoder blocks. We introduce a Newton's approximation-scaled weight initialization, which ensures norm preservation for all layers at the start of training. Empirically, we show that our Hyperbolic U-Nets strongly outperform standard Euclidean U-Nets across multiple medical image datasets for Gaussian, Speckle, Poisson, and Rician noise, as well as to brightness and contrast shift. We conclude that a fully Hyperbolic U-Net is highly robust to out-of-the-box noise, without the need for denoising or additional objectives, highlighting the potential of hyperbolic geometry for medical imaging. All code and models will be made publicly available.

Primary Subject Area: Segmentation

Secondary Subject Area: Geometric Deep Learning

Registration Requirement: Yes

Reproducibility: https://github.com/swastishreya/Hyperbolic-U-Net

Visa & Travel: Yes

Read CFP & Author Instructions: Yes

Originality Policy: Yes

Single-blind & Not Under Review Elsewhere: Yes

LLM Policy: Yes

Submission Number: 99