KHiM-Mamba: Injecting Pathology Knowledge into Mamba via Hidden-State Modulation for Whole Slide Image Analysis
Authors: Qixiang Zhang, Yi Li, Tianqi Xiang, Haonan Wang, Mengjiao Wei, Bo Xu, Xiaomeng Li
Organizations: Department of Electronic and Computer Engineering, The Hong Kong University of Science and Technology, Hong Kong, China · Center for Intelligent Oncology, Chongqing University Cancer Hospital and Chongqing University School of Medicine, and Chongqing Key Laboratory of Intelligent Oncology for Breast Cancer, Chongqing, 400030, China
Whole slide image analysis is commonly formulated as multiple instance learning (MIL), where instance features are contextually updated and aggregated into a slide representation, a process we term slide encoding dynamics. Recently, selective state-space models (SSM) have emerged as promising MIL architectures due to their long-sequence modeling capability and linear complexity. However, existing SSM-based MIL methods rely solely on visual features during MIL. Meanwhile, in large-scale WSIs, where sparse diagnostically decisive regions are surrounded by abundant irrelevant information, such purely vision-driven selective dynamics can misallocate state updates and readouts, causing the evolving SSM state to accumulate task-irrelevant evidence and dilute critical diagnostic cues over long scan trajectories. In this work, we propose the Knowledge-Aware Hidden-State Modulation architecture (KHiM-Mamba), which innovatively regulates Mamba's core selective state-space mechanism with explicit knowledge priors, steering slide encoding dynamics toward diagnostically meaningful evidence accumulation. Specifically, we redesign the original SSM layer to perform knowledge modulation operations during the evolution of hidden states, thereby guiding what visual evidence is accumulated and retrieved from the hidden state at each encoding step. Furthermore, we additionally introduce a local-adaptive vocabulary retrieval module that uses large language models to assign each patch fine-grained, tissue-specific semantic descriptions, enabling precise modulation across diverse tasks. Experiments on 11 public benchmarks across 4 tasks show that KHiM-Mamba consistently achieves state-of-the-art performance.
Figures & tables
Figure 1: (a) MIL pipeline: a WSI is split into patches, encoded as patch features, and processed by slide encoding dynamics ( i.e. , contextual interaction and aggregation). (b) Example of diagnostically irrelevant visual patterns in WSIs: under a certain stain, tumor slides appear redder while normal slides appear purpler. (c)(d) Comparison of SSM layer in prior Mamba-based MIL methods and ours.
Figure 2: (a) Existing VLM-MIL methods typically inject language priors either before slide encoding [ 10 , 21 ] to guide patch selection via image-text similarity, or after slide encoding as auxiliary supervision [ 24 ] . (b) KHiM-Mamba injects knowledge priors into the core slide encoding dynamics by modulating two key matrix: input-to-state Bt , which governs how each patch’s information is written into the hidden state, and state-to-output Ct , which governs how accumulated context is read out to enrich patch representations. (c) An illustrative example of language-guided contextual interaction across patches. In invasive ductal carcinoma, ‘‘conspicuous cytologic atypia and variation in nuclear size’’ guides the model to relate distinct patches with complementary nuclear morphologies, thereby grounding the slide-level representation in pathology-relevant evidence. (d) Histograms of patch response scores for background and tumor-related patches, with and without our Knowledge-aware Hidden-state Modulation (KHiM). KHiM markedly reduces the overlap between the two distributions, indicating stronger patch selectivity and clearer semantic separation; More quantitative and qualitative results to support this statement can be found in Tab. 5 and Fig. 5 respectively
Figure 3: Overview of KHiM-Mamba . Our framework consists of two main blocks: Local-adaptive vocabulary retrieval module (Sec. 3.3 ) provides precise language priors across diverse downstream tasks; and K nowledge-aware Hi dden-state M odulation (KHiM-Mamba, Sec. 3.2 ) which leverages language priors to modulate the Mamba-based slide encoding dynamics.
BRACS-7 (n=547, p<0.05 )
BRACS-3 (n=547, p<0.01 )
TCGA-BRCA (n=952, p<0.05 )
Method
ACC. ( ± std)
AUC ( ± std)
RANK
ACC. ( ± std)
AUC ( ± std)
RANK
ACC. ( ± std)
AUC ( ± std)
RANK
Mean Pooling
0.5378±0.0575
0.8341±0.0261
15.5
0.7746±0.0442
0.8769±0.0248
15
0.8970±0.0402
0.9204±0.0385
17
Max Pooling
0.5315±0.0426
0.8084±0.0588
19.5
0.7463±0.0424
0.8782±0.0203
19.5
0.8812±0.0399
0.9042±0.0449
22
ABMIL [ 14 ] [ICML’21]
0.5378±0.0354
0.8268±0.0280
16.5
0.7571±0.0665
0.8730±0.0410
19
0.9028±0.0275
0.9354±0.0424
11
Table 1: Comparison Results of Cancer Sub-typing Task using CONCH [ 18 ] Features. p denotes paired t-tests between our KHiM-Mamba and the second-best counterpart on the corresponding dataset.
Method
BLCA (n=437)
BRCA (n=1023)
KIRC (n=498)
KIRP (n=261)
LUAD (n=455)
STAD (n=363)
MEAN
RANK
Mean-Pooling
0.6758±0.0333
0.6274±0.0368
0.7287±0.0348
0.8179±0.0426
0.6334±0.0897
0.5917±0.0412
0.6792
19
Max-Pooling
0.6417±0.0446
0.6095±0.1011
0.7316±0.0557
0.8164±0.0743
0.6245±0.0671
0.5554±0.0423
0.6632
21
ABMIL [ 14 ] [ICML’21]
0.6604±0.0183
0.6621±0.0263
0.7185±0.0609
0.8201±0.0480
0.6433±0.0766
0.6053±0.0329
0.6850
17
TransMIL [ 23 ] [NIPS’21]
0.6705±0.0244
0.6592±0.0593
0.7174±0.0708
0.8357±0.0477
0.6185±0.1102
0.6597±0.0374
0.6935
11
Table 2: Comparison Results of Survival Analysis Task on Six TCGA Subsets.
Figure 4: Kaplan–Meier survival curves of second-best method GMMamba [ 40 ] and our proposed KHiM-Mamba.
Method
CENTER 1
CENTER 2
CENTER 3
CENTER 4
CENTER 5
MEAN
RANK
Mean-Pooling
0.6760±0.1237
0.3780±0.0183
0.7850±0.0838
0.6758±0.1063
0.6647±0.0301
0.6359
16
Max-Pooling
0.9210±0.0164
0.7780±0.1454
0.9510±0.0094
0.9333±0.0112
0.8330±0.0200
0.8833
4
ABMIL [ 14 ] [ICML’21]
0.9120±0.0248
0.7450±0.1344
0.9480±0.0098
0.9323±0.0197
0.8460±0.0918
0.8767
7
Table 3: Domain Generalization Performance on CAMELYON16 to CAMELYON17 [ 6 ] .
8-SHOT
16-SHOT
MEAN
Method
ACC.
AUC
ACC.
AUC
ACC.
AUC
CLAM-SB [ 19 ]
59.96±9.44
73.68±6.72
78.52±4.16
87.03±3.78
69.11
80.36
CLAM-MB [ 19 ]
66.63±2.48
79.69±6.34
73.76±3.88
88.01±2.99
70.20
83.85
TransMIL [ 23 ]
75.82±7.44
83.69±7.05
73.88±2.58
88.53±3.17
74.85
86.11
ViLa-MIL [ 24 ]
73.51±9.22
84.20±7.09
71.28±3.56
88.30±4.12
72.40
86.25
CoOp [ 41 ]
69.48±5.93
80.11±5.91
71.81±4.35
82.10±3.72
70.65
81.11
Table 4: Few-shot performance on the TCGA-NSCLC [ 28 ] dataset.
Settings
ΔB
ΔC
BRACS-7
TCGA-BRCA
CAMELYON16
Original SSM
✗
✗
0.8277±0.0542
0.9046±0.0372
0.8573±0.0549
Version 1
✔
✗
\uline0.8333±0.0439
\uline0.9067±0.0538
0.8745±0.0576
Version 2
✗
✔
0.8325±0.0643
0.8976±0.0422
0.8620±0.0628
Table 5: Ablations on our core Knowledge-aware Hidden-state Modulation (KHiM) module upon cancer sub-typing datasets, i.e. , BRACS-7, TCGA-BRCA, and CAMELYON16 tumor region localization (patch-level classification) task.
Vocabulary
BRACS-7
TCGA-BRCA
MEAN
Class-label
0.8466±0.0451
0.9411±0.0675
0.8939
VilaMIL [ 24 ] Prompt
0.8671±0.0420
0.9509±0.0483
0.9090
TOP [ 21 ] Prompt
0.8644±0.0387
0.9521±0.0576
0.9083
Table 6: Ablations on our Local-Adaptive Vocabulary Retrieval module upon cancer sub-typing datasets.
Figure 5: Visualization of patch attention map using patch features with and without Knowledge-aware Hidden-state Modulation (KHiM). Major tumor areas are circled in yellow.
n=1,000
n=5,000
Method
Params.
Time
Thro.
Time
Thro.
Mean-Pooling
525.8 K
0.25 ms
4.02 K
0.53 ms
1.90 K
Max-Pooling
525.8 K
0.19 ms
5.22 K
0.51 ms
1.97 K
CLAM [ 19 ] [NAT.BME’20]
790.8 K
0.82 ms
1.22 K
0.94 ms
1.07 K
ABMIL [ 14 ] [ICML’21]
591.6 K
0.71 ms
1.41 K
0.80 ms
1.25 K
DSMIL [ 15 ] [CVPR’21]
872.7 K
1.48 ms
0.68 K
1.95 ms
0.51 K
Table 7: Efficacy evaluation with different numbers of patches per slide ( i.e. , 1,000 and 5,000). Parameter counts (Params.), inference time, and throughput (Thro.) are reported in thousands (K), milliseconds (ms), and thousands of patches per second (K patches/s), respectively.
Whole-slide image (WSI) analysis is challenging due to the gigapixel scale of slides and their inherent hierarchical multi-resolution structure. Existing multiple instance learning (MIL) approaches often model WSIs as unordered collections of patches, which limits their ability to capture structured dependencies between global tissue organization and local cellular patterns. Although recent State Space Models (SSMs) enable efficient modeling of long sequences, how to structure WSI tokens to fully exploit their spatial hierarchy remains an open problem.We propose MoEMambaMIL, a structure-aware SSM framework for WSI analysis that integrates region-nested selective scanning with mixture-of-experts (MoE) modeling. Leveraging multi-resolution preprocessing, MoEMambaMIL organizes patch tokens into region-aware sequences that preserve spatial containment across resolutions. On top of this structured sequence, we decouple resolution-aware encoding and region-adaptive contextual modeling via a combination of static, resolution-specific experts and dynamic sparse experts with learned routing. This design enables efficient long-sequence modeling while promoting expert specialization across heterogeneous diagnostic patterns. Experiments demonstrate that MoEMambaMIL achieves the best performance across 9 downstream tasks.
Whole Slide Image (WSI) analysis is pivotal in computational pathology, enabling cancer diagnosis by integrating morphological and architectural cues across magnifications. Multiple Instance Learning (MIL) serves as the standard framework for WSI analysis. Recently, Mamba has become a promising backbone for MIL, overtaking Transformers due to its efficiency and global context modeling capabilities originating from Natural Language Processing (NLP). However, existing Mamba-based MIL approaches face three critical challenges: (1) disruption of 2D spatial locality during 1D sequence flattening; (2) sub-optimal modeling of fine-grained local cellular structures; and (3) high memory peaks during inference on resource-constrained edge devices. Studies like MambaOut reveal that Mamba's SSM component is redundant for local feature extraction, where Gated CNNs suffice. Recognizing that WSI analysis demands both fine-grained local feature extraction akin to natural images, and global context modeling akin to NLP, we propose MambaBack, a novel hybrid architecture that harmonizes the strengths of Mamba and MambaOut. First, we propose the Hilbert sampling strategy to preserve the 2D spatial locality of tiles within 1D sequences, enhancing the model's spatial perception. Second, we design a hierarchical structure comprising a 1D Gated CNN block based on MambaOut to capture local cellular features, and a BiMamba2 block to aggregate global context, jointly enhancing multi-scale representation. Finally, we implement an asymmetric chunking design, allowing parallel processing during training and chunking-streaming accumulation during inference, minimizing peak memory usage for deployment. Experimental results on five datasets demonstrate that MambaBack outperforms seven state-of-the-art methods. Source code and datasets are publicly available.
Sicheng Chen, Chad Wong, Tianyi Zhang +3
Nhu Department of Electrical Engineering and Computer Science, University of California, Irvine · Department of Electrical & Computer Engineering, National University of Singapore · PuzzleLogic Pte Ltd, Singapore 229594, Singapore
Whole Slide Images (WSIs) exhibit hierarchical structure, where diagnostic information emerges from cellular morphology, regional tissue organization, and global context. Existing Computational Pathology (CPath) Multimodal Large Language Models (MLLMs) typically compress an entire WSI into a single embedding, which hinders fine-grained grounding and ignores how pathologists synthesize evidence across different scales. We introduce MLLM-HWSI, a Hierarchical WSI-level MLLM that aligns visual features with pathology language at four distinct scales, cell as word, patch as phrase, region as sentence, and WSI as paragraph to support interpretable evidence-grounded reasoning. MLLM-HWSI decomposes each WSI into multi-scale embeddings with scale-specific projectors and jointly enforces (i) a hierarchical contrastive objective and (ii) a cross-scale consistency loss, preserving semantic coherence from cells to the WSI. We compute diagnostically relevant patches and aggregate segmented cell embeddings into a compact cellular token per-patch using a lightweight Cell-Cell Attention Fusion (CCAF) transformer. The projected multi-scale tokens are fused with text tokens and fed to an instruction-tuned LLM for open-ended reasoning, VQA, report, and caption generation tasks. Trained in three stages, MLLM-HWSI achieves new SOTA results on 13 WSI-level benchmarks across six CPath tasks. By aligning language with multi-scale visual evidence, MLLM-HWSI provides accurate, interpretable outputs that mirror diagnostic workflows and advance holistic WSI understanding. Code is available at: GitHub.
Basit Alawode, Arif Mahmood, Muaz Khalifa Al-Radi +6
Department of Computer Science, Khalifa University of Science and Technology, UAE. · Information Technology University, Pakistan. · KAU, KSA. +1