SSMamba: A Self-Supervised Hybrid State Space Model for Pathological Image Classification

Enhui Chai, Sicheng Chen, Tianyi Zhang, Xingyu Li, Tianxiang Cui

View Original ↗
AI 導讀 technology AI 重要性 4/5

僅 25.3M 參數,SSMamba 於病理影像分類徹底擊敗 11 款主流大模型。

  • MAMIM 模組採 75% 遮蔽預訓練,有效克服病理跨倍率領域偏移。
  • 結合 DMS 與 LPR 模組,剔除座標編碼以抵抗 H&E 染色偽影干擾。
  • 以 25.3M 參數,在 16 個病理資料集擊敗所有主流基礎模型。

病理影像基礎模型雖具備強大表徵能力,但其龐大參數往往伴隨鉅額運算成本。西北大學與諾丁漢大學聯合提出的 SSMamba 框架,僅使用 25.3M 的輕量化參數,即在 10 個公開感興趣區域(ROI)資料集上擊敗了包含 Virchow2Prov-GigaPath 在內的 11 款主流模型。該研究證實,針對病理空間特徵的專屬架構優化,比單純盲目擴充參數更能有效克服跨倍率領域偏移與形態辨識挑戰。

突破 ViT 模型跨倍率與局部感知的限制

臨床病理診斷高度依賴數位玻片中的感興趣區域(ROI)作為證據。目前,基於視覺變換器(ViT,Vision Transformer)與大規模自監督學習(SSL)的基礎模型已成為業界主流。然而,這類模型多半在固定的單一倍率下進行預訓練(例如 UNI 模型採用 0.5mpp 解析度),導致應用於不同硬體或多樣化掃描條件時,極易產生跨倍率的領域偏移(domain shift),無法充分利用預訓練學到的病理知識。

分析架構層面,常規卷積神經網路(CNN)缺乏長距離建模能力,而主流的 ViT 模型在處理高解析度影像時又面臨運算複雜度過高的瓶頸。同時,傳統的自注意力機制偏好全局關係,缺乏對局部形態變化的歸納偏置(inductive bias)。這使得模型對腺體輪廓、細胞核多形性等細微且具決定性的診斷線索不夠敏感。雖然狀態空間模型(SSM,State Space Model)以線性時間複雜度帶來新解方,但既有 Mamba 架構的自迴歸特性,卻與病理組織非序列性的空間分佈存在嚴重衝突。

導入 75% 遮蔽率的 MAMIM 預訓練機制

為了解決跨倍率領域偏移,研究團隊設計出獨創的 Mamba 遮蔽影像建模(MAMIM) 模組。該機制針對輸入影像套用 75% 的高隨機遮蔽率,將特定區塊替換為可學習的遮蔽標記。有別於傳統基於 ViT 的自編碼器,MAMIM 完全採用客製化的 SSMamba 區塊進行特徵擷取與解碼。

這種以病理特徵為導向的重建任務,能有效引導模型將注意力集中於細胞形態與組織拓撲等結構性關鍵訊號。預訓練過程不需依賴通用的大規模外部圖像數據,而是直接在目標 ROI 資料集上進行領域內學習。透過保留未遮蔽區塊的上下文線索,此設計從根本上減少了因倍率不同而產生的特徵錯位,避免模型學習到特定倍率的偽影。

DMS 模組重構空間感知的雙向全局建模

標準 Mamba 架構的單向因果卷積設計,無法捕捉相鄰區域之間的相互作用,這對於辨識腫瘤與基質邊界等病理特徵極為不利。方向性多尺度(DMS) 模組為此徹底重構了核心路徑,將單向掃描替換為雙向深度可分離卷積(Depthwise Separable Convolution)。此舉不僅大幅降低參數負擔,更實現了以中心點向外擴展的雙向上下文聚合。

為彌補序列化處理對局部平行互動的忽視,DMS 模組額外引入了一組具備 SiLU 啟動函數的常規 1D 卷積平行分支。此對稱式架構能同步處理所有標記,顯著強化局部鑑別度,並提升面對視覺退化時的穩健性。最終,該模組透過通道拆分連接機制,完美調和了 Mamba 原本的自迴歸偏差與自監督預訓練框架,確保了訓練過程的高度穩定。

LPR 模組剔除座標編碼以抵抗染色偽影

病理影像分析面臨獨特的空間挑戰,常規的絕對位置編碼極易被 H&E 染色的偽影(如染料擴散或組織褶皺)破壞。局部感知殘差(LPR) 模組為此拋棄了傳統的座標編碼策略,轉而使用動態深度卷積來實現空間平移不變性(translation invariance)。

處理流程中,系統會先利用逐點卷積(Pointwise Convolution)壓縮通道數,在降低運算負載的同時保留細胞的細粒度特徵。接著套用空間共享權重的深度卷積,實現無涉於特定座標的局部感知。最後,透過殘差路徑恢復原始表徵,將多尺度局部特徵與全局組織拓撲無縫融合。此一設計徹底解耦了染色雜訊與診斷訊號,確保了空間感知的純粹性。

僅 25.3M 參數於 10 個 ROI 資料集奪冠

經過架構的深度重構,SSMamba 在涵蓋肺癌、大腸直腸癌等 10 個公開 ROI 資料集上展現出驚人效能。實驗數據顯示,該模型達到 95.56% 的平均 F1 分數、95.98% 的平均準確率與 95.02% 的平均 AUC。相較於表現次佳的大型模型 Virchow2,優勢分別達 1.13%、1.19% 與 1.88%。

在最具代表性的 CAM16 資料集上,SSMamba 的準確率與 F1 分數比具備十億參數的 Prov-GigaPath 高出 5.09%3.83%。此外,在跨資料集的泛化能力測試中,從 NCT 資料集訓練的模型遷移至 MHIST 與 CRC 資料集時,依然保持極為穩定的領先成績。這證明了相較於動輒數億參數的通用大模型,僅有 25.3M 參數的專屬架構更能精準捕捉疾病特定的診斷訊號。

於 6 個 WSI 下游任務擊敗 8 款主流算法

模型在感興趣區域取得成功後,其優勢也順利延伸至全玻片影像(WSI)等級的複雜診斷。研究團隊將其整合為 SSMambaMIL 架構,挑戰 PANDA、TCGA 等 6 個 WSI 資料集的下游任務,涵蓋前列腺癌侵襲性分級、HER2 表現狀態預測與整體存活期回歸等多重情境。

在與 ABMIL、CLAM 及 MambaMIL 等 8 款先進演算法的較量中,SSMambaMIL 在總計 17 項評估指標裡,囊括了 12 項第一。例如在 TCGA-LGG 的腫瘤分期任務中,其 AUC 領先次佳模型達 4.54%。這一系列結果證實,具有線性運算時間與先進序列建模能力的 SSMamba 框架,已在臨床應用的廣泛適應性上樹立了全新的標竿。

專為病理特徵客製化的輕量狀態空間模型,證明了底層架構的領域適應性,遠比單純擴展參數規模更能有效萃取複雜的醫療診斷訊號。

Abstract

Pathological diagnosis is highly reliant on image analysis, where Regions of Interest (ROIs) serve as the primary basis for diagnostic evidence, while whole-slide image (WSI)-level tasks primarily capture aggregated patterns. To extract these critical morphological features, ROI-level Foundation Models (FMs) based on Vision Transformers (ViTs) and large-scale self-supervised learning (SSL) have been widely adopted. However, three core limitations remain in their application to ROI analysis: (1) cross-magnification domain shift, as fixed-scale pretraining hinders adaptation to diverse clinical settings; (2) inadequate local-global relationship modeling, wherein the ViT backbone of FMs suffers from high computational overhead and imprecise local characterization; (3) insufficient fine-grained sensitivity, as traditional self-attention mechanisms tend to overlook subtle diagnostic cues. To address these challenges, we propose SSMamba, a hybrid SSL framework that enables effective fine-grained feature learning without relying on large external datasets. This framework incorporates three domain-adaptive components: Mamba Masked Image Modeling (MAMIM) for mitigating domain shift, a Directional Multi-scale (DMS) module for balanced local-global modeling, and a Local Perception Residual (LPR) module for enhanced fine-grained sensitivity. Employing a two-stage pipeline, SSL pretraining on target ROI datasets followed by supervised fine-tuning (SFT), SSMamba outperforms 11 state-of-the-art (SOTA) pathological FMs on 10 public ROI datasets and surpasses 8 SOTA methods on 6 public WSI datasets. These results validate the superiority of task-specific architectural designs for pathological image analysis.