[MICCAI2023]MDViT

论文代码

[siyi-wind/MDViT: MICCAI 2023] MDViT: Multi-domain Vision Transformer for Small Medical Image Segmentation Datasets (an official impleme ntation) (github.com)

摘要

MDViT是Multi-domain Vision Transformer的缩写，意为多领域视觉变换器。MDViT是专为医学图像分割（MIS）任务设计的视觉变换器模型。医学图像分割是一项具有挑战性的任务，因为医学图像具有复杂性和变异性。

近年来，视觉变换器（ViTs）在改善MIS方面显示出了潜力，但通常需要大量的训练数据集。为了解决数据稀缺的问题，提出了数据效率高的ViTs。然而，这些模型通常只在单一数据源上进行训练，忽视了从其他可用数据集中获取宝贵知识的可能性。

MDViT通过引入域适配器来克服这个限制，使模型能够自适应地利用多个小型数据资源或领域中的知识。通过整合域适配器，MDViT可以缓解不同领域之间存在显著领域异质性时可能发生的负面知识转移（NKT）。

此外，MDViT采用了互相知识蒸馏的方法来增强跨领域的表示学习。这种方法涉及在跨越所有领域的通用网络和辅助特定领域网络分支之间转移知识。

对于四个皮肤病变分割数据集的实验结果表明，MDViT在分割性能方面优于现有的算法，即使在添加更多领域的情况下也是如此。在推理时，MDViT保持固定的模型大小，同时实现卓越的分割性能。

解决问题

解决ViT（Vision transformers) data-hungry问题

以前的方法中一个值得注意的限制是它们并不是通用的，即它们依赖于针对每个数据集的单独训练，而不是整合相关领域的宝贵知识。因此，它们可能会带来额外的训练、推断和内存成本。

主要创新点

• 将多维度学习引入解决目标（见上）
• MDViT(multi-domain ViT)
• 4 skin lesion segmentation datasets 10.16% 提升in IOU

我们提出的MDViT在BASE的基础上进行了扩展，具体包括以下两个

方面：
1）在分解的多头自注意力（MHSA）内部添加了一个领域适配器（DA）模块，以使模型适应不同的领域（图1-b、c）；
2）采用相互知识蒸馏（MKD）策略，提取跨领域更稳健的表示（图1-d）。

方法

假设X ∈ RH×W×3是输入的RGB图像，Y ∈ {0, 1}H×W是其对应的分割掩码（groundtruth segmentation mask）。训练样本 {(X,Y)} 来自M个数据集，每个数据集代表一个领域。

BASE是一个基于U-Net 和金字塔ViTs 架构的U型ViT。它包含编码（前四个）和解码（后四个）的Transformer块，一个两层的CNN桥接层和跳跃连接。如[19]所述，第i个Transformer块包括一个卷积块嵌入层，采用3×3的补丁大小，和Li个使用线性复杂度的分解MHSA的Transformer层。前者将特征图Xi-1转换为补丁嵌入序列zi ∈ RNi×Ci，其中
是补丁的数量，Ci是通道维度。我们使用与[19]相同的位置嵌入和与[27]相同的跳跃连接。为了降低计算复杂性，根据[19]的方法，在Transformer块之前添加两个CNN层，在Transformer块之后添加一个CNN层，使得第一个Transformer块能够处理从较低分辨率开始的特征：$\frac{H}{4}+\frac{W}{4}$。在BASE中，我们没有使用集成和分层的CNN骨干网络，例如ResNet，如数据效率高的混合ViTs ，以明确评估多领域学习在缓解ViTs对数据需求方面的有效性。

MDViT网络结构

DA (DomainAdapter)

在多领域自适应训练中，一些方法会在主网络的旁路上构建领域特定的层。而我们在不添加领域特定层的情况下，利用ViTs中现有的并行结构，即MHSA，进行领域适应。MHSA的H个并行头模拟了人类从不同角度观察同一对象的方式。类似地，我们插入DA到MHSA的想法是使不同的头在不同领域中具有不同的观点。与手动指定每个头对应一个领域不同，MDViT通过领域标签的指导，在遇到一个领域时学习将注意力集中在不同头部的相应特征上。DA包含两个步骤：注意力生成和信息选择（图1-c）。
注意力生成为每个头生成注意力。我们首先将领域标签向量m（采用独热编码m ∈ RM，但也可以使用其他编码方式）通过一个带有ReLU激活函数的线性层，得到一个领域感知向量d ∈ $R^{\frac{K}{r}}$。这里K是头部特征的通道维度。我们将缩减比例r设置为2。之后，类似于[20]，我们计算每个头的注意力：
，其中ψ是对头部进行softmax操作。

信息选择自适应地从不同的头部选择信息。在从第h个头部获得特征
后，我们利用ah来校准沿着通道维度的信息：

MKD(Mutual Knowledge Distillation)

从领域特定网络中提取知识已被证明对于通用网络来学习更健壮的表示是有益的。此外，将知识在教师和学生之间相互传递的互相学习使得两者可以同时优化。为了实现这些优势，我们提出了互相知识蒸馏（MKD），在辅助对等网络和通用网络之间相互传递知识。在图1-d中，第m个辅助对等网络仅在第m个领域上进行训练，产生输出${\hat{Y}}^m$，而通用网络的输出是ˆY。类似于[21]，我们利用对称Dice损失
作为知识蒸馏损失。每个对等网络是某个特定领域的专家，引导通用网络学习领域特定的信息。通用网络经历所有的领域并掌握领域共享的知识，这对于对等网络的学习是有益的。
每个辅助对等网络在一个小的、特定于该对等网络的个体数据集上进行训练（图1-d）。为了实现快速训练过程并防止过拟合，特别是在处理众多训练数据集时，我们将一个轻量级的多层感知器（MLP）解码器（针对ViT编码器设计的）调整为我们对等网络的架构。具体而言，从编码Transformer块（图1-a)得到的多层级特征经过一个MLP层和一个上采样操作，将通道维度和分辨率统一为$\frac{H}{4}+\frac{W}{4}$，然后与通用网络最后一个Transformer块中包含领域共享信息的特征进行连接。最后，我们将融合后的特征传递到一个MLP层，并进行上采样以获得分割图。

损失函数

Lseg=Ldice+ Lbce

实验效果

4个皮肤病变分割数据集：ISIC 2018（ISIC）、Dermofit图像库（DMF）、皮肤癌检测（SCD）和PH2。

ST-seperate Training

JT-Joint Training

MAT-multi-domain adaptive training

ResNet-34作为BAT骨架，其他的一些数据集上好，一些不好

消融实验