[Paper Review] FANet: A Feedback Attention Network for Improved Biomedical Image Segmentation

Nikhil Kumar Tomar, Debesh Jha+6 more|arXiv (Cornell University)|Mar 31, 2021

Radiomics and Machine Learning in Medical Imaging4 citations

TL;DR

FANet proposes a feedback attention network that iteratively refines biomedical image segmentation masks by leveraging predictions from previous training epochs as hard attention guides. By unifying prior epoch masks with current feature maps, FANet improves segmentation accuracy across seven public datasets with minimal additional computation, outperforming SOTA methods without ensembling or data augmentation.

ABSTRACT

The increase of available large clinical and experimental datasets has contributed to a substantial amount of important contributions in the area of biomedical image analysis. Image segmentation, which is crucial for any quantitative analysis, has especially attracted attention. Recent hardware advancement has led to the success of deep learning approaches. However, although deep learning models are being trained on large datasets, existing methods do not use the information from different learning epochs effectively. In this work, we leverage the information of each training epoch to prune the prediction maps of the subsequent epochs. We propose a novel architecture called feedback attention network (FANet) that unifies the previous epoch mask with the feature map of the current training epoch. The previous epoch mask is then used to provide a hard attention to the learned feature maps at different convolutional layers. The network also allows to rectify the predictions in an iterative fashion during the test time. We show that our proposed extit{feedback attention} model provides a substantial improvement on most segmentation metrics tested on seven publicly available biomedical imaging datasets demonstrating the effectiveness of FANet. The source code is available at \url{https://github.com/nikhilroxtomar/FANet}.

Motivation & Objective

• To address the limitation of single-step deep learning models in capturing intra- and inter-class variability in biomedical images.
• To improve segmentation performance by leveraging information from previous training epochs during both training and inference.
• To enable self-rectification of predictions through a feedback mechanism that prunes and refines masks iteratively.
• To reduce reliance on large, diverse datasets and complex data augmentation or ensemble strategies.
• To develop a lightweight, end-to-end trainable network that maintains high accuracy with fewer training epochs.

Proposed method

• FANet introduces a feedback mechanism that uses the predicted mask from the previous training epoch as a hard attention map to guide feature learning in the current epoch.
• A MixPool block is used to unify the feature map and the previous epoch's mask, enabling spatial attention through element-wise multiplication.
• The network employs SE-Residual blocks to model channel-wise feature dependencies and enhance representation learning.
• Spatial attention is integrated via a contrastive attention mechanism that suppresses background clutter and emphasizes relevant regions.
• During inference, the model performs iterative mask refinement using the same feedback mechanism, with up to 10 iterations shown to be optimal.
• The feedback mechanism is applied during both training and test time, allowing the network to progressively improve predictions without retraining.

Experimental results

Research questions

• RQ1 Can iterative refinement of segmentation masks using feedback from previous predictions improve performance on diverse biomedical image datasets?
• RQ2 How does incorporating prior epoch predictions as hard attention affect feature representation and segmentation accuracy?
• RQ3 Does the feedback mechanism reduce the need for extensive data augmentation or ensemble methods in biomedical segmentation?
• RQ4 What is the optimal number of inference-time iterations for mask refinement in the FANet architecture?
• RQ5 How does FANet compare to SOTA methods in terms of accuracy, training efficiency, and generalization across multiple datasets?

Key findings

• FANet achieves substantial improvements across all seven benchmark biomedical image segmentation datasets, outperforming existing SOTA methods on multiple metrics.
• The optimal number of inference-time iterations for mask refinement is 10, though significant gains are observed even by iteration 2, with marginal gains beyond iteration 5.
• Ablation studies confirm that the MixPool block with feedback is critical, as removing it reduces performance despite identical parameter counts.
• The network achieves near-SOTA performance with fewer training epochs, indicating improved sample efficiency and faster convergence.
• Qualitative results show that FANet produces clearer, more distinguishable segmentation maps with reduced background noise compared to SOTA methods.
• The feedback mechanism acts as an effective hard attention mechanism, enhancing feature representation without requiring additional learnable parameters or complex transformations.