# PSA-MIL: A Probabilistic Spatial Attention-Based Multiple Instance Learning for Whole Slide Image Classification

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 59 references that are in the map.
- First posted on arXiv 20 Mar 2025.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Data Efficient and Weakly Supervised Computational Pathology on Whole Slide Images](https://mlsanity.com/paper/2004.09666/) (2020)
- [TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification](https://mlsanity.com/paper/2106.00908/) (2021)
- [Measuring and Relieving the Over-smoothing Problem for Graph Neural Networks from the Topological View](https://mlsanity.com/paper/1909.03211/) (2019)
- [Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning](https://mlsanity.com/paper/2011.08939/) (2020)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Scaling Vision Transformers to Gigapixel Images via Hierarchical Self-Supervised Learning](https://mlsanity.com/paper/2206.02647/) (2022)
- [DTFD-MIL: Double-Tier Feature Distillation Multiple Instance Learning for Histopathology Whole Slide Image Classification](https://mlsanity.com/paper/2203.12081/) (2022)
- [LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models](https://mlsanity.com/paper/2309.12307/) (2023)
- [FLatten Transformer: Vision Transformer using Focused Linear Attention](https://mlsanity.com/paper/2308.00442/) (2023)
- [Benchmarking Self-Supervised Learning on Diverse Pathology Datasets](https://mlsanity.com/paper/2212.04690/) (2022)
- [Scaling Laws of RoPE-based Extrapolation](https://mlsanity.com/paper/2310.05209/) (2023)
- [A Tutorial on Kernel Density Estimation and Recent Advances](https://mlsanity.com/paper/1704.03924/) (2017)
- [Interventional Bag Multi-Instance Learning On Whole-Slide Pathological Images](https://mlsanity.com/paper/2303.06873/) (2023)
- [A graph-transformer for whole slide image classification](https://mlsanity.com/paper/2205.09671/) (2022)
- [Multiple Instance Learning for Digital Pathology: A Review on the State-of-the-Art, Limitations & Future Potential](https://mlsanity.com/paper/2206.04425/) (2022)
- [Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers](https://mlsanity.com/paper/2006.03555/) (2020)
- [Improving Transformers with Probabilistic Attention Keys](https://mlsanity.com/paper/2110.08678/) (2021)
- [CAMIL: Context-Aware Multiple Instance Learning for Cancer Detection and Subtyping in Whole Slide Images](https://mlsanity.com/paper/2305.05314/) (2023)
- [Rethinking Transformer for Long Contextual Histopathology Whole Slide Image Analysis](https://mlsanity.com/paper/2410.14195/) (2024)
- [Sm: enhanced localization in Multiple Instance Learning for medical imaging classification](https://mlsanity.com/paper/2410.03276/) (2024)
- [RoFormer for Position Aware Multiple Instance Learning in Whole Slide Image Classification](https://mlsanity.com/paper/2310.01924/) (2023)
- [QuST: QuPath Extension for Integrative Whole Slide Image and Spatial Transcriptomics Analysis](https://mlsanity.com/paper/2406.01613/) (2024)

## Built on by (most cited first, 2 of 2)

- [Normal Guidance is what Attention Needs](https://mlsanity.com/paper/2605.27306/) (2026)
- [Predicting Metastatic Risk from Primary Cancer Tissue Architecture via Distance-Aware Spatial Modeling](https://mlsanity.com/paper/2606.28676/) (2026)

## Abstract (arXiv)

Whole Slide Images (WSIs) are high-resolution digital scans widely used in medical diagnostics. WSI classification is typically approached using Multiple Instance Learning (MIL), where the slide is partitioned into tiles treated as interconnected instances. While attention-based MIL methods aim to identify the most informative tiles, they often fail to fully exploit the spatial relationships among them, potentially overlooking intricate tissue structures crucial for accurate diagnosis. To address this limitation, we propose Probabilistic Spatial Attention MIL (PSA-MIL), a novel attention-based MIL framework that integrates spatial context into the attention mechanism through learnable distance-decayed priors, formulated within a probabilistic interpretation of self-attention as a posterior distribution. This formulation enables a dynamic inference of spatial relationships during training, eliminating the need for predefined assumptions often imposed by previous approaches. Additionally, we suggest a spatial pruning strategy for the posterior, effectively reducing self-attention's quadratic complexity. To further enhance spatial modeling, we introduce a diversity loss that encourages variation among attention heads, ensuring each captures distinct spatial representations. Together, PSA-MIL enables a more data-driven and adaptive integration of spatial context, moving beyond predefined constraints. We achieve state-of-the-art performance across both contextual and non-contextual baselines, while significantly reducing computational costs.