基于自注意力连接UNet的磁共振成像去吉布斯伪影算法

刘 阳，陆志扬，王 骏，施 俊

（上海大学 通信与信息工程学院，上海 200444）

0 引言

吉布斯伪影是磁共振成像（Magnetic Resonance Imaging，MRI）中常见的一种环状伪影，当磁共振扫描高频采样信号不足时，傅里叶变换成像后就会产生吉布斯伪影［1-2］。吉布斯伪影的出现会改变组织结构的强度、形状和解剖细节，影响医生对疾病的诊断［3］。

针对吉布斯伪影，传统方法采用小波变换、总变分方法、亚像素移位和自适应滤波等技术实现去伪影［4-7］。但滤波处理方法往往需要复杂的参数调整或预处理，导致其很难在医学图像处理中得到广泛应用［8］。亚像素移位可以在不使用部分傅里叶编码情况下减少吉布斯伪影，但是由于欠采样产生的模糊效应和傅里叶级数近似产生的吉布斯伪影之间相互作用，这些效应的组合很难明确建模［9］。

近年来，深度学习（Deep Learning，DL）作为一种机器学习（Machine Learning，ML）方法在磁共振成像处理中得到了广泛的关注，许多DL 网络被提出，用于去除MRI 中的噪声和伪影。由于高斯噪声的非相干性和位置无关性，常利用卷积神经网络（Convolutional Neural Network，CNN）自身小感受野的特性进行图像去噪。通过构建非线性模型灵活地去除图像中的吉布斯伪影和噪声，对含吉布斯伪影的MRI 进行并行化和高效处理［10-11］。例如，GAS-CNN（Gibbs-ringing Artifact Suppression Convolutional Neural Network）通过从自然图像到MRI 的知识转移来抑制吉布斯伪影，利用自然图像训练MRI进行微调［10］；GRACNN（Gibbs-Ringing Artifact reduction using CNN）通过学习图像的伪影部分来提取伪影信息，得到无伪影图像，获得了比传统算法更优的实验结果［11］。这种单尺度CNN 对改善噪声很有效，但对分布于整张图像的吉布斯伪影的去除，则需要更大感受野的网络来改善［12］。

UNet 通常用于图像分割，利用多尺度特征模块进行特征的提取和恢复，具有较大的感受野［13］。其中，跳连接可以帮助编码-解码过程恢复完整的空间分辨率，将编码端子网络特征与解码端子网络特征进行直接串联融合。许多基于UNet 的改进算法，通过改变跳连接的连接方式减小编码端和解码端特征之间的语义差距，提升了网络的性能。例如，UNet++使用了嵌套的密集跳跃连接代替普通跳跃连接，缩小了编码端和解码端子网络特征之间的语义差距［14］；UCTransNet（U-Net with a proposed CTrans module）设计了一个多尺度通道融合的子模块来引导融合的多尺度通道信息有效连接到解码器，来消除编解码之间的语义差异［15］。CNN卷积运算的固有局限性限制了感受野的扩大，而Transformer利用注意力机制模拟数据间的长期依赖关系，可以学习到特征之间明确的全局和远程语义交互信息［16］。因此考虑在UNet 跳连接处加入Transformer 模块［17-18］，通过自注意力连接的方式进行编解码端特征融合，进一步扩大网络的感受野，同时利用UNet 良好的编解码能力捕捉更多的伪影信息。

自蒸馏（Self-distillation）是一种无需教师模型的知识蒸馏方式，其中知识是在网络本身中蒸馏出来的［19］。例如，基于阶段的自蒸馏方法Metadistiller，将知识从深层阶段转移到浅层阶段［20］。目前自蒸馏在分类任务中可以替代过参数化的教师模型，实现模型压缩和快速推理［21］。而自蒸馏技术在图像重建任务中的应用还处于起步阶段。为了提升网络自身性能，增强网络在小样本条件下对图像边缘和精细结构的处理能力，提高UNet 解码端的重建精度，本文将自蒸馏训练与UNet 图像恢复模块相结合，通过自蒸馏训练，使得深层网络输出的更精细化特征指导浅层输出特征，提升网络在MRI 重建和去伪影方面的性能，获得更精确的无伪影输出图像和更高的评价指标。

针对MRI 去除吉布斯伪影任务，本文提出了一种基于自蒸馏训练的自注意力连接UNet（Self-attention connection UNet based on Self-Distillation training，SD-SacUNet）算法。首先设计了一种基于自注意力连接模块的UNet 以加强编码-解码端的特征融合，提高特征提取能力；然后设计了一种基于图像重建的自蒸馏训练方式，以提高UNet 解码端图像恢复精度，进一步提升去伪影效果。

本文的主要工作如下：

1）将自注意力连接模块引入经典的UNet，提出一种新的SD-SacUNet，在UNet 编解码网络两端之间加入自注意力连接模块进行自注意力机制的计算，将处理后的编码端特征与解码端特征合并，有效融合编解码两端不同的语义特征，利用Transformer 强大的全局交互能力，获取更精细的网络输出结果，得到清晰的无伪影输出图像。

2）提出了一种基于图像重建任务的自蒸馏训练方法，在SD-SacUNet 的编码器中间层特征输出之间计算损失函数，将深层特征的有用信息用于指导浅层特征的训练，获取更高的训练精度，提升网络整体的特征表达能力。

1 本文方法

1.1 总体流程

本文提出的SD-SacUNet 的总体结构如图1 所示。该算法将UNet 编码端每个下采样层的输出特征分别输入各自对应的自注意力连接模块进行自注意力机制的运算，而后与解码特征进行连接；在UNet 解码端加入瓶颈层用于调节特征尺度来进行自蒸馏训练，利用深层特征指导浅层特征的训练，编码器深层的精细化特征通过自蒸馏路径传递给浅层，以指导浅层网络的训练。

图1 SD-SacUNet总体结构Fig.1 Overall structure of SD-SacUNet

具体内容如下：

1）网络架构：在传统UNet 基础上加入了自注意力连接模块，将UNet 的编码端划分为不同的子块，通过下采样和特征映射，编码端特征被投影到更高维度。转换后的子块通过Transformer 模块以及子块合并操作进行特征学习和维度恢复，而后与解码特征合并，参与解码阶段的图像恢复过程。

2）自蒸馏训练：在UNet解码端图像恢复阶段进行自蒸馏训练，利用瓶颈层建立深层特征与浅层特征信息交互的通道，从而让深层重建阶段更精细的特征来指导浅层的特征训练，提高整体解码阶段的特征重建精度，获得更好的去伪影重建效果。自蒸馏训练时加入的瓶颈层将在测试阶段去除。

1.2 自注意力连接模块设计

本文所提出的SD-SacUNet 中的自注意力连接模块由子块嵌套操作、两个Swin Transformer（Hierarchical Vision Transformer using Shifted Windows）［17］块以及子块融合操作组成，替代了原有UNet 中的跳连接操作，目的是缩小编码特征和解码特征之间的语义差距，有利于在特征恢复阶段，解码端网络充分利用编码端特征进行特征重建与恢复。自注意力连接模块能够扩大感受野，有利于捕捉全局信息，帮助克服卷积操作感受野较小的局限性。对特征进行自注意力编码，网络根据特征感兴趣的区域进行特征融合，辅助解码端的特征重建。自注意力连接模块主要结构如图2 所示。

图2 自注意力连接模块Fig.2 Self-attention connection module

自注意力连接模块首先将输入的编码端特征进行子块嵌套操作，具体由两步组成：分别为对特征进行分块和将分块后的子块特征进行线性变化并嵌入位置向量信息。随后，处理好的子块被送入两个连续的Swin Transformer 块中进行注意力机制的运算。两个Swin Transformer 块分别计算普通窗口注意力和移位窗口注意力，以此达到将注意力的计算限制在每个窗口内的目的。移位窗口操作可以使不同窗口之间的特征能够进行信息交互，兼顾全局特征和控制计算的复杂度。最后通过子块融合操作，将子块进行合并和上采样，恢复子块分辨率与输入的编码特征相同，并与解码端特征进行合并操作。此处的合并操作保留了UNet 跳连接中的串联操作，将自注意力连接模块的输出特征与解码端特征进行直接的串联融合。通过网络的自我学习将有用信息进行传递，而不花费更多的计算成本［13］。

子块嵌套操作是将编码端输出的特征图进行不重叠的子块划分和位置编码。SD-SacUNet 编码端的输出共分为4层，其中第n层特征尺寸为为了方便阐述，将其定义为(Cn，hn，wn)。首先进行降采样操作来降低计算量，降采样因子设为2。因此输入Swin Transformer 模块的标记特征大小为

Swin Transformer 块最主要的特点是在传统Transformer的多头自注意力机制的基础上，增加局部关注和移位窗口机制。Swin Transformer 模块由层标准化（Layer Normalization，LN）层、多头自注意力（Multi-head Self-Attention，MSA）层、残差连接和多层感知器（MultiLayer Perceptron，MLP）模块组成，其中MLP 模块中含有两个GELU（Gaussian Error Linear Unit）非线性激活全连接层。使用的两个Swin Transformer 模块采用了不同的注意力模块。第一个模块采用了常规窗口多头自注意力模块（Multi-head Self-Attention modules with regular Windows，W-MSA），W-MSA 将输入特征划分为不重叠的窗口，在每个窗口内分别计算自注意力得分；第二个模块为移位窗口多头自注意力模块（Multi-head Self-Attention modules with Shifted Windows，SW-MSA），移位窗口每次移动，通过对移位窗口计算注意力，能获得特征图之间的远程依赖信息，使它兼顾了局部和全局的注意力。注意力得分定义为：

其中：Q、K、V∈分别代表查询矩阵、键矩阵和值矩阵；d代表查询矩阵和键矩阵的维度；B代表偏差矩阵。

子块融合操作是将前面嵌套操作降采样的特征图利用反卷积进行上采样，还原其原有分辨率，因此在经过自注意力连接模块处理后的特征与同级解码器特征尺寸相同，可以直接进行特征融合。经过注意力机制融合的特征保留了编码器提取特征的细节部分，并且在网络的训练过程中，注意力机制会调整编解码连接部分网络的权重，使后续特征能够得到有效融合，避免了编解码器两端特征由于语义上的差异带来的网络性能的下降。同时解码器在重建图像的过程中充分利用网络提取的有用信息，得到具有精准细节的重建图像，避免了过度平滑问题。

1.3 自蒸馏模块设计

重建任务的自蒸馏模块设计主要是为了从网络内部提供知识进行蒸馏训练，提升训练的效果。在重建任务中，网络在去除伪影的同时可以保留和恢复更多的图像细节，使得结果图像结构边缘清晰，具有清楚的轮廓结构。由于UNet解码端是一个图像分辨率逐渐扩大的过程，解码端深层的特征图尺寸更大，含有的细节信息更加丰富，利用这一特点来指导浅层特征的训练，可以使得浅层特征图进一步提升细节，从而使得整体网络输出结果保留更多图像细节。在网络进行训练时，自蒸馏方法可以将深层特征图信息传递给浅层。此外，在测试阶段，将用作自蒸馏训练的瓶颈层去除，因此不会对网络带来额外参数量的提升。

SD-SacUNet 的解码部分共有四个图像恢复模块。为了使得浅层网络更好地学习深层网络的知识，训练网络时在每个图像恢复模块后各增加一个瓶颈（Bottle Neck）层。添加瓶颈层可以调整特征大小与后一层输出特征相一致，减少各层输出之间的影响，通过计算与后一层特征之间的损失函数，使得深层解码特征的知识逐步过渡给浅层。在图像恢复过程中，特征恢复从浅层得以加强。在解码端输出至全连接层恢复完整无伪影图像的过程中，特征被进一步精炼，在上采样恢复时也排除了冗余噪声过多的干扰，更好地保留图像细节。自蒸馏训练的损失函数定义计算解码端相邻特征之间的欧氏距离，如式（2）所示：

其中：ni＇代表经过瓶颈层进行上采样的浅层网络特征；ni+1表示下一层解码端网络输出特征，通过计算二者之间的损失函数，将深层特征信息传递到浅层。

总的自蒸馏损失定义为各层损失函数之和，如式（3）所示：

其中：Lsd代表自蒸馏损失。按照本文网络架构，m=4，计算UNet 解码器端5 层中间输入的损失。

1.4 模型训练

如图1 所示，MRI 重建去伪影的训练目标是通过一个SD-SacUNet 编解码模型，学习输入的含吉布斯伪影图像与无伪影原图之间的映射关系。本文中用到的损失函数均为L2欧氏距离损失，以此来衡量预测输出的无伪影图与真实标签图之间的差异。通过计算反向梯度优化参数，指导图像重建模型的训练过程。

其中：Lres代表重建损失；yi表示网络的标签，即无伪影图像；为网络重建输出的预测；n为训练的样本数量。

训练网络的总损失定义为重建损失与自蒸馏损失的和，如公式（5）所示：

其中：α和β两个超参数用来调节损失函数的数量级，使得每种损失函数得到充分的训练和优化。

2 实验与结果分析

2.1 实验数据与预处理

本文采用CC359 数据库［22］验证所提出的SD-SacUNet 算法的有效性。CC359 数据库为单通道线圈数据，数据采集于临 床MRI 扫描仪（Discovery MR750，General Electric（CE）Healthcare，Waukesha，WI），共含有35 个全采样的T1 加权MRI 数据。该数据集为复数图像集，可以反向转换生成复数k空间样本。将数据集进行两个方向（切片编码和相位编码）上的欠采样时，可作为2D 数据集。

本文按照五折交叉验证的方法将数据集划分为训练集28 个病例，每个病例切片45 张，共1 440 张图像；测试集有7个病例，共315 张图像。由于吉布斯伪影是由高频数据缺失导致的，制作成对的输入和标签图像的时候，将全采样图像通过傅里叶变换到k空间域，此时为高频中心，通过频移得到低频中心傅里叶变换图，通过一个遮罩将外围部分高频数据置0，再进行傅里叶反变换得到的图像就含有吉布斯伪影，其中吉布斯伪影和噪声的影响程度与遮罩的大小有关。

2.2 实验设计

为了比较本文所提出的SD-SacUNet 算法的有效性，选取了如下经典的网络和目前在医学图像处理领域与本文工作类似的先进网络进行对比，本文对比实验设置如下：

市场需求一旦发生变化，订货量会随之变动，即q>q*或q

1）UNet［13］：作为本文算法的基准进行对比和消融实验。

2）UNet++［14］：对于UNet 中的跳连接进行了改进，在编解码器之间加入了嵌套的多尺度融合模块。

3）GRACNN［11］：深度学习去吉布斯伪影的经典方法，先提取伪影，再从原图中将伪影减去得到无伪影图像的方法。

4）UCTransNet［15］：一种利用Transformer 进行多尺度通道交叉融合代替UNet 跳连接的网络模型，应用于分割领域。

5）SwinIR（Image Restoration using Swin Transformer）［18］：基于Swin Transformer 模块进行图像恢复任务的算法。

6）Uformer［23］：一种利用Transformer 块构造的U 形网络，主要用于图像恢复任务。

此外，本文还对SD-SacUNet 算法本身进行了消融实验，分别验证了自蒸馏方法和SacUNet 结构的有效性。本文设置四组消融实验，具体如下：

1）UNet：采用常规训练方法训练UNet，不添加额外的损失函数。

2）SD-UNet：采用自蒸馏方法进行训练的UNet。

3）SacUNet：采用常规训练方法训练SacUNet，不添加额外损失函数。

4）SD-SacUNet：采用自蒸馏方法进行训练SacUNet。

所有算法的实验均采用五折交叉验证，取五折实验结果的平均值，并计算方差得到最终的实验数据。在网络层数和超参数的选择上尽可能地设置为与本文算法相同或近似，以防止参数不统一导致的结果差异。

实验结果定量评价主要使用峰值信噪比（Peak Signal-to-Noise Ratio，PSNR）和结构相似性（Structure Similarity Index Measure，SSIM）判断MRI 进行去伪影重建后的输出结果。其中：PSNR 旨在降低噪声，是评价图像重建质量的标准；SSIM 关注图像的可见结构，从亮度、对比度和结构三个方面衡量，是视觉质量的度量。利用这两个指标可以定量判断网络处理后的图像质量，数值越大，重建效果越好。实验结果定性评估使用可视化图像显示重建后的无伪影图像。

2.3 实现细节

在网络编码器部分，每个特征提取模块均由最大池化层，卷积层和线性整流函数（Rectified Linear Unit，ReLU）激活层构成，其中，卷积层和激活层将重复两次，输入特征经过特征提取模块的处理后，其通道数翻倍，特征面积不断缩小，感受野不断扩大。输入端的卷积层将图像通道数为1 变化为32，后接的4 个特征提取模块输出通道数分别为64，128，256，512。Swin Transformer 模块的窗口大小按照特征图由大到小设置为16，8，4，2。解码器端的图像恢复模块通过反卷积进行上采样，后接卷积层和激活层，在对自注意力连接块输出进行合并后，将再进行一次卷积和激活的操作。解码器后跟全连接层，将通道数由32 合并为1 通道的图像输出，即为无伪影重图像。自蒸馏瓶颈层在训练阶段参与网络的训练，并在测试阶段去除。

本文将输入图像裁剪为160×256 像素，去除了非组织区域。采用PyTorch 框架进行实验，利用Adam 优化器进行反向传播。

2.4 实验结果定量分析

表1 和表2 分别给出了对比实验和消融实验的定量结果。通过对比实验，可以看出：UNet++通过嵌套的网络结构对UNet 的跳连接进行了改进，获得了比UNet 更好的实验结果。经典的去吉布斯伪影网络GRACNN 通过设计的CNN 结构提取伪影信息，本文的SD-SacUNet 比GRACNN 在PSNR 指标上高出0.77 dB，SSIM 也提高了0.018 3。与其他两个基于UNet 与Transformer 的算法UCTransNet 和Uformer相比，在PSNR上分别高出了0.35 dB、0.10 dB，SSIM 分 别提升 了0.004 1、0.000 6。与基于Swin Transformer 模块的图像恢复算法SwinIR 相比，两项指标分别提高了0.14 dB、0.003 3。总的来说，本文算法结果在PSNR 和SSIM 这两个图像重建指标上均有所提升，图像重建效果相较于一些经典算法框架取得了一定进步。

表1 CC359数据集在不同算法中的重建结果Tab.1 Reconstruction results of different algorithms on CC359 dataset

表2 CC359数据集上的消融实验结果Tab.2 Ablation experimental results on CC359 dataset

表2 所作消融实验将SD-UNet 与UNet 进行对比，PSNR高出了0.19 dB，SSIM 提升了0.014，通过对比得出自蒸馏训练方法在图像重建任务可以在一定程度上提升精度。通过将SacUNet 和UNet 进行对比，两项指标分别提升了0.88 dB、0.026 8，由此可以得出基于自注意力连接的UNet 比传统的UNet 具有更高的重建水平。SD-SacUNet 在两项指标中均获得了明显的提升，说明基于自蒸馏训练的自注意力连接UNet 结构改进均有效，提升了网络整体对于MRI 去吉布斯伪影重建的效果。

2.5 实验结果定性分析

本文在处理网络输入图像的过程中去除了一部分的高频信息，因此，输入图像不仅含有吉布斯伪影，还含有一定程度影响图像重构细节的噪声。根据可视化结果图3，以及表1 和表2 定量结果进行综合分析，UNet、UNet++与GRACNN这类基于CNN 的方法在去除吉布斯伪影的同时，边缘区域存在一定的模糊，甚至是不正确的纹理。由于PSNR 指标并不高，图像中仍存在部分噪声，对于局部组织结构重建产生不良影响，并且重建后的图像过于平滑，可能损失一部分结构的细节信息。基于Transformer 经典模块的算法，如Uformer 和SwinIR 以及基于UNet 与Transformer 相结合的算法，如UCTransNet，这三个算法均取得了比较好的去伪影重建效果，这说明了此类利用CNN 和Transformer 相结合的算法拥有良好的特征提取能力，证明了二者结合的有效性。与Uformer、SwinIR 等基于多个堆叠的Transformer 基本模块的图像恢复算法相比，本文算法仅在UNet 编解码特征融合阶段加入了Transformer 子模块，通过将UNet 与Transformer模块相结合的方式，降低了参数量，节省了大量的运算空间，并且通过自蒸馏训练，提高了重建结果。

图3 不同算法在CC359数据集上的去伪影重建效果Fig.3 Reconstruction results of different algorithms with removal of Gibbs artifacts on CC359 dataset

通过细节比对（例如图3 中箭头所指区域），可以看出本文算法SD-SacUNet 可以很好地去除吉布斯伪影，恢复高频细节并减轻模糊瑕疵，在保证原始结构的同时，获得更清晰的结构和自然的纹理，并且不产生新的伪影，边缘清晰，对不平滑的边缘也有一定的理解能力，具有较高的对比度和清晰度。定量和定性分析证明了本文算法架构在MRI 去伪影重建的特征表示方面是高效的。

3 结语

本文提出一种利用自注意力连接模块替换UNet 原有跳连接的网络结构，采用了自蒸馏方法进行训练，应用于MRI去除吉布斯伪影的重建任务。通过UNet 编码结构完成对特征的提取，自注意力连接模块将编码端不同深度输出特征进行注意力机制的计算。完成自注意力建模的特征与解码器输出特征进行融合，在解码阶段帮助图像的重建。在图像恢复过程中，自蒸馏训练利用深层精细化特征指导浅层网络的训练，提升了训练效果。通过实验对比表明，本文算法在去除吉布斯伪影的同时，能够保留图像的细节，网络结构比较简单，没有采用复杂的嵌套结构对不同尺度的特征进行融合和分配，保留了同分辨率特征的细节信息。与传统图像处理方法相比，深度学习算法运行速度快，不需要复杂的先验知识和边缘检测等技术，具有更高的鲁棒性，在临床大规模应用也有很大优势，可以帮助医生辅助诊断。具有一定的研究意义和临床应用价值。

本文算法仍然需要进一步优化，通过改变注意力机制的运算方法，减小计算量，在目前UNet 与Transformer 相结合的研究热点中，进一步挖掘其融合的可行性。对于卷积神经网络的处理可以考虑与注意力机制相结合，实时调整网络节点权重，以便在UNet 特征提取阶段，在获取更高级语义特征的同时保留更多的细节传递到解码器，进一步优化重建图像细节。在子块融合时，可以考虑多分辨率信息的交互，或者基于注意力机制建立新型的特征融合方式，关注在特征融合方式上的多样化。