Advancing Multiple Instance Learning with Continual Learning for Whole Slide Imaging

TL;DR

This work targets continual learning for whole-slide image MIL, uncovering that forgetting in attention-based MIL is concentrated in the attention module rather than the classifier. It introduces Attention Knowledge Distillation (AKD) to preserve attention relationships across tasks and Pseudo-Bag Memory Pool (PMP) to enable memory-efficient replay by storing distilled patch subsets. The proposed framework yields significant gains in accuracy and memory efficiency across diverse WSI datasets, outperforming state-of-the-art CL methods and showing robust performance across MIL architectures. Together, AKD and PMP provide a practical, scalable path to adaptive, weakly supervised diagnostic models for gigapixel pathology images.

Abstract

Advances in medical imaging and deep learning have propelled progress in whole slide image (WSI) analysis, with multiple instance learning (MIL) showing promise for efficient and accurate diagnostics. However, conventional MIL models often lack adaptability to evolving datasets, as they rely on static training that cannot incorporate new information without extensive retraining. Applying continual learning (CL) to MIL models is a possible solution, but often sees limited improvements. In this paper, we analyze CL in the context of attention MIL models and find that the model forgetting is mainly concentrated in the attention layers of the MIL model. Using the results of this analysis we propose two components for improving CL on MIL: Attention Knowledge Distillation (AKD) and the Pseudo-Bag Memory Pool (PMP). AKD mitigates catastrophic forgetting by focusing on retaining attention layer knowledge between learning sessions, while PMP reduces the memory footprint by selectively storing only the most informative patches, or ``pseudo-bags'' from WSIs. Experimental evaluations demonstrate that our method significantly improves both accuracy and memory efficiency on diverse WSI datasets, outperforming current state-of-the-art CL methods. This work provides a foundation for CL in large-scale, weakly annotated clinical datasets, paving the way for more adaptable and resilient diagnostic models.

Figures (6)

• Figure 1: Example of the attention shift during CL for a Task 1 test image. (1) Input image with tumor area outlined in green. (2) The attention heat map for $t$=1 session. (3) and (4): the attention drift, relative to (2), after CL sessions $t \in \{2, 3\}$ for fine-tuning CL. The tumor areas are outlined in black. (5) and (6): the attention drift for our proposed method. Our method better preserves the attention much better than fine-tuning.
• Figure 2: Evolution of the distribution of gradient values of CLAM during CL for: (top) the attention network, (bottom) the classifier network. The data points are aggregated where each dot represents the mean of 50 consecutive training steps. The five colored dots at each aggregated step represent the minimum, maximum, and three quartiles (first, second, and third) of the gradient values. The vertical black dotted lines mark the transitions between tasks.
• Figure 3: Overview of our framework architecture. The black arrows denote the forward pass for current data, including patch-wise slide processing, feature encoding, and memory pool updates. Orange rectangles represent patch-level features, while deep orange rectangles indicate bag-level feature representations. The green paths indicate the knowledge distillation process between teacher and student networks, incorporating both attention-based and logit-based distillation losses alongside cross-entropy supervision.
• Figure 4: AACC performance of (left) CLAM and (right) TransMIL on Camelyon-TCGA as the number of tasks increases, i.e., the CL session $t$ increases. The gray dotted line indicates the AACC performance of joint training.
• Figure 5: The tradeoff between BWT and IM on Camelyon-TCGA for (left) CLAM and (right) TransMIL. BWT measures the amount of forgetting of previous tasks, with higher BWT indicating less forgetting. IM measures the ability to learn new tasks, with lower IM indicating better ability. Thus, better methods are closer to the upper-left corner. The points represent different methods, with our method achieving the best trade-off with both models.