TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging...

Area of Science:

• Medical Imaging
• Artificial Intelligence
• Computer Vision

Background:

• Transformers offer powerful global interactions in medical imaging but face computational challenges.
• Existing models like SwinUNETR integrate Transformers into UNets by tokenizing inputs, leading to quadratic scaling issues with attention mechanisms.
• The cubic scaling of 3D input resolution exacerbates computational demands, limiting deployment on standard hardware.

Purpose of the Study:

• To introduce TokenUNet, a family of 3D segmentation models designed for computationally constrained environments.
• To mitigate computational challenges associated with Transformers in 3D medical imaging.
• To develop a more efficient alternative to existing Transformer-based UNet models.

Main Methods:

• TokenUNet incorporates TokenLearner and TokenFuser modules within a UNet architecture.
• The convolutional encoder is maintained, with TokenLearner applied to 3D feature maps to pool tokens, decoupling token number from input size.
• The model was evaluated on the BraTS challenge dataset for glioma segmentation.

Main Results:

• Tokenization effectively encodes task-relevant information and produces interpretable attention maps.
• TokenUNet achieved a memory footprint of 38%, inference computation time of 10%, and parameter count of 17% compared to SwinUNETR.
• Statistically equivalent Dice score performance was observed with nnunetv2 5-fold cross-validation.

Conclusions:

• TokenUNet offers a computationally efficient solution for 3D medical image segmentation, suitable for limited hardware.
• The approach facilitates efficient model optimization, fine-tuning, and transfer learning in resource-constrained settings.
• This work accelerates the development and diversification of AI approaches in 3D medical imaging research.