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

Midl Latex Submission Checklist: Ensure no LaTeX errors during compilation., Replace NNN with your OpenReview submission ID., Includes \documentclass{midl}, \jmlryear{2026}, \jmlrworkshop, \jmlrvolume, \editors, and correct \bibliography command., Did not override options of the hyperref package., Did not use the times package., Use the correct spelling and format, avoid Unicode characters, and use LaTeX equivalents instead., Any math in the title and abstract must be enclosed within $...$., Did not override the bibliography style defined in midl.cls and did not use \begin{thebibliography} directly to insert references., Avoid using \scalebox; use \resizebox when needed., Included all necessary figures and removed *unused* files in the zip archive., Removed special formatting, visual annotations, and highlights used during rebuttal., All special characters in the paper and .bib file use LaTeX commands (e.g., \'e for é)., No separate supplementary PDF uploads., Acknowledgements, references, and appendix must start after the main content.

Submission Number: 99