[论文解读] TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

Jieneng Chen, Yongyi Lu+7 more|arXiv (Cornell University)|Feb 8, 2021

Advanced Neural Network Applications参考文献 16被引用 3,794

一句话总结

TransUNet 将基于 CNN 的高分辨率特征与 Transformer 全局上下文相结合，在多数据集的医学图像分割中达到最前沿，高于纯 CNN 与纯 Transformer 基线。

ABSTRACT

Medical image segmentation is an essential prerequisite for developing healthcare systems, especially for disease diagnosis and treatment planning. On various medical image segmentation tasks, the u-shaped architecture, also known as U-Net, has become the de-facto standard and achieved tremendous success. However, due to the intrinsic locality of convolution operations, U-Net generally demonstrates limitations in explicitly modeling long-range dependency. Transformers, designed for sequence-to-sequence prediction, have emerged as alternative architectures with innate global self-attention mechanisms, but can result in limited localization abilities due to insufficient low-level details. In this paper, we propose TransUNet, which merits both Transformers and U-Net, as a strong alternative for medical image segmentation. On one hand, the Transformer encodes tokenized image patches from a convolution neural network (CNN) feature map as the input sequence for extracting global contexts. On the other hand, the decoder upsamples the encoded features which are then combined with the high-resolution CNN feature maps to enable precise localization. We argue that Transformers can serve as strong encoders for medical image segmentation tasks, with the combination of U-Net to enhance finer details by recovering localized spatial information. TransUNet achieves superior performances to various competing methods on different medical applications including multi-organ segmentation and cardiac segmentation. Code and models are available at https://github.com/Beckschen/TransUNet.

研究动机与目标

• 说明为何卷积神经网络（U-Net）在医学分割中难以处理长程依赖。
• 提出一个混合的 CNN-Transformer 编码器，以同时利用高分辨率细节与全局上下文。
• 设计一个级联上采样解码器并具备跳跃连接，以恢复细粒度的空间细节。
• 在多种医学成像任务上展示相对于基于 CNN 与基于 Transformer 的基线的经验收益。

提出的方法

• 将图像分块进行标记并用 Transformer 编码以捕捉全局上下文。
• 使用 CNN 特征图为 Transformer 嵌入提供高分辨率的补丁（混合编码器）。
• 用级联上采样器（CUP）对 Transformer 特征进行上采样，并通过类似 U-Net 的跳跃连接进行融合。
• 使用标准 SGD 在预训练骨干网络上训练；输入为 224x224，默认补丁大小为 16。
• 比较“None”（朴素上采样）与 CUP 解码器以及不同的编码器选项。
• 对跳跃连接、分辨率、补丁大小和模型规模进行消融实验。

实验结果

研究问题

• RQ1 在结合 CNN 细节的情况下，Transformers 能否成为医学图像分割的强编码器？
• RQ2 混合的 CNN-Transformer 编码器加上级联上采样解码器是否在医学分割任务中优于纯 Transformer 或纯 CNN 基线？
• RQ3 跳跃连接、输入分辨率、补丁大小和模型规模对分割质量的影响是什么？
• RQ4 TransUNet 在 CT 多器官分割和心脏 MRI 分割数据集上的泛化性能如何？

主要发现

• 在 Synapse 多器官 CT 上，使用 R50-ViT-CUP 基线时，TransUNet 实现了 77.48% 的平均 Dice 分数（DSC），并在 ACDC 心脏 MRI 数据集达到 89.71 DSC（见表 5）。
• 消融研究表明在多个 CUP 分辨率添加跳跃连接可以提高性能，最佳结果出现在 1/2、1/4 和 1/8 尺度进行跳跃连接时。
• 混合编码器（CNN + ViT）优于纯 ViT 和纯 CNN 基线，证明了将高分辨率 CNN 特征与全局 Transformer 上下文结合的好处。
• CUP 解码器显著优于朴素上采样，且更大模型尺寸在测试中带来更好性能（Base 与 Large）。
• 更高的输入分辨率（512x512）在计算成本增加的前提下将平均 DSC 提高约 6.88%；补丁大小为 16（序列长度 196）比更大补丁表现更好。
• 定性结果显示，与仅 CNN 和其他基于 Transformer 的模型相比，TransUNet 产生更少的假阳性并保留细小的器官边界。