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

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

Medical Image Segmentation Techniques被引用 9

一句话总结

该论文提出VT-UNet，一种用于3D肿瘤分割的纯体素Transformer架构，无需切片即可处理完整的3D MRI体积，编码器采用分层自注意力机制，解码器采用并行自注意力与交叉注意力机制，并结合傅里叶位置编码。该方法在BraTS数据集上实现了最先进性能，同时对MRI伪影具有更强的鲁棒性，且计算复杂度较低。

ABSTRACT

We propose a Transformer architecture for volumetric segmentation, a challenging task that requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volume. Encoder of the proposed design benefits from self-attention mechanism to simultaneously encode local and global cues, while the decoder employs a parallel self and cross attention formulation to capture fine details for boundary refinement. Empirically, we show that the proposed design choices result in a computationally efficient model, with competitive and promising results on the Medical Segmentation Decathlon (MSD) brain tumor segmentation (BraTS) Task. We further show that the representations learned by our model are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

研究动机与目标

• 为解决在3D医学图像分割中避免基于2D切片处理而保留完整体素空间上下文的挑战。
• 设计一种纯Transformer-based的UNet架构，在保持计算效率的同时捕捉所有3D轴向的长距离依赖关系。
• 通过解码器中并行注意力机制增强特征精炼，改善肿瘤分割的边界勾勒。
• 提升模型对常见MRI伪影（如运动伪影、鬼影伪影和脉冲伪影）的鲁棒性。
• 证明3D自注意力机制在体素肿瘤分割任务中可实现比CNN-based和混合模型更高的准确率与鲁棒性。

提出的方法

• 编码器采用分层、移位窗口方式的连续两层自注意力机制，联合捕捉3D体素中的局部与全局上下文线索。
• 解码器采用并行的自注意力与交叉注意力机制，并共享查询投影，以在上采样过程中保持全局上下文信息。
• 将傅里叶位置编码注入解码过程，以增强空间归纳偏置并改善特征表示。
• 采用凸组合策略融合并行注意力头的特征，提升表示质量。
• 使用AdamW优化器进行训练，初始峰值学习率为1e-4，结合数据增强（旋转、噪声、模糊、伽马变换），并使用ImageNet-22K预训练的Swin-T权重进行模型初始化。
• 在MSD BraTS数据集的完整3D MRI体积上评估该架构，未进行任何2D切片分解。

实验结果

研究问题

• RQ1 纯体素Transformer架构是否能在保持低计算成本的同时，优于基于CNN和混合模型的3D肿瘤分割方法？
• RQ2 解码器中并行的自注意力与交叉注意力机制在多大程度上提升了3D分割的边界准确率与特征精炼能力？
• RQ3 傅里叶位置编码在3D Transformer-based分割网络中在多大程度上增强了空间建模能力？
• RQ4 所提出的架构是否在运动、鬼影和脉冲伪影等常见MRI伪影下表现出更强的鲁棒性？
• RQ5 完全基于注意力机制的3D UNet设计是否能在不依赖卷积归纳偏置的情况下实现SOTA性能？

主要发现

• VT-UNet-B在Whole Tumor（WT）类别上取得了91.9的最高Dice相似系数（DSC），优于BraTS数据集上的所有SOTA方法。
• 该模型在WT类别上的Hausdorff距离（HD95）为3.43，是所有对比方法中最低的，表明边界准确率更优。
• VT-UNet对MRI伪影表现出卓越的鲁棒性：在脉冲伪影上DSC达到86.6，在运动伪影上达到85.8，优于nnFormer在所有伪影类型下的表现。
• 消融实验证实，傅里叶位置编码、凸组合策略与并行注意力机制的组合显著提升性能，当移除所有组件时DSC下降1.78点。
• VT-UNet-S与VT-UNet-B的FLOPs分别为3.84和3.43 GFLOPs，低于大多数SOTA模型（包括nnFormer的4.05 GFLOPs），同时保持更高准确率。
• 定性结果表明，VT-UNet能生成更清晰的肿瘤边界，并更准确地分割复杂结构（如瘤周水肿和增强肿瘤）