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

TL;DR

TokenSeg tackles the high computational cost of 3D medical image segmentation by introducing a boundary-aware sparse token framework. It combines a four-level hierarchical encoder to generate 400 candidate tokens, a boundary-aware VQ-VAE-based tokenizer to keep 100 boundary-focused tokens, and a sparse-to-dense decoder that reprojects anchors and progressively reconstructs full-resolution masks with cross-level fusion. The approach achieves state-of-the-art Dice and IoU on a large private breast DCE-MRI dataset ($94.49\%$ Dice, $89.61\%$ IoU) and demonstrates strong cross-dataset generalization to MSD brain and cardiac MRI, while delivering major efficiency gains (e.g., $\sim64\%$ memory reduction, $\sim68\%$ latency reduction). The results validate that anatomically informed sparse representations can deliver accurate and efficient 3D segmentation, with significant implications for real-time clinical deployment and multi-center studies.

Abstract

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 \textbf{TokenSeg}, a boundary-aware sparse token representation framework for efficient 3D medical volume segmentation. Specifically, (1) we design a \emph{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 \emph{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 \emph{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.

Figures (6)

• Figure 1: Performance comparison between TokenSeg and state-of-the-art methods on 3D medical volume segmentation. TokenSeg-100 achieves 94.49% Dice score, outperforming the best baseline (nnU-Net) by +4.29% Dice while being 5.3× faster (48ms vs 256ms) with 54.5% fewer parameters (23.8M vs 52.3M).
• Figure 2: The architecture of TokenSeg for the DCE-MRI Breast Cancer segmentation
• Figure 3: Qualitative visualizations of segmentation results on DCE-MRI and MSD b33. The results presented from rows one to three correspond,in order, to breast tumors, brain tumors, and cardiac tumors. We present the visualizations on other datasets in the supplemental material.
• Figure 4: Qualitative comparison on a representative breast DCE-MRI slice. From left to right: input scan with ROI, zoomed view, predictions from 3D U-Net b8, V-Net b9, nnU-Net b10, Swin UNETR b12, our TokenSeg, and the ground truth.
• Figure SM5: Qualitative comparison of different segmentation models across three datasets: (a) Pancreas, (b) Hepatic Vessel, and (c) Lung. In each subfigure, the columns from left to right display: the input CT scan with ROI showing the target organ, a zoomed view of the ROI, predictions from 3D U-Net, V-Net, nnU-Net, Swin UNETR, our TokenSeg, and the ground truth segmentation.