論文鏈接：https://arxiv.org/pdf/2402.03302

Code： https://github.com/JiarunLiu/Swin-UMamba

摘要：

醫學圖像分割的挑戰在於應該同時整合多尺度信息，從局部細節到全局依賴關係。但現有方法難以有效建模長距離的全局信息：

  • 卷積神經網絡（CNN）由於受限於局部感受野，難以捕獲遠距離的全局信息。
  • 視覺變換器（ViTs）雖然能處理全局信息，但其注意力機制在高分辨率圖像上具有高二次複雜度，計算成本較大。

最近，基於Mamba的模型基於其在長序列建模方面表現出色，越來越受到關注。多項研究已顯示，這些模型在多種任務中的表現優於流行的視覺模型，具有更高的準確率、較低的內存需求和較少的計算資源。

目前大多數Mamba模型都是從頭訓練（隨機初始化）而不是利用預訓練，而預訓練已被證明能顯著提升素材有限情況下的醫學圖像分析效果。

為此，本文提出了一個新穎的Mamba基礎模型——Swin-UMamba，專為醫學圖像分割設計，結合了基於ImageNet的預訓練優勢。實驗顯示，ImageNet預訓練在提升Mamba模型性能方面起到了關鍵作用。

Swin-UMamba在多個數據集上（如腹部MRI、內窺鏡、顯微鏡）都表現優異，明顯優於CNN、ViTs及其他最新的Mamba模型，平均提升了2.72%的性能。

1. 引言

醫學圖像分割的作用與挑戰

  • 作用：醫學圖像分割在診斷、治療方案制定和療法實施中發揮關鍵作用。它能夠幫助醫生自動識別和劃分醫學影像中的不同結構，提升診斷效率和準確性。
  • 現有方法的侷限：目前大部分分割依賴經驗豐富的醫生進行手工操作，這不僅耗費大量人力和時間，而且不同醫生之間的判讀可能存在差異，影響結果的一致性。

深度學習在醫學圖像分割中的應用

  • 進步：近年來，深度學習已極大推動醫學圖像分割的發展，但仍然存在一些難題。
  • 需求：高效準確的分割需要結合局部特徵（細節信息）和全局依賴（整體結構、關係），這是一個複雜的任務。

模型的優勢與侷限

  • 卷積神經網絡（CNNs）：如SegResNet、U-Net、nnU-Net，擅長提取局部特徵（如邊緣、紋理），但在捕捉長距離全局信息方面有限，因為它們的感受野有限。
  • 視覺Transformers（ViTs）：能處理全局上下文，適合捕獲長程依賴，但基於基於注意力機制，計算複雜度很高，隨着輸入分辨率增加會呈指數增長（quadratic complexity），在高分辨率醫療圖像（如全-slide病理切片或高分辨率MRI/CT）中面臨挑戰。此外，ViTs對數據量也很敏感，容易過擬合，尤其是在數據有限的醫學領域。

結構化狀態空間模型（SSMs）與Mamba模型

  • 近期研究表明，結構化狀態空間模型（比如Mamba）在長序列建模方面十分高效，能夠線性或接近線性擴展、同時捕獲長距離依賴。
  • 在視覺任務中，相關研究（如Vim、VMamba）展示了這些模型的潛力。它們優於傳統模型，具備更高的性能、更低的存儲和計算需求。
  • 具體應用在醫學圖像分割上，包括基於nnU-Net和Swin-UNETR的專用架構（U-Mamba、SegMamba），已取得不錯的效果。

關鍵問題與挑戰

  • 預訓練的缺乏：目前大多數Mamba模型是從零開始訓練的，缺乏藉助ImageNet等大規模預訓練模型的經驗，導致模型可能難以充分利用已有的知識。
  • 遷移學習的重要性：預訓練已在CNN和ViT中被證明能提升數據有限情況下的效果，探索其在Mamba模型中的作用尤為關鍵。
  • 模型結構調整的需求：由於Mamba塊在視覺領域的應用尚屬新穎，現有模型需重新設計以更好地利用預訓練模型。
  • 實用性與效率：在資源有限的臨牀環境中，模型的規模和運算效率至關重要，需要確保模型具有良好的可擴展性和部署能力。

模型設計與目的：

  • Swin-UMamba：這是一個結合了通用預訓練視覺模型（如通過ImageNet預訓練的模型）與專門為醫學圖像分割設計的解碼器（decoder）的網絡。它利用預訓練模型的強大特徵提取能力，從而提高醫學圖像分割的性能。
  • Swin-UMamba†：這是Swin-UMamba的變體，採用了基於Mamba的解碼器，具有更少的參數和較低的計算量（FLOPs），從而實現更高效的應用。

創新點和貢獻：

  • 首創探索預訓練Mamba網絡在醫學圖像分割中的作用至關重要的。就是：這是首次嘗試評估預訓練的Mamba模型對醫學圖像分割任務的影響，實驗結果顯示，基於ImageNet的預訓練在提高模型性能方面起到關鍵作用，有時候甚至
  • 提出新型網絡結構：
  • 設計了Swin-UMamba，旨在充分發揮預訓練模型的優勢，同時保持模型架構的統一性。
  • 還提出了參數更少、計算更低（低FLOPs）的Swin-UMamba†變體，以適應資源有限的應用場景，同時保持了良好的性能。

性能表現：

  • 結果顯示，Swin-UMamba和Swin-UMamba†均超過了之前的多種模型，包括傳統的卷積神經網絡（CNNs）、視覺變換器（ViTs）以及最新的基於Mamba的模型，在多個醫學圖像數據集上都表現出明顯優勢。
  • 這一表現充分驗證了預訓練（特別是ImageNet預訓練）和所提出架構在醫學圖像分割中的有效性。

2. 方法

2020年人工智能論文總結 - deephub的個人空間 -_數據集

整體架構：由三個主要部分組成：

  1. Mamba基礎的編碼器：在大規模材料集（如ImageNet）預訓練，用以多尺度提取特徵。
  2. 解碼器：利用多個上採樣塊逐步恢復圖像的空間信息，以生成分割結果。
  3. 跨級跳躍連接：連接不同尺度的特徵，彌補低層細節與高層語義之間的差距。

2.1 Mamba基礎的VSS塊

Mamba模型的背景：

  • Mamba在自然語言處理中的應用利用空間狀態序列模型（SSMs）將注意機制的複雜度從二次（quadratic）降低到線性（linear），提升了模型的效率和性能。
  • 在計算機視覺中的應用優勢包括更高的準確性、較低的計算成本和較少的內存消耗。

挑戰：

  • 2D視覺數據與1D語言序列的差異：視覺上的空間關係至關重要，但直接將Mamba應用到扁平化的圖像中會導致“感受野”受限，難以捕捉未掃描區域的關係。

VSS塊的設計：

  • 目的：消除二維圖像數據的處理難題，確保全面捕獲空間關係。
  • 方法：
  • 將輸入圖像塊沿四個不同的方向展開（如上下左右），形成四個不同的序列（2020年人工智能論文總結 - deephub的個人空間 -_圖像分割_02，其中2020年人工智能論文總結 - deephub的個人空間 -_數據_03表示方向）。
  • 每個序列經過空間狀態序列模型（S6）進行處理，獲取長距離依賴關係。
  • 最後，將四個方向的輸出序列合併，形成完整的二維特徵圖。
  • S6（空間狀態序列模型）
  • 核心作用：允許序列中的每個元素都能與之前掃描過的樣本進行交互，這通過壓縮的隱狀態實現，從而高效捕獲長距離依賴關係。
  • 原理：模仿文本處理中的序列關係，從而在圖像或特徵序列中建模空間依賴。

2.2 基於ImageNet預訓練的集成

Swin-UMamba的編碼器由五個階段組成：

  • 第一個階段是“幹線（stem）”階段，採用一個7×7卷積層進行2×下采樣，配合padding和步長為2，之後加上2D實例歸一化。不同於VMamba的設計，作者在此階段採用了漸進式的2×下采樣策略，以保留低層細節。
  • 第二個階段採用2×2的patch embedding層，保持特徵分辨率為原始圖像的1/4，與VMamba中的特徵保持一致。
  • 後續階段則遵循VMamba-Tiny的設計格式，每個階段由patch merging層進行2×下采樣，配合若干VSS塊（空間狀態序列模型塊）提取高層次特徵。VSS塊的數量在各階段分別為2020年人工智能論文總結 - deephub的個人空間 -_數據_04。這些階段的特徵維度隨着階段遞增，最終達到2020年人工智能論文總結 - deephub的個人空間 -_數據_05。

預訓練方面，作者將從ImageNet預訓練所得的權重初始化VSS塊和patch merging層，以提升模型性能。然而，由於patch embedding層的patch大小和輸入通道的差異，其未採用預訓練的權重。

2.3 主幹解碼器

  • 引入了帶殘差連接的額外卷積塊，用於處理跳躍連接的特徵。
  • 在每個尺度層加入深層監督的分割頭，以提升訓練效果。
  • 在上採樣階段
  • 利用拼接操作（Cat）結合上一尺度和當前尺度的特徵
  • 通過兩個殘差卷積塊（Res）處理後，經過轉置卷積（DeConv）和1×1卷積生成最終的分割圖。

2.4 Swin-UMamba†

目標與背景

  • 目的：探索Mamba模型在醫學語義分割中的潛力。
  • 方法：通過將Swin-UMamba的解碼器替換為基於Mamba的解碼器，得到Swin-UMamba†，希望其在性能上具有競爭力，同時擁有更低的參數量和計算成本。
  • 發現：大規模預訓練在醫學圖像分割中具有重要作用，無論解碼器結構如何。

主要改動

  • 替換上採樣塊：原Swin-UMamba中的上採樣塊被替換為patch expanding和兩個VSS塊。Patch expanding是直接將輸入圖像從原始尺寸2020年人工智能論文總結 - deephub的個人空間 -_數據集_06通過4×4的補丁嵌入層投影到特徵圖2020年人工智能論文總結 - deephub的個人空間 -_圖像分割_07，從而簡化了結構。
  • 減少參數：原模型參數為60M，FLOPs為68.0G；改造後，參數減至28M，FLOPs降低至18.9G。
  • 移除部分連接：去除了源自輸入圖像和較低分辨率特徵的跳躍連接，因為這些在解碼過程中沒有對應的解碼塊。
  • 深度監督：在尺度為1×、1/4×、1/8×和1/16×的特徵層上加入了深度監督（通過1×1卷積映射到K類），增強模型的監督信號。

結構示意：Swin-UMamba†的結構由多次4×上採樣和VSS塊組成，整體架構簡化，更加高效，結構細節如圖所示。

2020年人工智能論文總結 - deephub的個人空間 -_數據集_08

3. 實驗

3.1 數據集：

Abdomen MRI（腹部MRI）：該數據集由MICCAI 2022的AMOS挑戰提供，目標是分割13個腹部器官（如肝臟、脾臟、胰腺、左右腎、胃、膽囊、食管、主動脈、下腔靜脈、左右腎上腺和十二指腸）。訓練集包括60個MRI掃描（共5615個切片），測試集包括50個掃描（共3357個切片）。圖像被裁剪成尺寸為（320, 320）的塊，利用nnU-Net框架進行訓練和測試。

Endoscopy（內窺鏡）：該數據集來自MICCAI 2017的EndoVis挑戰，目標是分割7種內窺鏡儀器（如針驅動器、夾鉗、彎曲剪刀、卡迪爾鉗、雙極鉗、血管封堵器和超聲探頭）。訓練集有1800幀圖像，測試集1200幀。圖像裁剪尺寸為（384, 640），且該數據集的圖像具有特殊的長寬比例，展示出不同的影像特性。

Microscopy（顯微鏡）：來自NeurIPS 2022細胞分割挑戰，用於細胞實例分割。訓練集具備1000張圖像，評估集有101張。圖像裁剪為（512, 512）尺寸。

3.2 實現細節：

基礎框架：模型在成熟的nnU-Net框架上構建。nnU-Net具有自適應配置的特性、。

損失函數：結合Dice損失和交叉熵損失的加和，用於訓練。並在不同尺度上引入深層監督，以提升模型的梯度傳播效果。

優化器：採用AdamW優化器，權重衰減（weight decay）設置為0.05，以改善正則化效果。

學習率策略：使用 餘弦退火（cosine decay）的學習率調整策略，初始學習率設為0.0001。

預訓練模型的使用：利用預訓練的VMamba-Tiny模型來初始化Swin-UMamba的參數，確保模型在訓練 少量數據和加快收斂。

參數凍結：在訓練的前10個epoch中，凍結預訓練模型中的參數，以使其他模塊（例如新加的層）更好地適應特定任務。

訓練週期差異：

  • AbdomenMRI：訓練100個epoch。
  • Endoscopy：訓練350個epoch。
  • Microscopy：訓練450個epoch。

3.3 基線方式和評估指標：

  • 對比模型：使用三類不同的模型作為基線，以確保評估的全面性：
  • CNN-based：nnU-Net、SegResNet；
  • Transformer-based：UNETR、Swin-UNETR、nnFormer；
  • Mamba-based：最新的U-Mamba，包含兩個變體：
  • U-Mamba_Bot：只在“瓶頸”位置引入Mamba塊；
  • U-Mamba_Enc：在每個編碼器階段都引入Mamba塊。
  • 指標：
  • Dice 相似係數（DSC）：用於評價AbdomenMRI和Endoscopy數據集的全局分割效果。
  • 歸一化表面距離（NSD）：衡量邊界的精確性。
  • F1分數：對於顯微鏡數據集（實例分割任務）。
  • 模型複雜度：
  • 計算網絡參數數（#param）；
  • 浮點運算量（FLOPs），均用fvcore包衡量。

3.4 在AbdomenMRI數據集上的表現：

Swin-UMamba和 Swin-UMamba†都優於所有基線技巧（包括CNN、Transformer和其他Mamba模型），證明了它們在醫學圖像分割中的潛力。

這些模型在Dice相似係數（DSC）和歸一化表面距離（NSD）指標上都比CNN和Transformer手段高至少1%，且Swin-UMamba在DSC上比前一狀態最優模型（U-Mamba_Enc）提升了1.34%。

ImageNet預訓練在性能提升中起到了關鍵作用。例如，使用預訓練後，Swin-UMamba的DSC和NSD分別提升了7.06%和7.74%，且預訓練模型使訓練更快、更穩定，只需訓練的迭代次數的十分之一。

Swin-UMamba†在沒有預訓練時，未能良好收斂，即使延長訓練時間也如此，但在利用預訓練權重後，表現明顯優於所有基線，且參數量和計算量都少於前一狀態最優模型U-Mamba。

2020年人工智能論文總結 - deephub的個人空間 -_數據_09

圖像顯示Swin-UMamba能更準確識別目標器官的形狀和類型，而基線模型則未能全部正確識別。

2020年人工智能論文總結 - deephub的個人空間 -_圖像分割_10

3.5 在Endoscopy（內窺鏡）數據集上的表現：

  • Swin-UMamba†在DSC和NSD指標上比U-Mamba_Bot分別提升2.43%和2.41%。
  • 預訓練帶來的性能提升非常明顯（分別提升了12.84%和12.90%），尤其是在內容較少、易過擬合的小數據集（如該內窺鏡素材集）上，預訓練顯得尤為重要，有助於緩解過擬合。
  • Swin-UMamba†的參數較少，更易避免過擬合，表現優於未加入預訓練的模型。

2020年人工智能論文總結 - deephub的個人空間 -_數據集_11

3.6 在Microscopy（顯微鏡）細胞圖像數據集上的表現：

  • Swin-UMamba和Swin-UMamba†繼續超越所有基線模型，提升範圍從1.99%到20.15%。
  • 由於Microscopy數據集具有更高的圖像分辨率、樣本更少且差異更大，因此對模型的長距離依賴建模能力和數據效能提出了更高要求。
  • 預訓練再一次表現出顯著優勢（提升12.45%和7.96%），即使在這種更復雜的高分辨率、有限樣本的場景下，預訓練依然十分關鍵。

2020年人工智能論文總結 - deephub的個人空間 -_圖像分割_12

4. 結論

本研究旨在探討基於ImageNet預訓練對Mamba模型在二維醫學圖像分割中的影響。研究團隊提出了一種新穎的基於Mamba的模型，叫做Swin-UMamba，以及它的變體Swin-UMamba†，這兩種模型都能充分利用預訓練模型的能力來改善分割任務。在多個醫學圖像分割內容集上的實驗結果顯示，採用ImageNet預訓練的手段帶來了多方面的優勢，包括更高的分割準確率、訓練過程更穩定、減少了過擬合問題、提高了數據利用效率，以及降低了模型的計算資源需求。