论文研读：TransMorph—用于无监督医学图像配准的Transformer

论文地址：TransMorph: Transformer for unsupervised medical image registration

摘要

在过去，卷积神经网络(ConvNets)一直是医学图像分析领域的研究热点。但是它的缺点是很少关注图像中的远程空间关系。最近，人们提出Transformer来解决ConvNet的缺点，并在许多医学成像应用中产生了最先进的性能，因为其大得多的感受野能够更精确地理解运动图像和固定图像之间的空间对应关系。在本文中作者提出了Transmorph—用于无监督医学图像配准的Transformer。

本文还给出了变形的微分同胚型和贝叶斯变型：

微分同胚型变型保证了变形的拓扑性
贝叶斯变型产生了校准良好的配准不确定性估计。

一、介绍

作者提出了一种混合Transformer和ConvNet的框架Transmorph，用于3D医学图像配准。在这种方法中，Swin Transformer作为编码器，捕获输入运动图像和固定图像之间的空间对应关系。然后，ConvNet解码器将Transformer编码器提供的信息处理成密集的位移场。部署了长跳跃连接以维持编码器和解码器之间的定位信息流。我们还引入了变形的微分同胚变量，以确保平滑和保持拓扑的变形。此外，我们对变形参数进行了变分推理，得到了一个贝叶斯模型，该模型基于给定的图像对来预测配准不确定性。

这项工作的主要贡献概括如下：

基于Transformer的图像配准模型：本文介绍了使用Transformer进行图像配准的开创性工作。提出了一种新的基于Transformer的神经网络用于仿射和变形图像配准。
结构分析：本文的实验表明，Positional Embedding是传统的Transformer中常用的元素，对于所提出的Transformer-ConvNet混合模型来说，位置嵌入是不必要的。其次，作者证明了基于Transformer的模型比基于ConvNets的模型具有更大的有效感受场。此外，作者还证明了TransMorph促进了一种更平坦的配准损失。
微分同胚配准：我们证明了TransMorph可以很容易地集成到两个现有的框架中，作为配准主干来提供微分同胚配准。
不确定度量化：本文还提供了TransMorph的贝叶斯不确定度变量，它产生了transformer的不确定度和经过完美校准的外观不确定度估计。
State-of-the-art results

二、相关工作

2.1 图像配准

可变形图像配准(DIR)通过优化能量函数来建立两个图像之间的空间对应关系：

其中，Im和If分别表示运动图像和固定图像， $\Phi$ 表示扭曲运动图像的变形场(即，Im◦φ)，R(φ)施加变形场的平滑度，并且λ是确定图像相似性和变形场规则性之间的权衡的正则化超参数。最优翘曲ˆ $\hat{\Phi }$ 是通过最小化该能量函数给出的：

$E_{sim}$ 测量变形的运动图像Im◦φ和固定图像If之间的对准程度。 $E_{sim}$ 的一些常见选择是均方误差(MSE)，归一化互相关(NCC)，结构相似指数(SSIM)和相互信息(MI)。

正则化项R(φ)对变形场施加空间光滑性。在大多数应用中，一个常见的假设是运动图像和固定图像中都存在相似的结构。因此，需要一个连续和可逆的变形场(即，微分同胚)来保持拓扑，而正则化R(φ)旨在加强或鼓励这一点。

2.1.1 基于深度神经网络的图像配准

基于DNN(深度神经网络)的方法优化用于训练数据集的能量函数，从而学习图像配准的全局表示，其使得能够配准未知的图像。DNN方法通常分为监督和非监督两类，前者需要ground-truth真实形变场进行训练，而后者仅依赖于图像数据集。

2.1.2 微分同胚图像配准

微分同胚图像配准由于其具有拓扑保持性和变换可逆性等特殊性质，在许多医学图像应用中具有重要的意义。微分同胚变换是具有可逆导数(即非零雅可比行列式)的光滑连续的一对一映射。这种转换可以通过依赖于时间的时间积分来实现或时间固定速度场(SVF)。

在依赖时间的设置中和SYN，一个微分同胚变换 $\Phi$ 是通过积分足够的光滑时变速度场 $v^{(t)}$ 得到的，即 $\frac{}{}$ $\frac{d}{dt}\Phi ^{(t)}=v^{(t)}(\Phi ^{(t)})$ ,其中 $\Phi ^{(t)}=id$ 。

另一方面，在静止速度场(SVF)环境中，假设速度场随时间是平稳的，即 $\frac{d}{dt}\Phi ^{(t)}=v(\Phi ^{(t)})$ 。Dalca等人首先在深度学习模型中采用微分同胚公式，使用具有有效缩放和平方方法的SVF设置。在缩放和平方方法中，变形场被表示为李代数成员，该李代数成员被指数化以生成时间1变形 $\Phi^{1}$ ，其是李群的成员： $\Phi ^{1}=exp(v)$ 。这意味着指数化的流场迫使使用相同的流场的映射是微分同胚的和可逆的。

从初始变形场开始

其中p表示空间位置。φ(1)可以使用递归公式获得：

因此。

在实际过程中，神经网络首先生成位移场，然后按1/2T缩放以产生初始变形场 $\Phi ^{(1/2t)}$ ,随后，递归公式通过空间变换函数递归地应用于 $\Phi ^{(1/2t)}$ T次，得到最终的微分同胚形变场φ(1)。尽管微分同胚在理论上被保证是可逆的，但内插误差会导致可逆误差，该误差随着内插步数的增加而线性增加。

2.2 自我注意机制与Transformer

Transformer使用自注意力机制，该机制通过Query-Key-Value(QKV)模型估计一个输入序列和另一个输入序列的相关性。输入序列通常来自图像的patches。

设x是在3D空间域上定义的图像体，（ $x\in \mathbb{R}^{H\times W\times L}$ ）。

首先将图像体分成N个平面化的3D块 $x_{p}\in \mathbb{R}^{N\times P^{3}}$ ,(HWL)是原始图像的大小，（P,P,P）是每个图像块的大小， $N=\frac{HWL}{P^{3}}$ 。

随后将可学习的linear embedding $E$ 应用于 $x_{p}$ ，它将每个patch投影到 $D\times 1$ 向量：

$\hat{x_{e}}=[x_{p}^{1};x_{p}^{2};...;x_{p}^{N}E], E\in \mathbb{R}^{p^{3}\times D}$

其中，维度D是用户定义的超参数。然后，将可学习的位置嵌入添加到 $\hat{x_{e}}$ ，使得patch可以保留它们的位置信息，即 $x_{e}=\hat{x_{e}}+E_{pos}$ ，其中 $E_{pos}\in \mathbb{R}^{N\times D}$ 。这些通常被称为tokens的矢量表示随后被用作自注意计算的输入。
自注意力机制。为了自我注意力(SA)， $x_{e}\in \mathbb{R}^{N\times D}$ 通过线性层编码为三个矩阵表示 $Q\in \mathbb{R}^{N\times D_{k}}$ ， $K\in \mathbb{R}^{N\times D_{k}}$ 和 $V\in \mathbb{R}^{N\times D_{v}}$ ，点积注意力由以下公式得出：

2.3 贝叶斯深度学习

我们先回顾一下贝叶斯公式

其中，𝑝(𝑧|𝑥) 被称为后验概率（posterior），𝑝(𝑥,𝑧) 被称为联合概率，𝑝(𝑥|𝑧) 被称为似然（likelihood），𝑝(𝑧) 被称为先验概率（prior），𝑝(𝑥) 被称为 evidence。

如果再引入全概率公式 𝑝(𝑥)=∫𝑝(𝑥|𝑧)𝑝(𝑧)𝑑𝑧，上式可以再变成如下形式：

如果 𝑧 是离散型变量，则将式中分母积分符号 ∫ 改成求和符号 ∑ 即可。

什么是贝叶斯深度学习？

在深度学习中，𝑤𝑖,(𝑖=1,...,𝑛) 和 𝑏 都是一个确定的值，例如 𝑤1=0.1,𝑏=0.2。即使我们通过梯度下降（gradient decent）更新 𝑤𝑖=𝑤𝑖−𝛼⋅∂𝐽/∂𝑤𝑖，我们仍未改变 “𝑤𝑖 和 𝑏 都是一个确定的值” 这一事实。

那什么是贝叶斯深度学习？将 𝑤𝑖 和 𝑏 由确定的值变成分布（distributions），这就是贝叶斯深度学习。

三、方法

图像配准的传统范例

图像配准的常规范例在上图中示出。分别表示为 $\hat{I}_{m}$ 和 $I_{f}$ 的移动图像和固定图像首先被仿射变换到单个坐标系中。

所得到的仿射对准的运动图像被表示为 $I_{m}$ 。随后，使用由DIR算法生成的变形场Φ将 $I_{m}$ 扭曲到 $I_{f}$ （即， $I_{m}o\phi$ ）。

下图呈现了作者所提出的方法的概述。这里，仿射变换和可变形配准都是使用基于Transformer的神经网络来执行的。Transformer将 $\hat{I}_{m}$ 和 $I_{f}$ 作为输入，并计算仿射变换参数的集合（例如，旋转角度、平移等）。这些参数用于经由仿射变换函数将 $\hat{I}_{m}$ 与 $I_{f}$ 仿射对准，从而产生对准图像 $I_{m}$ 。然后，在给定 $I_{m}$ 和 $I_{f}$ 的情况下，变形网络计算变形场φ，其使用空间变换函数。在训练期间，所述训练网络可以有选择地包含补充信息（例如，解剖分割）。

3.1 仿射变换网络

仿射变换通常用作图像配准的初始阶段，因为它有助于优化后续更复杂的仿射过程。仿射网络全局地检查一对移动图像和固定图像，并产生一组变换参数，该变换参数将移动图像与固定图像对齐。这里，所提出的基于transformer的仿射网络的架构是修改的Swin Transformer，其将两个3D体积作为输入（即，If和Im），并生成12个仿射参数：三个旋转角度，三个平移参数，三个缩放参数和三个剪切参数。由于仿射配准相对简单，我们减少了原始Swin Transformer中的参数数量。Transformer的体系结构和参数设置的细节将在后续部分中介绍。

3.2 可变形配准网络

上图显示了作者所提出的TransMorph的网络架构。

网络的编码器首先将输入的移动体积和固定体积分成不重叠的3D Patch，每个Patch的大小为2×P×P×P，其中P通常设置为4。我们将第i片表示为 $X_{p}^{i}$ ，其中 $i\in\left \{ 1,... N\right \}$ ， $N=\frac{H}{P}\times\frac{W}{P}\times\frac{L}{P}$ 是曲面片的总数。每个Patch被展平并被视为一个“token”，然后使用线性投影层将每个token投影到任意维度的特征表示（表示为C）：

其中 $E\in \mathbb{R}^{2P^{3}\times C}$ 表示线性投影，输出z0的维数为N ×C。

由于线性投影对图像块进行操作，并且不保持标记相对于图像的整体位置，因此先前的基于Transformer的模型通常向线性投影添加位置嵌入，以便将位置信息集成到标记中，即 $z_{0}+E_{pos}$ 。然而这种方法通常应用于图像分类操作，在图像配准中，对于像图像配准这样的像素级任务，网络通常包括一个解码器，该解码器生成与输入或目标图像具有相同分辨率的密集预测。通过使用损失函数将输出与目标图像进行比较来加强输出图像中的体素之间的空间对应性。输出和目标之间的任何空间不匹配都将导致丢失并反向传播到Transformer编码器中。因此，Transformer应该固有地捕获令牌的位置信息。

在这项工作中，作者观察到，位置嵌入对于图像配准是不必要的，它只会向网络添加额外的参数，而不会提高性能。

在线性投影层之后，对 $z_{0}$ 进行多个连续阶段的Patch合并和Swin Transformer块。

Swin Transformer块输出与输入相同数量的token，而Patch合并层连接每组2 × 2 × 2相邻token的特征，因此它们将令牌的数量减少了2 × 2 × 2 = 8倍（例如, $H\times W\times L\times C\rightarrow \frac{H}{2}\times\frac{W}{2}\times\frac{L}{2}\times8C$ 。然后，在8C维级联特征上应用线性层以产生每个2C维的特征。在Swin Transformer块的四个阶段和Transformer阶段之间的补丁合并的三个阶段（即，图中的橙子框），编码器最后一级的输出尺寸为 $\frac{H}{32}\times\frac{W}{32}\times\frac{L}{32}\times8C$ 。解码器由连续上采样和卷积层组成，内核大小为3 × 3。

解码阶段中的每个上采样特征图通过跳过连接与来自编码路径的对应特征图级联，然后是两个连续的卷积层。

Transformer编码器只能提供高达HP × WP × LP分辨率的特征图。因此，Transformer可能无法提供高分辨率的特征图并在较低层聚合本地信息。为了解决这个缺点，我们采用了两个卷积层，使用原始和下采样图像对作为输入来捕获本地信息并生成高分辨率特征图。这些层的输出与解码器中的特征图连接以产生变形场。输出变形场φ由16个3 × 3卷积产生。除了最后一个卷积层，每个卷积层后面都有一个Leaky Rectified Linear Unit激活。最后，使用空间变换函数，利用网络提供的变形场φ（或位移场u）对运动图像Im进行非线性扭曲。

3.2.1 3D Swin Transformer Block

Swin Transformer可以通过使用Patch合并层来生成各种分辨率的分层特征图，使其非常适合用作像素级任务（如图像配准和分割）的通用骨干。

（a）Swin Transformer通过合并图像块来创建分层特征图。自我注意力在每个局部3D窗口（红色框）内计算。在每个分辨率下生成的特征图被发送到ConvNet解码器以产生输出。

（b）局部窗口的3D循环移位，用于基于移位窗口的自注意计算。

3.2.2 损失函数

损失函数由两部分组成：一部分计算变形的运动图像和固定图像之间的相似性，另一部分对变形场进行正则化以使其平滑：

其中Lsim表示图像相似度度量，并且R表示变形场正则化。

当If和Im用于器官分割时，TransMorph可以在训练期间利用该辅助信息来改善 $I_{m}o\phi$ 和 $I_{f}$ 之间的解剖映射。量化分割重叠的损失函数Lseg被添加到总体损失函数：

其中， $s_{f}$ 和 $s_{m}$ 分别表示 $I_{f}$ 和 $I_{m}$ 的器官分割，并且γ是控制Lseg的强度的加权参数。在图像配准领域，通常使用Dice得分作为品质因数来量化配准性能。因此，我们直接最小化 $s_{f}^{k}$ 和 $s_{m}^{k}$ 之间的Dice损失，其中k表示第k个结构/器官：

3.3 概率和B样条变量

待更新

3.4 贝叶斯不确定性变量

待更新

四、总结

在本文中，作者提出了TransMorph，无监督的可变形图像配准。TransMorph是建立在Transformer，建立图像体素之间的长距离空间对应。提出了两种变体的transMorph，提供拓扑保持变形。此外，我们还将贝叶斯深度学习引入到TransMorph的Transformer编码器中，从而在不降低配准性能的情况下实现变形不确定性估计。与传统方法和基于学习的方法相比，该方法具有较高的配准精度，证明了其在医学图像配准中的有效性。

标签： transformer 深度学习人工智能

本文转载自: https://blog.csdn.net/weixin_42937740/article/details/140448663
版权归原作者 俺不是文盲 所有，如有侵权，请联系我们删除。