Source-linked AI summary
Attention U-Net: Learning Where to Look for the Pancreas
Ozan Oktay, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori, Steven McDonagh, Nils Y Hammerla, Bernhard Kainz, Ben Glocker, Daniel Rueckert
TL;DR
Variable organ shape and size can require cascaded CNN localization for segmentation. Attention U-Net adds attention gates that suppress irrelevant regions, improving segmentation accuracy across datasets and training sizes while avoiding multiple CNN models.
Problem
U-Net-based segmentation may require cascaded CNNs to localize organs with substantial inter-patient variation in shape and size.
Method
Attention U-Net integrates attention gates that generate soft region proposals and suppress irrelevant feature activations during dense prediction.
Results
Attention gates consistently improve prediction accuracy across two abdominal CT datasets and different training sizes while preserving computational efficiency.
Takeaways & Limitations
The approach can remove the necessity for external organ localization, particularly benefiting identification of variable small organs such as the pancreas.
Takeaways & Limitations
Residual connections did not significantly improve performance in the reported experiments, leaving training-behavior improvements for future work.
Abstract
from arXiv · showhide
We propose a novel attention gate (AG) model for medical imaging that automatically learns to focus on target structures of varying shapes and sizes. Models trained with AGs implicitly learn to suppress irrelevant regions in an input image while highlighting salient features useful for a specific task. This enables us to eliminate the necessity of using explicit external tissue/organ localisation modules of cascaded convolutional neural networks (CNNs). AGs can be easily integrated into standard CNN architectures such as the U-Net model with minimal computational overhead while increasing the model sensitivity and prediction accuracy. The proposed Attention U-Net architecture is evaluated on two large CT abdominal datasets for multi-class image segmentation. Experimental results show that AGs consistently improve the prediction performance of U-Net across different datasets and training sizes while preserving computational efficiency. The code for the proposed architecture is publicly available.
1 Introduction
The paper introduces self-attention gates for CNN-based dense prediction, integrating them into U-Net to suppress irrelevant regions and improve segmentation sensitivity and accuracy. It evaluates Attention U-Net on challenging pancreas CT segmentation across two abdominal CT benchmarks.
- 1 Introduction: Automated segmentation addresses the tedious and error-prone burden of dense manual labelling while supporting faster clinical workflows and quantitative decision-making.CNNs provide high representation power, fast inference, and filter sharing, making them standard architectures for image segmentation.
- 1 Introduction: Self-attention gates generate task-specific soft region proposals that highlight salient features and suppress irrelevant activations without significant computational overhead.The gates operate without additional supervision and avoid the parameter burden of multi-model frameworks.
- 1 Introduction: Grid-based gating makes attention coefficients specific to local regions, improving over global feature-vector gating while supporting dense predictions without adaptive pooling.The approach filters features propagated through U-Net skip connections using contextual information from coarser scales.
- 1 Introduction: Attention gates can replace hard-attention methods and external organ-localisation models in medical image segmentation frameworks.This provides an integrated alternative to explicit localisation modules within feed-forward CNN-based segmentation systems.
- 1 Introduction: The proposed Attention U-Net extends standard U-Net to improve sensitivity to foreground pixels without complicated heuristics, with consistent accuracy improvements across imaging datasets.The model is applied to pancreas segmentation, a difficult CT task because of low tissue contrast and substantial variation in organ shape and size.
2 Methodology
The methodology integrates additive attention gates into a standard U-Net to identify salient target regions and suppress irrelevant responses. Gating uses contextual coarse-scale information while preserving computational efficiency through lightweight transformations and resolution matching.
- U-Net integration: The method builds attention gates on a standard U-Net, whose multi-scale feature extraction supports accurate segmentation with efficient GPU-memory use.Coarse feature maps provide contextual information about foreground category and location.
- Attention gates: Attention coefficients α_i ∈ [0, 1] identify salient regions and element-wise scale input feature maps, preserving task-relevant activations while pruning others.The gates suppress irrelevant feature responses before they are passed onward.
- Attention gates: Each gate uses a coarse-scale gating vector containing contextual information to determine focus regions, with multi-dimensional coefficients supporting multiple semantic classes.This allows individual gates to focus on subsets of target structures.
- Attention gates: The gates use sigmoid rather than softmax activation because sequential softmax produces sparser outputs, while sigmoid experimentally improves convergence of attention-gate parameters.The formulation includes channel-wise 1x1x1 convolutions for input transformations.
- U-Net integration: In U-Net, gates filter skip-connection activations before concatenation, using coarse-scale information to disambiguate irrelevant and noisy responses.This highlights salient features passed through the skip connections during the forward pass.
- Efficient implementation: The implementation reduces trainable parameters and computational complexity with spatially unsupported 1x1x1 convolutions and downsampling to the gating-signal resolution.Attention coefficients are resampled on the grid using trilinear interpolation.
3 Experiments and Results
Experiments on two abdominal CT datasets show that Attention U-Net improves pancreas-focused segmentation over standard U-Net while retaining computational efficiency. Attention gates progressively localize organ boundaries and provide benefits beyond uniformly increasing network capacity.
- Segmentation Experiments: Attention U-Net was evaluated against standard 3D U-Net on challenging multi-class abdominal CT segmentation, including pancreas boundary delineation.The comparison covered segmentation performance, model capacity, computation time, and memory requirements.
- Attention Map Analysis: Across training epochs, attention gates changed from broadly distributed responses to localized organ-boundary focus, with coarse outlines refined at finer resolutions.Visualizations showed the model focusing on the pancreas, kidney, and spleen.
- Segmentation Experiments: Attention gates increased pancreas recall on CT-150 by helping the model localize foreground pixels, while reduced training size caused a smaller spleen-DSC drop than with U-Net.Kidney accuracy was similar between models because kidney tissue contrast was higher.
- Capacity and Efficiency: Adding attention gates contributed more to performance than uniformly distributing additional filters across U-Net layers.The standard U-Net gained 2–3% DSC with 8% extra capacity, but higher-capacity comparisons still favored the smaller attention model.
- Comparison to State-of-the-Art: On the public CT-82 pancreas benchmark, attention gates improved segmentation accuracy and lowered pancreas surface distances, including when models were trained from scratch.The reported improvements were attributed to increased pancreas-pixel recall and compared with state-of-the-art CNN approaches without post-processing.
4 Discussion and Conclusion
The paper presents a novel, generic, modular attention-gate model for medical image segmentation that eliminates the need for an external object-localisation model. Attention-gate training can benefit from transfer learning and multi-stage schemes.
- Contributions: The proposed attention-gate approach eliminates the necessity of applying an external object localisation model in medical image segmentation.The model is presented as generic and modular, allowing application to image classification and regression problems.
- Training considerations: Attention-gate training can benefit from transfer learning and multi-stage training schemes.Pre-trained U-Net weights can initialise the attention network before gates are trained during fine-tuning.