Title: 3D UX-Net: a Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation

Author: Ho Hin Lee et al. (范德堡大學(xué))

Paper: https://arxiv.org/abs/2209.15076

Github: https://github.com/MASILab/3DUX-Net

引言

眾所周知，大多數(shù)醫(yī)學(xué)圖像如 MRI 和 CT 是屬于 volumetric data 類型。因此，為了更加充分的利用體素信息，近幾年已經(jīng)提出了不少 3D CNNs 的模型，如 SwinUNETR、UNETR以及筆者前段時(shí)間分享過(guò)的UNETR++等。

整體來(lái)說(shuō)，這些模型性能是越來(lái)越高，在幾個(gè)主流的 3D 數(shù)據(jù)基準(zhǔn)測(cè)試中也實(shí)現(xiàn)了大大小小的 SOTA，特別是 3D 醫(yī)學(xué)圖像分割這塊。當(dāng)然，時(shí)代在進(jìn)步，作為一名高科技前沿從業(yè)者本身也是需要不斷汲取新的知識(shí)營(yíng)養(yǎng)才能不被輕易的淘汰。今天小編就帶大家解讀下 ICLR 2023 新鮮出爐的 3D 醫(yī)學(xué)圖像分割之星——3D UX-Net。

3D UXNet 是一種輕量級(jí)3D卷積神經(jīng)網(wǎng)絡(luò)，其使用 ConvNet 模塊調(diào)整分層 Transformer 以實(shí)現(xiàn)穩(wěn)健的體素分割，在三個(gè)具有挑戰(zhàn)性的腦體積和腹部成像公共數(shù)據(jù)集與當(dāng)前的 SOTA 模型如 SwinUNETR 對(duì)比，同時(shí)在以下三大主流數(shù)據(jù)集性能均達(dá)到了最優(yōu)：

• MICCAI Challenge 2021 FLARE
• MICCAI Challenge 2021 FeTA
• MICCAI Challenge 2022 AMOS

具體的，相比于 SwinUNETR，3D UX-Net 將 Dice 從 0.929 提高到 0.938 (FLARE2021)，Dice 從 0.867 提高到 0.874 (Feta2021)。此外，為了進(jìn)一步評(píng)估 3D UX-Net 的遷移學(xué)習(xí)能力，作者在 AMOS2022 數(shù)據(jù)集上訓(xùn)練得到的模型在另一個(gè)數(shù)據(jù)集上取得了 2.27% Dice 的提升(0.880 → 0.900)。

動(dòng)機(jī)

先來(lái)看看最近提出的一眾基于 Transformer 架構(gòu)的 ViT 模型有什么優(yōu)劣勢(shì)。以 SwinUNETR 為例，其將 ConvNet 的一些先驗(yàn)引入到了 Swin Transformer 分層架構(gòu)中，進(jìn)一步增強(qiáng)了在 3D 醫(yī)學(xué)數(shù)據(jù)集中調(diào)整體素分割的實(shí)際可行性。此類 "Conv+Transforemr" 的組合拳的有效性在很大程度上歸功于以下兩個(gè)因素：

• 非局部自注意力所帶來(lái)的大感受野
• 大量的模型參數(shù)

為此，本文作者想到了應(yīng)用深度卷積以更少的模型參數(shù)來(lái)模擬這些方法讓網(wǎng)絡(luò)學(xué)會(huì)如何捕獲更大感受野的行為。這一點(diǎn)筆者在前文《關(guān)于語(yǔ)義分割的億點(diǎn)思考》中也提過(guò)，文中對(duì)整個(gè)任務(wù)進(jìn)行了深度剖析，有興趣的讀者可以自行翻閱公眾號(hào)歷史文章：

語(yǔ)義分割任務(wù)的核心思想是如何高效建模上下文信息，它是提升語(yǔ)義分割性能最為重要的因素之一，而有效感受野則大致決定了網(wǎng)絡(luò)能夠利用到多少上下文信息。

回到正文，本文的設(shè)計(jì)思路主要是受 ConvNeXt 啟發(fā)，還沒(méi)了解過(guò)的同學(xué)也可自行移動(dòng)到公眾號(hào)翻閱歷史文章，筆者之前對(duì)該系列也講解過(guò)了，此處不再詳述。3D UX-Net 核心理念是設(shè)計(jì)出一種簡(jiǎn)單、高效和輕量化的 網(wǎng)絡(luò)，其適用于 hierarchical transformers 的能力同時(shí)保留使用 ConvNet 模塊的優(yōu)勢(shì)，如歸納偏置。具體地，其編碼器模塊設(shè)計(jì)的基本思想可分為: 1) block-wise(微觀層面) 及 2) layer-wise(宏觀層面)【可以類比下 ConvNeXt】。

block-wise

首先，我們先從下列三個(gè)不同的視角進(jìn)行討論。

Patch-wise Features Projection

對(duì)比 ConvNets 和 ViTs 之間的相似性，兩個(gè)網(wǎng)絡(luò)都使用一個(gè)共同的基礎(chǔ)模塊將特征表示縮小到特定的尺寸。以往的方法大都沒(méi)有將圖像塊展平為具有線性層的順序輸入，因此作者采用具備大卷積核的投影層來(lái)提取 patch-wise 特征作為編碼器的輸入。

Volumetric Depth-wise Convolution with LKs

Swin transformer 的內(nèi)在特性之一是用于計(jì)算非局部 MSA 的滑動(dòng)窗口策略?？偟膩?lái)說(shuō)，有兩種分層方法來(lái)計(jì)算 MSA：基于窗口的 MSA (W-MSA) 和 移動(dòng)窗口 MSA (SW-MSA)。這兩種方式都生成了跨層的全局感受野，并進(jìn)一步細(xì)化了非重疊窗口之間的特征對(duì)應(yīng)關(guān)系。

受深度卷積思想的啟發(fā)，作者發(fā)現(xiàn)自注意力中的加權(quán)和方法與每通道卷積基礎(chǔ)之間的相似之處，其認(rèn)為使用 LK 大小的深度卷積可以在提取類似于 MSA 塊的特征時(shí)提供大的感受野。因此，本文建議采用 LK 大?。ɡ?，從 7 × 7 × 7 開(kāi)始）通過(guò)深度卷積壓縮 Swin transformer 的窗口移動(dòng)特性。如此一來(lái)便可以保證每個(gè)卷積核通道與相應(yīng)的輸入通道進(jìn)行卷積運(yùn)算，使得輸出特征與輸入具備相同的通道維度。

Inverted Bottleneck with Depthwise Convolutional Scaling

Swin transformer 的另一個(gè)固有結(jié)構(gòu)是，它們被設(shè)計(jì)為 MLP 塊的隱藏層維度比輸入維度寬四倍，如下圖所示。有趣的是，這種設(shè)計(jì)與 ResNet 塊中的擴(kuò)張率相關(guān)。因此，我們利用 ResNet 塊中的類似設(shè)計(jì)并向上移動(dòng)深度卷積來(lái)計(jì)算特征。此外，通過(guò)引入了具有 1 × 1 × 1 卷積核大小的深度卷積縮放(DCS)，以獨(dú)立地線性縮放每個(gè)通道特征。通過(guò)獨(dú)立擴(kuò)展和壓縮每個(gè)通道來(lái)豐富特征表示，可以最小化跨通道上下文產(chǎn)生的冗余信息，同時(shí)在每個(gè)階段增強(qiáng)了與下采樣塊的跨通道特征對(duì)應(yīng)。最后，通過(guò)使用 DCS，可以進(jìn)一步將模型復(fù)雜度降低 5%，并展示了與使用 MLP 模型的架構(gòu)相當(dāng)?shù)慕Y(jié)果。

layer-wise

介紹完微觀層面的設(shè)計(jì)思想，再讓我們從宏觀層面出發(fā)，以另外三個(gè)嶄新的視角去理解作者的動(dòng)機(jī)。

Applying Residual Connections

從上圖左上角可以看到，標(biāo)準(zhǔn)的 3D U-Net 模型內(nèi)嵌的模塊為 2 個(gè) 3 x 3 x 3 卷積的堆疊，其展示了使用小卷積核提取具有增加通道的局部表示的樸素方法；而其右手邊的 SegResNet 則應(yīng)用了類似 3D 版本的瓶頸層，先降維再升維最后再接殘差表示；緊接著左下角的便是 Swin Transformer，其基于窗口注意力+MLP層的組合；最后右下角便是本文所提出的模塊，其在最后一個(gè)縮放層之后應(yīng)用輸入和提取特征之間的殘差連接。此外，在殘差求和前后并沒(méi)有應(yīng)用到任何的歸一化層和激活層。

Adapting Layer Normalization

我們知道，在卷積神經(jīng)網(wǎng)絡(luò)中，BN 是一種常用策略，它對(duì)卷積表示進(jìn)行歸一化以增強(qiáng)收斂性并減少過(guò)擬合。然而，之前的工作已經(jīng)證明 BN 會(huì)對(duì)模型的泛化能力產(chǎn)生不利影響。因此作者這里跟 ConvNeXt 一致，將 BN 替換為 LN。

Using GELU as the Activation Layer

ReLU 是個(gè)好東西，幾乎是現(xiàn)代 CNNs 模型的首選激活函數(shù)。作者在這里提倡使用 GELU，這是一種基于高斯誤差的線性變換單元，相對(duì) ReLU 更加平滑，也是其中一種變體，解決 ReLU 因負(fù)梯度被硬截?cái)喽鴮?dǎo)致的神經(jīng)元失活問(wèn)題。

方法

可以看出，整體的設(shè)計(jì)思路非常簡(jiǎn)潔，框架也是一目了然，就一個(gè)標(biāo)準(zhǔn)的 3D U-Net 架構(gòu)，由編碼器-解碼器組成，同時(shí)結(jié)合長(zhǎng)跳躍連接操作幫助網(wǎng)絡(luò)更好的恢復(fù)空間細(xì)節(jié)的定位。也沒(méi)啥好分析的，下面就簡(jiǎn)單拆開(kāi)來(lái)看看，帶大家快速的過(guò)一遍就行。

DEPTH-WISE CONVOLUTION ENCODER

首先，輸入部分應(yīng)用一個(gè)大卷積核將原始圖像映射到一個(gè)低維空間分辨率的潛在空間特征表示，一來(lái)降低顯存參數(shù)量計(jì)算量等，二來(lái)可以增大網(wǎng)絡(luò)感受野，順便再處理成編碼器輸入所需的格式，總之好處多多，不過(guò)都是基操。

處理完之后就輸入到編碼器中進(jìn)行主要的特征提取，該編碼器共4個(gè)stage，也是標(biāo)準(zhǔn)的16倍下采樣。每個(gè)stage由多個(gè)不同的 3D UX-Net Block 構(gòu)成，具體長(zhǎng)啥樣自己看下左圖一點(diǎn)就通。主要的特點(diǎn)就是主打輕量化和大感受野。

DECODER

編碼器中每個(gè)階段的多尺度輸出通過(guò)長(zhǎng)跳躍連接鏈接到基于 ConvNet 的解碼器，并形成一個(gè)類似U形的網(wǎng)絡(luò)以用于下游的分割任務(wù)。這一塊就跟標(biāo)準(zhǔn)的 3D U-Net 幾乎沒(méi)啥兩樣了。

實(shí)驗(yàn)

消融實(shí)驗(yàn)

可視化效果

可以看到，與 GT 相比，3D UX-Net器官和組織的形態(tài)得到了很好的保存。

與 SOTA 方法的對(duì)比

從上述表格可以看出，3D UX-Net 在所有分割任務(wù)中均展示出最佳性能，并且 Dice 分?jǐn)?shù)有了顯著提高（FeTA2021：0.870 到 0.874，F(xiàn)LARE2021：0.929 到 0.934）。

總結(jié)

本文為大家介紹了 3D UX-Net，這是第一個(gè)將分層 Transformer 的特征與用于醫(yī)學(xué)圖像分割的純 ConvNet 模塊相適應(yīng)的3D網(wǎng)絡(luò)架構(gòu)。具體地，本文重新設(shè)計(jì)了具有深度卷積的編碼器塊，以更低的代價(jià)實(shí)現(xiàn)與 Transformer 想媲美的能力。最后，通過(guò)在三個(gè)具有挑戰(zhàn)性的公共數(shù)據(jù)集上進(jìn)行的廣泛實(shí)驗(yàn)表明所提方法的高效性。