[小样本分割] Few Shot Medical Image Segmentation with Cross Attention Transformer

Few Shot Medical Image Segmentation with Cross Attention Transformer

摘要

1. 在深度学习医学图像分割领域，训练一个性能强，可以大规模部署落地的模型，往往需要大量手动标注的数据进行监督训练，其中花费的成本是非常高的。
2. 为了解决这一挑战，少样本学习（few-shot）技术有潜力从有限的几个sample中学习新类别的能力。本文提出了一种基于交叉掩码注意力Transformer的少样本医学图像分割新框架CAT-Net：通过挖掘support和query图像之间的相关性，并限制模型仅关注有用的前景信息，来提高support和query特征的表达能力。同时，本文还进一步设计了一个迭代细化训练框架来优化查询query图像分割。
3. 作者在三个公共数据集（Abd-CT，Abd-MRI和Card-MRI）上验证了所提出方法的有效性。

1 引言

​ 图像分割在医学影像中是一项非常常见的任务。在工业界，绝大部分医学图像分割任务还是基于全监督进行训练（或许只因为它稳定且精度高，只要公司不断的堆高质量标注数据就可以了，这个成本公司还是愿意花的）；但医学图像的标注是一项耗时长，成本高的任务，比如对于一个3D的volume图像（以CT和MRI居多），图像标注则更具有挑战性：第一，标注者需要浏览每个3D扫描的数百个2D切片进行标注，量非常大；第二，有些标注外行人还真的做不好，这往往需要一些专业医生动用专业知识进行标注；第三，这些专业医生往往还有其它很多事情要做，你让医生去标注大量数据还真得考虑可行性。

​ 基于这个出发点，为了解决手动标注所带来的挑战，开辟了各种研究方向：例如自监督学习，半监督学习和弱监督学习。尽管利用未标记或弱标记数据的信息，这些技术仍然需要大量的训练数据，这对于医学领域中仅有有限样本的新类别可能不可行。于是，研究学者开辟了few-shot这么一个研究方向：few-shot 学习范式旨在从少量标记数据（称为support）中学习模型，然后将其应用于仅有少量标记数据的新任务（称为query），而无需重新训练。

​ 考虑到人体内的数百个器官和无数的疾病，few-shot学习为各种医学图像分割任务带来了巨大的潜力，可以在数据高效的情况下轻松地研究新任务。

​ 大多数few-shot分割方法都在学习如何学习（旨在学习元学习器），根据support图像及其相应的分割标签的知识预测query图像的分割，而这里的核心是：如何有效地将知识从support图像传递到query图像。现有的少样本分割方法主要集中在以下两个方面：

1. 如何学习一个元学习器
2. 如何更好地将知识从support图像传递到query图像

尽管基于原型的方法效果已经不错，但它们通常忽略了训练过程中support和query特征之间的交互。

​ 因此，本文提出了一种名为CAT-Net的新型网络结构，其基于交叉注意力Transformer，可以更好地捕捉support图像和query图像之间的相关性，促进support和query特征之间的相互作用，同时减少无用像素信息，提高特征表达能力和分割性能；此外，本文还提出了一个迭代训练框架，将先前的support分割结果反馈到注意力Transformer中，以有效增强并细化特征和分割结果。作者在三个公共数据集上验证了CAT-Net的有效性和性能优越性。

2 方法

2.1 问题定义

少样本分割（Few-shot segmentation，FSS）的目的是通过只有少量标注的样本来分割新类别。在FSS中，数据集被分为训练集Dtrain和测试集Dtest，其中训练集包含基类别Ctrain，测试集包含新类别Ctest，且Ctrain和Ctest没有交集。为了获得用于FSS的分割模型，采用了通常使用的episode训练方法。每个训练 / 测试episodes(Si,Qi)实例化一个N-way, K-shot分割学习任务。具体而言: support集包含N个类别的K个样本，而query集包Qi含同一类别的一个样本。FSS模型通过episode训练以预测query图像的新类别。在模型推理测试时，模型直接在Dtest上进行评估，无需重新训练。

2.2 方法

如上图1展示了CAT-Net网络框架图，主要由三部分组成：

1. 带有mask的特征提取MIFE子网络，用于提取初始query和support特征以及query mask。
2. 交叉mask注意力Transformer模块CMAT，其中query和support特征相互促进，从而提高query预测的准确性。
3. 迭代细化框架，顺序应用CMAT模块以持续促进分割性能，整个框架以端到端的方式进行训练。

2.3 Mask Incorporated Feature Extraction

​ CAT-Net中的Mask Incorporate Feature Extraction (MIFE)子网络。MIFE子网络接收查询和支持图像作为输入，生成它们各自的特征，同时集成支持掩膜。然后，使用一个简单的分类器来预测查询图像的分割结果。具体地，首先使用一个特征提取器网络（即ResNet-50）将查询和支持图像对Iq和Is映射到特征空间中，分别产生查询图像的多层特征图Fq和支持图像的特征图Fs。接下来，将支持掩膜与Fs进行池化，然后将其扩展并与Fq和Fs进行连接。此外，还将一个先验掩膜进一步与查询特征进行连接，通过像素级相似度图来增强查询和支持特征之间的相关性。最后，使用一个简单的分类器来处理查询特征，得到查询掩膜。关于MIFE架构的更多细节可以在补充材料中找到。

2.4 Cross Masked Attention Transformer

​ 如图1b所示，CMAT模块包括三个主要组成部分：自注意力模块、交叉掩码注意力模块，和原型分割模块。其中，自注意力模块用于提取查询query特征和支持support特征中的全局信息；交叉掩码注意力模块用于在传递前景信息的同时消除冗余的背景信息；原型分割模块用于生成查询图像的最终预测结果。

自注意力模块：

​ 自注意力模块首先将查询特征F₀^q和支持特征F₀^s展平为1D序列，然后输入到两个相同的自注意力模块中。每个自注意力模块由一个多头注意力层和一个多层感知器MLP层组成。给定一个输入序列S，MHA层首先使用不同的权重将序列投影为三个序列Q,K 和 V。然后计算注意力矩阵 A，公式为：

​ 其中， d是输入序列的维度。注意力矩阵通过 softmax 函数归一化，并乘以值序列 以获得输出序列V 。MLP层是一个简单的1×1卷积层，将输出序列O映射到与输入序列S相同的维度。最终，将输出序列O添加到输入序列S中，并使用层归一化（LN）对其进行规范化，以获得最终的输出序列X。自注意力对齐编码器的输出特征序列分别表示为X^q∈R^HW×D和X^q∈R^HW×D ，分别对应于查询和支持特征。

交叉掩码注意力模块：

​ 用于将查询特征和支持特征按照它们的前景信息结合起来。在attention矩阵中，通过支持和查询的掩码来限制注意力区域。具体来说，给定查询特征X^q和来自自注意力模块的支持特征X^s，首先使用不同的权重将输入序列投影到三个序列K，Q和V中，从而得到K^q、Q^q、V^q和K^v、Q^v、V^v。以查询特征为例，交叉注意力矩阵通过下面的公式计算得到：

其中，d表示查询特征的维度。这里使用的是点积注意力的形式，通过K^q和Q^s的点积计算查询和支持之间的相关性。通过√d来缩放点积，防止在较高维度时点积的大小对注意力分布的影响过大。

原型分割模块：

首先，通过一个“masked average pooling”的方法，建立每个类别的原型（prototype）p_c，用于表示该类别的特征分布。

其中，K是支持集中图像的数量，m^s_(k,x,y,c)是一个二进制掩模，(x,y)表示位置在支持特征k中是否属于类别c，F₁^s是支持特征。具体来说，对于每个类别，该原型是在所有支持图像中该类别对应位置的特征平均值，这样可以得到每个类别的原型p_c。

接着使用非参数度量学习方法进行分割。原型网络计算查询特征向量与原型P=P_c|c∈C之间的距离。对所有类别应用softmax函数，生成查询分割结果：

其中cos(·)表示余弦距离，α是一个缩放因子，有助于在训练中反向传播梯度，其中α设置为20。

Iterative Refinement framework:

该模块的设计目的是优化查询和支持特征以及查询分割掩模。因此可通过迭代优化的思路进行精细化分割，第i次迭代后的结果由以下公式给出：

每个步骤的细分可表示如下：

其中CMA(·)表示自注意力和交叉掩码注意力模块，Proto(·)代表原型分割模块，该公式表示通过多次迭代应用CMA和Proto模块，来获得增强的特征和优化的分割结果。

3 实验

3.1 数据集和评价指标

​ 我们在三个公共数据集上评估了所提出的方法，即Abd-CT ， Abd-MRI和Card-MRI。Abd-CT包含30次腹部CT扫描，标注左、右肾(LK、RK)、脾(Spl)、肝(Liv)。Abd-MRI包含20次腹部MRI扫描，与Abd-CT相同的器官标注。Card-MRI包括35次心脏MRI扫描，并标注左心室血池(LV-B)、左心室心肌(LV-M)和右心室(RV)。我们使用Dice分数作为评估指标如下。

​ 为了保证比较的公平性，所有的实验都是在1-way 1-shot场景下进行的，使用了5次交叉验证。我们遵循[16]在训练期间删除所有包含测试类的片，以确保在验证期间测试类都是不可见的。在每一次折叠中，我们遵循[16,5,22]，以最后一个患者作为支持图像，其余患者作为查询(设置I)。我们进一步提出了一种新的验证设置(设置II)，将每个折叠中的每张图像交替作为支持图像，其他图像作为查询。报告每次折叠的平均结果。通过减少支持图像选择的影响来评价模型的泛化能力。

3.2 实现细节

所提出的方法是使用PyTorch实现的。每个3D扫描被切成2D切片，并重新塑造为256×256像素。训练数据采用常用的三维图像预处理技术，如强度归一化和重采样。我们应用2万次迭代的插曲训练。采用SGD优化器，学习率为0.001，批大小为1。使用单个NVIDIA RTX 3090 GPU，每集训练大约需要4个小时。

3.3 与最新方法的比较

我们将提出的CAT-Net与最先进的(SOTA)方法进行了比较，包括SE-Net ， PANet ， ALP-Net ， AD-Net和Q-Net。PANet是自然图像域FSS方法的典型原型，SE-Net、ALP-Net、AD-Net和Q-Net是医学FSS任务中最具代表性的工作。实验结果见表1证明在设置I和设置II下，所提出的方法在所有三个数据集上都优于SOTAs。在设置I下，提出的CAT-Net在Abd-CT上的Dice达到66.59%，在Abd-MRI上的Dice达到75.18%，在Card-MRI上的Dice达到79.03%，分别优于SOTAs 1.76%， 0.75%和0.45%。在设置II下，CAT-Net在Abd-CT上的Dice达到70.88%，在Abd-MRI上的Dice达到75.22%，在Card-MRI上的Dice达到79.36%，分别比SOTAs高出2.56%、2.02%和1.32%。在三个数据集和两种评估设置下，我们的方法相对于SOTAs的一致优势表明了所提出的CAT-Net的有效性和泛化能力。

在 Abd-CT 和 Abd-MRI数据集上，相比于之前的最先进方法（SOTAs），这个提出的方法能够生成更加准确和详细的分割结果

3.4 消融实验

我们进行了一项消融实验，以调查CAT-Net中每个组件的有效性。所有消融研究均在设置II的Abd-MRI上进行

CMAT Block的有效性:

为了证明我们提出的CAT-Net在缩小查询和支持图像之间的信息差距和获得增强特征方面的重要性，我们进行了消融实验。具体地说，我们比较了仅从支持图像(S→Q)或查询图像(Q→S)学习前景信息的结果，并得到单个增强特征而不是两个增强特征(S↔Q)。可以观察到，使用增强型查询特性(S→Q)在Dice中达到66.72%，仅使用增强型支持特性(Q→S)的性能优于0.74%。使用我们的CMAT块，相互增强支持和查询特性(S↔Q)可以将Dice提高1.90%。迭代精化框架在Dice中对上述三种变化分别提高了0.96%、0.56%和2.26%。

迭代掩码优化块的影响:

为了确定迭代细化CMAT块的最佳数量，我们对不同数量的块进行了实验。在图3中，我们观察到增加块的数量可以提高性能，当使用5个块时，Dice的性能最大提高了2.26%。考虑到使用4个和5个CMAT块之间的性能增益不显著，因此我们选择在最终模型中使用4个CMAT块，以在效率和性能之间取得平衡。

4 结论

本文提出了一种用于few-shot医学图像分割的交叉注意力Transformer网络CAT-Net。通过交叉掩码注意力模块实现了查询和支持特征之间的交互，增强了特征表达能力。此外，所提出的CMAT模块可以通过迭代优化的方式以持续提高分割性能，实验结果表明了每个模块的有效性以及模型相对于SOTA方法的卓越性能。其中论文中的各个组件属于即插即用模块，可很好的嵌入到few shot任务中，以提高少样本分割的性能。