# TokenSeg: Efficient 3D Medical Image Segmentation via Hierarchical Visual Token Compression

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 39 references that are in the map.
- First posted on arXiv 8 Jan 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Feature Pyramid Networks for Object Detection](https://mlsanity.com/paper/1612.03144/) (2016)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [CBAM: Convolutional Block Attention Module](https://mlsanity.com/paper/1807.06521/) (2018)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [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)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [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)
- [3D Semantic Segmentation with Submanifold Sparse Convolutional Networks](https://mlsanity.com/paper/1711.10275/) (2017)
- [Vector-quantized Image Modeling with Improved VQGAN](https://mlsanity.com/paper/2110.04627/) (2021)
- [The Medical Segmentation Decathlon](https://mlsanity.com/paper/2106.05735/) (2021)
- [Holistically-Nested Edge Detection](https://mlsanity.com/paper/1504.06375/) (2015)

## Built on by (most cited first, 3 of 3)

- [BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens](https://mlsanity.com/paper/2607.26829/) (2026)
- [SAMRI-3D: Adapting SAM2 for 3D MRI Segmentation with Global Volume Tokens](https://mlsanity.com/paper/2607.18014/) (2026)
- [Adaptive Computation with Conformal Certification](https://mlsanity.com/paper/2610.19189/) (2026)

## Abstract (arXiv)

Three-dimensional medical image segmentation is a fundamental yet computationally demanding task due to the cubic growth of voxel processing and the redundant computation on homogeneous regions. To address these limitations, we propose TokenSeg, a boundary-aware sparse token representation framework for efficient 3D medical volume segmentation. Specifically, (1) we design a multi-scale hierarchical encoder that extracts 400 candidate tokens across four resolution levels to capture both global anatomical context and fine boundary details; (2) we introduce a boundary-aware tokenizer that combines VQ-VAE quantization with importance scoring to select 100 salient tokens, over 60% of which lie near tumor boundaries; and (3) we develop a sparse-to-dense decoder that reconstructs full-resolution masks through token reprojection, progressive upsampling, and skip connections. Extensive experiments on a 3D breast DCE-MRI dataset comprising 960 cases demonstrate that TokenSeg achieves state-of-the-art performance with 94.49% Dice and 89.61% IoU, while reducing GPU memory and inference latency by 64% and 68%, respectively. To verify the generalization capability, our evaluations on MSD cardiac and brain MRI benchmark datasets demonstrate that TokenSeg consistently delivers optimal performance across heterogeneous anatomical structures. These results highlight the effectiveness of anatomically informed sparse representation for accurate and efficient 3D medical image segmentation.