Source-linked AI summary
Shallow Attention Network for Polyp Segmentation
Jun Wei, Yiwen Hu, Ruimao Zhang, Zhen Li, S. Kevin Zhou, Shuguang Cui
TL;DR
Polyp segmentation is important for colorectal cancer diagnosis but faces color variation, small-polyps degradation, and foreground-background imbalance. SANet addresses these challenges with color exchange, shallow attention, and inference-time probability correction. Across five challenging benchmarks, it reports superior segmentation performance and about 72FPS.
Problem
Polyp segmentation is hindered by inconsistent colors, repeated downsampling that degrades small polyps, and severe foreground-background pixel imbalance.
Method
SANet uses color exchange to decouple image content and color, shallow attention to filter shallow-feature background noise, and PCS to correct biased predictions during inference.
Results
SANet achieves state-of-the-art performance on five public benchmarks and runs at about 72FPS.
Takeaways & Limitations
Color exchange, shallow attention, and inference-time probability correction together reduce interference from color, background noise, and pixel imbalance in polyp segmentation.
Abstract
from arXiv · showhide
Accurate polyp segmentation is of great importance for colorectal cancer diagnosis. However, even with a powerful deep neural network, there still exists three big challenges that impede the development of polyp segmentation. (i) Samples collected under different conditions show inconsistent colors, causing the feature distribution gap and overfitting issue; (ii) Due to repeated feature downsampling, small polyps are easily degraded; (iii) Foreground and background pixels are imbalanced, leading to a biased training. To address the above issues, we propose the Shallow Attention Network (SANet) for polyp segmentation. Specifically, to eliminate the effects of color, we design the color exchange operation to decouple the image contents and colors, and force the model to focus more on the target shape and structure. Furthermore, to enhance the segmentation quality of small polyps, we propose the shallow attention module to filter out the background noise of shallow features. Thanks to the high resolution of shallow features, small polyps can be preserved correctly. In addition, to ease the severe pixel imbalance for small polyps, we propose a probability correction strategy (PCS) during the inference phase. Note that even though PCS is not involved in the training phase, it can still work well on a biased model and consistently improve the segmentation performance. Quantitative and qualitative experimental results on five challenging benchmarks confirm that our proposed SANet outperforms previous state-of-the-art methods by a large margin and achieves a speed about 72FPS.
1 Introduction
Polyp segmentation supports colorectal cancer diagnosis but is challenged by color variation, small targets, and limited data. SANet addresses these issues with color exchange, shallow attention, and probability correction.
- Colorectal cancer is a serious health threat, and polyps can transform into malignant lesions, making early diagnosis and treatment important.
- Different acquisition conditions create color distributions that can encourage models to overfit color rather than polyp structure.
- Color exchange randomly transfers colors between images to decouple image content and color and reduce this overfitting.
- SANet combines shallow attention and probability correction to improve small-polyp segmentation and balance biased predictions.
- Experiments on five public benchmarks report state-of-the-art performance for SANet.
2 Related Work
Polyp segmentation progressed from handcrafted features to fully convolutional networks and specialized boundary modeling. SANet extends this line of work to address persistent challenges under limited data.
- Handcrafted features such as color and texture struggle to capture global context and remain robust in complex scenarios.
- U-Net, SegNet, U-Net++, and ResUNet++ use encoder-decoder, pooling-index, dense-connection, or pretrained-backbone designs to improve segmentation.
- Psi-Net and SFA incorporate boundary information or area-boundary constraints, while PraNet uses reverse attention to refine object boundaries.
- Color Exchange transforms paired images between RGB and LAB spaces, computes channel statistics, and transfers color statistics before conversion back to RGB.
- SANet is proposed to further improve polyp segmentation because existing methods still face challenges associated with limited data.
3 Method
SANet combines color exchange, shallow attention, and inference-time probability correction to address color overfitting, small-polyp information loss, and foreground–background imbalance in polyp segmentation.
- Architecture: Res2Net supplies features from the last three encoder blocks for SANet’s subsequent processing.The selected blocks are f_i for i ∈ (3, 4, 5).
- Color Exchange: Color exchange randomly transfers another image’s color to each input, reducing the association between color and polyp segmentation.This operation is intended to prevent the model from overfitting color-related spurious correlations.
- Shallow Attention Module: Small polyps suffer information loss during repeated CNN downsampling, motivating greater use of high-resolution shallow features.Shallow features preserve clear object boundaries, whereas deep features provide cleaner backgrounds but coarser boundaries.
- Shallow Attention Module: The shallow attention module filters background noise from shallow features and dynamically weights features from different blocks.The resulting shallow features provide cleaner cues for segmenting small polyps.
- Probability Correction Strategy: PCS corrects biased predictions by normalizing positive and negative logits according to their sample proportions during inference.It targets severe foreground–background imbalance, especially for small polyps, without participating in training.
- Probability Correction Strategy: Predictions with PCS have clearer boundaries, while the method adds very little computation during inference.The paper presents visual comparisons between predictions with and without PCS.
- Loss Function: SANet uses binary cross entropy and Dice losses for supervision, with both weighting coefficients set to 1.The loss combines BCE(P, G) and Dice(P, G).
4 Experiments
SANet is evaluated on five polyp datasets against six state-of-the-art models using quantitative, threshold-based, visual, and ablation comparisons. It achieves the best reported performance across datasets, maintains an advantage across thresholds, and runs at about 72FPS.
- Experimental Setup: Experiments use five polyp datasets and compare SANet with six state-of-the-art segmentation models under the same training and testing division.Input images are uniformly resized to 352×352.
- Quantitative Comparison: SANet achieves the best scores across five datasets on both mIoU and mDice.The comparison includes quantitative results against six state-of-the-art models.
- Quantitative Comparison: SANet consistently outperforms other models under different thresholds on three polyp datasets.The threshold analysis is presented through Dice curves.
- Quantitative Comparison: 72FPS on an RTX 2080Ti GPU exceeds the 64FPS reported for PraNet.The speed comparison is reported alongside the segmentation results.
- Visual Comparison: Visual comparisons show that SANet highlights polyp regions, suppresses background noise, and generates accurate masks in challenging scenarios.The qualitative comparison includes the proposed method and four state-of-the-art methods.
- Ablation Study: Controlled experiments on ColonDB and Kvasir find all proposed modules or strategies necessary for the final predictions.Combining the methods yields new state-of-the-art performance.
5 Conclusion
The conclusion addresses overfitting in limited-data polyp segmentation through color-content decoupling and shallow attention for small-polyps. It identifies incorporating more prior knowledge into robust features as future work.
- Conclusion: Limited datasets make polyp segmentation models vulnerable to overfitting.The conclusion frames this as the central problem addressed by the paper.
- Conclusion: Color exchange decouples image color and content to address false color causality.This is one of the two aspects used to alleviate overfitting.
- Conclusion: Shallow attention reduces data noise for difficult small-polyp segmentation.The conclusion presents it as the second aspect of the proposed solution.
- Conclusion: Future work will combine more prior knowledge to design robust features that remove interference from independent factors.This is the stated direction for improving feature robustness.