将Transformer融入3D脑影像分割模型,显著降低计算开销且保持高精度。
Token-UNet: A New Case for Transformers Integration in Efficient and Interpretable 3D UNets for Brain Imaging Segmentation
- 用TokenLearner和TokenFuser模块在UNet中轻量集成Transformer,保留卷积编码器
- 在3D医学图像上实现87.21%的Dice分数,推理时间仅为SwinUNETR的10%
- 适合算力受限场景,支持快速训练与模型迁移,利于医疗研究落地
我们提出Token-UNet,通过引入TokenLearner与TokenFuser模块,将Transformer高效融入UNet架构。尽管Transformer在医学图像中可实现全局交互,但其计算开销大,难以在普通硬件部署。现有模型如(Swin)UNETR采用(Swin)Transformer编码器处理每块$8^3$体素的输入片段,导致注意力机制随令牌数呈二次增长,与3D输入分辨率的立方增长相关联。本文重新审视卷积与注意力的作用,设计出可在资源受限环境运行的Token-UNets系列3D分割模型。为缓解计算压力,保留UNet类模型的卷积编码器,并对3D特征图应用TokenLearner模块,从局部与全局结构中池化固定数量的令牌。实验表明,该方式有效编码任务相关特征,生成自然可解释的注意力图。最重模型的内存占用、推理时间与参数量分别降至SwinUNETR的33%、10%和35%,且平均性能更优(SwinUNETR: 86.75% ± 0.19%,本方法: 87.21% ± 0.35%)。该工作为算力有限场景下的高效训练铺平道路,助力3D医学影像领域的快速迭代与广泛应用。
原文摘要 · Abstract (English)
We present Token-UNet, adopting the TokenLearner and TokenFuser modules to encase Transformers into UNets. While Transformers have enabled global interactions among input elements in medical imaging, current computational challenges hinder their deployment on common hardware. Models like (Swin)UNETR adapt the UNet architecture by incorporating (Swin)Transformer encoders, which process tokens that each represent small subvolumes ($8^3$ voxels) of the input. The Transformer attention mechanism scales quadratically with the number of tokens, which is tied to the cubic scaling of 3D input resolution. This work reconsiders the role of convolution and attention, introducing Token-UNets, a family of 3D segmentation models that can operate in constrained computational environments and time frames. To mitigate computational demands, our approach maintains the convolutional encoder of UNet-like models, and applies TokenLearner to 3D feature maps. This module pools a preset number of tokens from local and global structures. Our results show this tokenization effectively encodes task-relevant information, yielding naturally interpretable attention maps. The memory footprint, computation times at inference, and parameter counts of our heaviest model are reduced to 33\%, 10\%, and 35\% of the SwinUNETR values, with better average performance (86.75\% $\pm 0.19\%$ Dice score for SwinUNETR vs our 87.21\% $\pm 0.35\%$). This work opens the way to more efficient trainings in contexts with limited computational resources, such as 3D medical imaging. Easing model optimization, fine-tuning, and transfer-learning in limited hardware settings can accelerate and diversify the development of approaches, for the benefit of the research community.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。