[Paper Review] Swin Deformable Attention Hybrid U-Net for Medical Image Segmentation

Lichao Wang, Jiahao Huang+2 more|arXiv (Cornell University)|Feb 28, 2023

Radiomics and Machine Learning in Medical ImagingMedicine3 citations

TL;DR

This paper proposes SDAH-UNet, a hybrid U-Net architecture integrating Swin Deformable Multi-Head Self-Attention (SDMSA) with parallel convolutional branches to enhance medical image segmentation accuracy and interpretability. By enabling dynamic, deformable attention that focuses precisely on salient regions like myocardium or tumors, the model reduces computation redundancy and provides visual explanations through deformation points, achieving state-of-the-art performance on ACDC and BraTS2020 datasets with a 92.23% DSC on LV segmentation.

ABSTRACT

Medical image segmentation is a crucial task in the field of medical image analysis. Harmonizing the convolution and multi-head self-attention mechanism is a recent research focus in this field, with various combination methods proposed. However, the lack of interpretability of these hybrid models remains a common pitfall, limiting their practical application in clinical scenarios. To address this issue, we propose to incorporate the Shifted Window (Swin) Deformable Attention into a hybrid architecture to improve segmentation performance while ensuring explainability. Our proposed Swin Deformable Attention Hybrid UNet (SDAH-UNet) demonstrates state-of-the-art performance on both anatomical and lesion segmentation tasks. Moreover, we provide a direct and visual explanation of the model focalization and how the model forms it, enabling clinicians to better understand and trust the decision of the model. Our approach could be a promising solution to the challenge of developing accurate and interpretable medical image segmentation models.

Motivation & Objective

• To address the lack of interpretability in hybrid CNN-Transformer models for medical image segmentation.
• To reduce computational redundancy in standard multi-head self-attention by introducing deformable sampling in the Swin Transformer framework.
• To enhance segmentation accuracy by combining global structural modeling via SDMSA with local texture learning via parallel convolutional branches.
• To provide intrinsic, visual explanations of model attention through deformation points, enabling clinicians to interpret model decisions.
• To develop a model that balances high performance with clinical trustworthiness through explainable focalization.

Proposed method

• The proposed SDAH-UNet uses a U-Net architecture with three encoder and decoder blocks, each incorporating a novel SDAPC block combining deformable attention and convolutional features.
• The SDAPC block integrates a Swin Deformable Multi-Head Self-Attention (SDMSA) module that learns dynamic, non-uniform sampling points to focus on irregularly shaped anatomical structures like the left ventricle wall.
• A parallel convolutional branch preserves local texture details and provides inductive bias, complementing the global attention mechanism.
• The model employs a convolutional embedding module with small kernels and GELU activation to preserve spatial resolution and receptive field.
• Deconvolutional expanding modules upsample features for mask prediction, maintaining high-resolution output.
• Model interpretability is enhanced via deformation points, attention score heatmaps, and gradient-based visualization (SEG-GRAD-CAM), with deformation points showing precise focalization on targets.

Experimental results

Research questions

• RQ1 Can deformable attention in a Swin Transformer-based architecture improve segmentation accuracy on medical images with complex, irregular structures?
• RQ2 Does the integration of deformable attention with convolutional branches reduce computation redundancy while enhancing feature learning?
• RQ3 Can deformation points in the attention mechanism provide a more accurate and robust explanation of model focalization than traditional gradient-based methods?
• RQ4 Does the proposed hybrid architecture outperform existing state-of-the-art models on both anatomical and lesion segmentation tasks?
• RQ5 Can the model’s attention mechanism be visually interpreted in a way that supports clinical trust and understanding?

Key findings

• SDAH-UNet achieved a 92.23% Dice Similarity Coefficient (DSC) on left ventricle segmentation in the ACDC dataset, outperforming state-of-the-art models.
• The model achieved a 96.91% DSC on the left ventricle in the ACDC dataset when using four SDAPC blocks, with HD95 of 1.22 mm.
• Ablation studies showed that replacing standard blocks with SDAPC blocks improved DSC by nearly 1.2% on average.
• Removing either the SDMSA or convolutional branch in the SDAPC block reduced DSC by approximately 2%, confirming the contribution of both components.
• Deformation points provided clearer and more robust visual explanations of model focalization than attention score heatmaps or gradient-based methods, especially in capturing tumor boundaries and enhanced regions.
• The model demonstrated superior performance on both anatomical (ACDC) and lesion (BraTS2020) segmentation tasks, confirming its generalization capability.