[Paper Review] SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation

Jesse Sun, Fatemeh Darbehani+2 more|arXiv (Cornell University)|Jan 21, 2020

Advanced Neural Network ApplicationsComputer Science39 references22 citations

TL;DR

SAUNet proposes a Shape Attentive U-Net architecture that enhances medical image segmentation interpretability and robustness by integrating a parallel shape stream to learn structural features alongside traditional texture-based features. The model achieves state-of-the-art performance on SUN09 and AC17 cardiac MRI datasets while generating multi-resolution saliency maps during inference, eliminating the need for post-hoc gradient-based interpretability methods.

ABSTRACT

Medical image segmentation is a difficult but important task for many clinical operations such as cardiac bi-ventricular volume estimation. More recently, there has been a shift to utilizing deep learning and fully convolutional neural networks (CNNs) to perform image segmentation that has yielded state-of-the-art results in many public benchmark datasets. Despite the progress of deep learning in medical image segmentation, standard CNNs are still not fully adopted in clinical settings as they lack robustness and interpretability. Shapes are generally more meaningful features than solely textures of images, which are features regular CNNs learn, causing a lack of robustness. Likewise, previous works surrounding model interpretability have been focused on post hoc gradient-based saliency methods. However, gradient-based saliency methods typically require additional computations post hoc and have been shown to be unreliable for interpretability. Thus, we present a new architecture called Shape Attentive U-Net (SAUNet) which focuses on model interpretability and robustness. The proposed architecture attempts to address these limitations by the use of a secondary shape stream that captures rich shape-dependent information in parallel with the regular texture stream. Furthermore, we suggest multi-resolution saliency maps can be learned using our dual-attention decoder module which allows for multi-level interpretability and mitigates the need for additional computations post hoc. Our method also achieves state-of-the-art results on the two large public cardiac MRI image segmentation datasets of SUN09 and AC17.

Motivation & Objective

• To address the lack of interpretability and robustness in deep learning-based medical image segmentation models.
• To improve model generalization by learning shape features, which are more robust than texture-based features.
• To eliminate reliance on computationally expensive post-hoc saliency methods like SmoothGrad by embedding multi-resolution interpretability directly into the model.
• To achieve state-of-the-art performance on large public cardiac MRI segmentation benchmarks (SUN09 and AC17).
• To develop a clinically adoptable segmentation framework that supports transparency and reliability in diagnostic applications.

Proposed method

• Introduces a dual-stream encoder-decoder architecture with a primary texture stream and a secondary gated shape stream to learn shape-dependent features in parallel.
• Employs a dual-attention decoder module that generates multi-resolution spatial and shape attention maps during inference, enabling built-in interpretability.
• Uses a dual-task loss objective that jointly optimizes for segmentation accuracy and shape feature learning, improving robustness.
• Applies skip connections between encoder and decoder paths to preserve spatial context while enabling feature fusion across scales.
• Generates saliency maps via attention maps at different decoder levels, avoiding the need for additional forward-backward passes required by methods like SmoothGrad.
• Utilizes RAdam optimizer with a learning rate of 5e-4 and batch size of 4 for training on SUN09 and AC17 datasets.

Experimental results

Research questions

• RQ1 Can a dual-stream U-Net architecture that explicitly learns shape features improve segmentation robustness compared to standard CNNs?
• RQ2 Does integrating a shape stream reduce the performance drop when generalizing from training data (SUN09) to a different test set (AC17)?
• RQ3 Can built-in attention maps provide multi-resolution interpretability without relying on post-hoc gradient-based methods such as SmoothGrad?
• RQ4 To what extent do shape-aware features improve segmentation accuracy on bi-ventricular cardiac MRI segmentation tasks?
• RQ5 Can the proposed method achieve state-of-the-art performance while maintaining inherent interpretability?

Key findings

• The model with the shape stream achieved a 90.84% mIoU on the SUN09 training set, outperforming the baseline without the shape stream (90.20%).
• On the AC17 validation set, the model with the shape stream achieved an 81.73% mIoU, showing a 1.76% smaller performance drop compared to the baseline (79.33%).
• The shape stream improved generalization, as evidenced by a 1.76% smaller drop in mIoU when transferring from SUN09 to AC17, indicating greater robustness.
• The proposed method generated all multi-resolution saliency maps in 20 seconds for the entire AC17 validation set, compared to 24 minutes using SmoothGrad, demonstrating significant computational efficiency.
• Attention maps from the shape stream clearly localized on the anatomical shapes of the left and right ventricles, confirming that the model learned meaningful structural features.
• The dual-attention decoder enabled hierarchical interpretability, with higher-resolution blocks focusing on fine details and lower-resolution blocks capturing global shape priors.