基於狀態空間模型的遞迴神經網路 Mamba 在三維醫學影像語義分割中的應用
核心概念
本文提出將基於狀態空間模型 (SSM) 的新型遞迴神經網路 (RNN) 架構 Mamba 應用於三維醫學影像語義分割任務,並探討如何解決 Mamba 架構在處理三維數據時遇到的方向性問題,最終在多個公開數據集上取得優於現有基於 CNN 和 Transformer 模型的性能表現。
Taming Mambas for Voxel Level 3D Medical Image Segmentation
研究背景
醫學影像語義分割是醫學影像分析的關鍵步驟,近年來基於卷積神經網路 (CNN) 的方法,特別是 U-Net 架構及其變體,在該領域取得了主導地位。
然而,CNN 受限於局部感受野,難以捕捉全局信息。Transformer 的出現試圖解決這個問題,但其平方級的内存占用和數據需求使其難以應用於大規模三維醫學影像。
近期出現的基於狀態空間模型 (SSM) 的遞迴神經網路 (RNN) 架構 Mamba,在自然語言處理和基因組學等長序列建模任務中展現出優於 Transformer 的性能,且具有線性複雜度,為高效處理三維醫學影像提供了新的可能性。
研究方法
本文提出將 Mamba 應用於三維醫學影像語義分割,並針對其方向性問題提出解決方案。
提出 SegMamba、SegMambaSkip、BiSegMamba 和 MultiSegMamba 四種不同的 Mamba 架構變體,將 Mamba 層整合到 U-Net 架構的不同階段,並分別利用單向、雙向和多方向的實現方式。
這些架構結合了卷積和狀態空間模型的優勢,利用卷積編碼精確的空間信息,同時利用 Mamba 建模長距離的體素級交互。
實驗結果
在 MSD BrainTumour、Synapse Multi-organ 和 ACDC 三個公開數據集上進行實驗,結果表明,與 nnU-Net 和其他 Transformer 模型相比,所提出的 Mamba 模型在 HD95 和 DSC 指標上均有顯著提升。
其中,MultiSegMamba 模型在多個數據集和類別上均取得最佳性能,證明了多方向建模的重要性。
結論
本文證明了 Mamba 在三維醫學影像語義分割任務中的有效性,並提出解決其方向性問題的方案。
Mamba 架構結合了卷積和狀態空間模型的優勢,為醫學影像語義分割提供了新的思路。
統計資料
MultiSegMamba 在 Synapse Abdomen 數據集上分割膽囊的 Dice 指標比 SegMamba 提高了 10 個百分點。
在 BrainTumor 數據集上,基於 Mamba 的架構比性能最佳的基於 Transformer 的架構的 Dice 指標高出 3 個百分點以上,比 nnU-Net 高出 1 個百分點。
所提出的 Mamba 模型的參數數量約為 60M,是 nnU-Net 的兩倍,與 nnU-Net ResEnc 相當,遠低於基於 Transformer 的模型(從 TransUNet 和 UNETR 的約 95M 到 nnFormer 的 150M)。
深入探究
Mamba 架構如何應用於其他醫學影像分析任務,例如影像分類、目標檢測等?
Mamba 架構作為一種高效的長序列建模工具,除了在三維醫學影像分割任務中展現出優勢,其應用也可以拓展至其他醫學影像分析任務,例如影像分類和目標檢測:
1. 影像分類:
特徵提取: Mamba 可以作為特徵提取器,將三維醫學影像視為長序列數據,提取具有全局上下文信息的影像表徵。具體而言,可以將 Mamba 層整合到現有的 CNN 分類模型中,例如 ResNet 或 DenseNet,取代部分卷積層或作為額外的分支網路,以增強模型對全局信息的捕捉能力。
序列分類: 對於某些醫學影像數據,例如時間序列的醫學影像或切片堆疊的影像,可以直接將其視為長序列數據,利用 Mamba 進行分類。
2. 目標檢測:
整合至檢測器: 類似於影像分類,可以將 Mamba 層整合到現有的目標檢測器中,例如 Faster R-CNN 或 YOLO,以提升模型對目標與其周圍環境之間的長距離依賴關係的建模能力。例如,可以用 Mamba 層替換目標檢測器中用於特徵融合的模塊,或作為一個額外的分支來增強特徵表示。
基於序列的目標檢測: 可以將目標檢測任務轉化為序列預測問題,例如將影像分割成一系列的圖像塊或像素序列,利用 Mamba 預測每個圖像塊或像素屬於目標的概率。
需要注意的是,將 Mamba 應用於影像分類和目標檢測任務時,需要根據具體任務和數據特性對模型進行調整,例如輸入數據的預處理、模型結構的設計以及損失函數的選擇等。
如何進一步優化 Mamba 架構,以更好地處理三維醫學影像數據,例如探索更有效的掃描路徑或體積劃分方法?
雖然 Mamba 架構在處理三維醫學影像數據方面已展現出顯著的優勢,但仍有進一步優化的空間,特別是在處理三維數據固有的空間特性方面:
1. 探索更有效的掃描路徑:
基於影像內容的掃描路徑: 目前的 Mamba 模型採用固定的掃描路徑將三維數據展平成序列,可以探索根據影像內容自適應地調整掃描路徑的方法,例如利用注意力機制學習更有效的掃描順序,使得模型能夠更有效地捕捉重要的空間信息。
分層級掃描: 可以借鉴多尺度分析的思想,設計分層級的 Mamba 模型,在不同尺度上提取影像特徵,並利用更有效的掃描路徑整合不同尺度的信息。
2. 更有效的體積劃分方法:
重疊劃分: 可以將三維影像體積劃分為若干個重疊的子區域,並利用 Mamba 模型分別處理每個子區域,最後將各個子區域的結果融合,以減少邊界效應,提高分割精度。
自適應劃分: 可以根據影像內容自適應地調整子區域的大小和形狀,例如利用注意力機制或其他學習方法,將具有相似特征的區域劃分到一起,以提高模型的效率和性能。
3. 其他優化方向:
結合局部信息: Mamba 模型擅長捕捉長距離依賴關係,但可能會忽略局部細節信息。可以結合卷積操作或其他局部操作,以增強模型對局部信息的捕捉能力。
模型壓縮和加速: 可以探索模型壓縮和加速方法,例如模型剪枝、量化和知識蒸餾等,以降低模型的計算複雜度和内存占用,使其更易於部署到資源受限的設備上。
除了醫學影像分析領域,Mamba 架構還可以用於哪些其他需要長序列建模的領域?
Mamba 架構作為一種高效的長序列建模工具,其應用遠不止於醫學影像分析領域,還可以拓展至其他需要處理長序列數據的領域,例如:
1. 自然語言處理 (NLP):
文本生成: Mamba 在處理長文本序列方面具有優勢,可以應用於文本生成任務,例如機器翻譯、文本摘要和对话系统等。
語義理解: Mamba 可以捕捉長文本中的語義依賴關係,應用於語義理解任務,例如情感分析、問答系統和信息抽取等。
2. 時間序列分析:
語音識別: 語音信號可以視為時間序列數據,Mamba 可以用於語音識別任務,提取語音特徵並進行語音轉文本。
預測和預警: Mamba 可以用於分析時間序列數據,例如金融數據、氣象數據和交通數據等,進行預測和預警,例如股票預測、天氣預報和交通流量預測等。
3. 生物信息學:
基因組學: 基因序列可以視為長序列數據,Mamba 可以用於分析基因序列,例如基因預測、基因功能註釋和疾病診斷等。
蛋白質結構預測: 蛋白質序列也可以視為長序列數據,Mamba 可以用於預測蛋白質的三維結構,這對於理解蛋白質功能和設計新藥具有重要意義。
總之,Mamba 架構作為一種通用的長序列建模工具,具有廣泛的應用前景,可以預期在未來會在更多領域發揮重要作用。