# Token-UNet: A New Case for Transformers Integration in Efficient and Interpretable 3D UNets for Brain Imaging Segmentation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 46 references that are in the map.
- First posted on arXiv 23 Feb 2026.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [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)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [Sanity Checks for Saliency Maps](https://mlsanity.com/paper/1810.03292/) (2018)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [nnU-Net Revisited: A Call for Rigorous Validation in 3D Medical Image Segmentation](https://mlsanity.com/paper/2404.09556/) (2024)
- [U-Net and its variants for medical image segmentation: theory and applications](https://mlsanity.com/paper/2011.01118/) (2020)
- [On The Computational Complexity of Self-Attention](https://mlsanity.com/paper/2209.04881/) (2022)
- [The Information Bottleneck Problem and Its Applications in Machine Learning](https://mlsanity.com/paper/2004.14941/) (2020)
- [How we won BraTS 2023 Adult Glioma challenge? Just faking it! Enhanced Synthetic Data Augmentation and Model Ensemble for brain tumour segmentation](https://mlsanity.com/paper/2402.17317/) (2024)
- [Nesterov's Accelerated Gradient and Momentum as approximations to Regularised Update Descent](https://mlsanity.com/paper/1607.01981/) (2016)
- [A Recent Survey of Vision Transformers for Medical Image Segmentation](https://mlsanity.com/paper/2312.00634/) (2023)
- [Vision Transformers in Medical Computer Vision -- A Contemplative Retrospection](https://mlsanity.com/paper/2203.15269/) (2022)
- [Improving Pre-trained Adult Glioma Segmentation Models Using only Post-processing Techniques](https://mlsanity.com/paper/2512.14937/) (2025)

## Abstract (arXiv)

We present Token-UNet, adopting the TokenLearner and TokenFuser modules to encase Transformers into UNets. While Transformers have enabled global interactions among input elements in medical imaging, current computational challenges hinder their deployment on common hardware. Models like (Swin)UNETR adapt the UNet architecture by incorporating (Swin)Transformer encoders, which process tokens that each represent small subvolumes (83 voxels) of the input. The Transformer attention mechanism scales quadratically with the number of tokens, which is tied to the cubic scaling of 3D input resolution. This work reconsiders the role of convolution and attention, introducing Token-UNets, a family of 3D segmentation models that can operate in constrained computational environments and time frames. To mitigate computational demands, our approach maintains the convolutional encoder of UNet-like models, and applies TokenLearner to 3D feature maps. This module pools a preset number of tokens from local and global structures. Our results show this tokenization effectively encodes task-relevant information, yielding naturally interpretable attention maps. The memory footprint, computation times at inference, and parameter counts of our heaviest model are reduced to 33%, 10%, and 35% of the SwinUNETR values, with better average performance (86.75% ± 0.19% Dice score for SwinUNETR vs our 87.21% ± 0.35%). This work opens the way to more efficient trainings in contexts with limited computational resources, such as 3D medical imaging. Easing model optimization, fine-tuning, and transfer-learning in limited hardware settings can accelerate and diversify the development of approaches, for the benefit of the research community.