鑽石舞台

很久之前讀了這篇接收在 MICCAI 2021 上的文章，復現調試了代碼之後還沒有及時整理一篇閱讀筆記。由於在 MICCAI 上，這篇文章同樣沒有大量的實驗對比工作，但是提出的方法思路清晰易懂，值得借鑑。arXiv: https://arxiv.org/abs/2107.00781 。作為前置依賴，本篇閱讀筆記首先介紹了 Transformer Architecture 和在醫學圖像分割上的應用；其次，分析了論文中提出的 UTNet 架構（主幹 U-Net，混合 Transformer 等 module）；最後，看了下論文的實驗結果。

在語義分割上，FCN 這類卷積的編碼器-解碼器架構衍生出的模型在過去幾年取得了實質性進展，但這類模型存在兩個局限。第一，卷積僅能從鄰域像素收集信息，缺乏提取明確全局依賴性特徵的能力；第二，卷積核的大小和形狀往往是固定的，因此它們不能靈活適應輸入的圖像或其他內容。相反，Transformer architecture 由於自注意力機制具有捕獲全局依賴特徵的能力，且允許網絡根據輸入內容動態收集相關特徵。

值得注意的是，Transformer 架構的訓練需要的更大的開銷，因為自注意力機制（self-attention）在序列長度方面具有的時間和空間複雜度。基於此，標準的 self-attention 大多數以 patch-wise 方式應用到模型中，比如使用 16 × 16 這種小扁平圖像塊作為輸入序列，或者在來自 CNN 主幹的特徵圖之上對圖像進行編碼，這些特徵圖一般是下採樣後的低分辨率圖像。這裡問題就出來了，對於醫學圖像分割任務目標位置敏感的特殊性，一些欠分割或者過分割的區域都在目標周圍，往往需要高分辨率特徵。此外，有些實驗論證，在 ImageNet 上進行了預訓練，Transformer 要比 ResNet 差，Transformer 真的能適應醫學圖像這種小數據集分割任務嗎？

為了解決上面的問題，文章中提出的 U-Net 混合 Transformer 網絡：UTNet，它整合了卷積和自注意力策略用於醫學圖像分割任務。應用卷積層來提取局部強度特徵，以避免對 Transformer 進行大規模的預訓練，同時使用自注意力來捕獲全局特徵。為了提高分割質量，還提出了一種 efficient self-attention，在時間和空間上將整體複雜度從 O(n2) 顯着降低到接近 O(n)。此外，在 self-attention 模塊中使用相對位置編碼來學習醫學圖像中的內容-位置關係。

上面我們對 Transformer 在醫學圖像分割上的應用和局限做了概述，這裡我們簡單分析下 Transformer 架構，對這部分很了解的讀者可以跳過。

Transformer 建立在多頭自注意機制 (MHSA) 模塊上，MHSA 是由多個 Self-Attention 組成的。下圖是 Self-Attention 的結構，在計算的時候需要用到矩陣 **Q(查詢)，K(鍵值)，V(值)**。在實際中，Self-Attention 接收的是輸入(單詞的表示向量 x 組成的矩陣 X) 或者上一個 Encoder block 的輸出。而 Q,K,V 是通過 Self-Attention 的輸入進行線性變換得到的。

得到矩陣 Q, K, V 之後就可以計算出 Self-Attention 的輸出了，計算的公式如下。其中 d 是 Q,K 矩陣的列數（向量維度），公式中計算矩陣 Q 和 K 每一行向量的內積，為了防止內積過大，因此除以 d 的平方根。將 Q, K, V 展平並轉置為大小為 n × d 的序列，其中 n = HW。P ∈ Rn×n 被命名為上下文聚合矩陣，用作權重以收集上下文信息。

通過這種方式，self-attention 本質上具有全局感受野，擅長捕捉全局依賴。此外，上下文聚合矩陣可以適應輸入內容，以實現更好的特徵聚合。關於更細節的內容，這裡就不多介紹了。需要關注的是，n×d 矩陣的點乘會導致 O(n2d) 複雜度。通常，當特徵圖的分辨率很大時，n 遠大於 d，因此序列長度 n 在自注意力計算中占主導地位，這使得高分辨率特徵圖中應用自注意力是不可行的，例如對於 16 × 16 特徵圖，n = 256，對於 128 × 128 特徵圖，n = 16384。這一點在本篇筆記的概述中提到過。

那麼如何降低 O(n2) 的複雜度呢？由於醫學圖像是高度結構化的數據，除了邊界區域外，局部像素的高分辨率特徵和圖中的其他像素特徵存在相似性，因此，所有像素之間的成對注意力計算往往是低效和冗餘的。所以需要一種高效的自我注意機制，計算如下。

觀察與 2.1 中公式不同之處，主要思想很簡單，是將 K 和 V 矩陣做了一個 Low-dimension Embedding 達到減小計算量的目的，對應的上下文聚合矩陣 P 的 size 也會被修改。通過這樣做，可以把複雜度降低到 O(nkd)。相當於，這裡我們可以做任意下採樣操作，例如平均/最大池化或帶步長卷積操作。比如，使用 1×1 卷積，然後使用雙線性插值對特徵圖進行下採樣，縮小後的 size 是可以推斷的。

基於 MHSA 的 Encoder 和 Decoder block 分別如下圖所示，關於這兩部分如何插入到 U-Net 中，在後面會有解釋。

Transformer 的相對位置編碼大概作用是讓像素間保持空間位置關係（從ViT的實驗看，去掉 position embedding 後，性能會下降3個點以上），對於圖像就是保持二維信息，它對於圖像任務來講很重要。像目前的相對位置編碼設計，都算是將位置編碼整合到了 Attention 的定義中去了，沒有顯式地使用絕對位置編碼。讀者們應該可以發現 2.2 的 MHSA 圖裡，在 softmax 之前的 pair-wise attention logit 計算中使用了像素 i 和 j 的位置編碼，具體的計算如下。

相應的，self-attention 的計算公式又需要做修正了，和之前的區別就是相對寬度和高度是在 low-dimensional projection （低維投影）之後計算的。對 2.2 的高效 self-attention 修正後的計算方式如下。

上一部分我們對 UTNet 中的 Transformer 架構做了清晰的介紹，這一部分會整體解釋下 UTNet 結構，也就是如何把 Transformer 的 encoder 和 decoder 合理加入到主幹 U-Net 中。

如上圖所示 UTNet 結構圖，整體上還是保持 U 型。(b) 是一個經典的殘差塊，傳統的 U-Net 改進方法也是這麼做的，這樣也可以提高分割任務的準確率，避免網絡深度帶來的梯度爆炸和梯度消失等問題，這些都是老生常談了，我們不重點關注。(c) 是一個標準的 Transformer Decoder 設計。可以發現，遵循了 U-Net 的標準設計，但將每個構建塊的最後一個卷積（最高的除外）替換為 2.2 的 Transformer 模塊。此外，低三層的跨層連接也被替換為了 Transformer Decoder，我認為還是很好理解的哈，關於這樣做的目的在筆記前面陸陸續續都說了一些，下面總結下。

這種混合架構可以利用卷積圖像的歸納偏差來避免大規模預訓練，以及 Transformer 捕獲全局特徵關係的能力。由於錯誤分割的區域通常位於感興趣區域的邊界，高分辨率的上下文信息可以在分割中發揮至關重要的作用。因此，重點放在了自我注意模塊上，這使得有效處理大尺寸特徵圖成為可能。沒有將自注意力模塊簡單地集成到來自 CNN 主幹的特徵圖之上，而是將 Transformer 模塊應用於編碼器和解碼器的每個級別，以從多個尺度收集長期依賴關係。請注意，沒有在原始分辨率上應用 Transformer，因為在網絡的非常淺層中添加 Transformer 模塊對實驗沒有幫助，但會引入額外的計算。一個可能的原因是網絡的淺層更多地關注詳細的紋理，其中收集全局上下文特徵效果肯定不理想。

數據集選用 Multi-Centre, Multi-Vendor & Multi-Disease Cardiac Image Segmentation Challenge。

下圖是消融實驗的結果。下表除了參數量和預測時間比 U-Net 等大，其他均是 SOTA。(a) 為不同自注意力位置的影響消融；(b) 為不同高效 Self-attention 的 projection 效果（對應 2.2）；（c）為 Transformer 編碼器、Transformer 解碼器和相對位置編碼的影響消融。

更多的實驗對比和結果可視化如下所示，大多數指標比 CBAM 要優秀。

這篇閱讀筆記大多為個人理解，代碼復現我後面也許還會更新一篇文章，由於一些醫學圖像處理任務數據集的特殊性，vit 在醫學圖像上的應用還需要不斷優化，最近有一些不錯的想法，也歡迎交流和糾正！