【论文笔记】双重监督的脑部图像配准网络BIRNet

本文是论文 BIRNet：Brain image registration using dual-supervised fully convolutional networks 的阅读笔记。

一、简介

BIRNet 是 Brain Image Registration Networks 的缩写，该网络采用了双重监督的方式，一个监督是 ground-truth，即真实的形变场，另一个监督是图像的差异性，直接测量两幅图像在配准之后图像灰度值之间的差异性。并且采用了以下措施来提高网络的准确率和效率：

• 使用了分层监督（hierarchical supervision），引入了分层的损失函数，以使得类 U-Net 网络的前些层（即编码器/收缩路径部分）更容易收敛；
• 间隙填充（gap filling）是在类 U-Net 网络的收缩路径和扩张路径之间插入了额外的卷积层；
• 多通道输入（multi-channel inputs）是将图像、两幅图像之间的差异图、图像的梯度图同时作为输入；
• 数据增强（data augmentation）来扩大训练样本。

上图是整个网络的结构示意图。

一些记号：形变场记作 $\varphi$，subject image（即浮动图像）记作 $S$，template image（即固定图像/模版图像）记作 $T$。

二、网络结构

1. 双重监督

该网络采用了双重监督的方式，一个监督是 ground-truth，即真实的形变场，另一个监督是图像的差异性，直接测量两幅图像在配准之后图像灰度值之间的差异性。ground-truth 可以让模型快速的学习图像的形变以及让形变场具有正则化（即光滑性），图像的差异性可以避免对估计得到的 ground-truth 的过度依赖。

上图为 BIRNet 的损失函数示意图。

因为有两个监督，所以损失函数也包括两部分，一部分是预测的形变场和 ground-truth 形变场（通过其他已有的模型或方法得到的）之间的损失函数 $los{s}_{\varphi }$，具体使用的是两个形变场之间的欧氏距离 $los{s}_{\varphi }=\frac{1}{N}||\varphi -{\varphi }_g|{|}_2^2$，N 是体素的个数。另一部分是固定图像和经过形变之后的浮动图像之间的差异损失函数 $los{s}_M$，其计算公式为 $los{s}_M=\frac{1}{N}{\sum }_u||S(u+\varphi (u))-T(u)|{|}_2^2$，其中 $u$ 是体素的坐标，$\varphi (u)$ 是体素 $u$ 的位移。

$los{s}_{\varphi }$ 的取值范围为 [-30, 30]，而 $los{s}_M$ 的取值范围为 [-255, 255]，所以为了让两个损失函数保持平衡，在计算时 $los{s}_M$ 会先乘以 0.1，总的损失函数是两者的结合，可以表示为：
$$loss=\alpha \cdot los{s}_{\varphi }+\beta \cdot los{s}_M,\alpha ,\beta \ge 0,\alpha +\beta =1$$
在训练的初试阶段，为了加速训练的收敛并让预测的形变场光滑，可以采用较大 $\alpha$ 值，也就是多从 ground-truth 中学习，在训练后期，为了对预测的形变场进行微调和完善，可以采用较大的 $\beta$ 值，也就是多从两幅图像之间的差异性中学习。在具体操作中，在训练的前 5 个 epochs 会采用 $\alpha =0.8,\beta =0.2$ 的设置，在训练的后 5 个 epochs 会采用 $\alpha =0.5,\beta =0.5$ 的设定。

我个人的理解是如果只用 ground-truth 作为监督标签，那么训练结果的上限就是 ground-truth 了，所以还要设定其他的标签作为监督的方式，以突破 ground-truth 的局限。

2. 分层监督

在传统的 U-Net 中，由于在网络最终得到输出之后才计算损失函数，然后反向传播更新网络参数，所以网络的前半部分比后半部分收敛的要慢。我个人的理解是网络太深了之后容易导致梯度消失，所以反向传播时更新不了那么深或者深层的（前面的）网络参数更新的幅度要小，所以收敛的慢。

为了解决这个问题，在网络的每一个 level 都加入了个一个损失函数，用来直接监督网络前半部分的训练。网络有三层，高分辨率层（最上面一层）的输出是 $24\times 24\times 24$ 大小的形变场块（patch）${\varphi }_g^{high}$，中分辨率层（中间一层）的输出是 $14\times 14\times 14$ 大小的形变场块（patch）${\varphi }_g^{mid}$，低分辨率层（最下面一层）的输出是 $9\times 9\times 9$ 大小的形变场块（patch）${\varphi }_g^{low}$。最后总的损失函数 $los{s}_{\varphi }$ 为：
$$los{s}_{\varphi }=los{s}_{\varphi }^{high}+los{s}_{\varphi }^{mid}+los{s}_{\varphi }^{low}$$

3. 间隙填充

在 U-Net 的收缩路径中使用的是两个 $3\times 3\times 3$ 的卷积层，后面跟着 ReLU 激活函数和 batch normalization，然后是一个 $2\times 2\times 2$ 的最大池化层，其步长为2。在扩张路径使用的是一个 $2\times 2\times 2$ 是反卷积层，后面是两个 $3\times 3\times 3$ 的卷积层，最后一层是一个 $1\times 1\times 1$ 的卷积层用来实现全连接。

上图为间隙填充示意图，网络黑色的部分是传统的 U-Net，绿色的部分是间隙填充的卷积操作，个人感觉特征图 A 的虚线少了一个从 Conv 指向图片的箭头。

特征图 A 和 B 分别是收缩路径和扩张路径的中间结果，两者存着着较大的差异，前者像原图，后者像形变场，并且特征图 B 有点不连续。为了解决这个问题，文章在收缩路径和扩张路径的同一 level 的层之间加入了额外的卷积层。进行过间隙填充得到的特征图 C 就和特征图 B 比较相像了。

4. 多通道输入

输入包括原图、两幅图像的差异图和梯度图。两幅图像的差异图是通过计算图像 $S$ 和 $T$ 之间的差异得到的，梯度图提供了图像的边界信息。上图分别是原图、两幅图像的差异图和梯度图。

5. 其他设置

优化器采用的是 Adam 优化器，学习率在训练的前期为 0.001，在训练的后期为 1e-8，输入的 patch 之间是有重叠的，其步长为 24，和输出的 patch 大小相同，这样就保证了输出正好组成一个完整的没重叠的形变场。

三、总结

上图是几种模型的对比，黄色框中的部分是几种方法差异比较明显的部分，可以发现 BIRNet 达到了很好的配准效果。

上图是网络模型的具体配置。