Search papers, labs, and topics across Lattice.
This paper introduces BATS, a novel 3D medical image segmentation architecture that optimizes resource efficiency by employing a boundary-aware mixed-resolution token strategy. By focusing fine-resolution processing on predicted class boundaries and utilizing a dense boundary predictor, BATS significantly reduces peak GPU memory usage by over 53% compared to the leading dense baseline, MedNeXt-L, while achieving competitive segmentation accuracy across multiple datasets. The method's ability to adaptively refine segmentation based on boundary density allows for faster inference times and substantial memory savings without sacrificing performance.
BATS achieves a remarkable 53% reduction in GPU memory usage while maintaining competitive segmentation accuracy, revolutionizing resource-efficient volumetric segmentation.
Many high-performing volumetric segmentation models maintain dense multi-scale feature maps, leading to high activation memory and inference cost. We present BATS (Boundary-Aware Token Selection), a 3D medical image segmentation architecture that concentrates fine-resolution processing near predicted class boundaries. A dense boundary predictor identifies where additional resolution is needed, while a fine-first context cascade constructs an input-dependent mixed-resolution hierarchy. Homogeneous regions are represented coarsely, with finer tokens retained around boundaries, thin structures, and small targets. The sparse hierarchy is refined and rasterised into a dense segmentation. BATS predicts boundary relevance independently at every resolution level, preventing an erroneous coarse-scale decision from suppressing fine-scale evidence. Parent cluster attention further injects hierarchical ancestor tokens into local attention neighbourhoods, providing cross-scale context without dense multi-scale feature maps or cross-scale neighbour search. We evaluate BATS on five public CT and MRI datasets using the standardised nnU-Net Revisited protocol. BATS achieves the highest LiTS Dice among the compared methods and averages within 0.37 Dice points of the strongest dense baseline, MedNeXt-L, across the five datasets. Relative to MedNeXt-L, it reduces peak allocated GPU memory by more than 53% on KiTS, LiTS, and BraTS. Inference is up to 30% faster on KiTS and LiTS, which retain fewer tokens, but slower on the more token-dense BraTS. Mixed-resolution processing therefore provides consistent memory savings, while runtime and accuracy gains depend on dataset boundary density.
