摘要

得益于Transformer模型的成功，近期研究開始探索其在3D醫學分割任務中的適用性。在Transformer模型中，自注意力機制是核心構建模塊之一，與基于局部卷積的設計相比，它致力于捕捉長距離依賴關系。然而，自注意力操作存在二次復雜度問題，這已成為計算瓶頸，尤其在體醫學成像中——此類輸入為3D形式且包含大量切片。本文提出一種3D醫學圖像分割方法，名為UNETR++，該方法既能生成高質量分割掩碼，又在參數數量、計算成本和推理速度方面具備高效性。我們設計的核心是引入一種新型高效配對注意力（EPA）模塊：該模塊通過一對基于空間注意力和通道注意力的互依賴分支，高效學習空間和通道維度的判別性特征。我們提出的空間注意力公式具有高效性，其復雜度與輸入呈線性關系。為實現空間聚焦分支與通道聚焦分支之間的信息交互，我們共享查詢和鍵映射函數的權重——這不僅帶來互補優勢（配對注意力），還降低了整體復雜度。我們在五個基準數據集（Synapse、BTCV、ACDC、BraTS和Decathlon-Lung）上進行了大量評估，結果證實了所提方法在效率和準確性上的優勢。在Synapse數據集上，UNETR++的Dice評分達到87.2%，刷新了當前最優性能；同時與文獻中最佳方法相比，其參數數量和計算量（FLOPs）均減少了71%以上。我們的代碼和模型可通過以下鏈接獲取：https://tinyurl.com/2p87x5xn。

關鍵詞——深度學習、高效注意力、混合架構、醫學圖像分割?

INTRODUCTION

體素級（3D）分割是醫學成像領域的基礎性問題，具有眾多應用場景，包括用于診斷目的的腫瘤識別和器官定位等[1][2]。該任務通常通過類U-Net[3]的編碼器-解碼器架構來解決：編碼器生成3D圖像的分層低維表征，解碼器則將這種學習到的表征映射為體素級分割結果。早期基于CNN（卷積神經網絡）的方法在編碼器和解碼器中分別使用卷積和反卷積操作，但難以獲得精確的分割結果。近期有研究[4]嘗試解決基于CNN方法的局限性——通過擴大CNN的感受野，增強其對上下文表征的建模能力。盡管在編碼更優上下文表征方面取得了進展，但感受野有限、局部連接以及權重固定等問題，仍是基于CNN方法面臨的挑戰。這些挑戰可能影響其捕捉廣泛全局依賴關系的效果。此外，雖然基于3D CNN的方法通過權重共享實現了參數層面的高效性，但由于操作數量較多，其計算量（FLOPs）相應增加，推理速度也隨之變慢，這是一種性能上的權衡。相比之下，基于Transformer的方法本身具有全局性，近期已展現出有競爭力的性能，但代價是模型復雜度提高。具體而言，這類方法依賴全局自注意力機制，而該機制的復雜度與輸入呈二次關系。而且，在輸入為3D體數據的體素級醫學圖像分割任務中，這種復雜度問題更為突出。

為解決這些挑戰，探索融合CNN和Transformer優勢的混合架構的研究熱度日益上升。近期已有多項研究[1][5][6]嘗試設計混合架構，以結合局部卷積和全局注意力的優點。其中，部分方法[1]采用基于Transformer的編碼器搭配卷積解碼器，另一些方法[5][6]則致力于為編碼器和解碼器子網絡設計混合模塊。然而，這些研究主要聚焦于提高分割精度，這反過來導致模型在參數數量和計算量（FLOPs）上大幅增加，穩健性不盡如人意。我們認為，這種不理想的穩健性可能源于其低效的自注意力設計——這在體素級醫學圖像分割任務中問題更為明顯。此外，現有方法未能捕捉空間特征與通道特征之間的顯式依賴關系，而這種依賴關系本可以提升分割質量。在本研究中，我們旨在通過一個統一框架，同時提高分割精度和模型效率。

A. Contributions

我們提出了一種用于3D醫學圖像分割的高效混合分層架構，名為UNETR++。該架構力求在分割精度以及參數數量、計算量（FLOPs）、GPU內存消耗和推理速度等效率相關方面均實現提升。我們提出的UNETR++分層方法以近期的UNETR框架[1]為基礎，引入了一種新型高效配對注意力（EPA）模塊。該模塊通過在兩個分支中分別應用空間注意力和通道注意力，高效捕捉豐富的、互相關聯的空間特征與通道特征。EPA中的空間注意力將鍵（keys）和值（values）投影到固定的低維空間，使得自注意力計算的復雜度與輸入標記（tokens）數量呈線性關系。另一方面，EPA中的通道注意力通過在通道維度上對查詢（queries）和鍵執行點積運算，著重捕捉通道特征圖之間的依賴關系。此外，為了強化空間特征與通道特征之間的關聯性，兩個分支共享查詢和鍵的權重——這同時還有助于控制網絡的參數數量。與之相對，值的權重保持獨立，以確保兩個分支能夠學習到互補的特征。

UNETR++相比現有方法具有顯著優勢：它大幅降低了GPU內存消耗，同時提升了推理速度并減少了計算量。這一點對于3D分割任務至關重要，因為此類任務本身復雜度較高。例如，與nnFormer[5]相比，UNETR++不僅性能更優，GPU推理速度提升2.4倍，GPU內存消耗減少5倍，計算量減少6倍；與Swin-UNETR[6]相比，UNETR++的推理速度提升3.6倍，GPU內存消耗和計算量均減少8倍；此外，相較于高效的基于CNN的方法nnUNet[2]，UNETR++的內存消耗降低35%，計算量減少6倍。我們優先考慮高效的資源利用，從而在CPU和GPU上均實現了更快的推理速度。而且我們認為，上述這些混合方法難以有效捕捉特征通道之間的相互依賴關系，因此無法獲得同時編碼空間信息和通道間特征依賴的豐富特征表征。在本研究中，我們致力于在一個統一的混合分割框架中集中解決上述問題。

我們進行了全面評估，以檢驗UNETR++在五個廣泛使用的基準數據集上的性能，包括Synapse[7]、BTCV[7]、ACDC[8]、BraTS[9]和Decathlon-Lung[10]。評估結果證實了我們所提方法在效率和精度上的有效性，也體現出UNETR++在不同模態（包括CT掃描和MRI）的多樣化數據集上具有出色的泛化能力。在Synapse數據集上，UNETR++的Dice評分為87.2%，與現有最優方法[5]相比，參數數量和計算量均顯著減少71%以上；在ACDC數據集上，UNETR++的平均Dice評分為93.83%，比近期提出的MexNeXt[4]高出1.4%；在BraTS和Decathlon-Lung分割任務中，UNETR++在三項評估指標上均優于當前最優方法，在實現更優分割性能的同時，還具備更快的推理速度，且對GPU內存的需求顯著降低。

II. RELATED WORK

A. CNN-Based Segmentation Methods

自U-Net設計[3]提出以來，多種基于CNN的方法[11]、[12]、[13]、[14]對標準U-Net架構進行了擴展，以適用于各類醫學圖像分割任務。在3D醫學圖像分割領域[15]、[16]、[17]、[18]、[19]，完整的體素圖像通常被當作一系列2D切片來處理。已有多項研究探索了用于捕捉上下文信息的分層框架。Milletari等人[18]提出對體素圖像進行下采樣以降低分辨率，通過這種3D表征方式保留有益的圖像特征。?i?ek等人[17]將U-Net架構擴展到體素分割任務中，用3D操作替代2D操作，并從稀疏標注的體素圖像中學習特征。Isensee等人[2]提出了名為nnUNet的通用分割框架，該框架能自動配置架構，以提取多尺度特征。Roth等人[20]設計了多尺度3D全卷積網絡，用于從不同分辨率的圖像中學習表征，實現多器官分割。此外，文獻中已有多項研究嘗試在基于CNN的框架中編碼整體上下文信息，例如借助圖像金字塔[21]、大卷積核[22]、空洞卷積[23]和可變形卷積[24]等方式。近期，Roy等人[4]提出了名為MedNeXt的全卷積3D編解碼器網絡。該架構是ConNeXt框架[25]的擴展版本，融合了自適應核大小與殘差連接，旨在數據有限的3D醫學成像場景中提升分割精度。盡管MedNeXt-M-K3[4]展現出了良好的精度表現，但其代價是GPU內存消耗較nnUNet[2]增加了4.1倍，導致GPU和CPU的推理速度顯著變慢。

B. Transformers-Based Segmentation Methods

視覺Transformer（ViT）近年來備受關注，這得益于其編碼長距離依賴關系的能力——該能力使其在包括分類[26]和檢測[27]在內的各類視覺任務中均取得了良好效果。Transformer架構的核心構建模塊之一是自注意力操作，它能對圖像塊序列間的交互進行建模，進而學習全局關聯。近期已有少量研究[28]、[29]、[30]、[31]嘗試在Transformer框架中緩解標準自注意力操作的復雜度問題。然而，這些研究大多聚焦于分類任務，尚未在密集預測任務中展開探索。

在醫學圖像分割領域，近期僅有少數研究[32]、[33]對純Transformer設計進行了探索。Karimi等人[32]提出將體素圖像劃分為3D塊，隨后將這些塊展平以構建1D嵌入，并將其輸入主干網絡以獲取全局表征。Cao等人[33]則針對2D醫學圖像分割提出了一種含移位窗口的架構：先將圖像劃分為塊，再將其輸入U形編碼器-解碼器以進行局部-全局表征學習。

C. Hybrid Segmentation Methods

除了純CNN或基于Transformer的設計之外，近期多項研究[1]、[5]、[19]、[34]、[35]、[36]探索了混合架構，通過結合卷積與自注意力操作來實現更優的分割效果。TransFuse[34]提出了一種并行的CNN-Transformer架構，該架構配備BiFusion模塊，用于在編碼器中融合多尺度特征。MedT[19]在自注意力中引入門控位置敏感軸向注意力機制，以控制編碼器中的位置嵌入信息，同時解碼器中的ConvNet模塊生成分割模型。TransUNet[35]將Transformer與U-Net架構相結合——Transformer對來自卷積特征的嵌入圖像塊進行編碼，解碼器則將上采樣后的編碼特征與高分辨率CNN特征融合以實現定位。Ds-transunet[36]采用基于Swin Transformer[37]的雙尺度編碼器，用于處理多尺度輸入，并通過自注意力從不同語義尺度編碼局部和全局特征表征。Hatamizadeh等人[1]提出了3D混合模型UNETR，該模型將Transformer的長距離空間依賴能力與CNN的歸納偏置融入“U形”編碼器-解碼器架構。UNETR中的Transformer塊主要用于編碼器，以提取固定的全局表征，隨后在多個分辨率下與基于CNN的解碼器融合。Zhou等人[5]提出了名為nnFormer的方法，該方法對SwinUNet[33]架構進行了適配：卷積層將輸入掃描圖像轉換為3D塊，并引入基于體數據的自注意力模塊來構建分層特征金字塔。盡管nnFormer取得了良好的性能，但與UNETR及其他混合方法相比，其計算復雜度顯著更高。CoTr[38]是一種混合架構，由CNN編碼器和高效可變形Transformer組成——卷積編碼器提取局部特征圖，可變形Transformer僅關注少數關鍵位置，并對提取到的特征表征中的部分依賴關系進行編碼。

如上所述，UNETR[1]、nnFormer[5]等多數近期混合方法相較于純CNN和純Transformer方法，分割性能有所提升。然而我們注意到，這些混合方法為提高分割精度，付出了模型規模（從參數數量和計算量（FLOPs）來看）大幅增加的代價，這可能進一步導致穩健性不佳。例如，UNETR雖取得了理想的精度，但與現有最優的基于CNN的nnUNet[2]相比，其參數數量多出2.5倍；而nnFormer雖性能優于UNETR，但參數數量進一步增加1.6倍，計算量（FLOPs）增加2.8倍。

D. Efficient Attention Methods

近年來，為2D視覺應用設計高效注意力模塊受到了廣泛關注。CBAM[39]是一種基于卷積神經網絡的高效注意力模塊，該模塊對特征圖沿空間和通道維度進行高效處理，隨后將生成的注意力圖與輸入特征圖相乘，從而實現自適應特征優化。盡管我們提出的EPA模塊也對通道和空間信息進行編碼，但EPA模塊與CBAM存在重大差異：（1）CBAM通過兩個串行子模塊對空間表征和通道表征進行編碼，而EPA模塊則采用并行方式對二者進行編碼；（2）EPA模塊中的通道注意力和空間注意力基于注意力機制構建，而CBAM的注意力則基于池化和卷積操作；（3）CBAM中的通道注意力和空間注意力擁有獨立的權重矩陣，而在EPA中，我們共享查詢（Q）和鍵（K）矩陣的權重——這種共享機制通過僅學習互補特征，將參數數量減少了25%，同時將性能提升了0.23%。

Transformer-CBAM[40]是對CBAM[39]的增強版本，它通過集成多尺度Transformer模塊，實現了對不同尺度上下文信息的建模，因此在遙感圖像變化檢測任務中尤為有效。Squeeze-and-Excitation（SE）[41]則聚焦于通道特征圖之間的關系，其提出的“擠壓-激勵”（SE）模塊通過顯式捕捉通道間的相互依賴關系，對特定通道的特征圖進行重新調整。Attention Gated U-Net[42]對U-Net[3]架構進行了擴展，為醫學成像場景引入了注意力門控（AG）模塊——該AG模塊通過不同的門控機制學習對不同形狀和大小的目標結構進行優先級排序，且計算開銷極小。

盡管上述多數方法在效率和準確性之間取得了良好的平衡，但它們主要是為2D視覺任務設計的。為驗證所提EPA模塊在3D醫學分割場景中的有效性，我們通過將EPA模塊替換為這些高效方法的3D對應版本進行了實驗，相關細節詳見表IX。

III. METHOD

A. Overall Architecture

圖2展示了我們的UNETR++架構，該架構包含分層的編碼器-解碼器結構。我們的UNETR++框架以近期提出的UNETR[1]為基礎，在編碼器和解碼器之間設有跳躍連接，后續通過卷積模塊（ConvBlocks）生成預測掩碼。與在整個編碼器中使用固定特征分辨率的設計不同，UNETR++采用分層設計——在每個階段，特征分辨率會以2倍的比例逐步降低。

在UNETR++框架中，編碼器包含4個階段，這4個階段的通道數分別為[C1, C2, C3, C4]。第一階段由補丁嵌入（patch embedding）和我們提出的新型高效配對注意力（EPA）模塊組成：補丁嵌入會將體素輸入分割為3D補丁，具體而言，將3D輸入（體數據）x ∈ R^H×W×D分割為非重疊的補丁xu ∈ R^N×(P1,P2,P3)，其中（P1, P2, P3）為每個補丁的分辨率，N = (H/P1 × W/P2 × D/P3)表示序列長度；隨后，這些補丁被投影到C個通道維度，生成尺寸為H/P1 × W/P2 × D/P3 × C的特征圖。我們采用與文獻[5]相同的補丁分辨率（4, 4, 2）。對于編碼器的其余每個階段，我們先通過非重疊卷積下采樣層將分辨率降低2倍，再接入EPA模塊。

在我們提出的UNETR++中，每個EPA模塊包含兩個注意力子模塊。它們通過共享鍵-查詢（keys-queries）機制對空間和通道維度的信息進行編碼，從而高效學習豐富的空間-通道特征表征。編碼器各階段通過跳躍連接與解碼器各階段相連，以融合不同分辨率下的輸出——這能夠恢復下采樣操作中丟失的空間信息，進而得到更精確的預測結果。

與編碼器類似，解碼器也包含4個階段。每個解碼器階段先通過反卷積上采樣層將特征圖分辨率提高2倍，再接入EPA模塊（最后一個解碼器階段除外）。相鄰兩個解碼器階段之間的通道數以2倍的比例減少。因此，最后一個解碼器的輸出會與卷積特征圖融合，以恢復空間信息并增強特征表征能力。融合后的輸出隨后被輸入至3×3×3和1×1×1卷積模塊，生成體素級的最終掩碼預測。接下來，我們將詳細介紹EPA模塊。

B. Efficient Paired-Attention Block

如前所述，大多數現有混合方法所采用的自注意力操作，其復雜度與標記（tokens）數量呈二次關系。這在體素分割任務中計算成本極高，而當在混合設計中交錯使用窗口注意力和卷積組件時，問題會更為突出。此外，通過將鍵（keys）和值（values）的空間矩陣投影到低維空間，可高效學習空間注意力信息。有效結合空間維度的交互關系與通道特征間的依賴關系，能夠得到豐富的上下文空間-通道特征表征，進而提升掩碼預測效果。

所提出的EPA模塊可執行高效全局注意力計算，并有效捕捉豐富的空間-通道特征表征。該模塊包含空間注意力子模塊和通道注意力子模塊：空間注意力子模塊將自注意力的復雜度從二次降至線性；通道注意力子模塊則高效學習通道特征圖之間的依賴關系。EPA模塊通過在兩個注意力子模塊之間共享鍵-查詢（keys-queries）實現信息互通，從而生成更優且高效的特征表征。這可能得益于通過共享鍵和查詢、同時使用不同值（value）層來學習互補特征。為增強訓練穩定性，每個EPA模塊的起始處均采用了層歸一化（LayerNorm）操作。

如圖2（右側）所示，輸入特征圖x被送入EPA模塊的通道注意力子模塊和空間注意力子模塊。查詢（Q）和鍵（K）線性層的權重在兩個注意力子模塊之間共享，且每個注意力子模塊使用不同的值（V）層。兩個注意力子模塊的計算方式如下：

其中，X??、X?c、SA?和SAc分別表示空間注意力圖、通道注意力圖、空間注意力子模塊和通道注意力子模塊。Q?????d、K?????d、V???????和Vc??????分別為共享查詢矩陣、共享鍵矩陣、空間值層矩陣和通道值層矩陣。?

1) 空間注意力：在該子模塊中，我們致力于通過將復雜度從O(n2)降至O(np)來高效學習空間信息（其中n為標記數量，p為投影向量的維度，且p遠小于n）。給定形狀為hwd×C的歸一化張量X，我們通過三個線性層計算Q?????d（共享查詢）、K?????d（共享鍵）和V???????（空間值）的投影，得到Q?????d=W?X、K?????d=W?X和V???????=W?X（維度均為hwd×C）——其中W?、W?和W?分別為Q?????d、K?????d和V???????的投影權重。

之后，我們執行三個步驟：第一步，將K?????d和V???????層從hwd×C投影到形狀為p×C的低維矩陣；第二步，通過將Q?????d層與投影后的K?????d的轉置相乘，再經過softmax運算計算空間注意力圖，以此衡量每個特征與其他空間特征之間的相似度；第三步，將這些相似度與投影后的V???????層相乘，生成形狀為hwd×C的最終空間注意力圖。空間注意力的定義如下：

其中，Q?????d（共享查詢）、K????（投影后的共享鍵）和????????（投影后的空間值層）分別表示共享查詢矩陣、投影后的共享鍵矩陣和投影后的空間值層矩陣，d為每個向量的維度。

2) 通道注意力：該子模塊通過在通道維度上對通道值層與通道注意力圖執行點積運算，捕捉特征通道之間的依賴關系。我們使用與空間注意力子模塊相同的Q?????d（共享查詢）和K?????d（共享鍵），并通過線性層計算通道的 值層以學習互補特征，得到V???????=W?X（維度為hwd×C）——其中W?為V???????（通道值層）的投影權重。通道注意力的定義如下：?

其中，V???????（通道值層）、Q?????d（共享查詢）和K?????d（共享鍵）分別表示通道值層、共享查詢矩陣和共享鍵矩陣。最后，我們對兩個注意力子模塊的輸出進行求和融合，并通過卷積模塊對融合結果進行轉換，以獲得豐富的特征表征。EPA模塊的最終輸出X?計算如下：?

其中，X??和X?c分別表示空間注意力圖和通道注意力圖，Conv1和Conv3分別為1×1×1和3×3×3卷積模塊。

C. Loss Function

遵循基線模型UNETR [1]和nnFormer [5]，我們的損失函數基于常用的軟Dice損失[18]與交叉熵損失的總和，以同時利用這兩種互補損失函數的優勢。其定義如下：

其中，I表示類別數量，V表示體素數量，Yv,i和Pv,i分別表示體素v處類別i的真實標簽和輸出概率。

IV. EXPERIMENTS

V. CONCLUSION AND DISCUSSION

我們提出了一種名為UNETR++的分層方法，用于3D醫學分割。UNETR++引入了高效配對注意力（EPA）模塊，通過空間注意力和通道注意力對豐富的、互依賴的空間特征與通道特征進行編碼。在EPA模塊中，我們共享查詢和鍵映射函數的權重，以促進空間分支與通道分支之間的信息交互，這不僅帶來了互補優勢，還減少了參數數量。

我們的UNETR++在五個數據集（Synapse、ACDC、BTCV、BraTS和Decathlon-Lung）上取得了優異的分割結果，同時與現有最佳方法相比，顯著降低了模型復雜度（從參數數量、計算量（FLOPs）來看）。此外，我們的研究表明，UNETR++的GPU內存消耗更低——這是3D分割任務的關鍵因素，且在GPU和CPU平臺上均能以更快的推理速度運行。因此，UNETR++提供了一種更通用且資源效率更高的解決方案，提升了在移動平臺部署醫學分割模型以進行實時醫學圖像分析的可行性。

所提出的EPA模塊具有通用性，可應用于其他研究中。為驗證這一點，我們在TransBTS[43]的框架中，用所提出的EPA模塊替換了其自注意力模塊，并在Synapse數據集上對更新后的模型進行了評估。結果顯示，Dice相似系數（DSC）提升了1.2%（從83.28%升至84.47%），歸一化表面距離（NSD）提升了0.8%，95%豪斯多夫距離（HD95）從12.34顯著降至7.92，平均對稱表面距離（MASD）從3.65降至3.11。

為探究UNETR++可能存在的局限性，我們對不同的異常案例進行了分析。盡管我們的預測結果優于現有方法，且更接近真實標簽，但我們發現，在少數案例中，我們的模型以及大多數現有方法都難以對某些器官進行分割。當少數切片中器官的幾何形狀異常（表現為邊界纖細）時，我們的模型和大多數現有模型都難以對其進行精確分割。原因可能是，與正常形狀的樣本相比，這種異常形狀的訓練樣本數量有限。在某些案例中，與最新的基于CNN的方法MedNeXt[4]相比，我們的模型在歸一化表面距離（NSD）上表現較低，這一定量結果體現了上述定位誤差。其原因可能在于基于CNN的方法具有歸納偏置，擅長捕捉空間層次結構和局部模式。我們計劃通過在預處理階段應用幾何數據增強技術來解決這一問題。