PseMix: Pseudo-Bag Mixup Augmentation for MIL-Based Whole Slide Image Classification

IEEE Transaction on Medical Imaging, 2024

Abstract: Multiple instance learning (MIL) has become one of the most important frameworks for gigapixel Whole Slide Images (WSIs). 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. To address these problems, this paper proposes a new Pseudo-bag Mixup (PseMix) data augmentation scheme, inspired by the basic idea of Mixup. Cooperated by pseudo-bags, this scheme fulfills the critical size alignment and semantic alignment in Mixup. Moreover, it is efficient and plugin-and-play, neither involving time-consuming operations nor relying on model predictions. Experimental results show that PseMix could often improve the performance of state-of-the-art MIL networks. Most importantly, it could also boost the generalization performance of MIL models in special test scenarios, and promote their robustness to patch occlusion and label noise.

<div align="center"> <a href="https://"><img width="100%" height="auto" src="./docs/procedure-psemix.png"></a> </div>

📚 Recent updates:

25/02/07: upload the patch features used in the experiment (36.42G in files); you can access them from here.
24/05/27: 🔥 new SOTA performance with PseMix using CONCH features
24/05/19: add data_split and Notebook - Generate Data Split
24/05/09: add two PseMix Walkthrough notebooks: Notebook - Pseudo-bag Generation and Notebook - Pseudo-bag Mixup
24/05/08: add link to a detailed tutorial of WSI preprocessing
24/02/27: add missing codes regrading the module of optim
24/02/22: add useful research papers involving pseudo-bags

Why PseMix?

Adopting PseMix in training MIL networks could

(1) improve performance with minimal costs:

| Network | BRCA | NSCLC | RCC | Average AUC | |---------|------|-------|-----| :--------------------------------------: | | ABMIL | 87.05 | 92.23 | 97.36 | 92.21 | | ABMIL w/ PseMix | 89.49 | 93.01 | 98.02 | 93.51 | | DSMIL | 87.73 | 92.99 | 97.65 | 92.79 | | DSMIL w/ PseMix | 89.65 | 93.92 | 97.89 | 93.82 | | TransMIL | 88.83 | 92.14 | 97.88 | 92.95 | | TransMIL w/ PseMix | 90.40 | 93.47 | 97.76 | 93.88 |

(2) obtain better generalization and robustness:

Training curves, showing the AUC performance on training and test data (exported from wandb), are given below.

(3) often obtain improvements when using the SOTA feature extractor CONCH:

🔥 New SOTA results on TCGA-BRCA, tested with the patch features extracted by CONCH:

| Network | Prototype | AUC | ACC | F1-Score | |---------|-----------|-----|-----|----------| | ABMIL | - | 94.48 | 92.34 | 93.44 | | ABMIL w/ PseMix | PANTHER | 94.77 | 92.33 | 94.07 | | ABMIL w/ PseMix | ProDiv | 94.63 | 92.55 | 94.28 | | DSMIL | - | 92.16 | 89.51 | 90.91 | | DSMIL w/ PseMix | PANTHER | 93.46 | 91.60 | 93.46 | | DSMIL w/ PseMix | ProDiv | 93.90 | 91.49 | 93.63 | | TransMIL | - | 93.67 | 91.92 | 94.40 | | TransMIL w/ PseMix | PANTHER | 94.46 | 91.81 | 94.79 | | TransMIL w/ PseMix | ProDiv | 94.64 | 90.65 | 93.80 |

There are some fun observations from the table above:

Strong feature extractor could boost the performance by a large margin
ABMIL is a very strong baseline, outperforming DSMIL and TransMIL when using CONCH features
Training with PseMix often leads to SOTA performance

PseMix Walkthrough

PseMix contains two key steps: i) pseudo-bag generation and ii) pseudo-bag mixup. Here are two alternative means that could help you quickly understand the two key steps.

Option 1: Notebooks

Get started with PseMix:

Option 2: Overview and pseudo-codes

An overall description of the two key steps is as follows:

Step 1. Pseudo-bag Generation

Pseudo-bag generation contains two sub-steps:

Phenotype clustering by
- bag-prototype-based instance clustering
- cluster fine-tuning
Phenotype-stratified sampling

Its implementation details can be found via the function generate_pseudo_bags. This predefined function will be directly used and called in the next step, pseudo-bag mixup.

Step 2. Pseudo-bag Mixup

Below is the pseudo-code of pseudo-bag mixup:

# generate_pseudo_bags: function for dividing WSI bags into pseudo-bags
# ALPHA: the hyper-parameter of Beta distribution
# N: the number of pseudo-bags in each WSI bag
# PROB_MIXUP: random mixing parameter for determining the proportion of mixed bags. 
for (X, y) in loader: # load a minibatch 
    n_batch = X.shape[0] # with `n_batch` WSI bags (samples)

    # 1. dividing each bag into `N` pseudo-bags
    X = generate_pseudo_bags(X)

    new_idxs = torch.randperm(n_batch)
    # draw a mixing scale from Beta distribution
    lam = numpy.random.beta(ALPHA, ALPHA) 
    lam = min(lam, 1.0 - 1e-5) # avoid numerical overflow when transforming it into discrete ones
    lam_discrete = int(lam * (N + 1)) # transform into discrete values

    # 2. pseudo-bag-level Mixup generates samples (new_X, new_y)
    new_X, new_y = [], []
    for i in range(n_batch):
    	# randomly select pseudo-bags according to `lam_discrete`
        masked_bag_A = select_pseudo_bags(X[i], lam_discrete) # select `lam_discrete` pseudo-bags
        masked_bag_B = select_pseudo_bags(X[new_idxs[i]], N - lam_discrete) # select `n-lam_discrete` pseudo-bags

        # random-mixing mechanism for two purposes: more data diversity and efficient learning on mixed samples.
        if np.random.rand() <= PROB_MIXUP:
            mixed_bag = torch.cat([masked_bag_A, masked_bag_B], dim=0) # instance-axis concat
            new_X.append(mixed_bag)
            mix_ratio = lam_discrete / N
        else:
            masked_bag = masked_bag_A 
            new_X.append(masked_bag)
            mix_ratio = 1.0

        # target-level mixing
        new_y.append(mix_ratio * y[i] + (1 - mix_ratio) * y[new_idxs[i]]) 

    # 3. minibatch training
    minibatch_training(new_X, new_y)

More details can be found at

pseudo-bag-level Mixup.
training with mixed labels.
weighted loss for mixed samples, following the implementation of Mixup.

WSI Preprocessing

The procedure of WSI preprocessing is elaborated in Pipeline-Processing-TCGA-Slides-for-MIL. You can

move to that repo for a detailed tutorial and try it on your own dataset;
or download the patc

PseMix

Install / Use

README