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

2026-04-08

AI Summary

This work addresses the limitations of existing multiple instance learning (MIL) approaches in whole-slide image classification, which struggle to model pathological heterogeneity due to shared aggregation pathways, and conventional mixture-of-experts (MoE) routing schemes that often suffer from imbalanced expert utilization. To overcome these issues, the authors propose the ROAM framework, which constructs a spatial region graph and employs an entropy-regularized optimal transport mechanism with capacity constraints to achieve balanced and locally consistent assignment of region tokens to expert poolers. By integrating graph-regularized Sinkhorn iterations, ROAM ensures both expert load balancing and neighborhood routing consistency without requiring additional balancing losses. The method demonstrates strong performance across four whole-slide benchmarks, achieving an external AUC of 0.845 ± 0.019 on the NSCLC generalization task (TCGA–CPTAC), significantly outperforming current MIL and MoE approaches.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Intelligent Robots: Learning & Optimization for ROB

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

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.

Problem

• Multiple Instance Learning
• Mixture-of-Experts
• Whole-Slide Image Classification
• Expert Imbalance
• Pathological Heterogeneity

Innovation

• Mixture-of-Experts
• Optimal Transport
• Spatial Graph Regularization
• Multiple Instance Learning
• Whole-Slide Image Classification