BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens

Volumetric medical image segmentation typically relies on dense multi-scale feature grids, which are computationally expensive and memory-intensive due to the cubic growth of 3D spatial data. The authors investigate whether it is possible to achieve high-performance segmentation by dynamically allocating fine-resolution processing only to regions that require it—such as boundaries, thin structures, and small targets—while representing homogeneous regions more coarsely.

BATS (Boundary-Aware Token Selection) introduces a three-part architecture:

1. Dense Boundary Predictor: A CNN that evaluates potential boundary regions across multiple scales in parallel to identify where high resolution is necessary.
2. Fine-First Context Cascade: A mechanism that constructs a sparse, mixed-resolution hierarchy based on boundary predictions. It ensures that fine-scale evidence is retained even if coarse-scale predictions are negative, preventing the loss of small structures.
3. Parent Cluster Attention: A cross-scale attention mechanism that injects hierarchical ancestor tokens into local attention neighborhoods. This allows tokens at different resolutions to exchange information without requiring dense multi-scale feature maps or expensive nearest-neighbor searches.

BATS was evaluated against the nnU-Net Revisited protocol on five public CT and MRI datasets (ACDC, LiTS, BraTS, KiTS, and AMOS). The model achieved the highest Dice score on the LiTS dataset and performed within 0.37 Dice points of the strong MedNeXt-L baseline on average. Notably, BATS reduced peak GPU memory allocation by over 53% on KiTS, LiTS, and BraTS compared to MedNeXt-L. While inference speed improved on datasets with sparse structures, the model's runtime performance varies based on the density of the required token hierarchy.