# Geometry-Aware State Space Model: A New Paradigm for Whole-Slide Image Representation

arXiv: https://arxiv.org/abs/2605.05164 · HTML: https://mlsanity.com/paper/2605.05164/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 33 references that are in the map.
- First posted on arXiv 6 May 2026.

## Builds on (most cited first, 19 of 19)

- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification](https://mlsanity.com/paper/2106.00908/) (2021)
- [Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning](https://mlsanity.com/paper/2011.08939/) (2020)
- [DTFD-MIL: Double-Tier Feature Distillation Multiple Instance Learning for Histopathology Whole Slide Image Classification](https://mlsanity.com/paper/2203.12081/) (2022)
- [Virchow2: Scaling Self-Supervised Mixed Magnification Models in Pathology](https://mlsanity.com/paper/2408.00738/) (2024)
- [Whole Slide Images are 2D Point Clouds: Context-Aware Survival Prediction using Patch-based Graph Convolutional Networks](https://mlsanity.com/paper/2107.13048/) (2021)
- [Hyperbolic Vision Transformers: Combining Improvements in Metric Learning](https://mlsanity.com/paper/2203.10833/) (2022)
- [MambaMIL: Enhancing Long Sequence Modeling with Sequence Reordering in Computational Pathology](https://mlsanity.com/paper/2403.06800/) (2024)
- [Structured State Space Models for Multiple Instance Learning in Digital Pathology](https://mlsanity.com/paper/2306.15789/) (2023)
- [HIGT: Hierarchical Interaction Graph-Transformer for Whole Slide Image Analysis](https://mlsanity.com/paper/2309.07400/) (2023)
- [MamMIL: Multiple Instance Learning for Whole Slide Images with State Space Models](https://mlsanity.com/paper/2403.05160/) (2024)
- [Improving Visual Recognition with Hyperbolical Visual Hierarchy Mapping](https://mlsanity.com/paper/2404.00974/) (2024)
- [UnPuzzle: A Unified Framework for Pathology Image Analysis](https://mlsanity.com/paper/2503.03152/) (2025)
- [HyperPath: Knowledge-Guided Hyperbolic Semantic Hierarchy Modeling for WSI Analysis](https://mlsanity.com/paper/2506.16398/) (2025)

## Abstract (arXiv)

Accurate analysis of histopathological images is critical for disease diagnosis and treatment planning. Whole-slide images (WSIs), which digitize tissue specimens at gigapixel resolution, are fundamental to this process but require aggregating thousands of patches for slide-level predictions. Multiple Instance Learning (MIL) tackles this challenge with a two-stage paradigm, decoupling tile-level embedding and slide-level prediction. However, most existing methods implicitly embed patch representations in homogeneous Euclidean spaces, overlooking the hierarchical organization and regional heterogeneity of pathological tissues. This limits current models' ability to capture global tissue architecture and fine-grained cellular morphology. To address this limitation, we introduce a hybrid hyperbolic-Euclidean representation that embeds WSI features in dual geometric spaces, enabling complementary modeling of hierarchical tissue structures and local morphological details. Building on this formulation, we develop BatMIL, a WSI classification framework that leverages both geometric spaces. To model long-range dependencies among thousands of patches, we employ a structured state space sequence model (S4) backbone that encodes patch sequences with linear computational complexity. Furthermore, to account for regional heterogeneity, we introduce a chunk-level mixture-of-experts (MoE) module that groups patches into regions and dynamically routes them to specialized subnetworks, improving representational capacity while reducing redundant computation. Extensive experiments on seven WSI datasets spanning six cancer types demonstrate that BatMIL consistently outperforms state-of-the-art MIL approaches in slide-level classification tasks. These results indicate that geometry-aware representation learning offers a promising direction for next-generation computational pathology.