EU-Nets: Enhanced, Explainable and Parsimonious U-Nets

TL;DR

EU-Nets address the need for transparent and reliable segmentation in medical imaging by extending U-Nets with the MHEX+ framework. They introduce the Equivalent Convolutional Kernel to improve interpretability and a collaboration gradient approach for uncertainty, achieving an average accuracy gain of $1.389\%$ and a variance reduction of $0.83\%$ with fewer than $0.1$M parameters. The framework is architecture-agnostic, enabling EU-Nets to be plugged into U-Net, U-Net++, U-Net Transformer, and AHF-U-Net. Salience maps from MHEX+ offer multi-scale, decoding-stage insights beyond simple Grad-CAM outputs, and the gradient-based uncertainty correlates with Deep Ensembles while remaining computationally efficient. Overall, EU-Nets provide a practical path to more trustworthy segmentation in diverse medical imaging tasks.

Abstract

In this study, we propose MHEX+, a framework adaptable to any U-Net architecture. Built upon MHEX+, we introduce novel U-Net variants, EU-Nets, which enhance explainability and uncertainty estimation, addressing the limitations of traditional U-Net models while improving performance and stability. A key innovation is the Equivalent Convolutional Kernel, which unifies consecutive convolutional layers, boosting interpretability. For uncertainty estimation, we propose the collaboration gradient approach, measuring gradient consistency across decoder layers. Notably, EU-Nets achieve an average accuracy improvement of 1.389\% and a variance reduction of 0.83\% across all networks and datasets in our experiments, requiring fewer than 0.1M parameters.

Figures (6)

• Figure 1: The left side illustrates the working mechanism of the MHEX+ module, while the right side represents the Equivalent Convolutional Kernel
• Figure 2: Example of MHEX+ integration with U-Net Transformer (EU-Net Transformer)
• Figure 3: Comparison of segmentation performance using Dice coefficient (%). The most suitable loss function was chosen for each dataset: $^D$ denotes the use of Dice loss, $^C$ indicates cross-entropy loss, and $^*$ signifies cases where cross-entropy loss was used due to training failure.
• Figure 4: Comparison of different models on the MSD-Heart dataset. From top to bottom, the rows represent predictions, confidence maps (pixel-wise probabilities), explanation maps (Grad-CAM for U-Nets, MHEX+ for EU-Nets), and uncertainty maps generated by MHEX+'s gradient.
• Figure 5: Right: Grad-CAM salience map. Left: MHEX+ salience map, where deeper decisions, lightly activated, reveal decision traces and provide deeper insights. These regions are marked by dashed lines.