本文是关于论文《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 z。然后把 moving image x x x 和隐含变量 z z z 输入到空间转换网络(STN)中产生最终的配准后的图像。
(
F
x
i
,
F
y
i
)
(F_x^i,F_y^i)
(Fxi,Fyi) 表示两个CNN第
i
i
i 层产生的特征图,
F
z
i
F_z^i
Fzi 表示概率模型产生的第
i
i
i 层的隐含变量, 它实际是让
F
x
i
F_x^i
Fxi 对齐到
F
y
i
F_y^i
Fyi 的 STN 的参数,或者说形变场。在已知
F
x
i
,
F
y
i
F_x^i,F_y^i
Fxi,Fyi 的情况下,可以通过最大化后验概率
p
(
F
z
i
∣
F
x
i
;
F
y
i
)
p(F_z^i|F_x^i;F_y^i)
p(Fzi∣Fxi;Fyi) 来得到最优的
F
z
i
F_z^i
Fzi。具体的,引入了一个近似后验概率
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
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(F_z^i|F_x^i;F_y^i)
p(Fzi∣Fxi;Fyi) 和
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
q_\psi(F_z^i|F_x^i;F_y^i)
qψ(Fzi∣Fxi;Fyi) 之间的KL散度来使得两个分布尽可能的相似,该过程可以用下式表示:
min
ψ
K
L
[
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
∥
p
(
F
z
i
∣
F
x
i
;
F
y
i
)
]
=
min
ψ
K
L
[
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
∥
p
(
F
z
i
)
]
−
E
q
log
p
(
F
y
i
∣
F
z
i
;
F
x
i
)
\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}
=ψminKL[qψ(Fzi∣Fxi;Fyi)∥p(Fzi∣Fxi;Fyi)]ψminKL[qψ(Fzi∣Fxi;Fyi)∥p(Fzi)]−Eqlogp(Fyi∣Fzi;Fxi)
其中
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
q_\psi(F_z^i|F_x^i;F_y^i)
qψ(Fzi∣Fxi;Fyi) 来自于多元正态分布:
q
ψ
(
F
z
i
∣
F
x
i
;
F
y
i
)
=
N
(
z
;
μ
F
z
i
∣
F
x
i
;
F
y
i
,
σ
F
z
i
∣
F
x
i
,
F
y
i
2
)
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,Fyi2)
其中
μ
F
z
i
∣
F
x
i
;
F
y
i
\mu_{F_{z}^{i} | F_{x}^{i} ; F_{y}^{i}}
μFzi∣Fxi;Fyi 是分布的均值,
σ
F
z
i
∣
F
x
i
,
F
y
i
2
\sigma_{F_{z}^{i} | F_{x}^{i}, F_{y}^{i}}^{2}
σFzi∣Fxi,Fyi2 是分布的方差,它们是通过概率模型得到的(如图1(b))。
p
(
F
z
i
)
p(F_z^i)
p(Fzi) 和
p
(
F
z
i
∣
F
x
i
;
F
y
i
)
p(F_z^i|F_x^i;F_y^i)
p(Fzi∣Fxi;Fyi) 符合以下多元正态分布:
p
(
F
z
i
)
=
N
(
F
z
i
;
0
,
σ
F
z
i
2
)
p\left(F_{z}^{i}\right)=\mathcal{N}\left(F_{z}^{i} ; 0, \sigma_{F_{z}^{i}}^{2}\right)
p(Fzi)=N(Fzi;0,σFzi2)
p ( F y i ∣ F z i ; F x i ) = N ( F y i ; F x i ∘ ϕ F z i , σ F i 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)
其中 σ F z i 2 \sigma_{F_z^i}^2 σFzi2 是分布的方差, F x i ∘ ϕ F z i F_{x}^{i} \circ \phi_{F_{z}^{i}} Fxi∘ϕFzi 是噪音, σ F i 2 \sigma_{F^i}^2 σFi2 是噪音项的方差。
在CNN浅层的特征图具有较高的分辨率并且具有丰富的细节信息,而CNN深层的特征图具有较低的分辨率并且具有高层次的语义信息。高层语义信息可以帮助全局配准,但是忽略了很多细节。而细节信息则是捕获了局部的配准信息。所以将浅层到深层的特征图 F z i F_z^i Fzi 混合得到最终的隐含变量 z z z,然后输入到 STN 中,对 moving image 进行变形。
模型总的损失为:
D
total
=
L
(
z
;
x
,
y
)
+
∑
i
=
1
n
w
i
L
(
F
z
i
;
F
x
i
,
F
y
i
)
\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∑nwiL(Fzi;Fxi,Fyi)
其中,
L
(
z
;
x
,
y
)
\mathcal{L}(z;x,y)
L(z;x,y) 表示从输入图像
x
x
x 和
y
y
y 到输出的配准后的图像
z
z
z 的KL散度,
L
(
F
z
i
;
F
x
i
,
F
y
i
)
\mathcal{L}\left(F_{z}^{i} ; F_{x}^{i}, F_{y}^{i}\right)
L(Fzi;Fxi,Fyi) 是从输入特征图
F
x
i
F_x^i
Fxi 和
F
y
i
F_y^i
Fyi 到输出配准转换变量
F
z
i
F_z^i
Fzi 的KL散度。
n
n
n 是CNN的层数,
w
i
w_i
wi 是第
i
i
i 层损失的权重。通常设置
n
=
4
,
w
i
=
1
n=4,w_i=1
n=4,wi=1。基于KL散度的损失为:
L
(
Z
;
X
,
Y
)
=
1
2
σ
Z
∣
X
;
Y
2
∥
Y
−
X
∘
ϕ
Z
∥
2
+
1
2
[
tr
(
σ
Z
∣
X
;
Y
2
)
+
∥
μ
Z
∣
X
;
Y
∥
−
log
det
(
σ
Z
∣
X
;
Y
2
)
]
\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;Y21∥Y−X∘ϕZ∥2+21[tr(σZ∣X;Y2)+∥∥μZ∣X;Y∥∥−logdet(σZ∣X;Y2)]
其中第一项是使得配准后的图像
X
∘
ϕ
Z
X\circ\phi_Z
X∘ϕZ 与图像
Y
Y
Y 相似的重建损失,第二项是公式1第一项的近似,它可以让
q
ψ
(
Z
∣
X
;
Y
)
q_\psi(Z|X;Y)
qψ(Z∣X;Y) 与
p
(
Z
)
p(Z)
p(Z) 相似;
μ
Z
∣
X
;
Y
\mu_{Z|X;Y}
μZ∣X;Y 和
σ
Z
∣
X
;
Y
\sigma_{Z|X;Y}
σZ∣X;Y 分别是分布
q
ψ
(
Z
∣
X
;
Y
)
q_\psi(Z|X;Y)
qψ(Z∣X;Y) 的均值和标准差。
初始学习率为 1 e − 4 1e^{-4} 1e−4,并且周期性的减少(乘以0.1),一共有100个epoch,使用Adam优化器,优化器的第一个动量为0.9,第二个动量为0.999,衰减权重为0.0001。
下图是实验的结果对比图。