# AFTer-UNet: Axial Fusion Transformer UNet for Medical Image Segmentation

arXiv: https://arxiv.org/abs/2110.10403 · HTML: https://mlsanity.com/paper/2110.10403/

## Key facts

- Cited by 48 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 56 references that are in the map.
- First posted on arXiv 20 Oct 2021.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://mlsanity.com/paper/1506.01497/) (2015)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Is Space-Time Attention All You Need for Video Understanding?](https://mlsanity.com/paper/2102.05095/) (2021)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [UNet++: A Nested U-Net Architecture for Medical Image Segmentation](https://mlsanity.com/paper/1807.10165/) (2018)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [VoxelMorph: A Learning Framework for Deformable Medical Image Registration](https://mlsanity.com/paper/1809.05231/) (2018)
- [A large annotated medical image dataset for the development and evaluation of segmentation algorithms](https://mlsanity.com/paper/1902.09063/) (2019)
- [UNet 3+: A Full-Scale Connected UNet for Medical Image Segmentation](https://mlsanity.com/paper/2004.08790/) (2020)
- [ResUNet-a: a deep learning framework for semantic segmentation of remotely sensed data](https://mlsanity.com/paper/1904.00592/) (2019)
- [CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.03024/) (2021)
- [Recurrent Mask Refinement for Few-Shot Medical Image Segmentation](https://mlsanity.com/paper/2108.00622/) (2021)
- [Evaluate the Malignancy of Pulmonary Nodules Using the 3D Deep Leaky Noisy-or Network](https://mlsanity.com/paper/1711.08324/) (2017)
- [Momentum Contrastive Voxel-wise Representation Learning for Semi-supervised Volumetric Medical Image Segmentation](https://mlsanity.com/paper/2105.07059/) (2021)
- [Closing the Gap between Deep and Conventional Image Registration using Probabilistic Dense Displacement Networks](https://mlsanity.com/paper/1907.10931/) (2019)
- [Predicting Camera Viewpoint Improves Cross-dataset Generalization for 3D Human Pose Estimation](https://mlsanity.com/paper/2004.03143/) (2020)
- [TransFusion: Cross-view Fusion with Transformer for 3D Human Pose Estimation](https://mlsanity.com/paper/2110.09554/) (2021)
- [Unsupervised Wasserstein Distance Guided Domain Adaptation for 3D Multi-Domain Liver Segmentation](https://mlsanity.com/paper/2009.02831/) (2020)

## Built on by (most cited first, 30 of 48)

- [Transformers in Medical Imaging: A Survey](https://mlsanity.com/paper/2201.09873/) (2022)
- [Medical Image Segmentation Review: The success of U-Net](https://mlsanity.com/paper/2211.14830/) (2022)
- [STAR-Transformer: A Spatio-temporal Cross Attention Transformer for Human Action Recognition](https://mlsanity.com/paper/2210.07503/) (2022)
- [Transformers in Medical Image Analysis: A Review](https://mlsanity.com/paper/2202.12165/) (2022)
- [Topology-Preserving Shape Reconstruction and Registration via Neural Diffeomorphic Flow](https://mlsanity.com/paper/2203.08652/) (2022)
- [MedGen3D: A Deep Generative Framework for Paired 3D Image and Mask Generation](https://mlsanity.com/paper/2304.04106/) (2023)
- [Transforming medical imaging with Transformers? A comparative review of key properties, current progresses, and future perspectives](https://mlsanity.com/paper/2206.01136/) (2022)
- [Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation](https://mlsanity.com/paper/2402.10887/) (2024)
- [PPT: token-Pruned Pose Transformer for monocular and multi-view human pose estimation](https://mlsanity.com/paper/2209.08194/) (2022)
- [CVTHead: One-shot Controllable Head Avatar with Vertex-feature Transformer](https://mlsanity.com/paper/2311.06443/) (2023)
- [A Recent Survey of Vision Transformers for Medical Image Segmentation](https://mlsanity.com/paper/2312.00634/) (2023)
- [Alternate Diverse Teaching for Semi-supervised Medical Image Segmentation](https://mlsanity.com/paper/2311.17325/) (2023)
- [3D Vision with Transformers: A Survey](https://mlsanity.com/paper/2208.04309/) (2022)
- [Diffeomorphic Image Registration with Neural Velocity Field](https://mlsanity.com/paper/2202.12498/) (2022)
- [Transformers in Healthcare: A Survey](https://mlsanity.com/paper/2307.00067/) (2023)
- [A survey on attention mechanisms for medical applications: are we moving towards better algorithms?](https://mlsanity.com/paper/2204.12406/) (2022)
- [MedSegFactory: Text-Guided Generation of Medical Image-Mask Pairs](https://mlsanity.com/paper/2504.06897/) (2025)
- [SPPNet: A Single-Point Prompt Network for Nuclei Image Segmentation](https://mlsanity.com/paper/2308.12231/) (2023)
- [Identity-Aware Hand Mesh Estimation and Personalization from RGB Images](https://mlsanity.com/paper/2209.10840/) (2022)
- [Recent Progress in Transformer-based Medical Image Analysis](https://mlsanity.com/paper/2208.06643/) (2022)
- [Attention Mechanisms in Medical Image Segmentation: A Survey](https://mlsanity.com/paper/2305.17937/) (2023)
- [TBConvL-Net: A Hybrid Deep Learning Architecture for Robust Medical Image Segmentation](https://mlsanity.com/paper/2409.03367/) (2024)
- [Transformer based Models for Unsupervised Anomaly Segmentation in Brain MR Images](https://mlsanity.com/paper/2207.02059/) (2022)
- [CSAM: A 2.5D Cross-Slice Attention Module for Anisotropic Volumetric Medical Image Segmentation](https://mlsanity.com/paper/2311.04942/) (2023)
- [A Refreshed Similarity-based Upsampler for Direct High-Ratio Feature Upsampling](https://mlsanity.com/paper/2407.02283/) (2024)
- [Transformers Meet Visual Learning Understanding: A Comprehensive Review](https://mlsanity.com/paper/2203.12944/) (2022)
- [Understanding the Tricks of Deep Learning in Medical Image Segmentation: Challenges and Future Directions](https://mlsanity.com/paper/2209.10307/) (2022)
- [MCANet: Medical Image Segmentation with Multi-Scale Cross-Axis Attention](https://mlsanity.com/paper/2312.08866/) (2023)
- [Advancing 3D Medical Image Analysis with Variable Dimension Transform based Supervised 3D Pre-training](https://mlsanity.com/paper/2201.01426/) (2022)
- [ConvTransSeg: A Multi-resolution Convolution-Transformer Network for Medical Image Segmentation](https://mlsanity.com/paper/2210.07072/) (2022)

## Abstract (arXiv)

Recent advances in transformer-based models have drawn attention to exploring these techniques in medical image segmentation, especially in conjunction with the U-Net model (or its variants), which has shown great success in medical image segmentation, under both 2D and 3D settings. Current 2D based methods either directly replace convolutional layers with pure transformers or consider a transformer as an additional intermediate encoder between the encoder and decoder of U-Net. However, these approaches only consider the attention encoding within one single slice and do not utilize the axial-axis information naturally provided by a 3D volume. In the 3D setting, convolution on volumetric data and transformers both consume large GPU memory. One has to either downsample the image or use cropped local patches to reduce GPU memory usage, which limits its performance. In this paper, we propose Axial Fusion Transformer UNet (AFTer-UNet), which takes both advantages of convolutional layers' capability of extracting detailed features and transformers' strength on long sequence modeling. It considers both intra-slice and inter-slice long-range cues to guide the segmentation. Meanwhile, it has fewer parameters and takes less GPU memory to train than the previous transformer-based models. Extensive experiments on three multi-organ segmentation datasets demonstrate that our method outperforms current state-of-the-art methods.