南大提出VM-UNet v2 | 在VM-UNet基礎(chǔ)上細(xì)節(jié)&上下文信息都不放過(guò)，更小更快性能更強(qiáng)

在醫(yī)學(xué)圖像分割領(lǐng)域，基于CNN和Transformer的模型已經(jīng)得到了深入研究。然而，CNN在建模長(zhǎng)距離依賴方面存在局限性，這使得充分利用圖像內(nèi)的語(yǔ)義信息變得具有挑戰(zhàn)性。另一方面，二次計(jì)算復(fù)雜度為Transformers提出了挑戰(zhàn)。近來(lái)，狀態(tài)空間模型（SSMs），如Mamba，被認(rèn)為是一種有前景的方法。它們不僅在建模長(zhǎng)距離交互方面表現(xiàn)出卓越的性能，同時(shí)保持了線性的計(jì)算復(fù)雜度。

受到Mamba架構(gòu)的啟發(fā)，作者提出了Vison Mamba-UNetV2，引入了視覺狀態(tài)空間（VSS）塊來(lái)捕捉廣泛的上下文信息，引入了語(yǔ)義與細(xì)節(jié)注入（SDI）來(lái)增強(qiáng)低層次和高層次特征的融合。

作者在ISIC17、ISIC18、CVC-300、CVC-ClinicDB、Kvasir、CVC-ColonDB和ETIS-LaribPolyDB等公共數(shù)據(jù)集上進(jìn)行了全面實(shí)驗(yàn)。結(jié)果表明，VM-UNetV2在醫(yī)學(xué)圖像分割任務(wù)中表現(xiàn)出競(jìng)爭(zhēng)力。

代碼：https://github.com/nobodyplayer1/VM-UNetV2

1 Introduction

隨著醫(yī)學(xué)成像技術(shù)的不斷發(fā)展，醫(yī)學(xué)影像已成為診斷疾病和規(guī)劃治療的關(guān)鍵工具。在醫(yī)學(xué)圖像分析的基本和關(guān)鍵技術(shù)中，醫(yī)學(xué)圖像分割占有重要位置。這個(gè)過(guò)程涉及到在醫(yī)學(xué)圖像中區(qū)分器官或病變的像素，例如CT掃描和內(nèi)窺鏡視頻。醫(yī)學(xué)圖像分割是醫(yī)學(xué)圖像分析中最困難的任務(wù)之一，其目標(biāo)是為這些器官或組織的形狀和體積提供并提取關(guān)鍵信息。近年來(lái)，深度學(xué)習(xí)技術(shù)已被用于提高醫(yī)學(xué)圖像分割的性能。這些模型從圖像中提取有用信息，提高準(zhǔn)確性，并適應(yīng)不同的數(shù)據(jù)集和任務(wù)。

在語(yǔ)義圖像分割中，常用的方法是采用帶有跳躍連接的編碼器-解碼器網(wǎng)絡(luò)。在這個(gè)框架中，編碼器從輸入圖像中捕獲分層和抽象的特征。另一方面，解碼器使用編碼器生成的特征圖來(lái)構(gòu)建像素級(jí)的分割 Mask 或地圖，為輸入圖像中的每個(gè)像素賦予一個(gè)類別標(biāo)簽。已經(jīng)進(jìn)行了許多研究，以將全局信息整合到特征圖中并增強(qiáng)多尺度特征，這導(dǎo)致了分割性能的顯著提升。

U-Net 是一種關(guān)鍵架構(gòu)，因其平衡的編碼器-解碼器設(shè)計(jì)和跳層連接的融合而備受贊譽(yù)。這種結(jié)構(gòu)通過(guò)其各種編碼器和解碼器允許在多個(gè)層級(jí)上提取特征信息。此外，跳層連接有效地促進(jìn)了這一特征信息的轉(zhuǎn)換。對(duì)U-Net的眾多研究主要關(guān)注以下方面：編碼器部分 - 用不同的 Backbone 網(wǎng)來(lái)獲得不同層級(jí)的特征圖；跳層連接 - 采用各種通道注意力機(jī)制，并交換不同的連接部分；解碼器部分 - 使用不同的采樣和特征融合方案。

基于卷積神經(jīng)網(wǎng)絡(luò)（CNN）的模型由于其局部感受野，難以捕捉長(zhǎng)距離信息，這可能導(dǎo)致特征提取不充分和分割結(jié)果不理想。基于Transformer的模型在全球建模方面表現(xiàn)出色，但它們自注意力機(jī)制的雙曲復(fù)雜性造成了高昂的計(jì)算成本，特別是在像醫(yī)學(xué)圖像分割這類需要密集預(yù)測(cè)的任務(wù)中。這些局限性使得作者需要一種新的架構(gòu)來(lái)進(jìn)行醫(yī)學(xué)圖像分割，這種架構(gòu)能夠高效捕捉長(zhǎng)距離信息同時(shí)保持線性計(jì)算復(fù)雜性。

近期在狀態(tài)空間模型（SSMs）方面的進(jìn)展，尤其是結(jié)構(gòu)化SSMs（S4），由于它們?cè)谔幚黹L(zhǎng)序列方面的專長(zhǎng)，提供了一個(gè)有效的解決方案。例如，Mamba。Mamba模型通過(guò)增加選擇機(jī)制和硬件優(yōu)化來(lái)增強(qiáng)S4，在密集數(shù)據(jù)領(lǐng)域展示了卓越的性能。在視覺狀態(tài)空間模型（VMDaba）中融入交叉掃描模塊（CSM）進(jìn)一步提升了Mamba適用于計(jì)算機(jī)視覺任務(wù)的能力。它通過(guò)促進(jìn)空間域的遍歷，將非因果視覺圖像轉(zhuǎn)換為有序的斑塊序列來(lái)實(shí)現(xiàn)這一點(diǎn)。

受到VMDaba 在圖像分類任務(wù)和VM-Unet在醫(yī)學(xué)圖像分割中成功的啟發(fā)。遵循UNetV2的框架，本文提出了視覺Mamba-UNetV2（VM-UNetV2），作者重新整合了低 Level 和高 Level 的特征，將語(yǔ)義信息融入到低 Level 特征中，同時(shí)使用更詳細(xì)的信息來(lái)細(xì)化高級(jí)特征。

作者在與胃腸病學(xué)語(yǔ)義分割任務(wù)以及皮膚病變分割相關(guān)的任務(wù)上進(jìn)行了詳盡的實(shí)驗(yàn)，以展示純基于SSM的模型在醫(yī)學(xué)圖像分割領(lǐng)域的性能。特別是對(duì)ISIC17、ISIC18、CVC-300、CVC-ClinicDB、Kvasir、CVC-ColonDB和ETIS-LaribPolyDB等公開數(shù)據(jù)集以及作者的私有數(shù)據(jù)集ZD-LCI-EGGIM進(jìn)行了廣泛的測(cè)試。結(jié)果表明，VM-UNetV2能夠提供具有競(jìng)爭(zhēng)力的分割結(jié)果。

本研究的主要貢獻(xiàn)可以概括為以下幾點(diǎn)：

1. 作者提出了VM-UnetV2，在基于SSM的醫(yī)療圖像分割算法探索方面開創(chuàng)了先河。
2. 在七個(gè)數(shù)據(jù)集上進(jìn)行了詳盡的實(shí)驗(yàn)，結(jié)果證明了VM-UNetV2具有顯著競(jìng)爭(zhēng)力。
3. 作者正在開創(chuàng)將基于SSM的方法與Unet變體結(jié)合的探索，推動(dòng)更高效、更有效的基于SSM的分割算法的發(fā)展。

2 Methods

Preliminaries

在當(dāng)前的基于SSM的模型中，即結(jié)構(gòu)化狀態(tài)空間序列模型（S4）和Mamba，都依賴于一個(gè)傳統(tǒng)的連續(xù)系統(tǒng)，這個(gè)系統(tǒng)將一維輸入函數(shù)或序列，表示為 ，通過(guò)中間隱含狀態(tài) 映射到一個(gè)輸出 。這個(gè)過(guò)程可以描述為一個(gè)線性常微分方程（ODE）：

在這里， 代表狀態(tài)矩陣，而 和 分別表示投影參數(shù)。

S4和Mamba通過(guò)離散化這個(gè)連續(xù)系統(tǒng)，使其更適合深度學(xué)習(xí)環(huán)境。具體來(lái)說(shuō)，它們引入了一個(gè)時(shí)間尺度參數(shù) ，并使用一致的離散化規(guī)則將 和 轉(zhuǎn)換為離散參數(shù) 和 。零階保持（ZOH）通常被用作離散化規(guī)則，可以概述如下：

在離散化之后，基于SSM的模型可以通過(guò)兩種不同的方法進(jìn)行計(jì)算：線性遞推或全局卷積，分別表示為方程(3)和方程(4)。

其中 表示一個(gè)結(jié)構(gòu)化的卷積核，而 指的是輸入序列 的長(zhǎng)度。

VM-UNetV2 Architecture

視覺Mamba UNetV2的全面結(jié)構(gòu)如圖1所示。它由三個(gè)主要模塊組成：編碼器、SDI（語(yǔ)義和細(xì)節(jié)注入）模塊和解碼器。給定一個(gè)輸入圖像I，其中 ，編碼器在M個(gè)層次上生成特征。作者將第 級(jí)的特征表示為 ，其中 。這些累積的特征， 隨后被轉(zhuǎn)發(fā)到SDI模塊進(jìn)行進(jìn)一步增強(qiáng)。

如圖所示，編碼器輸出通道 的大小為 ， 共同輸入到SDI模塊進(jìn)行特征融合，并且 對(duì)應(yīng)于 作為第 階段的輸出。 的特征大小為 。在VM-UNetV2中，作者使用深度監(jiān)督來(lái)計(jì)算 和 特征的損失。

在本文中，作者在編碼器的四個(gè)階段使用了 VSS塊，每個(gè)階段的通道數(shù)分別為 。根據(jù)作者在VMamba [13]中的觀察， 和 的不同值是區(qū)分Tiny、Small和Base框架規(guī)格的重要因素。遵循VMamba的規(guī)格，作者讓C取值為96， 和 各自取值為2， 取值從集合中選取。這表示作者有意使用VMamba的Tiny和Small模型作為作者消融實(shí)驗(yàn)的主干。

VSS And SDI Block

VSS塊源自V Mamba，作為VMTUeV2編碼器的基礎(chǔ)部分，VSS塊的結(jié)構(gòu)如圖2 所示。輸入首先通過(guò)一個(gè)初始的線性嵌入層進(jìn)行處理，之后分為兩個(gè)獨(dú)立的信息流。一個(gè)信息流通過(guò)一個(gè) 深度卷積[11]層，然后經(jīng)過(guò)一個(gè)Silu激活[19]函數(shù)，再進(jìn)入主要的2D-Selective-Scan(SS2D)模塊。SS2D的輸出接著通過(guò)一個(gè)層歸一化層，并與另一個(gè)信息流的輸出合并，該信息流也經(jīng)過(guò)Silu激活處理。這個(gè)合并后的輸出構(gòu)成了VSS塊的最終結(jié)果。

SDI模塊，如圖2（b）所示。利用編碼器生成的分層特征圖 ，其中 表示第  Level 。

不同的注意力機(jī)制可以用于SDI模塊中，以計(jì)算空間和通道的注意力得分。遵循UNetV2 [15]中提到的方法，作者使用CBAM來(lái)實(shí)現(xiàn)空間和時(shí)間的注意力。計(jì)算公式如下，其中 表示第 注意力計(jì)算：

然后作者使用 卷積將 的通道對(duì)齊到 ，得到的結(jié)果特征圖表示為 。

在SDI解碼器的第 階段， 表示目標(biāo)參考。然后作者調(diào)整每個(gè) 級(jí)特征圖的大小，以匹配 的大小，如下公式所示：

在公式6中， 、 和 分別代表自適應(yīng)平均池化、身份映射和雙線性插值。在公式7中， 表示平滑卷積的參數(shù)， 是第 個(gè)層次上的第 個(gè)平滑特征圖。這里， 表示哈達(dá)瑪積（Hadamard product）。隨后， 會(huì)被送入第 個(gè)層次的解碼器，進(jìn)行進(jìn)一步的分辨率重建和分割。

Loss function

對(duì)于作者的醫(yī)學(xué)圖像分割任務(wù)，作者主要采用基本的交叉熵和Dice損失作為損失函數(shù)，因?yàn)樽髡叩乃袛?shù)據(jù)集 Mask 都包含兩個(gè)類別：?jiǎn)我荒繕?biāo)和背景。

（ ）是常數(shù)，通常選擇（1,1）作為默認(rèn)參數(shù)。

3 Experiments and results

Datasets

作者使用了三種類型的數(shù)據(jù)集來(lái)驗(yàn)證作者框架的有效性。第一種類型是開源的皮膚病數(shù)據(jù)集，包括ISIC 2017和ISIC 2018，作者將皮膚數(shù)據(jù)集以7:3的比例劃分為訓(xùn)練集和測(cè)試集。第二種是開源的胃腸道息肉數(shù)據(jù)集，其中包括Kvasir-SEG、ClinicDB、ColonDB、Endoscene和ETIS，在這一類數(shù)據(jù)集中，作者遵循PraNet中的實(shí)驗(yàn)設(shè)置。對(duì)于這些數(shù)據(jù)集，作者提供了在幾個(gè)指標(biāo)上的詳細(xì)評(píng)估，包括平均交并比(mIoU)、Dice相似系數(shù)(DSC)、準(zhǔn)確率(Acc)、敏感性(Sen)和特異性(Spe)。

Experimental setup

在V Mamba工作的基礎(chǔ)上，作者將所有數(shù)據(jù)集中的圖像尺寸調(diào)整為256×256像素。為了抑制過(guò)擬合，作者還引入了數(shù)據(jù)增強(qiáng)方法，比如隨機(jī)翻轉(zhuǎn)和隨機(jī)旋轉(zhuǎn)。在操作參數(shù)方面，作者將批處理大小設(shè)置為80，采用AdamW優(yōu)化器，初始學(xué)習(xí)率為1e-3。作者使用CosineAnnealingLR作為調(diào)度器，其操作最多持續(xù)50個(gè)迭代，學(xué)習(xí)率最低降至1e-5。作者的訓(xùn)練歷時(shí)300個(gè)周期。

對(duì)于VM-UNetV2，編碼器單元的權(quán)重最初設(shè)置為與V Mamba-S相匹配。實(shí)施工作在一個(gè)Ubuntu 20.04系統(tǒng)上進(jìn)行，使用Python3.9.12、PyTorch2.0.1和CUDA11.7，所有實(shí)驗(yàn)都在單個(gè)NVIDIA RTX V100 GPU上完成。

Results

作者對(duì)比了VM-UNetV2與一些最先進(jìn)模型的性能，并在表1和表2中展示了實(shí)驗(yàn)結(jié)果。

對(duì)于ISIC數(shù)據(jù)集，VM-UNetV2在mIoU、DSC和Acc指標(biāo)上均優(yōu)于其他模型。在息肉相關(guān)數(shù)據(jù)集中，VM-UNetV2在所有指標(biāo)上也超過(guò)了最先進(jìn)的UNetV2模型，在mIoU參數(shù)上提高了高達(dá)7%。

Ablation studies

在本節(jié)中，作者針對(duì)VM-UNetV2編碼器的初始化和解碼器中的深度監(jiān)督操作，使用息肉數(shù)據(jù)集進(jìn)行消融實(shí)驗(yàn)。正如VMamba論文中所指出的，編碼器的深度和特征圖中的通道數(shù)決定了VMamba的規(guī)模。在本文中，所提出的VM-UNetV2僅在編碼器部分使用在ImageNet-1k上預(yù)訓(xùn)練的VMamba權(quán)重。

因此，在進(jìn)行本研究中的模型規(guī)模消融實(shí)驗(yàn)時(shí)，作者僅改變編碼器的深度，如表4所示。對(duì)于輸出特征，作者采用深度監(jiān)督機(jī)制，融合了兩層輸出特征，然后與真實(shí)標(biāo)簽進(jìn)行比較以計(jì)算損失。

如表4和表5所示，當(dāng)編碼器的深度設(shè)置為 時(shí)，分割評(píng)估指標(biāo)相對(duì)較好。因此，在使用VM-UNetV2時(shí)，無(wú)需選擇特別大的深度。在大多數(shù)使用深度監(jiān)督機(jī)制的情況下，分割評(píng)估指標(biāo)相對(duì)較好，但這不是一個(gè)決定性因素。對(duì)于不同的數(shù)據(jù)集，需要分別進(jìn)行消融實(shí)驗(yàn)以確定是否采用深度監(jiān)督機(jī)制。

4 Conclusions

在本文中，作者提出了一種基于SSM的UNet類型醫(yī)療圖像分割模型VM-UNetV2，充分利用了基于SSM模型的性能。作者分別使用VSS塊和SDI來(lái)處理編碼器和解碼器連接部分。使用VMamba的預(yù)訓(xùn)練權(quán)重來(lái)初始化VM-UNetV2的編碼器部分，并采用深度監(jiān)督機(jī)制來(lái)監(jiān)督多個(gè)輸出特征。

VM-UNetV2已經(jīng)在皮膚病和息肉數(shù)據(jù)集上進(jìn)行了廣泛測(cè)試。結(jié)果顯示，在分割任務(wù)中，VM-UNetV2具有很高的競(jìng)爭(zhēng)力。復(fù)雜性分析表明，VM-UNetV2在FLOPs、Params和FPS方面也是高效的。

參考

[1].VM-UNet-V2: Rethinking Vision Mamba UNet for Medical Image Segmentation.