【论文学习】Non-local U-Nets for Biomedical Image Segmentation

组会要分享的文章，顺便翻译一下

1.摘要

深度学习在不同的生物图像分割上表现出色。现存的模型基本上是在U-Net基础上进行的改进，并且采用encoder-decoder的架构，这种架构通过堆积local operators来逐渐汇集长范围的信息。然而，使用local operators限制了效率和其有效性。在我们的工作中，我们提出了non-local U-Nets，采用了具有弹性的global aggregation blocks来进行生物图像分割。这些block可以被插入进Unet网络来进行size-preserving 的过程，下采样层和上采样层都可以用到。我们通过在3d多模态同强度的婴儿脑部MR图像分割任务上来评估我们的non-local Unet网络。结果显示，我们的方法相比于其他方法性能更高，参数更少且计算更快。

2.介绍

近年来，深度学习的方法，比如FCN,U-Net，Deeplab，RefineNet都不断的创建了图像分割上的新纪录。比方说residual deconvolutional network和residual symmetric Unet都采用了基于Unet的网络模型并加入了额外的short-range residual connections来进行2d电子显微镜图像的分割。除此之外，Unet从2D图像拓展到了3D立体图像，产生了像3D-Unet，V-Net，3D-FCN等网络模型。
除了这些研究的成功以外，我们对这些基于Unet的网络进行了研究，发现两点不足：
1.这些网络的encoder部分通常采用交替堆积size-preserving卷积层和下采样层的方式，来逐渐减少feature map的空间尺寸。卷积层和下采样层都是传统的local operators，采用小卷积核来扫描输入图像并且提取局部信息。将他们像瀑布一样堆积在一起来产生一个大而有效的卷积核，使其能够汇集长范围信息。但医学图像分割通常需要很大程度的上下文信息，之前的很多模型都有一个很深的encoder部分，也就是堆积了很多local operators。这将导致大量的训练参数，导致降低模型的效率，特别是在采用大量的下采样操作时，因为每次下采样操作过后，feature map的数量通常会加倍。除此之外，更多的下采样会导致在encoding的过程中损失更多的对医学图像来说十分重要的空间信息。
2.decoder部分和encoder部分采用同样的方式，只是用上采样层来代替了下采样层。常用的上采样层，像deconvolution和unpooling层也都是局部操作。然而上采样部分需要恢复图像的全局信息，如果不考虑 全局信息就会很难。
总的来说，对于基于unet的模型，采用能够进行non-local信息汇聚的operator将会提高它们的效率和效力。正如unet既有保持尺寸不变的卷积层，也有上采样和下采样层，这个新的operator也应该能够适应这些情况。
在本文中我们提出了unet的两种缺点和用于医学图像分割的non-local unet网络。为了解决第一个问题，我们提出了基于自注意力机制的global aggregation block，它能够不通过深层次的encoder来汇聚全局信息。这个block被拓展应用到了up-sampling global aggregation block，这将减轻第二个问题。据我们所知，我们是第一个做出这种延伸的。我们探索了这些具有弹性的global aggregation blocks在unet上的应用。实验证明我们提出的non-local unet表现最佳并且参数更少，计算更快。

3. Non-local U-Net

我们首先介绍了我们所提出来的unet网络架构。在此基础上我们的模型由不同的size-preserving，down-sampling和up-sampling block组成。我们描述了每个block并且提出我们的global aggregation blocks来组成non-local unets。

U-Net Framework

我们通过3d的案例来阐释我们的non-local unets。低维或高维的可以很容易的由此衍生出来。结构如图所示：
输入将会首先通过input block，将会提取出低维特征。两个下采样模块将会用来减少空间尺寸并且用来提取高维特征。在每次下采样过后通道的数量加倍，然后一个bottom block用来汇聚全局信息并且产生encoder的输出。相对应的，decoder用两个上采样block来恢复空间信息，每次上采样过后，通道数量减半。
上采样部分，通过skip connection来将encoder部分的下采样的feature map copy到decoder部分。与其他模型不同的是，这里的copied feature map通过加法与decoder部分的feature map相结合，而不是concatenation。采用加法有两点好处：
1.首先加法不会增加feature map的数量，这样就能够减少训练的参数。
2.其次，加法的skip connection可以看作是长范围的残差连接，可以促进网络的训练。
得到decoder部分的输出，output block产生分割的概率图。具体来说，对于每个体素，输出每个体素对应的分割类别。最终的分割图可以通过在probability map上进行argmax操作获得。

Residual Blocks

残差连接被证明可以促进深度学习模型的训练并表现得更好。我们unet网络中使用的加法连接相当于长范围的残差链接。为了更好地提升unet网络，Lee提出了加入短程残差链接的想法，然而这些想法并没有被应用于下采样和上采样块的残差链接。带有残差链接的下采样模块在ResNet里面已经有所提及。我们研究了基于我们提出的up-sampling global aggregation block的up-sampling block。
在我们提出的模型中，用四种不同的残差块来形成一个完整的残差网络。如图所示。

值得注意的是，他们都采用了pre-activation的方式。
（a）展示了一个常见的有着两个连续卷积层的residual block。在每个卷积层之前使用了bn和relu6 激活函数。这个模块在我们的网络中作为input block。output block则是在该模块后面加入了一个111，stride为1的卷积层。此外，我们在每个加法的skip connection过后也加入一个这个模块。（b）是一个下采样的残差模块。使用一个步长为2的111的卷积来代替原先的残差链接，是为了调整feature map的空间尺寸。(c )描述了我们的bottom block 。主要采用global aggregation block。（d）和下采样模块(b)相似，只不过换成了步长为2的333的deconvolution，另外一个分支换成了up-sampling global aggregation block。我们将这个模块作为up-sampling block。

Global Aggregation Block

为了融合全局信息，输出的feature map的每个位置，都应该取决于输入的feature map的每个位置。这样的操作和像卷积和反卷积一样的局部操作不同，它们每个输出的位置都在输入上有一个局部的感受野。实际上，全连接网络有这样的全局性质，然而这很容易过拟合并且实际中表现并不好。Transformer（论文）通过考虑输入的每一个位置来计算输出的一个位置。后来也有研究提出了non-local networks 用于视频识别。两项研究都采用self-attention blocks来捕捉序列的长期依赖信息，因此我们指出图像feature map的全局信息可以通过self-attention blocks来进行汇集。
基于这个观点，我们提出了global aggression block可以从任意尺寸的feature map来获取全局信息。我们将其推广到下采样和上采样的过程中，让它变成一个在深度学习模型任意位置都可以使用的block。
用X来表示输入，Y表示输出。用Conv_1N来表示步长为1，通道数为N的111的卷积，它不会改变空间大小，第一步是生成Q(query),K(key),V(value)三个矩阵。

Unfold把DHWC的tensor变为(DHW)C的矩阵
QueryTransformCk是产生Ck个feature map的任意操作，Ck，Cv是表示keys和values维度的超参数。假设X的维度是DHWC.那么K和V的维度就是（DHW)Ck和（DHW)Cv。Q的维度是 (DQ ×HQ ×WQ)×CK,DQ,HQ,WQ 都取决于 QueryTransform。
Q,K,V的每一行都是一个query vector，key vector和value vector。query vector的维度和key vector相同。同时，key vector的数量和value vector的数量相同，呈现一个一对一的关系。
第二步，在Q,K,V上面采用注意力机制，定义为

注意力权重矩阵 A 的维度是 (DQ × HQ × WQ) × (D × H × W) 输出矩阵 O 的维度是 (DQ × HQ × WQ) × CV . 为了展示它是如何运作的，我们从Q里面选择一个query vector来举例。在注意力机制中，query vector和所有的key vector相互作用，通过点乘的方式来产生value vector所对应的scalar weight。query vector的输出是一个有权重的value vector的总和，通过softmax来进行归一化。对所有的query vector都采用这个过程，最后生成 (DQ ×HQ ×WQ) CV 的向量。注意图中的dropout可以被用在A上来防止过拟合。
最后一步，
Fold的Unfold相对的操作，Co代表输出的维度。最终，Y的输出是 DQ×HQ×WQ×CO。值得注意的是，输出Y的空间大小是由Q矩阵来决定的，也就是通过QueryTransformCK这个操作，也因此global aggregation block可以灵活的适用于size-preserving，down-sampling和up-sampling。在我们提出的non-local U-Nets中，我们令CK = CV = CO 并且采用两种不同的QueryTransformCK函数，对于(c )而言，QueryTransformCK是Conv_1Ck。对于up-sampling global aggregation block (d)而言,QueryTransformCK是一个步长为2的33*3的反卷积层。这些block的使用，缓解了通过单独的反卷积层来上采样会丢失信息的问题。通过考虑全局信息，上采样模块可以恢复更多的精确信息。

4.结果与讨论

我们在3D多模态同强度的婴儿脑图像上进行我们的实验。任务是将MR图像分割成CSF,GM,WM三种区域。我们首先介绍了我们的baseline model和实验评估方法，然后描述了训练和推理过程，我们提供了效率和效力的对比结果，并且阐释了我们网络中每个global aggregation block是如何提升我们的效果的。除此之外，我们还探索了基于不同的overlapping step size的推理速度与准确度之间的平衡，并分析了patch size的影响。

Experimental Setup(略译重点)

1.采用CC-3D-FCN作为baseline，该方法被证明优于之前的深度学习模型比如2d，3d cnns，还有deepMedic，3d-unet等网络。
2.采用Dice ratio(DR) 和3d modified Hausdorff distance来作为评价指标。这两种方法通常被应用于二分类问题，为了实现我们的四分类问题，我们需要做一些改变。也就是对于每种分类，都应该建立一个3D 二分类的分割图，1表示这个位置的体素属于这个类别，而0则相反。在我们的实验中，我们直接从4-class segmentation maps中得到我们的binary segmentation map。CSF,GM,WM的分割都是在binary segmentation map上进行。
P代表某一类的预测binary segmentation map，L表示对应的label。

|·|表示segmentation中1的数量，|P ∩L|表示P和L的交集，DR在[0,1]之间，越大越准确。
The modified Hausdorff distance (MHD)是用来计算两个样本之间的相似程度的，在这里每个样本是一个点的集合，每一个点由一个向量表示。

Training and Inference Strategies

1.每个global aggregation block 和最终的111卷积前的output block 使用0.5的dropout。
2.2e-6的weight decay
3.采用随机的cropped small patches进行训练。这样可以有足够的训练数据，并且对于内存的需求降低。不用进行额外的数据增强。
4.Adam梯度下降，0.001

Comparison with the Baseline

Ablation studies of Different Modules

结果展示了不同的global aggregation blocks是如何提高表现的。

Impact of the Overlapping Step Size

Impact of the Patch Size

Conclusion

本论文提出了用于医学图像的non-local Unets。正如本文提到的，之前的基于Unet的模型并没有一个有效的汇集全局信息的方式，而仅仅是通过堆积局部操作来进行，这限制了它们的表现。为了解决这个问题，我们提出了可以灵活应用于unet网络中 size-preserving, down-sampling ， up-sampling 过程中的 global aggregation block 。实验证明我们的non-local unet网络优于其他模型，并且参数更少，计算更快。