Uncertainty-Guided Pixel Contrastive Learning for Semi-Supervised Medical Image Segmentation

用于半监督医学图像分割的不确定性引导像素对比学习

Paper：https://www.ijcai.org/proceedings/2022/201
Code：https://github.com/taovv/UGPCL

Abstract

最近，对比学习在医学图像分割中显示出巨大的潜力。然而，由于缺乏专家注释，在半监督场景中应用对比学习具有挑战性。

为了解决这个问题，我们提出了一种新颖的用于半监督医学图像分割的不确定性引导像素对比学习方法。具体来说，我们为每个未标记图像构建不确定性图，然后去除不确定性图中的不确定性区域以减少噪声采样的可能性。不确定性图是由精心设计的一致性学习机制确定的，该机制通过鼓励两个不同解码器的一致网络输出来生成对未标记数据的全面预测。

此外，我们建议图像编码器学习的有效全局表示应该与不同的几何变换等效。为此，我们构造了一个等变对比损失来增强编码器的全局表示学习能力。在流行的医学图像基准上进行的大量实验表明，所提出的方法比最先进的方法实现了更好的分割性能。

1 Introduction

医学图像分割在计算机辅助诊断系统中发挥着重要作用。基于深度学习的监督学习方法已经取得了很好的性能[Ronneberger et al, 2015;曹等人，2021； Li et al, 2021]依赖大量标记数据。然而，由于专业临床知识的要求以及数据收集和标记的时间消耗，很难获得大规模的医学图像注释。半监督学习可以同时利用标记数据和未标记数据，这大大减少了对注释的依赖。半监督学习旨在探索未标记数据的内部信息以提高模型的性能。一些流行的半监督学习策略包括使用伪标签进行自我训练[Qiao et al, 2018; Chen 等人，2021a]，自集成 [Tarvainen 和 Valpola，2017； Yu et al, 2019]、熵最小化 [Vu et al, 2019] 和一致性正则化 [Laine and Aila, 2016； Ouali 等人，2020]。上述方法在训练阶段利用未标记数据，通过构造可信标签或在输入数据遇到干扰时强制预测一致性。然而，这些方法使每个像素的分类相互独立，忽略了图像像素（或特征）之间的内部相关性。

为了加强像素之间的联系，最近的一些工作将对比学习应用于分割任务[Hu et al, 2021]。对比学习方法在自然图像的自监督表示学习中取得了优异的性能。对比学习的核心思想是相似样本的表示应该是相似的，不同种类的样本的表示应该是不同的。 如何定义相似样本是对比学习的关键。图像级对比学习将相似样本定义为同一图像的不同变换，而来自不同图像的样本则定义为相异性。然而，相似的像素在分割任务中是密集分布的，因此不同样本的定义不适合像素级对比学习。 为了解决这个问题，[1][Wang et al, 2021]使用分割标签来构建用于监督分割任务的对比样本。对于未标记的数据，[2][Chen et al,2021b]使用预测的伪标签来确定样本类别。[3][Zhong et al, 2021]利用弱增强图像的空间一致性来构造相似样本，并通过简单的跨图像和伪标签加权启发式构造不相似样本。

事实上，使用伪标签构建样本很可能与实际的语义类别不一致，这可能会导致对比学习上的噪声采样。另外，像素对比学习仅建立局部像素的关联，忽略了全局表示信息的学习。

[1] Exploring Cross-Image Pixel Contrast for Semantic Segmentation
[2] Bootstrap Representation Learning for Segmentation on Medical Volumes and Sequences
[3] Pixel Contrastive-Consistent Semi-Supervised Semantic Segmentation

在本文中，我们的目标是：1）解决使用伪标签的对比学习的噪声采样问题，2）增强编码器的全局表示学习能力。

为了实现这些目标，我们提出了一种基于不确定性的对比学习方法。图 1 显示了我们方法的核心思想。对于未标记的数据，我们使用不确定性图来指导伪标签采样的区域，减少错误样本的数量。然后计算样本对比损失来优化网络并减少预测的不确定性区域。为了获得更好的不确定性图，我们设计了一种带有 CNN 解码器和 Transformer 解码器的一致性学习策略，可以利用两个解码器之间的结构差异从不同的角度获得准确的预测。此外，分割模型应该具有识别几何变换的能力。在此基础上，我们定义了一个等变对比损失，通过在表示学习阶段添加变换类别预测来迫使网络学习几何变换的识别信息。总而言之，我们的贡献主要包括：

• 我们提出了一种新颖的不确定性引导对比学习方法，可以有效减轻来自未标记数据的伪标签的噪声采样。
• 设计了基于CNN和Transformer的异构解码器的一致性学习策略，通过对未标记数据进行一致性训练，可以获得可靠的预测结果和不确定性图。
• 我们为全局表示学习定义了等变对比损失，它使模型具有区分图像不同几何变换的能力。

2 Related Work

2.1 Semi-Supervised Medical Image Segmentation

由于不需要大规模标记数据，半监督学习在医学图像分割中备受关注。现有的半监督医学图像分割方法主要涉及熵最小化、伪标签自训练、协作训练和一致性学习。熵最小化[Vu et al, 2019]认为高质量的预测结果应该具有较低的熵，因此它通过最小化预测概率分布的信息熵来进行模型学习。伪标签自训练 [Chen et al, 2021a] 通过预测未标记数据的伪标签来执行类监督学习。协同训练 [Qiao et al, 2018] 假设存在多个包含互补信息的决策视图，并设计不同的分类器来学习不同的视图以提高分割性能。一致性学习 [Verma et al, 2019;莱恩和艾拉，2016；塔尔瓦宁和瓦尔波拉，2017； Ouali et al, 2020]假设即使图像样本遇到一些干扰，例如输入干扰或模型干扰，样本的预测结果也不应该改变。受这种直觉的启发，这些方法通过鼓励对未标记的干扰样本进行一致预测来进行模型训练。受协作训练和一致性学习的启发，我们提出利用 CNN 和 Transformer 之间的结构差异来表征不同视图数据的互补信息，并应用一致性约束来训练模型。

2.2 Contrastive Learning

在图像级表示学习中，对比学习可以充分利用未标记数据来学习有效的视觉表示，其核心思想是通过缩小相似对（正）和分离不相似对（负）来加强学习到的视觉表示的辨别力。 ）基于一些相似性约束。图像级对比学习的关键是如何构造对比样本。 [He et al, 2020]提出了一种可行的解决方案，通过引入记忆库和动量对比来增加对比样本的数量。

最近，一些作品[Chaitanya et al, 2020;王等人，2021；钟等人，2021； Hu et al, 2021] 已提出将对比学习从图像级扩展到像素级以进行图像分割。像素级对比学习的主要思想是借助分割标签构建像素样本对。对于未标记的数据，通过使用伪标签或空间结构来构造样本对。然而，这些方法在构造样本对的过程中可能会遇到噪声采样的问题。为了缓解这个问题，我们建议使用预测不确定性来指导样本采样并减少噪声样本的数量。此外，像素级对比学习缺乏全局表示的捕获能力，这促使我们在分割任务的表示学习中施加先验知识的约束。

2.3 Uncertainly Estimation

在半监督学习中，不确定性可用于评估模型预测的质量，以便更好地利用未标记数据。估计不确定性的度量方法主要包括1）利用预测概率分布的信息熵，2）利用同一输入在不同扰动下的多次预测结果的偏差[Yu et al, 2019]，3）计算方差相同输入的不同预测结果[Zheng和Yang，2021]。然而，这些方法既耗时又缺乏可靠性。在我们的方法中，我们通过计算不同预测变量获得的平均概率分布的熵来估计不确定性，以克服这些问题。

3 Methodology

给定一个标签数据集 D L = { ( x i , y i ) , i = 1 , . . . , N } D_L = \left\{(x_i , y_i), i = 1,..., N\right\} DL​={(xi​,yi​),i=1,...,N} 和未标记的数据集 D U = { x j , j = 1 , . . . , M } D_U = \left\{x_j , j = 1 ,..., M\right\} DU​={xj​,j=1,...,M}，其中 $M\gg N$，$D_L$和$D_U$中的图像首先经过几何变换，然后输入编码器网络以提取多尺度特征。之后，这些特征将被发送到以下三个分支，包括一致性学习分支、不确定性引导对比学习分支和等变对比学习分支。

对于一致性学习分支，我们提出了一个异构一致性网络来预测分割结果，该网络由监督损失 $l_{sup}$ 和一致性损失 $l_{con}$ 驱动。 $l_{sup}$ 通过 $D_L$ 的 Ground Truth 计算，$l_{con}$ 通过 $D_U$ 的预测一致性计算。

对于不确定性引导的对比学习分支，我们构建并维护一个内存队列来保留足够的样本用于对比学习。内存队列中样本的选择取决于 $D_L$ 的标签和 $D_U$ 的确定性标签。对于选定的样本，我们施加像素级对比损失 $l_c$，使同一类的像素彼此靠近，不同类的像素彼此远离。

对于等变对比学习分支，我们对所有标记和未标记数据进行几何变换类别预测，并设计等变对比损失 $l_e$ 以迫使编码器对几何变换具有鲁棒性。

为了便于理解，图 2给出了我们提出的方法的整体架构和训练过程的说明。总之，我们方法的总体目标是：

$$l=l_{sup}+{\lambda }_t{l}_{con}+{\lambda }_1{l}_c+{\lambda }_2{l}_e\text{\hspace{1em}(1)}$$

在本文中，我们设置 ${\lambda }_1={\lambda }_2=0.1$，${\lambda }_t$ 是一个从 0 增加到 0.01 的温度参数。下面详细展示上述三个分支。

3.1 Consistency Learning Between Decoders

在一致性学习分支中，我们设计了一个简单而有效的网络结构来实现以下两个目标：1）使用未标记的数据来促进分割网络的学习，2）从网络输出中获得可靠的不确定性估计。证明使用协同训练策略可以获得更好的分割性能，其核心思想是从不同的角度做出不同的分类预测，然后将预测的差异作为不确定性估计的衡量标准。受 [Luo et al, 2021] 的启发，我们没有使用需要添加一些干扰来进行协同训练的相同架构，而是采用了一种简单而有效的方案，该方案利用了 Transformer 解码器和 CNN 解码器之间的先天差异。具体来说，我们构造一个异构预测器来约束两个解码器生成一致的预测。然后使用平均预测的熵来估计不确定性图。

Patch Embedding and Position Encoding.（补丁嵌入和位置编码）
我们从 Swin-UNet 和 UNet 中选择两个不同的解码器 $f_{\theta }^t(\cdot )$ 和 $f_{\theta }^c(\cdot )$。通过编码器，我们可以获得一组特征 { f i , i = 0 , … , 3 } \left\{f_i , i = 0,…, 3\right\} {fi​,i=0,…,3}。在输入到 $f_{\theta }^t(\cdot )$ 之前，我们需要将特征 $f_i\in {\mathbb{R}}^{H_i\times W_i\times C_i}$ 重塑为一系列扁平补丁 $x_p^i\in {\mathbb{R}}^{P_i^2\times C_i}$ 并使用可学习的线性投影 $E$ 将它们嵌入到 $D$ 维空间中。为了保留空间信息，我们将绝对位置编码添加到嵌入的补丁中：

$$PE(x_i)=[x_p^{i1}E;x_p^{i2}E;\cdots ;x_p^{i{P}_i^2}E]+E_{pos}\text{\hspace{1em}(2)}$$

其中 $E=\in {\mathbb{R}}^{P_i^2\times C_i\times D}$ 是块嵌入投影，$E_{pos}\in {\mathbb{R}}^{N\times D}$ 表示绝对位置编码。

Consistency Learning.（一致性学习）
给定输入图像，我们可以从两个解码器获得两个预测概率分布 $p_t$ 和 $p_c$。对于标记数据，我们使用 Ground Truth 来计算监督分割损失：

$$l_{sup}={\mathcal{L}}_{seg}(p_c,y)+\alpha {\mathcal{L}}_{seg}(p_t,y)\text{\hspace{1em}(3)}$$

$${\mathcal{L}}_{seg}=\frac{1}{2}({\mathcal{L}}_{CE}+{\mathcal{L}}_{Dice})\text{\hspace{1em}(4)}$$

其中 ${\mathcal{L}}_{CE}$、${\mathcal{L}}_{Dice}$ 是交叉熵损失和 Dice 损失，$y$ 是标记数据的真实值。我们使用 $f_{\theta }^c(\cdot )$ 作为主要预测变量，因此 $\alpha$ 设置为 0.4。在推理阶段，CNN分支的预测就是最终结果。对于未标记的数据，一致性损失计算如下：

$$l_{con}={\mathcal{L}}_{dis}(p_c,p_t)\text{\hspace{1em}(5)}$$

其中 ${\mathcal{L}}_{dis}$ 是两个输出概率分布之间的距离度量。在这项工作中，我们选择使用均方误差（MSE）作为距离度量。

3.2 Uncertainty-Guided Contrastive Learning

图像分割任务通常涉及交叉熵损失来进行像素级分类。然而，这种损失使得每个像素的分类是独立的，因此忽略了像素之间的关系。

为了解决这个问题，我们设计了一种像素级对比学习机制，将同一类别（语义标签）的像素分类为正样本，将不同类别的像素分类为负样本。通过在嵌入空间中减小正样本之间的距离并扩大负样本之间的距离来建立像素之间的关系。 为了有效地利用未标记数据进行像素级对比学习，我们估计未标记像素的不确定性，并选择确定性较高的像素作为对比学习的锚点。

Mask Uncertainty Region. （屏蔽不确定性区域）
我们选择预测熵作为近似不确定性的度量。具体来说，我们首先计算预测结果的平均概率分布 $\hat{p}=(p_c+p_t)/2$，然后计算每个像素在通道维度上的概率分布的熵。可以概括为：

$$u=-\sum _c{\hat{p}}_{c}log({\hat{p}}_c+ϵ)\text{\hspace{1em}(6)}$$

其中 $ϵ$ 是一个非常小的常数，以避免奇异性。我们认为，大熵的预测在类别上是不确定的。在计算伪标签时，将那些不确定性预测作为非采样区域去除，然后得到确定的伪标签：

$$y_p=Argmax(\hat{p}{|}_{u<H})\text{\hspace{1em}(7)}$$

其中 $H$ 是掩盖不确定标签的阈值，$y_p$ 是最终的确定性伪标签。

Anchor Sampling.（锚定样本）
我们使用标记图像的标签和未标记图像的确定性伪标签作为使用对比样本的基础。

由于原始图像分辨率太大，在原始图像大小上进行对比学习的成本昂贵，并且像素的原型向量包含的语义信息较少。因此，我们在低分辨率的特征空间中使用对比学习。

首先，从编码器提取的特征将被嵌入到 $D$ 维空间中，其中每个 $D$ 维特征向量代表像素的原型向量。然后，将标签下采样到相同的分辨率，为每个原型向量指定类别，并且不对不确定区域中的向量进行采样。我们采用随机抽样策略，每个类别都有固定数量的样本。如果同一类别的样本数量较少，我们将从其他类别中采样anchor。对比负样本的数量极大地影响了对比学习的性能，但是大量的负样本会产生很大的开销。更好的解决方案是使用固定大小的外部存储来存储采样的样本，并随着训练更新存储内容。在我们的方法中，我们设置了一个内存队列来存储收集到的样本。在每次迭代中，随机选择的样本被用作anchor来计算对比损失，然后将它们更新到内存队列中。

Pixel Contrastive Loss.（像素对比损失）
原型向量及其像素类别保存在样本队列中。我们使用流行的 InfoNCE [van den Oord et al, 2018] 损失函数来计算对比损失。在每次迭代中，我们随机采样 $M$ 个锚并计算每个锚的对比损失。然后将所有锚点的损失平均作为整体对比损失。具体计算如下：

$$l_c^i=-\frac{1}{|P_i|}\sum _{v_i^{+}\in P_i}log\frac{e^{cos(v_i,v_i^{+})/\tau }}{e^{cos(v_i,v_i^{+})/\tau }+{\sum }_{v_i^{-}\in N_i}{e}^{cos(v_i,v_i^{-})/\tau }}\text{\hspace{1em}(8)}$$

$$l=\frac{1}{M}\sum _{i=1}^M{l}_c^i\text{\hspace{1em}(9)}$$

其中 $P_i$ 和 $N_i$ 表示像素 $i$ 的正样本和负样本的原型向量集合。 $v_i$ 是像素 $i$ 的原型向量，$v_i^{+}$ 是正原型向量，$v_i^{-}$ 是负向量，$\tau$ 是温度超参数。

3.3 Equivariant Contrastive Loss

为了进行对比学习，之前的一些工作通过同一图像的不同变换来构造正样本。然而，有些变换不符合分割任务的先验知识[Dangovski et al, 2021]，例如几何变换。在本文中，我们建议分割任务所需的有效特征表示对于不同的几何变换应该是等变的（或有区别的）。

基于上述，我们考虑将等变对比损失添加到分割模型的表示学习中来学习全局信息。具体来说，我们将分割模型定义为编码器-解码器形式：$f(x_i)=f_{\gamma }(f_{\theta }(x_i))$。对于一幅图像 $x_i$ ，当它经过某种几何变换 $G(\cdot )$ 时，相应的分割结果也会发生变化，即：

$$f(G(x_i))=G(f(x_i)))\text{\hspace{1em}(10)}$$

那么，我们可以推论：

$$f_{\theta }(G(x_i))\ne f_{\theta }(x_i)\text{\hspace{1em}(11)}$$

因此，我们可以明确地强化 $f_{\theta }(\cdot )$ 中这种几何变换信息的学习。我们添加一个分类预测器 $p_{\varphi }(\cdot )$ 来预测几何变换的判别结果。我们的等变对比损失函数如下：

$$l_e=\frac{1}{C}\sum _{i=0}^{C-1}{\mathcal{L}}_{CE}(p_{\varphi }(f_{\theta }(G^i(x))),i)\text{\hspace{1em}(12)}$$

其中几何变换 $G^i(\cdot )$ 在本文中表示四重旋转，所以$C=4$。

4 Experiments

4.1 Experimental Setup

Datasets and Metrics.（数据集和指标）
我们在两个公共数据集上验证了所提出的方法：

• ACDC 数据集 [Bernard et al, 2018] 包含来自 100 名患者的 200 张带注释的短轴心脏 MR 电影图像。我们将数据集按照7:3的比例进行划分，得到训练集和验证集。根据不同的半监督实验设置，分别对训练集中7名患者的136张图像和3名患者的68张图像进行标记。
• ISIC 数据集 [Codella et al, 2018] 包括 2594 个皮肤镜图像，我们使用 1815 个图像进行训练，779 个图像进行验证。在训练集中，5% (91) 和 10% (181) 图像被标记为不同的半监督实验设置。

两个数据集中的所有图像都调整为 224 × 224 以满足所提出方法的输入要求。我们使用标准数据增强来扩大训练集，包括随机裁剪、随机旋转、随机翻转和颜色抖动。在我们的方法中，将记录变换的类别来计算 $l_e$。为了评估我们方法的性能，我们选择Dice Coeffcient（记为Dice）和Jaccard Index（记为Jaccard）作为评估指标。

Implementation Detail.（实施细节）
为了公平比较，实验中使用的所有方法都选择 UNet 作为图像分割的基准架构。我们使用 ResNet-50 替换 UNet 的编码器部分，并使用 ImageNet 上预训练的权重初始化其参数。我们采用 SGD 作为优化器，权重衰减为 0.0005，动量为 0.9。初始学习率设置为 0.01，在训练过程中通过多项式调度策略将其降低到 0.001。我们使用 PyTorch 库实现这些方法，并在 NVIDIA RTX 2080Ti GPU 上训练它们。批量大小设置为 16，其中标记了 8 张图像。所有方法在训练期间都会执行 6000 次迭代。

4.2 Quantitative Comparison

Compared Methods.（比较方法）
我们将我们的方法与一些最近的半监督分割方法进行比较，包括：Meat-Teacher (MT) [Tarvainen and Valpola, 2017]、Entropy Minimization (EM) [Vu et al, 2019]、Uncertainty-Aware Mean Teacher (UA-MT)) [Yu et al, 2019]、深度协同训练 (DCT) [Qiao et al, 2018]、交叉一致性训练 (CCT) [Ouali et al, 2020] 和交叉伪监督 (CPS) [Chen et al, 2021a]。对于所有的比较方法，我们都采用官方的超参数设置。

Main Results.（主要结果）

表 1 显示了我们在ACDC和ISIC数据集上的定量比较实验结果。第一行表示仅使用标记数据训练的基线模型的性能。与基线模型相比，我们的方法可以有效地利用未标记的数据来实现出色的性能。在不同的数据集和不同的半监督设置中，我们提出的方法明显优于比较方法。特别是，当 ACDC 数据集中仅使用 68 个标记图像时，与其他方法相比，我们的方法将 Dice 提高了 3% 以上。

Visual Comparisons.（视觉比较）
图 3 显示了在 ACDC 数据集上使用 136 个标记图像时不同方法之间的一些视觉比较。我们选择了实验中表现较好的MT和CPS两种方法进行比较。与MT和CPS相比，我们的方法具有更好的预测结果和更少的错误预测。

4.3 Ablation Study

          

表 2 显示了我们的方法在具有 136 个标记图像的 ACDC 数据集上的消融实验结果。我们选择仅使用标记数据进行监督训练的 UNet 模型作为基线（第一行），并逐渐增加所提出的组件以证明其有效性。

此外，我们还添加了两个额外的比较设置，包括1）仅使用标记数据进行对比学习（$l_c^{sup}$）和 2）使用伪标签进行对比学习（$l_c^{pseudo}$）以证明我们提出的不确定性引导对比的有效性学习方法。实验结果表明，我们提出的方法的每个部分都具有积极的影响。对比学习的引入有效地建立了像素之间的关系，提高了模型的性能。与伪标签方案相比，该方法充分利用了未标签数据，因此带来了显着的性能提升（Dice 提高了约 1%）。

Contrast on Different Feature Scales.（不同特征尺度的对比）

          

不同特征的分辨率对对比学习样本的选择有重要影响。为了找到合适的特征尺度，我们探索了不同尺度对 ACDC 数据集进行对比学习的影响。表 3显示了四种不同特征尺度下的对比学习结果。正如我们所看到的，低分辨率（Conv4）下的对比学习性能很差，这可能是由于标签下采样的语义不一致造成的。更高分辨率的特征（Conv1）也会带来性能下降。潜在的原因是高分辨率像素向量包含较少的语义信息。因此，我们认为使用中级特征进行对比学习可以带来更好的分割性能。

Visualization of Features.（特征的可视化）
在图 4中，我们使用t-SNE算法来降低像素特征的维度以进行可视化。从左到右，它们是没有对比学习、使用伪标签的对比学习和所提出的方法的训练结果。与第一种方法相比，该方法可以使像素表示具有更好的类内紧凑性和类间可分离性，这表明了对比学习对于分割任务的有效性。与第二种方法相比，我们的方法具有更好的聚合结果，其潜在原因是我们的方法可以减少噪声采样的可能性。

Conclusion

我们提出了一种用于半监督医学图像分割的不确定性引导像素对比学习方法，该方法利用不确定性来解决像素对比学习中未标记数据的噪声采样问题。为了估计不确定性，基于CNN和Transformer的解码器精心设计了异构一致性学习策略。此外，我们构建了一个等变对比损失来增强我们模型的全局表示学习能力。大量的实验证明我们的方法可以实现最先进的性能。