【医学图像处理】用于精确生物医学图像分割的复合 Transformer

标题：Multi-compound Transformer for Accurate Biomedical Image Segmentation

作者：Yuanfeng Ji，香港大学；Ping Luo，商汤科技

来源：MICCAI 2021

代码：https://github.com/JiYuanFeng/MCTrans MCTrans-master.zip

主题：Transformer；注意力机制；医学图像分割

1. 引言

❓ 论文想要解决的问题是什么？

由于卷积核函数的局部性，传统的基于 CNN 的分割模型缺乏对长期依赖关系的建模能力。

Due to the local property of the convolutional kernels, the traditional CNN-based segmentation models (e.g. FCN) lack the ability for modeling long-term dependencies.

❓ 解决上述问题的现有方法有哪些？

为了解决上述问题，一些方法已经被用于进行强大的关系建模：

• 基于空间金字塔的方法

  • 基于空间金字塔的方法采用不同大小的卷积核将来自不同范围的上下文信息聚合到单个层中。
  • Zhao, Hengshuang, et al. “Pyramid scene parsing network.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.

• 基于 UNet 的编码器-解码器网络

  • 基于 UNet 的编码器-解码器网络通过应用跳跃连接将具有相同尺度的粗粒度深层特征和细粒度浅层特征合并。

• TransUNet

  • Chen, Jieneng, et al. “Transunet: Transformers make strong encoders for medical image segmentation.” arXiv preprint arXiv:2102.04306 (2021).
  • TransUNet 采用自注意力机制在最高级别的 CNN 特征上计算全局上下文，确保特定尺度下的各种范围依赖。
  • 
  • 然而，这样的设计对于医学图像分割仍然不是最佳的。首先，它只在单一尺度上使用自注意力机制进行上下文建模，而忽略了跨尺度的依赖性和一致性。后者通常在具有显著尺寸变化的病灶的分割中起关键作用。其次，它没有考虑到如何学习不同语义类别之间的相关性以及如何保证同一类别区域的特征的一致性。它们都是基于 CNN 的分割方案设计的关键。

    However, such a design is still sub-optimal for medical image segmentation for the following reasons. First, it only uses the self-attention mechanism for context modeling on a single scale but ignores the cross-scale dependency and consistency. The latter usually plays a critical role in the segmentation of lesions with dramatic size changes. Second, beyond the context modeling, how to learn the correlation between different semantic categories and how to ensure the feature consistency of the same category region are still not taken into account. But both of them have become critical for CNN-based segmentation scheme design.

❓ 论文的主要创新点是什么？

1. 本文提出了 Multi-Compound Transformer (MCTrans) 网络，它构建了跨尺度的上下文依赖关系，并挖掘了语义关系，用于准确的生物医学图像分割；

2. 引入了 Transformer-Self-Attention (TSA) 模块，通过自注意力机制实现跨尺度像素级上下文建模，从而实现不同尺度上的更全面的特征增强。

3. 开发了 Transformer-Cross-Attention (TCA) 模块，通过引入一种可学习的代理嵌入来自动学习不同语义类别的语义对应。然后进一步将这种代理嵌入通过交叉注意力机制与特征表示进行交互。通过为代理嵌入引入辅助损失，它可以有效地提高同一类别的特征相关性和不同类别之间的特征可辨别性。

We propose the Multi-Compound Transformer (MCTrans), which incorporates rich context modeling and semantic relationship mining for accurate biomedical image segmentation. MCTrans overcomes the limitations of conventional vision transformers by: (1) introducing the Transformer-Self-Attention (TSA) module to achieve cross-scale pixel-level contextual modeling via the self-attention mechanisms, leading to a more comprehensive feature enhancement for different scales. (2) developing the Transformer-Cross-Attention (TCA) to automatically learn the semantic correspondence of different semantic categories by introducing the proxy embedding. We further use such proxy embedding to interact with the feature representations via the cross-attention mechanism. By introducing auxiliary loss for the updated proxy embedding, we find that it could effectively improve feature correlations of the same category and the feature discriminability between different classes.

 

2. 方法

2.1 MCTrans 网络

如上图所示，本文在经典的 UNet 编码器和解码器架构之间引入了 MCTransformer，它由 Transformer-Self-Attention (TSA) 模块和 Transformer-Cross-Attention (TCA) 模块组成。前者用于对多个特征之间的上下文信息进行编码，从而产生丰富且一致的像素级上下文；后者引入了可学习的嵌入（embedding），为了语义关系建模并进一步增强特征表示。

The former is introduced to encode the contextual information between the multiple features, yielding rich and consistent pixel-level context. And the latter introduces learnable embedding for semantic relationship modeling and further enhances feature representations.

给定图像 $I\in {\mathbb{R}}^{H\times W}$，采用深度 CNN 提取不同尺度的多级特征 $\left\{X_i\in {\mathbb{R}}^{\frac{H}{2^i}\times \frac{W}{2^i}\times C_i}\right\}$。对于层级 $i$，特征以 $P\times P$ 的大小展开成各个块（patch），其中 $P$ 在本文中设置为 1，即第 $i$ 个特征图的每个位置都将被视为一个块，得到 $L_i=\frac{HW}{2^{2\ast i}\times P^2}$ 个块。接下来，将不同层级中的块输入具有相同输出特征维度 $C$ 的线性投影头（linear projections heads）（即 1×1 卷积层），得到嵌入标记（token） $T_i\in {\mathbb{R}}^{L_i\times C}$。然后，我们将 $i$ = 2, 3, 4 层级的特征进行拼接，形成整体标记 $T\in {\mathbb{R}}^{L\times C}$，其中 $L={\sum }_{i=2}^4{L}_i$。为了弥补丢失的位置信息，位置嵌入（positional embedding） $E_{pos}\in {\mathbb{R}}^{L\times C}$ 被添加到标记中，以提供关于特征在序列中的相对或绝对位置的信息，这样标记可以表示为 $T=T+E_{\text{pos }}$。接下来，我们将标记 $T$ 输入 TSA 模块以进行多尺度上下文建模。将输出的增强后的标记进一步输入 TCA 模块并与代理嵌入（proxy embedding）$E_{\text{pro }}\in {\mathbb{R}}^{M\times C}$ 进行交互，其中 $M$ 是数据集的类别数。最后，我们将编码后的标记折叠回金字塔特征图，并以自下而上的方式合并它们，以获得最终的特征图进行预测。

2.1.1 Transformer-Self-attention

将一维嵌入标记 $T$ 作为输入，TSA 模块用于学习多尺度特征之间的像素级上下文依赖关系。如上图所示，TSA 模块由 $K_s$ 层组成，每层由多头自注意力（multi-head self-attention，MSA）和前馈网络（feed forward networks，FFN）组成，在每个块之前应用层归一化（layer normalization，LN) ，在每个块之后应用残差连接（residual connection）。 FFN 包含两个带有 ReLU 激活的线性层。

对于第 $i$ 层，多头自注意力的输入是从输入 $T^{l-1}$ 计算得到的元组**（query, key, value）**：

$\text{ query }=T^{l-1}{\mathbf{W}}_Q^l,\text{ key }=T^{l-1}{\mathbf{W}}_K^l,\text{ value }=T^{l-1}{\mathbf{W}}_V^l$

其中 ${\mathbf{W}}_O^l\in {\mathbb{R}}^{C\times d_q}$，${\mathbf{W}}_K^l\in {\mathbb{R}}^{C\times d_k}$，${\mathbf{W}}_V^l\in {\mathbb{R}}^{C\times d_v}$ 是第 $i$ 层不同线性投影头的参数矩阵，$d_q$，$d_k$，$d_v$ 是三个输入的维度。

**自注意力（SA）**可以表示为：

$\mathrm{S}\mathrm{A}\left(T^{l-1}\right)=T^{l-1}+\mathrm{Softmax}\left(\frac{T^{l-1}{\mathbf{W}}_Q^l{\left(T^{l-1}{\mathbf{W}}_K^l\right)}^{\top }}{\sqrt{d_k}}\right)\left(T^{l-1}{\mathbf{W}}_V^l\right)$

**多头自注意力（MSA）**是具有 $h$ 个独立 SA 操作的扩展，将它们的输出进行拼接和投影：

$\mathrm{MSA}\left(T^{l-1}\right)=\mathrm{Concat}\left({\mathrm{S}\mathrm{A}}_1,\dots ,{\mathrm{S}\mathrm{A}}_{\mathrm{h}}\right)W_O^l$

其中 ${\mathbf{W}}_O\in {\mathbb{R}}^{h{d}_k\times C}$ 是输出线性投影头的参数矩阵。

在本文中，我们采用 $h$ = 8, $C$ = 128 ，$d_q$，$d_k$，$d_v$ 等于 $C/h$ = 32。

Transformer-Self-Attention 整个计算过程可以表示为：

$T^l=\mathrm{MSA}\left(T^{l-1}\right)+\mathrm{FFN}\left(\mathrm{MSA}\left(T^{l-1}\right)\right)\in {\mathbb{R}}^{L\times C}$

为简单起见，在上式中省略了层归一化 LN。

需要注意的是，由多尺度特征展开得到的标记 $T$ 的序列长度极长，MSA 的二次方计算复杂度使其无法处理。为此，在这个模块中，我们用**可变形自注意力（Deformable Self Attention，DSA）**机制来代替 SA。作为依赖于数据的稀疏注意力（不是全成对的），DSA 只关注整个序列中的一个稀疏元素集，而不管其序列长度如何，这在很大程度上降低了计算复杂度并允许多级特征图的交互。更多关于 Deformable Self Attention 的细节请参考 Deformable DETR: Deformable transformers for end-to-end object detection。

2.1.2 Transformer-Cross-attention

如上图所示，除了增强的标记 $T^{K_S}$ 之外，我们还提出了一种可学习的代理嵌入（proxy embedding）$E_{pro}$ 来学习类别之间（即类内/类间）的全局语义关系。在结构上，代理嵌入是具有特定维度（128）的张量。

与 TSA 模块一样，TCA 模块由 $K_c$ 层组成，但包含两个多头自注意力块（multi-head self-attention blocks）。

对于第 $j$ 层，代理嵌入 $E_{pro}^{j-1}$ 由不同的线性投影转换成第一个 MSA 块的输入（query, key, value）。在这里，MSA 块的自注意力机制与每一对类别连接并交互，从而对各种类别的语义对应进行建模。

Here, the MSA block’s self-attention mechanism connects and interacts with each pair of categories, thus modeling the semantic correspondence of various labels.

接下来，学习到的代理嵌入通过另一个 MSA 块中的交叉注意力机制提取输入标记 $T^{K_S}$ 的特征并与之交互，其中输入的 query 是代理嵌入，输入的 key 和 value 是标记 $T^{K_S}$。通过交叉注意力机制，标记的特征与学习到的全局语义关系进行通信，全面提高了特征表示的类内一致性和类间可区分性，产生了更新后的代理嵌入 $E_{pro}^j$。

Through the cross-attention, the features of tokens communicate with the learned global semantic relationship, comprehensively improving intra-class consistency and the inter-class discriminability of feature representation, yielding updated proxy embedding $E_{pro}^j$.

此外，我们引入了辅助损失 $Los{s}_{aux}$ 来促进代理嵌入的学习。TCA 模块的最后一层的输出 $E_{pro}^{K_c}$ 被进一步传递给线性投影头，并产生多类预测 $Pre{d}_{aux}\in {\mathbb{R}}^M$。基于真实分割掩膜，可以找到计算分类标签以进行监督的独特元素#疑问#。通过这种方式，代理嵌入被驱使学习适当的语义关系，并有助于提高同一类别的特征相关性和不同类别之间的特征可辨别性。

Base on the ground-truth segmentation mask, we find the unique elements to compute classification labels for supervision. In this way, the proxy embedding is driven to learn appropriate semantic relationship, and help to improve feature correlations of the same category and the feature discriminability between different categories.

最后，将编码后的标记 $T^{K_s}$ 折叠回二维特征图并附加未处理的特征以形成金字塔特征 $\left\{X_0,X_1,X_2^{\prime },X_3^{\prime },X_4^{\prime }\right\}$。我们使用 2 倍上采样层和 3×3 卷积以常规自下而上的方式逐步合并它们，以获得用于分割的最终特征图。

2.2 数据集

本文提出的 MCTrans 在三种类型的六个分割数据集上进行了评估。每个数据集都有不同的数据模式、数据大小和前景类，使其适合用来评估 MCTrans 的有效性和泛化性。

(1) 细胞分割：Pannuke 数据集（7904 例，6 类）；

（2）息肉分割： CVC-Clinic 数据集（612 例，2 类），CVC-ColonDB 数据集（380 例，2 类），ETIS-Larib 数据集（196 例，2 类），Kvasir 数据集（1000 例，2 类）；

（3）皮肤病变分割：ISIC2018 数据集（2594 例，2 类）。

 

3. 实验过程与结果

3.1 实施

• 硬件：NVIDIA Tesla V100 GPU
• 框架：Pytorch
• 优化器：Adam
• 学习率：初始学习率为 3e-4，在训练过程中线性衰减
• 损失函数：使用交叉熵损失和 Dice 损失来惩罚分割的训练误差，并使用权重为 0.1 的交叉熵损失来进行辅助监督。

3.2 消融实验

我们采用 VGG 网络作为主干网络（backbone）来提取多尺度特征表示。为了证明构建多尺度像素级依赖关系的有效性，我们在 UNet 的最高级别特征上使用非局部操作和 Transformer-Encoder 来实现单尺度的上下文传播，结果的精度远远落后于我们的方法。添加 TCA 后， MCTrans 的得分提高到了 68.40%，这表明学习语义关系以增强特征表示的有效性。去除辅助损失后，网络只对类别之间的语义关系进行隐式建模。该策略将性能降低到了 67.87%。

3.3 与 SOTA 方法的比较

我们在 Pannuke 数据集上比较了 MCTrans 和最先进的方法。我们采用传统 VGG 网络作为特征提取器。为了进行更全面的比较，在第二组中，我们采用了更强的特征提取器（例如 ResNet-34）。在下表中，MCTrans 以合理的计算开销为代价获得了更好的结果。与 UNet 相比，MCTrans 以几乎相同的参数和略微增加的计算量实现了 3.64% 的显著改善。请注意，其他先进方法（例如 UNet++）的计算量比 MCTrans 多很多，但性能较低。

下图提供了一些分割结果的示例。

 

4. 重要参考文献

• 可变形自注意力（Deformable Self Attention，DSA）

  Zhu, Xizhou, et al. “Deformable DETR: Deformable transformers for end-to-end object detection.” arXiv preprint arXiv:2010.04159 (2020).

 

5. 评审意见

官方评审意见 ：https://miccai2021.org/openaccess/paperlinks/2021/09/01/319-Paper0786.html

评分：6，7，8

5.1 优点

• Clearly listed the bottlenecks of current SOTA segmentation methods (both CNN based and transformer based) and introduced their method to make improvement.
• Deformable self-attention is implemented to limit the computational costs, which is an important issue in combined models;
• A novel learnable proxy embedding is introduced to enhance feature representation.
• The designed MCTrans can be plugged in and use in different models.

5.2 缺点

• The paper used many mathematical notations to explain their method, but didn’t provide the definition for all of them. Also some concepts might be self-explaining for the people who are familiar with the transformer, but not clear for the rest.

5.3 建议

• Transformers always consume much more computation resources. If you can extend your model to a 3D medical image, it will become a more great work.