一種基于卷積多層感知器（MLP）改進(jìn)U型架構的方法

隨著(zhù)醫學(xué)圖像的解決方案變得越來(lái)越適用，我們更需要關(guān)注使深度網(wǎng)絡(luò )輕量級、快速且高效的方法。具有高推理速度的輕量級網(wǎng)絡(luò )可以被部署在手機等設備上，例如 POCUS（point-of-care ultrasound）被用于檢測和診斷皮膚狀況。這就是 UNeXt 的動(dòng)機。

方法概述

之前我們解讀過(guò)基于 Transformer 的 U-Net 變體，近年來(lái)一直是領(lǐng)先的醫學(xué)圖像分割方法，但是參數量往往不樂(lè )觀(guān)，計算復雜，推理緩慢。這篇文章提出了基于卷積多層感知器（MLP）改進(jìn) U 型架構的方法，可以用于圖像分割。設計了一個(gè) tokenized MLP 塊有效地標記和投影卷積特征，使用 MLPs 來(lái)建模表示。這個(gè)結構被應用到 U 型架構的下兩層中（這里我們假設縱向一共五層）。

文章中提到，為了進(jìn)一步提高性能，建議在輸入到 MLP 的過(guò)程中改變輸入的通道，以便專(zhuān)注于學(xué)習局部依賴(lài)關(guān)系特征。還有額外的設計就是跳躍連接了，并不是我們主要關(guān)注的地方。最終，UNeXt 將參數數量減少了 72 倍，計算復雜度降低了 68 倍，推理速度提高了 10 倍，同時(shí)還獲得了更好的分割性能，如下圖所示。

UNeXt 架構

UNeXt 的設計如下圖所示?？v向來(lái)看，一共有兩個(gè)階段，普通的卷積和 Tokenized MLP 階段。其中，編碼器和解碼器分別設計兩個(gè) Tokenized MLP 塊。每個(gè)編碼器將分辨率降低兩倍，解碼器工作相反，還有跳躍連接結構。每個(gè)塊的通道數（C1-C5）被設計成超參數為了找到不掉點(diǎn)情況下最小參數量的網(wǎng)絡(luò )，對于使用 UNeXt 架構的實(shí)驗，遵循 C1 = 32、C2 = 64、C3 = 128、C4 = 160 和 C5 = 256。

TokMLP 設計思路

關(guān)于 Convolutional Stage 我們不做過(guò)多介紹了，在這一部分重點(diǎn)專(zhuān)注 Tokenized MLP Stage。從上一部分的圖中，可以看到 Shifted MLP 這一操作，其實(shí)思路類(lèi)似于 Swin transformer，引入基于窗口的注意力機制，向全局模型中添加更多的局域性。下圖的意思是，Tokenized MLP 塊有 2 個(gè) MLP，在一個(gè) MLP 中跨越寬度移動(dòng)特征，在另一個(gè) MLP 中跨越高度移動(dòng)特征，也就是說(shuō)，特征在高度和寬度上依次移位。

論文中是這么說(shuō)的：“我們將特征分成 h 個(gè)不同的分區，并根據指定的軸線(xiàn)將它們移到 j=5 的位置”。其實(shí)就是創(chuàng )建了隨機窗口，這個(gè)圖可以理解為灰色是特征塊的位置，白色是移動(dòng)之后的 padding。

解釋過(guò) Shifted MLP 后，我們再看另一部分：tokenized MLP block。首先，需要把特征轉換為 tokens（可以理解為 Patch Embedding 的過(guò)程）。為了實(shí)現 tokenized 化，使用 kernel size 為 3 的卷積，并將通道的數量改為 E，E 是 embadding 嵌入維度（ token 的數量），也是一個(gè)超參數。然后把這些 token 送到上面提到的第一個(gè)跨越寬度的 MLP 中。

這里會(huì )產(chǎn)生了一個(gè)疑問(wèn)，關(guān)于 kernel size 為 3 的卷積，使用的是什么樣的卷積層？答：這里還是普通的卷積，文章中提到了 DWConv（DepthWise Conv），是后面的特征通過(guò) DW-Conv 傳遞。使用 DWConv 有兩個(gè)原因：（1）它有助于對 MLP 特征的位置信息進(jìn)行編碼。MLP 塊中的卷積層足以編碼位置信息，它實(shí)際上比標準的位置編碼表現得更好。像 ViT 中的位置編碼技術(shù)，當測試和訓練的分辨率不一樣時(shí)，需要進(jìn)行插值，往往會(huì )導致性能下降。（2）DWConv 使用的參數數量較少。

這時(shí)我們得到了 DW-Conv 傳遞過(guò)來(lái)的特征，然后使用 GELU 完成激活。接下來(lái)，通過(guò)另一個(gè) MLP（跨越height）傳遞特征，該 MLP 把進(jìn)一步改變了特征尺寸。在這里還使用一個(gè)殘差連接，將原始 token 添加為殘差。然后我們利用 Layer Norm（LN），將輸出特征傳遞到下一個(gè)塊。LN 比 BN 更可取，因為它是沿著(zhù) token 進(jìn)行規范化，而不是在 Tokenized MLP 塊的整個(gè)批處理中進(jìn)行規范化。上面這些就是一個(gè) tokenized MLP block 的設計思路。

此外，文章中給出了 tokenized MLP block 涉及的計算公式：

其中 T 表示 tokens，H 表示高度，W 表示寬度。值得注意的是，所有這些計算都是在 embedding 維度 H 上進(jìn)行的，它明顯小于特征圖的維度 HN×HN，其中 N 取決于 block 大小。在下面的實(shí)驗部分，文章將 H 設置為 768。

實(shí)驗部分

實(shí)驗在 ISIC 和 BUSI 數據集上進(jìn)行，可以看到，在 GLOPs、性能和推理時(shí)間都上表現不錯。

下面是可視化和消融實(shí)驗的部分?？梢暬瘓D可以發(fā)現，UNeXt 處理的更加圓滑和接近真實(shí)標簽。

消融實(shí)驗可以發(fā)現，從原始的 UNet 開(kāi)始，然后只是減少過(guò)濾器的數量，發(fā)現性能下降，但參數并沒(méi)有減少太多。接下來(lái)，僅使用 3 層深度架構，既 UNeXt 的 Conv 階段。顯著(zhù)減少了參數的數量和復雜性，但性能降低了 4%。加入 tokenized MLP block 后，它顯著(zhù)提高了性能，同時(shí)將復雜度和參數量是一個(gè)最小值。

接下來(lái)，我們將 DWConv 添加到 positional embedding，性能又提高了。接下來(lái)，在 MLP 中添加 Shifted 操作，表明在標記化之前移位特征可以提高性能，但是不會(huì )增加任何參數或復雜性。注意：Shifted MLP 不會(huì )增加 GLOPs。

一些理解和總結

在這項工作中，提出了一種新的深度網(wǎng)絡(luò )架構 UNeXt，用于醫療圖像分割，專(zhuān)注于參數量的減小。UNeXt 是一種基于卷積和 MLP 的架構，其中有一個(gè)初始的 Conv 階段，然后是深層空間中的 MLP。具體來(lái)說(shuō)，提出了一個(gè)帶有移位 MLP 的標記化 MLP 塊。在多個(gè)數據集上驗證了 UNeXt，實(shí)現了更快的推理、更低的復雜性和更少的參數數量，同時(shí)還實(shí)現了最先進(jìn)的性能。

我在讀這篇論文的時(shí)候，直接注意到了它用的數據集。我認為 UNeXt 可能只適用于這種簡(jiǎn)單的醫學(xué)圖像分割任務(wù)，類(lèi)似的有 Optic Disc and Cup Seg，對于更復雜的，比如血管，軟骨，Liver Tumor，kidney Seg 這些，可能效果達不到這么好，因為運算量被極大的減少了，每個(gè) convolutional 階段只有一個(gè)卷積層。MLP 魔改 U-Net 也算是一個(gè)嘗試，在 Tokenized MLP block 中加入 DWConv 也是很合理的設計。

審核編輯：劉清