【论文笔记】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$ 。然后把 moving image $x$ 和隐含变量 $z$ 输入到空间转换网络（STN）中产生最终的配准后的图像。

$F_x^i,F_y^i)$ 表示两个CNN第 $i$ 层产生的特征图， $F_z^i$ 表示概率模型产生的第 $i$ 层的隐含变量，它实际是让 $F_x^i$ 对齐到 $F_y^i$ 的 STN 的参数，或者说形变场。在已知 $F_x^i,F_y^i$ 的情况下，可以通过最大化后验概率 $p(F_z^i|F_x^i;F_y^i)$ 来得到最优的 $F_z^i$ 。具体的，引入了一个近似后验概率 $q_\psi(F_z^i|F_x^i;F_y^i)$ ，然后最小化 $p(F_z^i|F_x^i;F_y^i)$ 和 $q_\psi(F_z^i|F_x^i;F_y^i)$ 之间的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}$
其中 $q_\psi(F_z^i|F_x^i;F_y^i)$ 来自于多元正态分布：
$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)$
其中 $\mu_{F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}}$ 是分布的均值， $\sigma_{F_{z}^{i} | F_{x}^{i}, F_{y}^{i}}^{2}$ 是分布的方差，它们是通过概率模型得到的（如图1(b)）。

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

$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)$

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

在CNN浅层的特征图具有较高的分辨率并且具有丰富的细节信息，而CNN深层的特征图具有较低的分辨率并且具有高层次的语义信息。高层语义信息可以帮助全局配准，但是忽略了很多细节。而细节信息则是捕获了局部的配准信息。所以将浅层到深层的特征图 $F_z^i$ 混合得到最终的隐含变量 $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)$
其中， $\mathcal{L}(z;x,y)$ 表示从输入图像 $x$ 和 $y$ 到输出的配准后的图像 $z$ 的KL散度， $\mathcal{L}\left(F_{z}^{i} ; F_{x}^{i}, F_{y}^{i}\right)$ 是从输入特征图 $F_x^i$ 和 $F_y^i$ 到输出配准转换变量 $F_z^i$ 的KL散度。 $n$ 是CNN的层数， $w_i$ 是第 $i$ 层损失的权重。通常设置 $n=4,w_i=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]$
其中第一项是使得配准后的图像 $X\circ\phi_Z$ 与图像 $Y$ 相似的重建损失，第二项是公式1第一项的近似，它可以让 $q_\psi(Z|X;Y)$ 与 $p (Z)$ 相似； $\mu_{Z|X;Y}$ 和 $\sigma_{Z|X;Y}$ 分别是分布 $q_\psi(Z|X;Y)$ 的均值和标准差。

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

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

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

猜你喜欢