【论文笔记】Prob-VoxelMorph：保证形变场微分同胚性并防止形变场重叠的医学图像配准模型

本文是论文《Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration》的阅读笔记。涉及数学的东西太多，没搞懂是怎么做的。

一、微分同胚

现有的基于学习的配准方法通常不能保证配准是微分同胚的，即保留拓扑性的（topology-preserving）。

本文在VoxelMorph模型的基础上加入了微分同胚积分层（diffeomorphic integration layer），以保证无监督的端到端的学习是微分同胚的。微分同胚是可微和可逆的，因此保留了拓扑性。形变场是通过下述常微分方程（OED）来定义的：
$$\frac{\partial {\varphi }^{(t)}}{\partial t}=\mathbf{v}\left({\varphi }^{(t)}\right)$$
其中，${\varphi }^{(0)}=Id$是恒等变换，$t$是时间，用静态速度场（stationary velocity field）$v$对$t=[0,1]$积分得到最终的配准场${\varphi }^{(1)}$。使用缩放和展平（squaring）来计算积分，静态ODE的积分可以表示为微分同胚的单参数子群，在群论中，$v$是李代数的一员，并且求幂得到${\varphi }^{(1)}$，${\varphi }^{(1)}$是李群${\varphi }^{(1)}=\exp (v)$的一员。从单参数子群的属性可知，对于任意常量$t$和$t^{\prime }$，有$\exp ((t+t^{\prime })v)=\exp (tv)\circ \exp (t^{\prime }v)$，其中$\circ$是和李群相关的成分图（composition map）。从${\varphi }^{\left(1/2^T\right)}=\mathbf{p}+\mathbf{v}(\mathbf{p})/2^T$开始，使用重现（recurrence）${\varphi }^{\left(1/2^{t-1}\right)}={\varphi }^{\left(1/2^t\right)}\circ {\varphi }^{\left(1/2^t\right)}$来获取${\varphi }^{(1)}={\varphi }^{(1/2)}\circ {\varphi }^{(1/2)}$，其中$p$是空间位置的映射，选择$T$使得$v/2^T\approx 0$。

二、Prob-VoxelMorph

让$x，y$表示3维MR图像，$z$表示转换函数${\varphi }_z$的隐变量，将$z$的先验概率表示为：
$$p(\mathbf{z})=\mathcal{N}\left(\mathbf{z};0,{\mathbf{\Sigma }}_z\right)$$
其中$\mathcal{N}\left(\mathbf{z};0,{\mathbf{\Sigma }}_z\right)$是多远正态分布，$\mu$和$\Sigma$分别是其均值和方差。$z$可以是形变场的低维嵌入，也可以是形变场本身。在本文中，让$z$表示静态形变场。让$L=D-A$表示定义在体素网格上的邻接图的拉普拉斯，其中$D$是图的度矩阵，$A$是体素邻接矩阵。通过让${\mathbf{\Sigma }}_z^{-1}={\mathbf{\Lambda }}_z=\lambda \mathbf{L}$来保证$z$的空间平滑，其中${\mathbf{\Lambda }}_z$是精密矩阵（precision matrix）。我们让$x$是扭曲图像$y$的噪声观测：
$$p(\mathbf{x}\mid \mathbf{z};\mathbf{y})=\mathcal{N}\left(\mathbf{x};\mathbf{y}\circ {\mathbf{\varphi }}_z,{\sigma }^2\mathbb{I}\right)$$
其中${\sigma }^2$是加性图像的方差。

我们的目标是估计后验配准概率$p(z|x;y)$，使用变分的方法，引入一个近似后验概率$q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})$，通过最小化以下KL散度（即变分下界的负）来计算：
$$\begin{array}{l}{\min }_{\psi }\mathrm{K}\mathrm{L}\left[q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})\Vert p(\mathbf{z}\mid \mathbf{x};\mathbf{y})\right]\\ ={\min }_{\psi }{\mathbf{E}}_q\left[\log q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})-\log p(\mathbf{z}\mid \mathbf{x};\mathbf{y})\right]\\ ={\min }_{\psi }{\mathbf{E}}_q\left[\log q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})-\log p(\mathbf{z},\mathbf{x},\mathbf{y})\right]+\log p(\mathbf{x};\mathbf{y})\\ ={\min }_{\psi }\mathrm{K}\mathrm{L}\left[q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})\Vert p(\mathbf{z})\right]-{\mathbf{E}}_q[\log p(\mathbf{x}\mid \mathbf{z};\mathbf{y})]\end{array}$$
近似后验概率$q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})$可以表示为：
$$q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})=\mathcal{N}\left(\mathbf{z};{\mathbf{\mu }}_{z\mid x,y},{\mathbf{\Sigma }}_{z\mid x,y}\right)$$
其中${\Sigma }_{z|x;y}$是对角矩阵。

使用参数为$\psi$的卷积神经网络$de{f}_{\psi }(x,y)$来估计${\mu }_{z|x,y}$和${\Sigma }_{z|x,y}$，可以通过随机梯度下降的方法优化变分下界来学习参数$\psi$，给定图像对$\{x,y\}$和样例${\mathbf{z}}_k\sim q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})$，通过以下损失来计算$y\circ {\varphi }_{zk}$：
$$\begin{array}{l}\mathcal{L}(\psi ;\mathbf{x},\mathbf{y})=-{\mathbf{E}}_q[\log p(\mathbf{x}\mid \mathbf{z};\mathbf{y})]+\mathrm{K}\mathrm{L}\left[q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})\Vert p(\mathbf{z})\right]\\ =\frac{1}{2{\sigma }^{2}K}{\sum }_k{\Vert \mathbf{x}-\mathbf{y}\circ {\mathbf{\varphi }}_{z_k}\Vert }^2+\frac{1}{2}\left[\mathrm{tr}\left(\lambda \mathbf{D}{\mathbf{\Sigma }}_{z\mid x;y}-\log \left|{\mathbf{\Sigma }}_{z\mid x;y}\right|\right)+{\mathbf{\mu }}_{z\mid x,y}^T{\mathbf{\Lambda }}_z{\mathbf{\mu }}_{z\mid x,y}\right]+\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}\end{array}$$
其中$K$是样例的个数，在实验中使用$K=1$。上式中第一项可以让配准后的图像$y\circ {\varphi }_{zk}$和$x$相似，第二项鼓励后现概率接近于先验概率$p(z)$。虽然变分协方差${\Sigma }_{z|x,y}$是对角矩阵，但是最后一项在空间上平滑了均值：${\mathbf{\mu }}_{z\mid x,y}^T{\mathbf{\Lambda }}_z{\mathbf{\mu }}_{z\mid x,y}=\frac{\lambda }{2}\sum {\sum }_{j\in N(I)}(\mathbf{\mu }[i]-\mathbf{\mu }[j]{)}^2$，其中$N(i)$是体素$i$的邻居，将${\sigma }^2$和$\lambda$看作是固定的超参数。

卷积神经网络$de{f}_{\psi }(x,y)$的输入是图像$x，y$，输出是${\mu }_{z|x,y}$和${\Sigma }_{z|x,y}$。神经网络包括一个卷积层（16通道），4个下采样层（32通道，步长为2）和3个上采样卷积层（32通道），每个卷积层后跟着Leaky ReLU激活函数，并且卷积核大小为$3\times 3$。模型的结构如上图所示。

为了使用公式(6)进行无监督的学习，使用了一个通过重参数化技巧来采样一个新的${\mathbf{z}}_k\sim \mathcal{N}\left({\mathbf{\mu }}_{z\mid x,y},{\mathbf{\Sigma }}_{z\mid x,y}\right)$。

文章提出了缩放和展平层来计算${\varphi }_{z_k}=\exp (z_k)$，给定两个3D向量场$a$和$b$，对于每个体素$p$，使用线性插值计算$(a\circ b)(p)=a(b(p))$，即$a$中一个非整数的体素位置$b(p)$。开始时有${\varphi }^{\left(1/2^T\right)}=\mathbf{p}+{\mathbf{z}}_k/2^T$，递归的计算${\mathbf{\varphi }}^{\left(1/2^{t+1}\right)}={\mathbf{\varphi }}^{\left(1/2^t\right)}\circ {\mathbf{\varphi }}^{\left(1/2^t\right)}$，使得${\varphi }^{(1)}\triangleq {\varphi }_{z_k}=\exp \left({\mathbf{z}}_k\right)$，在本文中$T=7$。

三、处理过程

总结一下就是，卷积神经网络$de{f}_{\psi }(x,y)$以$x，y$为输入，计算${\mu }_{z|x,y}$和${\Sigma }_{z|x,y}$，采样一个新的${\mathbf{z}}_k\sim \mathcal{N}\left({\mathbf{\mu }}_{z\mid x,y},{\mathbf{\Sigma }}_{z\mid x,y}\right)$，计算微分同胚${\varphi }_{z_k}$并将其作用在$y$上。

给定已经学习好的参数，使用${\varphi }_{{\hat{z}}_k}$来对图像对$(x,y)$进行近似配准，首先使用以下公式获取$\widehat{z_k}$：
$${\hat{\mathbf{z}}}_k=\arg \underset{z_k}{\max }p\left({\mathbf{z}}_k\mid \mathbf{x};\mathbf{y}\right)={\mathbf{\mu }}_{z\mid x;y}$$
然后使用缩放和展平层计算${\varphi }_{\widehat{z_k}}$，我们还得到${\Sigma }_{z|x,y}$，可以估计每个体素$j$处速度场$z$的不确定性：
$$H(\mathbf{z}[j])\approx \mathbf{E}\left[-\log q_{\psi }(\mathbf{z}\mid \mathbf{x},\mathbf{y})\right]=\frac{1}{2}\log 2\pi {\mathbf{\Sigma }}_{z\mid x;y}[j,j]$$
我们也估计形变场${\varphi }_z$的不确定性。采样几个表达${\mathbf{z}}_{k^{\prime }}\sim q_{\psi }(\mathbf{z}\mid \mathbf{x};\mathbf{y})$，通过微分同胚层传播它们来计算${\varphi }_{z_k^{\prime }}$和经验对角协方差（empirical diagonal covariance）${\hat{\Sigma }}_{{\varphi }_z}[j,j]$，不确定性为$H(\varphi [j])\approx \frac{1}{2}\log 2\pi {\hat{\Sigma }}_{{\varphi }_z}[j,j]$。

四、实验

实验部分进行了基于3D atlas的配准，使用的数据集有ADNI，OASIS，ABIDE，ADHD200，MCIC，PPMI，HABS和Harvard GSP。首先对数据进行标准预处理：将图像重采样成1mm的各向同性体素，并用FreeSurfer进行反射空间正则化和脑部提取（头骨去除），然后将图像裁剪到$160\times 192\times 224$大小。按照7329，250，250大小分别划分为训练集、验证集和测试集。

将配准得到的形变场作用在分割标签上，并计算其Dice score。为了验证微分同胚性，使用了雅克比矩阵$J_{\varphi }(p)=\nabla \varphi (p)\in R^{3\times 3}$，雅克比矩阵可以捕获形变场$\varphi$在体素$p$处的局部特性。只有在满足$|J_{\varphi }(p)|>0$的位置，局部形变才是微分同胚的，即可逆和方位保持（orientation-preserving）的。

baseline选用的是ANTs软件包中的SyN和VoxelMorph，后者不能保证形变场是微分同胚的或者是不确定的估计。上图两行分别是本文的方法和ANTs的配准结果对比，VoxelMorph的和ANTs的类似，没有画出。

由上图可知，本文所提出的模型不仅在Dice值上达到了最好的水平，而且运行时间更快，并且配准形变场是微分同胚的（几乎没有非负的雅克比体素）。

通过是实验发现${\sigma }^2\sim (0.035{)}^2$并且$\lambda \in (2w,10w)$时效果最好，本文中使用的是$\lambda =7w$。

上图是ANTs、VoxelMorph和本文配准结果在脑部各个结构的Dice值对比。

上图中的左图的速度场的不确定性$H(z)$表明了在结构边缘的低不确定性，如中间的折线图所示，这种相关性在最终的配准场的不确定性$H({\varphi }_z)$中并不明显，如右图所示。