【论文笔记】Probabilistic Multilayer Regularization Network for Unsupervised 3D Brain Image Registration

本文是关于论文《Probabilistic Multilayer Regularization Network for Unsupervised 3D Brain Image Registration》的阅读笔记。

文章提出了一个无监督的3D脑部图像配准网络，用来捕获 fixed image 和 moving image 之间特征级（feature-level）的信息。网络包括分别对 fixed image 和 moving image 进行处理的两个深度CNN，以及一个对以上两个CNN处理结果进行对齐的特征级概率网络。这两种网络实现了不同级别的特征提取。

传统的基于模型（深度学习）的配准网络都忽视了两张输入图像之间的特征级的转换关系，CNN的隐藏层学习到的特征对于隐含变量（latent variable）来说是透明的，所以在本文中首先使用了两个CNN，一个CNN用来 fixed image 中提取特征，另一个从 moving image 中提取特征。此外，还用一个概率网络在两个CNN对应的隐藏层之间捕获它们的转换关系。此外还在CNN的多个层中嵌入了正则项，以在不同层产生特征级的隐含变量。最后，通过把在所有层中预测得到的特正级隐含变量结合，得到最终的用于配准的隐含变量。

首先使用两个CNN产生两组具有不同分辨率的特征图集合，然后用一个特征级的概率推断模型来估计特征级的隐含变量，该隐含变量表示的是在两个CNN相同层的特征图之间的转换关系。然后把每一层产生的特征图扩大到相同的大小，将它们加起来产生最终的隐含变量 $z$z。然后把 moving image $x$x 和隐含变量 $z$z 输入到空间转换网络（STN）中产生最终的配准后的图像。

$(F_x^i,F_y^i)$(Fxi​,Fyi​) 表示两个CNN第 $i$i 层产生的特征图，$F_z^i$Fzi​ 表示概率模型产生的第 $i$i 层的隐含变量， 它实际是让 $F_x^i$Fxi​ 对齐到 $F_y^i$Fyi​ 的 STN 的参数，或者说形变场。在已知 $F_x^i,F_y^i$Fxi​,Fyi​ 的情况下，可以通过最大化后验概率 $p(F_z^i|F_x^i;F_y^i)$p(Fzi​∣Fxi​;Fyi​) 来得到最优的 $F_z^i$Fzi​。具体的，引入了一个近似后验概率 $q_\psi(F_z^i|F_x^i;F_y^i)$qψ​(Fzi​∣Fxi​;Fyi​) ，然后最小化 $p(F_z^i|F_x^i;F_y^i)$p(Fzi​∣Fxi​;Fyi​) 和 $q_\psi(F_z^i|F_x^i;F_y^i)$qψ​(Fzi​∣Fxi​;Fyi​) 之间的KL散度来使得两个分布尽可能的相似，该过程可以用下式表示：
$\begin{aligned} & \min _{\psi} K L\left[q_{\psi}\left(F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}\right) \| p\left(F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}\right)\right] \\ =& \min _{\psi} K L\left[q_{\psi}\left(F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}\right) \| p\left(F_{z}^{i}\right)\right]-E_{q} \log p\left(F_{y}^{i} | F_{z}^{i} ; F_{x}^{i}\right) \end{aligned}$=​ψmin​KL[qψ​(Fzi​∣Fxi​;Fyi​)∥p(Fzi​∣Fxi​;Fyi​)]ψmin​KL[qψ​(Fzi​∣Fxi​;Fyi​)∥p(Fzi​)]−Eq​logp(Fyi​∣Fzi​;Fxi​)​
其中 $q_\psi(F_z^i|F_x^i;F_y^i)$qψ​(Fzi​∣Fxi​;Fyi​) 来自于多元正态分布：
$q_{\psi}\left(F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}\right)=\mathcal{N}\left(z ; \mu_{F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}}, \sigma_{F_{z}^{i} | F_{x}^{i}, F_{y}^{i}}^{2}\right)$qψ​(Fzi​∣Fxi​;Fyi​)=N(z;μFzi​∣Fxi​;Fyi​​,σFzi​∣Fxi​,Fyi​2​)
其中 $\mu_{F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}}$μFzi​∣Fxi​;Fyi​​ 是分布的均值，$\sigma_{F_{z}^{i} | F_{x}^{i}, F_{y}^{i}}^{2}$σFzi​∣Fxi​,Fyi​2​ 是分布的方差，它们是通过概率模型得到的（如图1(b)）。

$p(F_z^i)$p(Fzi​) 和 $p(F_z^i|F_x^i;F_y^i)$p(Fzi​∣Fxi​;Fyi​) 符合以下多元正态分布：
$p\left(F_{z}^{i}\right)=\mathcal{N}\left(F_{z}^{i} ; 0, \sigma_{F_{z}^{i}}^{2}\right)$p(Fzi​)=N(Fzi​;0,σFzi​2​)

$p\left(F_{y}^{i} | F_{z}^{i} ; F_{x}^{i}\right)=\mathcal{N}\left(F_{y}^{i} ; F_{x}^{i} \circ \phi_{F_{z}^{i}}, \sigma_{F^{i}}^{2}\right)$p(Fyi​∣Fzi​;Fxi​)=N(Fyi​;Fxi​∘ϕFzi​​,σFi2​)

其中 $\sigma_{F_z^i}^2$σFzi​2​ 是分布的方差， $F_{x}^{i} \circ \phi_{F_{z}^{i}}$Fxi​∘ϕFzi​​ 是噪音，$\sigma_{F^i}^2$σFi2​ 是噪音项的方差。

在CNN浅层的特征图具有较高的分辨率并且具有丰富的细节信息，而CNN深层的特征图具有较低的分辨率并且具有高层次的语义信息。高层语义信息可以帮助全局配准，但是忽略了很多细节。而细节信息则是捕获了局部的配准信息。所以将浅层到深层的特征图 $F_z^i$Fzi​ 混合得到最终的隐含变量 $z$z，然后输入到 STN 中，对 moving image 进行变形。

模型总的损失为：
$\mathcal{D}_{\text {total}}=\mathcal{L}(z ; x, y)+\sum_{i=1}^n w_{i} \mathcal{L}\left(F_{z}^{i} ; F_{x}^{i}, F_{y}^{i}\right)$Dtotal​=L(z;x,y)+i=1∑n​wi​L(Fzi​;Fxi​,Fyi​)
其中，$\mathcal{L}(z;x,y)$L(z;x,y) 表示从输入图像 $x$x 和 $y$y 到输出的配准后的图像 $z$z 的KL散度，$\mathcal{L}\left(F_{z}^{i} ; F_{x}^{i}, F_{y}^{i}\right)$L(Fzi​;Fxi​,Fyi​) 是从输入特征图 $F_x^i$Fxi​ 和 $F_y^i$Fyi​ 到输出配准转换变量 $F_z^i$Fzi​ 的KL散度。$n$n 是CNN的层数，$w_i$wi​ 是第 $i$i 层损失的权重。通常设置$n=4,w_i=1$n=4,wi​=1。基于KL散度的损失为：
$\mathcal{L}(Z ; X, Y)=\frac{1}{2 \sigma_{Z | X ; Y}^{2}}\left\|Y-X \circ \phi_{Z}\right\|^{2}+\frac{1}{2}\left[\operatorname{tr}\left(\sigma_{Z | X ; Y}^{2}\right)+\left\|\mu_{Z | X ; Y}\right\|-\log \operatorname{det}\left(\sigma_{Z | X ; Y}^{2}\right)\right]$L(Z;X,Y)=2σZ∣X;Y2​1​∥Y−X∘ϕZ​∥2+21​[tr(σZ∣X;Y2​)+∥∥​μZ∣X;Y​∥∥​−logdet(σZ∣X;Y2​)]
其中第一项是使得配准后的图像 $X\circ\phi_Z$X∘ϕZ​ 与图像 $Y$Y 相似的重建损失，第二项是公式1第一项的近似，它可以让 $q_\psi(Z|X;Y)$qψ​(Z∣X;Y) 与 $p(Z)$p(Z) 相似；$\mu_{Z|X;Y}$μZ∣X;Y​ 和 $\sigma_{Z|X;Y}$σZ∣X;Y​ 分别是分布 $q_\psi(Z|X;Y)$qψ​(Z∣X;Y) 的均值和标准差。

初始学习率为 $1e^{-4}$1e−4，并且周期性的减少（乘以0.1），一共有100个epoch，使用Adam优化器，优化器的第一个动量为0.9，第二个动量为0.999，衰减权重为0.0001。

下图是实验的结果对比图。