[论文解读] A Volumetric Transformer for Accurate 3D Tumor Segmentation

Himashi Peiris, Munawar Hayat+3 more|arXiv (Cornell University)|Nov 26, 2021

Advanced Neural Network Applications参考文献 64被引用 25

一句话总结

本文提出VT-UNet，一种计算高效的3D体素Transformer，采用U形编码器-解码器架构，通过编码器中的双自注意力层捕捉局部与全局空间线索，并在解码器中引入基于并行移位窗口的自注意力与交叉注意力机制，结合傅里叶位置编码以优化肿瘤边界。该模型在BraTS 2021、胰腺和肝脏数据集上均达到最先进性能，展现出强大的泛化能力与对数据损坏的鲁棒性。

ABSTRACT

This paper presents a Transformer architecture for volumetric medical image segmentation. Designing a computationally efficient Transformer architecture for volumetric segmentation is a challenging task. It requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volumetric data. The proposed volumetric Transformer has a U-shaped encoder-decoder design that processes the input voxels in their entirety. Our encoder has two consecutive self-attention layers to simultaneously encode local and global cues, and our decoder has novel parallel shifted window based self and cross attention blocks to capture fine details for boundary refinement by subsuming Fourier position encoding. Our proposed design choices result in a computationally efficient architecture, which demonstrates promising results on Brain Tumor Segmentation (BraTS) 2021, and Medical Segmentation Decathlon (Pancreas and Liver) datasets for tumor segmentation. We further show that the representations learned by our model transfer better across-datasets and are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

研究动机与目标

• 为解决设计一种计算高效的3D体素医学图像分割Transformer的挑战。
• 在体素数据的所有三个维度上平衡局部与全局空间依赖关系的编码。
• 通过新颖的注意力机制改进肿瘤分割中的边界细化。
• 提升表示在不同医学影像数据集间的可迁移性与对数据损坏的鲁棒性。

提出的方法

• 该模型采用U形编码器-解码器架构，直接处理完整的3D体素输入，无需分块处理。
• 编码器使用两个连续的自注意力层，同时编码局部与全局空间上下文。
• 解码器引入并行的移位窗口自注意力与交叉注意力模块，以增强细粒度细节学习能力。
• 集成傅里叶位置编码，以在不增加计算成本的前提下改善位置建模。
• 架构设计旨在保持效率的同时，完整保留所有轴向的空间信息。
• 模型在BraTS 2021、胰腺和肝脏数据集上使用标准分割损失函数进行端到端训练。

实验结果

研究问题

• RQ1 基于Transformer的架构是否能在保持计算高效的同时，实现高精度的3D肿瘤分割？
• RQ2 该模型在体素医学图像中捕捉局部与全局空间依赖关系的效率如何？
• RQ3 所提出的注意力机制在肿瘤分割中对边界勾勒的改进程度如何？
• RQ4 所学习的表征在不同医学影像数据集之间以及面对数据损坏时的泛化能力如何？
• RQ5 傅里叶位置编码的引入是否在不增加计算开销的前提下提升了性能？

主要发现

• 所提出的VT-UNet在BraTS 2021数据集上实现了脑肿瘤分割的最先进性能。
• 该模型在胰腺和肝脏肿瘤分割任务中，于Medical Segmentation Decathlon数据集上表现出更强的泛化能力。
• 与基线模型相比，所学习的表征在不同数据集间展现出更优的可迁移性。
• 该模型对各类数据损坏表现出鲁棒性，表明其具备可靠的特征学习能力。
• 采用并行移位窗口注意力与傅里叶位置编码可在几乎不增加计算成本的前提下显著提升边界细化效果。
• 该架构在端到端处理完整3D体素的同时，保持了高效的计算性能。