Geometry-aware Multiple Instance Learning via a Shared Riemannian Metric for Few-Shot Whole-Slide Image Classification

The diagnostic paradigm of whole slide images (WSIs) involves a frozen pathomorphological model, providing patch features, while multiple instance learning (MIL) accomplishes aggregation under weak slide-level labels. Nevertheless, this paradigm has two issues in practice: feature manifold collapse and unstable attention. By analyzing them, we realize that both belong to the two ends of the same causal chain, where the disruption of the feature manifold causes oscillatory changes in the attention distribution that depends on the tangent-space component. Based on this finding, we propose geometry-aware multiple instance learning (GeoMIL), which simultaneously drives feature transformation and attention generation with a single learnable Riemannian metric. We design a tangent space low-rank projection (TSLP) module to perform task adaptation while maintaining feature geometry, and use geometric query attention (GQA) to anchor the attention distribution to the tangent space stabilized by TSLP. Through extensive experiments, we demonstrate that our method maintains superior few-shot WSI classification performance while achieving a more stable attention distribution.