From Correlation to Causation: Max-Pooling-Based Multi-Instance Learning Leads to More Robust Whole Slide Image Classification

TL;DR

This paper introduces FocusMIL, a novel MIL method that combines max-pooling with a variational information bottleneck to focus on causal features, improving robustness and localization in whole slide image classification.

Contribution

It provides a causal perspective on max-pooling in MIL, proposes FocusMIL with VIB and multi-bag training, and demonstrates improved out-of-distribution performance.

Findings

FocusMIL outperforms existing methods in OOD scenarios.

It enhances tumor region localization accuracy.

Theoretical analysis links max-pooling to causal feature focus.

Abstract

In whole slide images (WSIs) analysis, attention-based multi-instance learning (MIL) models are susceptible to spurious correlations and degrade under domain shift. These methods may assign high attention weights to non-tumor regions, such as staining biases or artifacts, leading to unreliable tumor region localization. In this paper, we revisit max-pooling-based MIL methods from a causal perspective. Under mild assumptions, our theoretical results demonstrate that max-pooling encourages the model to focus on causal factors while ignoring bias-related factors. Furthermore, we discover that existing max-pooling-based methods may overfit the training set through rote memorization of instance features and fail to learn meaningful patterns. To address these issues, we propose FocusMIL, which couples max-pooling with an instance-level variational information bottleneck (VIB) to learn compact, predictive latent representations, and employs a multi-bag mini-batch scheme to stabilize optimization. We conduct comprehensive experiments on three real-world datasets and one semi-synthetic dataset. The results show that, by capturing causal factors, FocusMIL exhibits significant advantages in out-of-distribution scenarios and instance-level tumor region localization tasks.