基于自注意力連接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 特征提取階段，在獲取更高級語義特征的同時保留更多的細節傳遞到解碼器，進一步優化重建圖像細節。在子塊融合時，可以考慮多分辨率信息的交互，或者基于注意力機制建立新型的特征融合方式，關注在特征融合方式上的多樣化。