🧬 biology

KHiM-Mamba: Injecting Pathology Knowledge into Mamba via Hidden-State Modulation for Whole Slide Image Analysis

KHiM-Mamba is a novel multiple instance learning architecture that enhances Whole Slide Image analysis by integrating pathology knowledge and LLM-derived semantic descriptions to modulate Mamba's hidden states, thereby preventing the accumulation of irrelevant information and achieving state-of-the-art performance across diverse diagnostic tasks.

Original authors: Qixiang Zhang, Yi Li, Tianqi Xiang, Haonan Wang, Mengjiao Wei, Bo Xu, Xiaomeng Li

Published 2026-08-18

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Qixiang Zhang, Yi Li, Tianqi Xiang, Haonan Wang, Mengjiao Wei, Bo Xu, Xiaomeng Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

In the world of modern medicine, a pathologist's most critical tool is a glass slide containing a tiny sample of human tissue, magnified until it reveals the cellular architecture of life and disease. Today, these slides are scanned into massive digital images, so large they contain billions of pixels, far too vast for a human eye to examine in a single glance. To help doctors navigate these digital landscapes, scientists have developed artificial intelligence systems that act as tireless assistants. These systems break the giant image into thousands of small, manageable pieces, analyze the texture and shape of cells in each piece, and then try to piece together a global diagnosis. The challenge lies in the sheer scale and the noise; a single slide might contain millions of cells, but only a few scattered clusters hold the key to a diagnosis, while the rest is healthy tissue or irrelevant background. If the computer gets distracted by the vast amount of unimportant information, it can miss the critical clues hidden in the corners.

Researchers have recently turned to a new type of artificial intelligence architecture, inspired by how the human brain processes sequences of information over time, to handle these massive digital slides. This technology is excellent at remembering what it saw earlier in a long sequence and deciding what to keep in memory as it moves forward. However, a team of scientists from Hong Kong and Chongqing discovered a flaw in how these systems were being used. They found that when these AI models scanned a slide, they relied entirely on what the pixels looked like. If a slide was stained with a slightly different color of dye—a common variation in labs—the model might get confused, treating a harmless color shift as a significant medical clue. This caused the system to accumulate irrelevant details in its memory, diluting the critical evidence needed for a diagnosis. To solve this, the researchers introduced a new method that injects actual medical knowledge directly into the AI's memory process, teaching the system to ignore the noise and focus on what truly matters.

The team, led by researchers at the Hong Kong University of Science and Technology and Chongqing University, developed a system they call KHiM-Mamba. Instead of letting the AI decide on its own what to remember as it scans a slide, they gave it a guidebook written in the language of pathology. This guidebook contains specific descriptions of what different types of tissue look like when they are healthy or diseased. As the AI moves from one small piece of the slide to the next, it consults this guidebook to decide what information is worth storing in its memory and what should be discarded. Imagine a librarian who, while sorting through a massive pile of books, is handed a specific list of titles to keep and a list of titles to ignore; the librarian no longer has to guess which books are important. In this case, the "librarian" is the AI, the "books" are the tiny image patches, and the "list" is the medical knowledge provided by a large language model.

The researchers found that previous attempts to use language to help AI diagnose slides were often too indirect. Those methods would use language to pick which slides to look at before the analysis began, or to check the final answer afterward. The new approach is different because it changes the core mechanism of how the AI thinks. It modifies the internal "write" and "read" processes of the memory system. When the AI writes a new memory, the medical knowledge tells it to focus on specific shapes and patterns associated with disease, rather than just the colors or textures that might be misleading. When the AI reads from its memory to make a decision, the knowledge guides it to retrieve the most relevant evidence accumulated so far. This ensures that the system builds a representation of the slide based on medically meaningful connections, rather than getting sidetracked by irrelevant visual variations.

To make this work, the team created a smart system that generates the right medical descriptions for each specific part of the slide. Since a single slide contains many different types of tissue, a generic description isn't enough. The system uses a powerful language model to look at a small patch of tissue, identify what kind of tissue it is, and then generate a precise, detailed description of what that tissue should look like if it were healthy or diseased. If the system is unsure about the tissue type, it falls back to a broader, safer description to avoid making mistakes. This dynamic process ensures that the AI receives the most accurate and specific guidance possible for every single piece of the puzzle it is solving.

The results of this approach were tested across eleven different public datasets involving various types of cancer and patient outcomes. The new system consistently outperformed the best existing methods. In tasks where the AI had to distinguish between different subtypes of breast cancer, it achieved higher accuracy than any previous model. When predicting patient survival times based on tissue samples, it provided more reliable risk assessments, better separating patients into high-risk and low-risk groups. Perhaps most importantly, the system proved to be robust when tested on slides from different hospitals that used different staining techniques. While other models struggled with these variations, often failing when the colors changed slightly, the new system maintained its high performance. This suggests that by grounding the AI in medical knowledge rather than just visual patterns, it learned to see the disease itself, not just the appearance of the slide.

The study also looked at how well the system works when it has very little data to learn from, a common problem in medicine where labeled examples are scarce. In tests where the AI was shown only eight or sixteen examples of a disease to learn from, the new method significantly outperformed its competitors. This indicates that the medical knowledge injected into the system acts as a strong foundation, allowing it to learn effectively even with minimal training data. The researchers also confirmed that adding this layer of knowledge did not slow the system down or make it too large to run on standard computer hardware. The system remained fast and efficient, capable of processing thousands of image patches in seconds.

By rethinking how artificial intelligence processes medical images, this work demonstrates that combining visual data with explicit medical knowledge creates a more reliable and accurate diagnostic tool. The researchers showed that the key to solving the problem of massive, noisy data is not just to build bigger computers, but to teach the computer what to look for. This approach offers a new path forward for computational pathology, where the goal is to assist doctors by filtering out the irrelevant and highlighting the critical, ensuring that the final diagnosis is based on the strongest possible evidence.

Technical Summary: KHiM-Mamba

Problem Statement

Whole-slide image (WSI) analysis in computational pathology is typically formulated as Multiple Instance Learning (MIL), where a slide is decomposed into patches, and a "slide encoder" aggregates patch features into a slide-level representation. This process, termed slide encoding dynamics, is critical because diagnostic evidence is often distributed across distant tissue regions.

Recently, Selective State Space Models (SSMs), specifically Mamba, have emerged as efficient architectures for this task due to their ability to model long sequences with linear complexity. However, the authors identify a fundamental limitation in existing Mamba-based MIL methods: they rely solely on visual features during the selective state transitions. In large-scale WSIs, where diagnostically decisive regions are surrounded by abundant irrelevant information (e.g., staining variations, background tissue), purely vision-driven dynamics can misallocate state updates. This causes the evolving hidden state to accumulate task-irrelevant evidence and dilute critical diagnostic cues over long scan trajectories.

While Vision-Language Model (VLM) based MIL methods exist, they typically treat language as an external add-on (e.g., for patch selection or auxiliary supervision) rather than integrating it into the core encoding dynamics. Consequently, the limitation of purely visual state transitions remains unresolved.

Methodology: KHiM-Mamba

The authors propose KHiM-Mamba (Knowledge-aware Hidden-state Modulation), a framework that injects explicit pathology knowledge priors directly into the core selective state-space mechanism of Mamba. The architecture consists of two primary components:

1. Knowledge-aware Hidden-state Modulation (KHiM)

Instead of using language as a peripheral signal, KHiM-Mamba redesigns the SSM layer to modulate the hidden state evolution using language-derived knowledge.

Mechanism: The standard Mamba recurrence involves input projection ( $B_t$ ) and output projection ( $C_t$ ) matrices that control how information is written into and read from the hidden state ( $h_t$ ).
Modulation: The authors introduce a low-rank modulation mechanism where language embeddings ( $e_t$ $e_{t}$ ) generate coefficient vectors ( $\alpha_B, \alpha_C$ $α_{B}, α_{C}$ ) via an MLP. These coefficients are used to create low-rank increments ( $\Delta B_t, \Delta C_t$ $Δ B_{t}, Δ C_{t}$ ) that are added to the original projection matrices.
- Write-in Modulation ( $\tilde{B}_t$ ): Guides what visual evidence is incorporated into the evolving hidden state, encouraging the accumulation of diagnosis-relevant features while attenuating irrelevant ones.
- Read-out Modulation ( $\tilde{C}_t$ ): Guides what accumulated contextual evidence is retrieved to enrich patch representations, ensuring the output emphasizes diagnostically meaningful patterns.
Result: This reshapes the "write-read" dynamics of the SSM without altering its temporal decay behavior, making the selective mechanism "diagnosis-aware."

2. Local-Adaptive Vocabulary Retrieval Module

To provide effective language priors for the modulation, the authors design a module that avoids fixed, global prompts.

Process:
1. Concept Generation: A Large Language Model (LLM) is prompted with the task definition to generate a set of task-relevant tissue concepts (e.g., specific tumor types) and their fine-grained morphological descriptions.
2. Confidence-Gated Routing: For each patch, a vision-language model (CONCH) computes the similarity between the patch image embedding and the concept text embeddings.
3. Adaptive Assignment: If the similarity to a specific tissue concept exceeds a threshold and the margin over other concepts is sufficient, the patch is assigned fine-grained, tissue-specific language. Otherwise, it falls back to a coarse-grained pathological description to avoid introducing noisy priors.
Outcome: This yields modulation signals that are semantically precise and robust across diverse tissue compositions.

Key Contributions

Identification of a Limitation: The paper highlights that existing Mamba-based MIL methods suffer from purely vision-driven state transitions that accumulate irrelevant evidence in long WSI sequences, and that existing VLM-MIL methods fail to address this because they do not intervene in the core encoding dynamics.
Architectural Innovation: KHiM-Mamba is presented as the first work to reshape Mamba's internal model dynamics for WSI analysis by injecting knowledge priors directly into the selective state transitions (write-in and read-out matrices), rather than using Mamba as a static backbone.
Local-Adaptive Vocabulary Retrieval: The proposal of a module that dynamically assigns fine-grained, tissue-specific language to patches via LLMs, providing precise modulation signals across diverse downstream tasks.
Empirical Performance: Extensive experiments demonstrate state-of-the-art (SOTA) performance across multiple tasks and datasets.

Experimental Results

The authors evaluated KHiM-Mamba on 11 public benchmarks across four tasks: cancer subtyping, survival prediction, domain adaptation, and few-shot learning.

Cancer Subtyping:
- BRACS Dataset: Achieved SOTA accuracy (0.5874) and AUC (0.8746) on the fine-grained 7-class task, and best results on the 3-class task, outperforming previous bests by ~1.77% and ~2.41% respectively.
- TCGA-BRCA: Achieved 0.9330 accuracy and 0.9559 AUC, improving upon the previous best accuracy by 1.97% and AUC by 1.01%.
Survival Prediction:
- Evaluated on six TCGA subsets (BLCA, BRCA, KIRC, KIRP, LUAD, STAD). KHiM-Mamba achieved a mean C-Index of 0.7394, surpassing the previous best (2DMamba) by 2.31%.
- Kaplan-Meier curves showed significantly sharper separation between high- and low-risk groups, with notably lower log-rank $p$ -values (e.g., $2.63 \times 10^{-6}$ on LUAD).
Domain Generalization (Cross-Center):
- Trained on CAMELYON16 and tested on CAMELYON17 (5 independent centers). KHiM-Mamba achieved a mean AUC of 0.9184, outperforming the previous best (TransMIL) by 2.27%, demonstrating robustness to staining and scanner variations.
Few-Shot Learning:
- On TCGA-NSCLC with 8-shot and 16-shot settings, KHiM-Mamba achieved 83.11% mean accuracy and 93.03% AUC, significantly outperforming other MIL and VLM-MIL methods, indicating strong inductive bias from language modulation.
Efficiency:
- The model maintains linear complexity with ~2.2M parameters, showing inference times comparable to or faster than other Mamba-based methods (e.g., PAM, 2DMamba) even with long sequences (5,000 patches).

Significance and Claims

The paper claims that KHiM-Mamba addresses a fundamental gap in computational pathology by moving beyond purely visual feature aggregation. By explicitly regulating the hidden-state evolution with diagnostic knowledge, the model steers the accumulation of evidence toward pathologically meaningful regions, effectively filtering out irrelevant background noise.

The authors emphasize that their approach is orthogonal to existing scanning strategies (e.g., grouping, hierarchical scanning) and can be combined with them. The significance lies in demonstrating that language-conditioned modulation of the core state-space mechanism is a powerful inductive bias that enhances generalization, robustness to domain shifts, and performance in data-scarce scenarios, offering a new perspective for designing knowledge-guided state-space models in medical imaging.