2.翻译和笔记Multi-Task Attention-Based Semi-Supervised Learning for Medical Image Segmentation

派大星的最爱海绵宝宝

已于 2022-10-29 17:34:13 修改

阅读量1.1k

点赞数

分类专栏：论文学习文章标签：深度学习计算机视觉人工智能

于 2022-09-30 11:52:38 首次发布

原文链接：https://paperswithcode.com/paper/multi-task-attention-based-semi-supervised

版权

论文学习专栏收录该内容

10 篇文章 0 订阅

订阅专栏

翻译

原文地址https://paperswithcode.com/paper/multi-task-attention-based-semi-supervised

摘要

我们提出了一种新型的半监督图像分割方法，它可以同时优化监督分割和无监督重建目标。重建目标使用了一种注意机制，该机制可以把不同类别对应的图像区域的重建分离开。在两个应用上对所提出的方法进行评估：大脑肿瘤和白质高强度分割。我们的模型在无标签和少量的有标签图像上训练，性能超过了在相同数量图像上训练的监督CNNs和在无标签数据上预训练的CNNs。在消融实验中，可以观察到我们所提出的注意力机制大大提高了分割性能。我们探讨了两种多任务训练策略：联合训练和交替训练。与联合训练相比，交替训练需要更少的超参数并且可达到更好、更稳定的性能。最后，我们分析了不同方法学习到的特征，并且发现注意力机制有助于学习在编码器深层次中的可判别特征(discriminative features)。
关键字：半监督学习、多任务学习、注意力、深度学习、分割、大脑肿瘤、白质高强度

1.介绍

半监督学习(SSL)使用无标签数据提高监督模型的泛化性能。这在医学图像分割中很有用，因为手动注释的生成昂贵又繁琐并且只能用于训练数据的一小部分。
半监督学习的一种方法是多任务学习，在这种学习中，除了额外使用有标签数据的目标（target objective）外，不需要使用有手动标记数据的辅助物体来训练网络。这可以通过包含额外自动编码器的目标实现，并且已经用于图像分类（例如引用[3，8]）。Sedai等人（[11]）首次将变分自动编码器引入到半监督分割中，并通过从其他预训练的重建自动编码器（reconstruction autoencoder）中学习到的编码嵌入（encoded embeddings）和重建分割掩码来训练分割自动编码器。但是，带有图像重建的多任务学习与主流图像分割架构（如U-Net及其变体[2,6]）结合并不容易，后者使用跳跃连接（skip-connection）来保存早期编码器层的高分辨信息。这些跳跃连接并不适合与自动编码器结合作为辅助任务（auxiliary task），因为它们允许网络从早期层中复制信息并且跳过自动编码器中的降维操作（dimensionality reduction）。
另一种半监督方法是为未标记的训练数据创建新的伪标签，例如自我训练[13]和co-training[12,15]，以获得更多可用的训练资源。但是，所创建的伪标签通常与用于目标分割（target segmentation objective）的ground truth的质量不一样，后者限制了它们从未标签数据中改进的潜力。
我们所提出的新型半监督方法叫基于注意的多任务半监督学习（MASSL），其中我们把自动编码器与类似于U-Net（U-Net-like）的网络结合了起来。我们不是训练它重建原始输入[11]，而是训练自动编码器重建由注意力机制产生的合成分割标签。这使得我们的模型可以从无标签图像中学习用于分割的可判别特征（discriminative features）。尽管注意力经常被应用于监督学习（例如[10]），但是据我们所知，注意力从未与半监督学习相结合过。我们的方法与自我训练[13]和co-training[12,15]有相似之处，后者可以动态的为无标签训练数据创建新标签。与这些方法相比，我们的模型为重建任务创建标签。这使得无监督辅助任务可以从无标签数据中学习到比从传统重建任务更具辨别力的潜在特征，后者不考虑类别差异。
我们的贡献概述如下。首先，我们提出了一个新的多任务半监督学习方法，并且研究了与两种训练策略结合后的性能。第二，我们在两种分割问题（大脑肿瘤和白质高强度）上评估了我们的方法，表明该模型性能超过了全监督CNN基线（fully-supervised CNN baseline）、两个预训练方法和没有所提出注意力机制的多任务学习。第三，我们研究了注意力机制是如何影响由编码器学习到的特征的，并且表明这有助于更深的层次学习更有辨别力的特征。

图一

在这里插入图片描述
图一：所提出的MASSL（基于注意的多任务半监督学习）框架。分割损失L1由soft分割预测和正确标记的数据（ground truth）计算。重建损失L2是由重建前景、背景预测和使用注意力机制产生的新标签计算。

2.方法

我们的半监督学习方法如图一所示。方法由共享同一个编码器的分割和重建网络，以及连接两个任务的注意力组成。

2.1架构和损失函数

类似于U-Net[9]，分割CNN架构具有跳跃连接，允许把来自于编码器更浅层次的精细细节（fine details）传输到解码器，并仅在标记图像上使用 Dice目标L1进行训练。
在我们模型的基线版本（baseline version）中，优化后的重建网络的输出用来预测输出图像。在本文的剩余部分中，我们把这种方法称为多任务SSL（MSSL）。
在我们方法的基于注意力版本中，我们称之为基于注意力的多任务SSL（MASSL），我们分别重建了图像的背景和前景部分，由从分割网络中获得的soft预测y~定义。前景和背景目标由各自的分割掩膜（msak）大小进行加权：
在这里插入图片描述
对于背景（b）和前景（f），y^和y~分别为重建和分割路径的预测；n是输入图像x的体素数（voxels）；⊙是元素乘积。值得注意的是，梯度不会通过y~传播到分割解码器。我们假设，将重建标签和分割预测相结合会导致在编码器的更深层次学到很好的特征，从而更好地分割。对客观条件进行加权（The objective terns are weighted），以防止过度强调前景重建的重要性。

2.2训练策略

MSSL和MASSL网络的两个任务可以联合优化或交替优化：
联合训练（joint training）：给定一个包含相同数量的小批次（minibatch）的标记样本xL和未标记样本xU，首次使用最新的分割网络参数对未标记样本xU进行分割，为重建任务创建前景和背景图像。其次，通过优化分割任务和重建任务的目标函数，更新整个网络的权重（weights）。损失Loss是由超参数γ∈[0,1]控制的分割和重建损失的线性组合：
在这里插入图片描述
交替训练（Alternating training）：对于每一个epoch，从相应的训练集中随机抽样相同数量（标记图像和未标记图像的较小数量）的标记图像和未标记图像。一个小批次（minibatch）包含标记样本xL或相同数量的未标记样本xU。这种两类型的batch在训练期间交替使用。分隔路径和重建路径的权重根据给定的batch类型和相应的损失loss分别更新：
在这里插入图片描述

3.实验

数据我们使用了BraTS 2018 挑战[5,1]和White Matter Hyperintensities 2017 挑战（ https://wmh.isi.uu.nl/）的公开数据。
BraTS18： 将220例带有高级别胶质瘤（high grade glioma）患者的MRI扫描随机分为120，50，50个扫描，分别用于训练，验证和测试，采用5倍蒙特卡罗交叉验证（with 5-fold Monte Carlo cross-validation）。为了简化不同分割任务之间的对比，我们采用了二元分类（binary classification），且只分割整个肿瘤，包括4种肿瘤类型，并且只使用FLAIR序列（FLAIR sequence）。
WMH17： 60例FLAIR MRI扫描提供了相应白质高强度（WMH）的手动分割。扫描需要3个site进行，每个site20次。在我们的实验中，我们使用30个扫描用于训练，10个用于验证和20个用于测试，确保三set中每个site的数量大致相同。我们采用5倍蒙特卡罗交叉验证。
网络和超参数 网络布局如图1所示。我们的网络受到UNet[9]结构的启发，但是有一些差异。网络的输入大小是128×128×32。在编码器和每个解码器中都会有5分辨率levels（5 resolution levels）。每个level包含2个3×3×3的的卷积层，使用0padding，instance normalization[14]和LeakyReLU激活函数，除了最后一层使用sigmod激活函数进行最后预测值的两个解码器。在每个level中间都有一个平均池化层/上采样层（average pooling/upsampling layer）。在第一个level中，特征通道（feature channel）的数量是16个，每次池化/上采样后都会加倍/减半，在最深的一个level中最大特征为256。特征上采样路径的特征映射通过跳跃连接与前面的特征映射接进行连接。重建网络与分割网络的结构除了没有跳跃连接外都相同。对于联合训练（joint training），我们使用一个Adam优化器优化公式2中的损失。对于交替训练（alternating training），我们使用两个单独的Adam优化器分别优化等式3中两种类型的损失。基于在验证集上的性能，我们把分割任务和重建任务的初始学习率分别设置为0.01和0.001。数据增强应用的是随机旋转，缩放和水平翻转。
特征分析 在每一个编码器level的最后一个层中，我们使用线性回归分析来评估特征区分前景与背景区域的能力。我们把每个体素（voxel）看作一个单独的样本，把它在每个特征映射（feature map）中的值作为回归变量。通过使用二元分割ground truth，然后使用平均池化下采样到所需要的分辨率（resolution），从而得到每个体素（voxel）的标签。

4.结果

分割结果如表一和表二所示。对于半监督的设置（前两个colomns），标记数据与未标记数据之间没有重叠。对于全监督设置（最后column），所有的图像被用作标记数据和未标记数据。对于Pretrain（Dec），我们先用未标记数据对重建网络进行预训练，然后再使用带标记数据对分割网络中的解码器路径进行训练，同时编码器部分保持固定，以确保分割任务仅使用从未标记图像学习到的特征。对于Pretrain（CNN），我们首先用未标记数据对重建网络进行预训练，然后再使用标记数据训练整个分割网络，必要的时候允许网络对编码器参数进行微调。MASSL和MSSL是所提出的带有注意力机制和不带注意力机制的多任务SSL方法，其中γ和alter分别表示联合训练（joint training）和交替训练（alter training）。对于联合训练（joint training），我们尝试γ=0.5，0.7，0.9，当γ=0.5时网络没有收敛。结果表明，在所有的方法中，MASSL（alter）达到了最好的分割性能。联合训练（joint training）策略的性能略低于交替训练（alter training），但在不同的标记/未标记数据分割之间差异很大，反映了联合训练（joint training）策略的不稳定性和调整γ的难度。

表一

在这里插入图片描述
表一 BraTS18结果
Dice similarity coefficient，所有交叉验证的平均值。最后一列除了CNN只使用带标签图像，所有标记图像也作为无标签图像使用。*：明显优于CNN（p<0.05）。：明显比MASSL（alter）要差（p<0.05）。P-值是由每一列的双面检验t（two-sided t-test）计算。

表二

在这里插入图片描述
表二 WMH17结果
Dice similarity coefficient，所有交叉验证的平均值。最后一列除了CNN只使用带标签图像，所有标记图像也作为无标签图像使用。*：明显优于CNN（p<0.05）。：明显比MASSL（alter）要差（p<0.05）。P-值是由每一列的双面检验t（two-sided t-test）计算。

特征分析的结果如表三所示。较高的R**2的值表明，在更深的level中，由MASSL学到的特征比CNN和MSSL的更有判别力（discriminative）。这支撑了我们的假设，即注意力机制可以使编码器的更深层次学到更有判别力的特征，同时还可以优化重建目标。

表三

在这里插入图片描述
表三编码器特征的判别能力
使用BRATS数据中5个folds的训练模型，50个标记数据/70未标记数据拆分。因为早期特征图的大小限制，从测试集中随机选择5个训练/测试数据并且用于所有模型。实验以不同的随机数据重复了5遍，报告了5次实验对所有5folds模型的平均R**2分数（方差）。

5.讨论和结论

在论文中，我们提出了一个叫MASSL的新的半监督学习方法，它通过多任务学习网络中的注意力机制将分割任务和重建任务结合起来。所提出的方法在两个应用上评估。对于两种应用，使用部分标记图像的MASSL优于使用相同数量标记图像、预训练＋微调方法和没有注意力的所提方法（MSSL）的全监督CNN基线（baseline）。当对所有图像使用分割和重建损失时，MASSL比baselineCNN有所改善，尽管这个差异仅对BRATS数据有统计意义。这主要是因为WMH数据中前景的稀疏分布，使得我们的注意力图（attention map）效率降低。
我们方法的改进主要来自于注意力机制，它把分割任务引入到重建任务中，并且把它们连接的比以前更好了。这个机制可以轻易的集合到任意的CNN结构并且可以推广到多类分割中。与联合训练（joint training）相比，交替训练（alternating training）是一种实用的策略，它允许学习率发生任务依赖的（task-dependent）的变化，且不需要微调λ，尽管它仍旧需要选择一个恰当的初始学习率。交替训练（alternating training）不能保证稳定，因为编码器的参数在两个人物之间不断变化。在实验期间，我们发现，为重建选一个比分割较小的初始学习率时，训练比较稳定，在多数情况下，交替优化（ alternating optimization）的性能远比联合优化（ joint optimization）的好。
与不同的多任务学习策略做比较时，我们做了一些简化。对于预训练方法，与Sedai等人不同[11]，本文中我们使用了常规（regular）自动编码器而不是变分自动编码器（VAE）。我们认为我们的SSL方法可以与VAE工作的很好并且能够更好的融合两个任务。在回归分析中我们使用了一个简单的回归模型，它只可以显示特征的线性判别（discriminative）能力。使用一个更复杂的非线性模型去展示非线性的有判别能力也会是一个很有趣的事。由于我们仅使用了一个MRI序列和一个扫描子集（ a subset of scans），我们在BraTS18和WMH17的性能比最新的技术要低。BraTS18和WMH7在测试集上的Dice性能分别为0.8839[7]和0.80[9]，并且首个工作也使用了变分编码器来提供更多与Ladder网络[8]和我们MSSL方法相似的正则效果。
综上所述，MASSL是一个很有前景的分割框架，对于简单高效的多任务学习，它能够在半监督和全监督的环境下实现强大的改进。