SemaMIL: Semantic-Aware Multiple Instance Learning with Retrieval-Guided State Space Modeling for Whole Slide Images

TL;DR

SemaMIL introduces a novel semantic-aware multiple instance learning framework for whole slide images, combining semantic reordering and retrieval-guided state space modeling to improve accuracy and efficiency in computational pathology.

Contribution

It proposes a new method integrating semantic reordering and retrieval-guided state space modeling to enhance interpretability and performance in WSI analysis.

Findings

Achieves state-of-the-art accuracy on four WSI datasets.

Uses fewer FLOPs and parameters than baseline methods.

Improves global modeling of histological features.

Abstract

Multiple instance learning (MIL) has become the leading approach for extracting discriminative features from whole slide images (WSIs) in computational pathology. Attention-based MIL methods can identify key patches but tend to overlook contextual relationships. Transformer models are able to model interactions but require quadratic computational cost and are prone to overfitting. State space models (SSMs) offer linear complexity, yet shuffling patch order disrupts histological meaning and reduces interpretability. In this work, we introduce SemaMIL, which integrates Semantic Reordering (SR), an adaptive method that clusters and arranges semantically similar patches in sequence through a reversible permutation, with a Semantic-guided Retrieval State Space Module (SRSM) that chooses a representative subset of queries to adjust state space parameters for improved global modeling. Evaluation on four WSI subtype datasets shows that, compared to strong baselines, SemaMIL achieves state-of-the-art accuracy with fewer FLOPs and parameters.

Tables

Table 1. Table 1 : Subtype classification results on four datasets. Bold : best performance, underline : second-best.

	EBRAINS		BRACS$\star$		BRACS		IPD		TCGA
	AUC	ACC	AUC	ACC	AUC	ACC	AUC	ACC	AUC	ACC
Max-Pooling	0.979$\pm$0.005	0.714$\pm$0.021	0.762	0.391	0.821$\pm$0.026	0.537$\pm$0.055	0.861$\pm$0.037	0.729$\pm$0.035	0.918$\pm$0.019	0.733$\pm$0.055
Mean-Pooling	0.978$\pm$0.003	0.714$\pm$0.035	0.711	0.345	0.827$\pm$0.023	0.556$\pm$0.045	0.880$\pm$0.038	0.756$\pm$0.056	0.914$\pm$0.013	0.699$\pm$0.035
ABMIL [7]	0.967$\pm$0.004	0.740$\pm$0.032	0.792	0.451	0.842$\pm$0.021	0.575$\pm$0.061	0.892$\pm$0.032	0.779$\pm$0.046	0.924$\pm$0.021	0.716$\pm$0.046
CLAM-MB [17]	0.976$\pm$0.005	0.722$\pm$0.033	0.796	0.461	0.836$\pm$0.027	0.573$\pm$0.040	0.885$\pm$0.037	0.760$\pm$0.049	0.921$\pm$0.020	0.717$\pm$0.045
DSMIL [8]	0.972$\pm$0.003	0.722$\pm$0.031	0.787	0.452	0.856$\pm$0.025	0.579$\pm$0.036	0.892$\pm$0.032	0.769$\pm$0.048	0.910$\pm$0.022	0.727$\pm$0.034
DTFDMIL [18]	0.970$\pm$0.003	0.741$\pm$0.023	0.802	0.462	0.853$\pm$0.028	0.576$\pm$0.037	0.895$\pm$0.036	0.781$\pm$0.047	0.926$\pm$0.010	0.749$\pm$0.037
TransMIL [19]	0.981$\pm$0.003	0.711$\pm$0.032	0.808	0.459	0.829$\pm$0.028	0.560$\pm$0.055	0.860$\pm$0.041	0.718$\pm$0.048	0.913$\pm$0.018	0.689$\pm$0.047
S4MIL [16]	0.965$\pm$0.003	0.715$\pm$0.029	0.778	0.425	0.834$\pm$0.021	0.573$\pm$0.036	0.897$\pm$0.046	0.762$\pm$0.077	0.927$\pm$0.019	0.735$\pm$0.039
MambaMIL [11]	0.982$\pm$0.004	0.744$\pm$0.034	0.820	0.467	0.855$\pm$0.027	0.586$\pm$0.059	0.910$\pm$0.034	0.781$\pm$0.056	0.924$\pm$0.023	0.736$\pm$0.048
Ours	0.984$\pm$0.004	0.751$\pm$0.024	0.821	0.472	0.861$\pm$0.025	0.603$\pm$0.036	0.918$\pm$0.025	0.797$\pm$0.047	0.933$\pm$0.023	0.755$\pm$0.047

Table 2. Table 2 : Ablation study for SR and SRSM.

Our Proposed		EBRAINS		BRACS
SR	SRSM	AUC	ACC	AUC	ACC
✗	✗	0.969$\pm$0.005	0.722$\pm$0.021	0.822$\pm$0.028	0.561$\pm$0.030
✗	✓	0.977$\pm$0.004	0.735$\pm$0.029	0.842$\pm$0.026	0.578$\pm$0.028
✓	✗	0.979$\pm$0.004	0.740$\pm$0.028	0.846$\pm$0.025	0.584$\pm$0.027
✓	✓	0.984$\pm$0.004	0.751$\pm$0.024	0.861$\pm$0.025	0.603$\pm$0.036

Table 3. Table 3 : Comparison of computational efficiency across different methods.

Methods	ABMIL	CLAM	DSMIL	DTFDMIL	TransMIL	S4MIL	MambaMIL	Ours
FLOPs(G)	0.352	0.354	0.624	0.318	0.502	0.706	0.255	0.248
Params(M)	0.467	0.662	0.476	0.586	0.542	0.924	0.470	0.464
ACC	0.740	0.722	0.722	0.741	0.711	0.715	0.744	0.751

Equations

$$h_i=GELU(W_1{x}_i),z_i=W_2{h}_i.$$

$$\pi =\arg sort(c_1,\cdots ,c_L)$$

$$y_i={y_{{\pi }^{-1}(i)}}^{\prime }.$$

$$\dot{h}(t)$$

$$y(t)$$

$$h_k$$

$$y_k$$

$$\begin{array}{rl}\Delta & =W_{\Delta }\mathrm{vec}(Q)+b_{\Delta },\\ B^{\prime }& =W_B\mathrm{vec}(Q)+b_B,\end{array}$$

$$\begin{array}{c}{A}_d(\Delta )=\exp (\Delta A_0),\\ B_d(\Delta ,B^{\prime })=(\Delta A_0{)}^{-1}(\exp (\Delta A_0)-I)B^{\prime },\end{array}$$

$$\begin{array}{c}{h}_k=A_d(\Delta )h_{k-1}+B_d(\Delta ,B^{\prime })c_k,\\ y_k=C{h}_k+D{c}_k,\end{array}$$

Full text

**Index Terms— ** Computational Pathology, Whole Slide Images, Multiple Instance Learning, Mamba

1 Introduction

The advent of digital pathology has positioned Whole Slide Images (WSIs) as a pivotal data modality for computational pathology, offering unprecedented opportunities for automated diagnosis and prognosis [20, 21, 22, 23, 24, 25, 26, 27]. Nonetheless, the gigapixel-scale resolution of WSIs, coupled with the scarcity of pixel-level annotations, poses substantial obstacles to the direct application of conventional deep learning techniques. Multiple Instance Learning (MIL) [1, 2, 3, 28, 29, 30, 31, 32] has therefore emerged as the prevailing paradigm to circumvent these challenges. By obviating the need for exhaustive, fine-grained annotations while still enabling effective exploitation of discriminative cues embedded in large-scale WSIs, MIL furnishes a principled framework that bridges state-of-the-art artificial intelligence methodologies with the stringent requirements of medical image analysis.

Current MIL pipelines typically compress each tissue patch into a low-dimensional embedding with a pretrained encoder [4, 5, 6, 57, 58, 59, 60, 61] and then aggregate these embeddings to form a bag-level representation for downstream tasks. This design turns WSI analysis into a long-sequence modeling problem: a model must capture both the relations among patches and the global context of the whole slide to extract truly discriminative cues. Attention‐based MIL methods [7, 17, 8, 39, 40, 41, 42, 43, 44], while effective in highlighting discriminative patches, typically treat each patch independently and neglect the contextual dependencies inherent to tissue architecture. Transformer‐based MIL models [19, 9, 10, 45, 46, 47, 48, 49, 50] can explicitly capture patch interactions, but their self‐attention mechanism incurs quadratic computation and memory costs, leading to prohibitive resource requirements and a propensity to overfit when annotation is scarce.

Recently, state‐space models (SSMs) [11, fang2024mammilren2024ultrapixel, 34, 35, 36, 37, 38] have recently emerged as a powerful alternative for long‐sequence modeling. By providing linear computational complexity alongside a global receptive field, SSMs can efficiently capture long‐range dependencies across thousands of tokens. These properties render SSMs especially attractive for WSI analysis, where sequences of patch embeddings can easily exceed tens of thousands of elements.

Nevertheless, applying SSMs to MIL directly in pathology encounters its own challenges. Existing pipelines typically randomize the ordering of patch embeddings before sequential processing, thereby discarding histological priors and separating semantically related regions in the sequence [51, 52, 53, 54, 55, 56]. This arbitrary reordering undermines the model’s ability to exploit tissue‐level context, diminishes interpretability of the learned interactions, and ultimately constrains classification performance [62, 63, 64].

To address these limitations, we propose SemaMIL, a novel Semantic-guided Multiple Instance Learning framework with the following contributions: (1) We introduce a semantic-aware patch ordering mechanism that arranges patches with higher semantic similarity closer together in the sequence, thereby enhancing interaction among histologically relevant regions. (2) We further design a Semantic-guided Retrieval State Space Module (SRSM) to reinforce longrange dependency modeling and data augmentation within the ordered sequence. (3) To evaluate the effectiveness of SemaMIL, we conduct comprehensive experiments on the subtype classification task of whole slide images across four challenging datasets. The results demonstrate that SemaMIL consistently outperforms state-of-the-art methods, further validating its superior performance and robustness.

2 Methods

To better exploit semantic relationships among tissue patches and further enhance long-sequence modeling in multiple instance learning, we propose SemaMIL, which replaces random instance reorderings with a semantically driven rearrangement that brings similar patches closer in the input sequence, and augments state-space sequential modeling with a cross-sequence querying mechanism to reinforce global context. As illustrated in Fig. 2 , semantic reordering strengthens interactions among related patches, while the augmented framework captures both local and global dependencies under linear computational complexity.

Specifically, given a WSI, we partition its tissue regions into a sequence of $L$ patches $\{{p_{1}},{p_{2}},\cdots,{p_{L}}\}$. A pretrained feature extractor maps these patches to embeddings $X\in\mathbb{R}^{L\times D}$, which are then projected to a lower dimension $d$ by a linear layer. The resulting features undergo the Semantic Reordering module, which computes pairwise similarities and permutes $X$ into a new sequence $X^{\prime}$ such that semantically similar patches are adjacent. The reordered embeddings $X^{\prime}$ are passed through a stack of state-space sequential modules enhanced with a global querying operation: each module alternates between scanning the reordered sequence to model instance interactions and querying across all positions to integrate distant contextual information. Finally, an aggregation head pools the refined sequence into a fixed-length bag representation for downstream subtype classification.

2.1 Semantic Reordering Module

In high‐resolution WSI patch sequences, semantically similar but spatially distant patches cannot directly interact within a single causal pass of the state‐space model. To overcome this limitation, we propose a semantic reordering mechanism that adaptively groups related patches in the processing sequence.

As shown in Fig. 2, given a sequence of patch embeddings $\{{x_{i}}\}_{i=1}^{L}$ from a WSI, the goal is to place semantically similar patches next to each other so that a single linear state-space pass can more effectively propagate discriminative context. a lightweight router which consists of two linear layers with GELU and produces a semantic score vector for each patch:

[TABLE]

We apply a softmax to obtain $p_{i}$ with $z_{i}$ and assign a hard semantic label $c_{i}$ with $argmax(p_{i})$, and then form a permutation:

[TABLE]

to reorder the sequence ${x_{\pi(i)}}$. After processing the reordered sequence with the causal state-space model to obtain outputs ${y_{i}}^{\prime}$, we invert the permutation to restore the original spatial order:

[TABLE]

This module shortens the effective interaction distance between histologically related regions such as dispersed tumor morphology while avoiding quadratic attention cost. When using soft assignments including Gumbel variants the process is fully differentiable and reversible and provides a semantically compressed input for subsequent retrieval enhanced state space modeling.

2.2 Semantic-guided Retrieval State Space Module

In the aligned and semantically reordered patch sequence $\{{x_{i}}\}_{i=1}^{N}$, we assign each patch an importance score via a lightweight linear projection and select the top $K$ scoring patches to form the query set $Q=\{{x_{{i_{j}}}}\}_{j=1}^{K}$, while the remaining $N-K$ patches constitute the context sequence $C=\{{c_{k}}\}_{k=1}^{N-K}$. This Patch Selector effectively suppresses background noise and directs model capacity toward tumor-relevant morphological features. Building on $Q$, we propose SRSM to jointly capture long-range dependencies and global semantic correlations within a unified state-space framework.

We begin with the continuous-time linear time-invariant system:

[TABLE]

and discretize it via zero-order hold into:

[TABLE]

To endow the system with dynamic adaptability, we derive its discrete parameters from $Q$ by:

[TABLE]

and then obtain:

[TABLE]

where $A_{0}$ is a fixed base transition matrix. The context sequence $\{c_{k}\}$ is then processed in one causal pass:

[TABLE]

which $A_{0}(\Delta)$ captures local tissue continuity, $B_{d}(\Delta,B^{\prime})$ gates global semantic querying and noise suppression under the guidance of $Q$. To fully exploit the two-dimensional structure of histopathological slides, we execute this SRSM in parallel along four scan directions, producing outputs $\{y_{k}^{(d)}\}_{d=1}^{4}$. Finally, global pooling of $y_{k}^{fused}$ yields a fixed‐dimensional representation, which combines fine‐grained morphological details with holistic semantic context and drives the downstream subtype classification head.

3 EXPERIMENTS

3.1 Datasets and Evaluation Metrics

To demonstrate the effectiveness of our proposed SemaMIL, we conduct extensive experiments on a single downstream task, subtype classification of pathology slides using features extracted by the TITAN model [5]. Comparative evaluations are performed on four challenging public datasets: EBRAINS [13], BRACS [14], IPD-Brain [15] and TCGA.

To ensure a robust assessment, we adopt ten-fold Monte Carlo cross-validation and partition each dataset into training, validation and test sets in the ratio of 76.5 percent to 13.5 percent to 10 percent. For fair comparison with prior work, we also evaluate the official BRACS split indicated by a star in Table 1. Following standard practice, we report the mean and standard deviation (std) of the area under the ROC curve (AUC) and accuracy (ACC), which together provide a reliable evaluation that mitigates the effects of class imbalance.

3.2 Implementation Details

We present the experimental results of our SemaMIL on four datasets compared with the following methods. These include traditional feature aggregation methods such as Mean Pooling and Max Pooling, ABMIL [7] and its three variants CLAM-MB [17], DSMIL [8], and DTFDMIL [18], the Transformer-based TransMIL [19], the state-space model-based S4MIL [16], and the recently proposed MambaMIL [11]. In accordance with standard experimental settings, we use the same data pre-processing pipeline as CLAM and apply a fixed learning rate of $5\times{10^{-5}}$ across all methods to ensure optimal performance and enable fair comparisons.

3.3 Comparison Results

Table 1 presents the experimental results across four datasets, covering both binary and multi-class classification tasks. Compared to state-of-the-art methods, our proposed SemaMIL consistently surpasses existing approaches in terms of accuracy, achieving an ACC of 75.1% on EBRAINS and 60.3% on BRACS, which represent notable improvements over ABMIL and its variants. As illustrated in Fig. 1(a), Mamba-based modeling paradigms adopt distinct patch sequence ordering strategies, and our semantic reordering yields a more coherent arrangement facilitating long-range contextual propagation. Fig. 1(b) shows that our method provides a strong solution to the WSI subtype classification task.

Furthermore, as shown in Table 3, SemaMIL demonstrates superior computational efficiency, with FLOPs of 0.2484G and only 0.4641M parameters, both lower than other competing methods. Despite this efficiency, SemaMIL still outperforms other methods in accuracy. Fig. 3 presents representative challenging slides misclassified by competing methods but correctly predicted by SemaMIL, serving as illustrative correct cases on difficult subtype instances.

3.4 Ablation Study

To assess the effectiveness of SR-Mamba, we conduct extensive experiments to evaluate the contributions of its constituent modules, as shown in Table 2. The baseline performs a single causal scan in the original spatial patch order. SR reorders patches according to learned semantic similarity and then applies the causal scan to the reordered sequence. SRSM preserves the original spatial order while selecting a salient subset of patches as queries to drive state space parameter updates during the full sequence scan. SR+SRSM applies semantic reordering followed by retrieval driven state space modeling. As shown in Table 2, these results indicate that both SR and SRSM are necessary and that their combination is the preferred configuration.

4 Conclusion

In this work, we introduce SemaMIL, a semantic‑guided multiple instance learning framework for the task of gigapixel whole slide image subtype classification. Its Semantic Reordering module places histologically related yet spatially distant patches contiguously within a reversible sequence, thereby enhancing information flow in a single causal scan. The Semantic‑guided Retrieval State Space Module (SRSM) further selects a salient query subset to dynamically modulate state space parameters, reinforcing long‑range dependency modeling, suppressing redundancy, and implicitly enriching contextual interactions with linear time complexity. Experiments on four challenging datasets demonstrate that SemaMIL consistently benefits from its semantic mechanisms and achieves state‑of‑the‑art performance across all evaluated metrics. We will further apply our method to a broader range of pathology tasks.