# Region-Graph Optimal Transport Routing for Mixture-of-Experts Whole-Slide Image Classification

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical and biomedical MoE, one of 12 directions in Mixture-of-experts.
- Builds on 12 of its 23 references that are in the map.
- First posted on arXiv 8 Apr 2026.

## Placed in

- Mixture-of-experts / [Medical and biomedical MoE](https://mlsanity.com/moe/medical-and-biomedical-moe/)

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

- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Data Efficient and Weakly Supervised Computational Pathology on Whole Slide Images](https://mlsanity.com/paper/2004.09666/) (2020)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning](https://mlsanity.com/paper/2011.08939/) (2020)
- [Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts](https://mlsanity.com/paper/2408.15664/) (2024)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [OTSurv: A Novel Multiple Instance Learning Framework for Survival Prediction with Heterogeneity-aware Optimal Transport](https://mlsanity.com/paper/2506.20741/) (2025)
- [SAM-MIL: A Spatial Contextual Aware Multiple Instance Learning Approach for Whole Slide Image Classification](https://mlsanity.com/paper/2407.17689/) (2024)
- [Improving Routing in Sparse Mixture of Experts with Graph of Tokens](https://mlsanity.com/paper/2505.00792/) (2025)
- [Mixture of Mini Experts: Overcoming the Linear Layer Bottleneck in Multiple Instance Learning](https://mlsanity.com/paper/2603.22198/) (2026)

## Abstract (arXiv)

Multiple Instance Learning (MIL) is the dominant framework for gigapixel whole-slide image (WSI) classification in computational pathology. However, current MIL aggregators route all instances through a shared pathway, constraining their capacity to specialise across the pathological heterogeneity inherent in each slide. Mixture-of-Experts (MoE) methods offer a natural remedy by partitioning instances across specialised expert subnetworks; yet unconstrained softmax routing may yield highly imbalanced utilisation, where one or a few experts absorb most routing mass, collapsing the mixture back to a near-single-pathway solution. To address these limitations, we propose ROAM (Region-graph OptimAl-transport Mixture-of-experts), a spatially aware MoE-MIL aggregator that routes region tokens to expert poolers via capacity-constrained entropic optimal transport, promoting balanced expert utilisation by construction. ROAM operates on spatial region tokens, obtained by compressing dense patch bags into spatially binned units that align routing with local tissue neighbourhoods and introduces two key mechanisms: (i) region-to-expert assignment formulated as entropic optimal transport (Sinkhorn) with explicit per slide capacity marginals, enforcing balanced expert utilisation without auxiliary load-balancing losses; and (ii) graph-regularised Sinkhorn iterations that diffuse routing assignments over the spatial region graph, encouraging neighbouring regions to coherently route to the same experts. Evaluated on four WSI benchmarks with frozen foundation-model patch embeddings, ROAM achieves performance competitive against strong MIL and MoE baselines, and on NSCLC generalisation (TCGA-CPTAC) reaches external AUC 0.845 +- 0.019.
