CAPRMIL: Context-Aware Patch Representations for Multiple Instance Learning

Andreas Lolos, Theofilos Christodoulou, Aris L. Moustakas, Stergios Christodoulidis, Maria Vakalopoulou

Keywords: Digital Pathology, Multiple Instance Learning, Context Aware Representations

TL;DR: We present an efficient MIL framework that performs on par with leading methods while using substantially fewer parameters and significantly lower computational cost.

Abstract:

In computational pathology, weak supervision has become the standard for deep learning due to the gigapixel scale of WSIs and the scarcity of pixel-level annotations, with Multiple Instance Learning (MIL) established as the principal framework for slide-level model training.
In this paper, we introduce a novel setting for MIL methods, inspired by proceedings in Neural Partial Differential Equation (PDE) Solvers. Instead of relying on complex attention-based aggregation, we propose an efficient, aggregator-agnostic framework that removes the complexity of correlation learning from the MIL aggregator. CAPRMIL produces rich context-aware patch embeddings that promote effective correlation learning on downstream tasks. By projecting patch features -extracted using a frozen patch encoder- into a small set of global context/morphology-aware tokens and utilizing multi-head self-attention, CAPRMIL injects global context with linear computational complexity with respect to the bag size. Paired with a simple Mean MIL aggregator, CAPRMIL matches state-of-the-art slide-level performance across multiple public pathology benchmarks, while reducing the total number of trainable parameters by $48$-$92.8$% versus SOTA MILs, lowering FLOPs during inference by $52$-$99$%, and ranking among the best models on GPU memory efficiency and training time. Our results indicate that learning rich, context-aware instance representations before aggregation is an effective and scalable alternative to complex pooling for whole-slide analysis. Our code is available at: https://github.com/mandlos/CAPRMIL

Primary Subject Area: Interpretability and Explainable AI

Secondary Subject Area: Detection and Diagnosis

Registration Requirement: Yes

Reproducibility: https://github.com/mandlos/CAPRMIL

Read CFP & Author Instructions: Yes

Originality Policy: Yes

Single-blind & Not Under Review Elsewhere: Yes

LLM Policy: Yes

Submission Number: 44