BATS通过边界感知的混合分辨率令牌,实现高效3D医学图像分割。
BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens

- 根据边界需求动态分配高分辨率处理,非均匀聚焦于关键区域。
- 在5个数据集上平均精度接近最强基线,内存占用减少超53%。
- 适合资源受限场景,尤其在边界稀疏的CT/MRI数据中表现优异。
许多高性能体素分割模型依赖密集的多尺度特征图,导致激活内存和推理成本高。本文提出BATS(边界感知令牌选择),一种3D医学图像分割架构,将细粒度处理集中在预测类别边界附近。密集边界预测器识别需更高分辨率的位置,细粒度优先的上下文级联构建输入相关的混合分辨率层次结构。同质区域以粗粒度表示,边界、细结构和小目标处保留更细粒度令牌。稀疏层次结构经细化并栅格化为密集分割结果。BATS在每个分辨率层级独立预测边界相关性,防止粗尺度错误决策抑制细尺度证据。父簇注意力将层级祖先令牌注入局部注意力邻域,提供跨尺度上下文,无需密集多尺度特征图或跨尺度邻居搜索。我们在五个公开的CT和MRI数据集上,采用标准化的nnU-Net Revisited协议评估BATS。BATS在LiTS上取得最高Dice分数,且在五个数据集上平均精度与最强基线MedNeXt-L相差不超过0.37个Dice点。相比MedNeXt-L,BATS在KiTS、LiTS和BraTS上峰值GPU内存降低超过53%。在KiTS和LiTS上推理速度最快提升30%,但因令牌密度更高,在BraTS上稍慢。混合分辨率处理带来一致的内存节省,而运行时间和精度增益取决于数据集边界密度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。