Pseudo-Bag Mixup Augmentation for Multiple Instance Learning-Based Whole Slide Image Classification

TL;DR

This paper tackles the data scarcity and memorization challenges in MIL-based Whole Slide Image (WSI) classification by introducing Pseudo-bag Mixup (PseMix), a data augmentation scheme that extends Mixup to WSIs through pseudo-bags to ensure size and semantic alignment. PseMix consists of phenotype-guided pseudo-bag division, bag-level masking and random mixing (r-mix), and target-level label mixing, designed to be plug-and-play and decoupled from MIL training. Empirical results across three TCGA datasets show consistent improvements over vanilla MIL and other mixing baselines, along with enhanced generalization, robustness to patch occlusion, and resilience to label noise. The method demonstrates strong versatility across multiple MIL architectures and feature extractors, suggesting practical impact for improving WSI classification in computational pathology.

Abstract

Given the special situation of modeling gigapixel images, multiple instance learning (MIL) has become one of the most important frameworks for Whole Slide Image (WSI) classification. In current practice, most MIL networks often face two unavoidable problems in training: i) insufficient WSI data and ii) the sample memorization inclination inherent in neural networks. These problems may hinder MIL models from adequate and efficient training, suppressing the continuous performance promotion of classification models on WSIs. Inspired by the basic idea of Mixup, this paper proposes a new Pseudo-bag Mixup (PseMix) data augmentation scheme to improve the training of MIL models. This scheme generalizes the Mixup strategy for general images to special WSIs via pseudo-bags so as to be applied in MIL-based WSI classification. Cooperated by pseudo-bags, our PseMix fulfills the critical size alignment and semantic alignment in Mixup strategy. Moreover, it is designed as an efficient and decoupled method, neither involving time-consuming operations nor relying on MIL model predictions. Comparative experiments and ablation studies are specially designed to evaluate the effectiveness and advantages of our PseMix. Experimental results show that PseMix could often assist state-of-the-art MIL networks to refresh their classification performance on WSIs. Besides, it could also boost the generalization performance of MIL models in special test scenarios, and promote their robustness to patch occlusion and label noise. Our source code is available at https://github.com/liupei101/PseMix.

Figures (6)

• Figure 1: Generalization gap of MIL models in training. It is measured by the model performance gap between training and test sets, to assess the generalization ability of models, following jiang2018predictingzhang2021how. Three state-of-the-art MIL models are trained on TCGA-BRCA WSIs. Vanilla models often show rapidly-growing gaps. The models trained with our PseMix could ease these without introducing extra complicated techniques.
• Figure 2: A conceptual framework of PseMix for MIL-based WSI classification. Size alignment means that two input representations (e.g., feature vectors and image matrices) are aligned in size at every dimension involving mixing. $X_{A}^{\tau}$ and $X_{B}^{\tau}$ denote pseudo-bags. $\mathbf{M}_{\lambda}$ is a binary mask. R-mix means random mixing. PseMix generalizes Mixup and fulfills its critical size alignment and semantic alignment via pseudo-bags.
• Figure 3: PseMix data augmentation for MIL-based WSI classification. (a) Classical MIL paradigm for weakly-supervised WSI analysis. (b) Illustration of PseMix. Two WSI bags $A$ and $B$ are taken as examples. Solid rectangular boxes give the t-SNE visualization of instances, where scatter points represents instances and each color indicates a specific phenotype. $A^1,A^2,...,B^3$ are pseudo-bag notations. R-mix means random mixing.
• Figure 4: Generalization gap (AUC and cross-entropy loss) between training and test set. Three MIL networks are trained on TCGA-BRCA.
• Figure 5: Cross-entropy loss of in-between training data and their corresponding soft-labels.