Source-linked AI summary
DAE-Former: Dual Attention-guided Efficient Transformer for Medical Image Segmentation
Reza Azad, René Arimond, Ehsan Khodapanah Aghdam, Amirhossein Kazerouni, Dorit Merhof
TL;DR
Medical image segmentation needs accurate models, but Transformer self-attention scales quadratically with token count and existing efficient designs may limit global context. DAE-Former redesigns attention to capture full spatial and channel relations efficiently and adds cross-attention to skip connections. It reports state-of-the-art performance on cardiac and skin-lesion datasets without pre-training weights.
Problem
Transformer self-attention has quadratic complexity in the number of tokens, while efficient variants can limit global context or local representation.
Method
DAE-Former combines efficient full-spatial attention, transpose attention for channel context, and skip-connection cross-attention in a hierarchical U-Net-like Transformer.
Results
DAE-Former achieves state-of-the-art results on Synapse and skin lesion segmentation datasets, including a 0.67% Dice increase over MISSFormer on Synapse.
Takeaways & Limitations
The model captures spatial and channel context across the full feature dimension while retaining the same number of parameters as previous architectures.
Takeaways & Limitations
Swin-Unet’s spatial context remains limited during processing, defining a scope boundary for prior window-based Transformer approaches.
Abstract
from arXiv · showhide
Transformers have recently gained attention in the computer vision domain due to their ability to model long-range dependencies. However, the self-attention mechanism, which is the core part of the Transformer model, usually suffers from quadratic computational complexity with respect to the number of tokens. Many architectures attempt to reduce model complexity by limiting the self-attention mechanism to local regions or by redesigning the tokenization process. In this paper, we propose DAE-Former, a novel method that seeks to provide an alternative perspective by efficiently designing the self-attention mechanism. More specifically, we reformulate the self-attention mechanism to capture both spatial and channel relations across the whole feature dimension while staying computationally efficient. Furthermore, we redesign the skip connection path by including the cross-attention module to ensure the feature reusability and enhance the localization power. Our method outperforms state-of-the-art methods on multi-organ cardiac and skin lesion segmentation datasets without requiring pre-training weights. The code is publicly available at https://github.com/mindflow-institue/DAEFormer.
1. Introduction
Medical image segmentation requires accurate algorithms, but deep models depend on extensive training data and costly pixel-level annotations. Transformers capture long-range dependencies yet face quadratic attention cost and weaker local representation, motivating efficient designs that preserve global context.
- Motivation: Accurate medical image segmentation supports disease monitoring, injury visualization, and treatment selection, but pixel-level annotation is tedious and expensive.Deep models also require large training volumes and many parameters to generalize well.
- Transformer limitations: ViT captures long-range dependencies and shape representations through self-attention, but requires substantial training data and has quadratic complexity in the token count.The naive ViT also performs worse than CNNs for local representation.
- Prior approaches: Efficient Transformer designs reduce cost by restricting attention locally, scaling down spatial dimensions, or using channel attention, but these approaches only partially capture global context.The cited strategies trade full spatial coverage for lower computational burden.
- Prior approaches: Swin-Unet uses shifted-window Transformer blocks and multi-scale representations, yet its spatial context remains limited during processing.Its U-Net-like structure attempts to recapture context from neighboring windows while reducing computational burden.
- Proposed approach: DAE-Former addresses these limitations with dual attention over the full spatial dimension and channel context, plus cross-attention in skip connections for multi-scale feature fusion.The proposed architecture is hierarchical, U-Net-like, and purely Transformer-based.
2. Proposed Method
DAE-Former is a convolution-free, U-Net-like hierarchical Transformer that combines efficient spatial attention with transpose-based channel attention. Its skip-connection cross-attention fuses encoder and decoder features to preserve information and recover fine-grained details.
- DAE-Former Architecture: DAE-Former uses overlapping 4 × 4 patch tokens, three encoder blocks, dual Transformer layers, and patch merging in a hierarchical encoder.The architecture is convolution-free and U-Net-like, with corresponding encoder and decoder pathways.
- Efficient Spatial Attention: Efficient attention captures spatial importance by reordering normalized query, key, and value operations instead of computing pairwise token similarities.This reduces the attention computation while preserving the stated representational behavior of dot-product attention.
- Transpose Channel Attention: Transpose attention captures channel context by basing attention weights on the cross-covariance between transposed keys and queries.Its temperature parameter improves training stability but reduces representational freedom.
- Computational Efficiency: Transpose attention scales as O(Nd^2/h), compared with O(N^2d) for standard self-attention, shifting the quadratic dependence from tokens to embedding dimension.The embedding dimension is usually smaller than the token count for larger images.
- Efficient Dual Attention: The dual attention block sequentially applies efficient attention and transpose attention, each followed by normalization and a feed-forward network.The block combines spatial and channel attention within the Transformer design.
- Skip Connection Cross Attention: SCCA cross-attends encoder features as queries with lower-decoder features as keys and values, preserving multi-level information for fine-grained mask recovery.The decoder feature is linearly projected to the encoder feature dimension before fusion.
3. Experimental Setup
DAE-Former is evaluated on Synapse multi-organ CT segmentation and ISIC2018 skin-lesion segmentation using an end-to-end training setup. It surpasses prior methods on reported segmentation results, with statistical and attention-map analyses providing additional evaluation evidence.
- Datasets and training: The model is trained end-to-end for 400 epochs on a single RTX 3090 GPU using cross-entropy and Dice losses.Training uses batch size 24, stochastic gradient descent, a base learning rate of 0.05, momentum 0.9, and weight decay 0.0001.
- Datasets and training: Evaluation uses the Synapse multi-organ dataset and the ISIC2018 skin-lesion segmentation dataset.Synapse contains 30 cases and 3779 axial abdominal clinical CT images; both evaluations follow previously published experimental settings.
- Quantitative and qualitative results: On Synapse, DAE-Former surpasses previous state-of-the-art methods in DSC and improves Dice by 0.67% over MISSFormer.The reported gains are especially observed for the gallbladder, kidney, liver, and spleen, while the model produces smooth maps for all organs, including small organs such as the pancreas.
- Quantitative and qualitative results: On ISIC2018, DAE-Former performs better than CNN- and Transformer-based methods in almost all reported metrics and exceeds TMU-Net.TMU-Net uses extra boundary and foreground-distribution information and has approximately 165.1M parameters versus 48.1M for DAE-Former.
- Statistical and attention analyses: Ten-run statistical analysis reports low variance for DAE-Former on several Synapse organs, while Grad-CAM visualizations assess attention across large and small organs.The statistical analysis covers comparisons with U-Net, TransUNet, Swin-Unet, HiFormer, and MISSFormer; the attention maps include liver, spleen, stomach, gallbladder, left kidney, and pancreas.
4. Conclusion
DAE-Former is a U-Net-like hierarchical pure Transformer that combines spatial and channel attention across the full feature dimension. It also uses skip-connection cross-attention to fuse multi-scale features and reports state-of-the-art results on Synapse and skin-lesion segmentation datasets.
- DAE-Former is a U-Net-like hierarchical pure Transformer using spatial and channel attention across the full feature dimension.
- The model retains the same parameter count as previous architectures while enriching the representational space through dual attention.
- Skip-connection cross-attention fuses multi-scale features, and the model reports state-of-the-art results on Synapse and skin-lesion segmentation datasets.
Appendix A. Attention in Detail
The appendix compares standard, efficient, and transpose attention and evaluates alternative ways to combine spatial and channel attention. Sequential fusion performs best while requiring the fewest parameters, so it is used throughout the proposed network.
- Attention mechanisms: Standard dot-product attention operates on tokens, whereas efficient attention forms a global context vector before combining it with queries.
- Attention mechanisms: Figure 6 compares standard dot-product, efficient, and transpose attention as alternative attention formulations.
- Dual-attention variants: The study evaluates sequential, simple additive, complex additive, and concatenation dual-attention structures.
- Dual-attention variants: Sequential fusion outperforms the other variants with both the lowest required parameter count and the best scores, so it is used in all network layers.
Appendix B. Input Resolution and Skip Connection Effect
The appendix examines skip-connection count and input resolution on Synapse. More skip connections improve performance, while higher image resolution also improves performance relative to lower resolution.
- Skip connections: Adding more skip connections improves performance, with two skip connections used as the base setting.
- Skip connections: Fine-grained information from higher layers is critical for fusing information through skip connections.
- Input resolution: The resolution study compares 128 × 128 and 288 × 288 inputs with the original 224 × 224 setting, and performance increases at higher resolution.
Appendix C. Further Visualization Results
Visual comparisons assess DAE-Former’s skin-lesion segmentation against other methods using ground-truth and predicted boundaries. The reported comparisons show strong performance across varied lesion shapes and patterns, particularly at lesion boundaries.
- Visual comparisons: Figure 8 compares visual skin-lesion segmentation results across methods on the ISIC2018 dataset.
- Visual comparisons: DAE-Former estimates lesion regions with varying shapes and patterns and reports better lesion-boundary predictions than the compared state-of-the-art approaches.
- Visual comparisons: Ground-truth lesion boundaries are shown in green, while predicted boundaries are shown in blue.
Appendix D. Computational Complexity
DAE-Former offers a favorable efficiency–performance trade-off, achieving stronger segmentation results than comparable approaches with fewer parameters.
- DAE-Former has a comparable parameter count to MISSFormer but slightly better performance.
- Figure 9 compares normalized performance and efficiency metrics, including DSC score, HD distance, and parameter count.HD distance and parameter count were reversed during normalization so higher values aligned with DSC.
- 48.1M parameters yielded a 4.95 higher Dice score and 14.23 lower HD distance than TransUNet's 105.28M parameters.The comparison reports substantially fewer parameters alongside better performance.