Token-UNet：用于脑成像分割的高效且可解释的 3D UNet 变换器集成的新方案
计算机视觉与模式识别
2026-02-24 v1
摘要
我们提出了 Token-UNet，通过引入 TokenLearner 和 TokenFuser 模块将变换器嵌入 UNet。尽管变换器在医学成像中实现了输入元素之间的全局交互，但当前的计算挑战阻碍了其在常规硬件上的部署。类如 (Swin)UNETR 的模型通过采用 (Swin)Transformer 编码器来改进 UNet 架构，这些编码器处理代表输入小体素子体积 ( 体素) 的 token。变换器注意力机制随 token 数量的平方增长，而这种数量与 3D 输入分辨率的立方增长相关。本 work 重新考虑卷积与注意力的作用，提出 Token-UNet 系列 3D 分割模型，可在受限的计算环境和时间限制下运行。为缓解计算需求，我们的方法保留了 UNet 类模型的卷积编码器，并对 3D 特征图应用 TokenLearner。该模块从局部和全局结构中抽取预设数量的 token。我们的结果表明，这种 tokenization 有效地编码了任务相关信息，生成了自然可解释的注意力图。我们最重型模型的内存占用、推理计算时间和参数数量分别降低至 SwinUNETR 的 33%、10% 和 35%，平均性能也更佳（SwinUNETR 为 86.75% Dice 评分，我方为 87.21% ）。本 work 为资源受限环境（如 3D 医学成像）中的更高效训练开辟了道路。简化模型优化、微调和迁移学习在受限硬件设置下的应用，可加速并丰富方法的发展，为研究社区的福祉。
引用
@article{arxiv.2602.20008,
title = {Token-UNet: A New Case for Transformers Integration in Efficient and Interpretable 3D UNets for Brain Imaging Segmentation},
author = {Louis Fabrice Tshimanga and Andrea Zanola and Federico Del Pup and Manfredo Atzori},
journal= {arXiv preprint arXiv:2602.20008},
year = {2026}
}