Multiple Instance Learning
By: Lori Sheng
Tags: Deep Learning, Image Classification, Digital Pathology Slide, Computer Vision, Convolutional Neural Net

Multiple Instance Learning

with MNIST dataset using Pytorch

When it comes to applying computer vision in the medical field, most tasks involve either 1) image classification for diagnosis or 2) segmentation to identify and separate lesions. However, in pathology cancer detection, this is not always possible. Obtaining labels is time-consuming and labor-intensive. Furthermore, pathology slides can be up to 200k x 100k pixels resolution, and they will not fit in memory for classification since for example, the ImageNet only uses 224 x 224 pixels for training. Downsampling normally is not an option because we are trying to detect a tiny area, such as a cancerous area varying from 300 x 300 pixels area (a few dots in Figure 1).

Figure 1: a slide from patient_004_node_004 (Camelyon17)

In this case, we can use Multiple Instance Learning, a weakly supervised learning method that takes a set of labeled bags containing many instances instead of receiving a set of labeled instances.

Suppose we have pathology slides and the label for each slide. Since we cannot train the classifier on the whole slide, we divide each slide into tiles and only process a few tiles at a time on GPU. However, we do not know labels for each tile, so we need Multiple Instance Learning. In the MIL framework, slides are the “bag” and tiles are “instances”. By using it, we are able to save the labeling effort and leverage weakly labeled data.

When we have pathology slides from patients and we want to predict if a large slide contains cancer cells or, zooming out if the patient has malignant cells, Multiple Instance Learning is a good option because doctors don’t need to segment individual cells or label each tile. Only the whole slide needs a label.

In general, Multiple Instance Learning can deal with classification problems, regression problems, ranking problems, and clustering problems, but we will mainly focus on classification problems here.

In this blog, I am going to explain how MIL works with a simple example based on the MNIST dataset. If you are not familiar with the MNIST dataset, here is a link to a Kaggle competition on the MNIST dataset and you can take a look.

In the next blog, I will talk in detail about MIL on the breast cancer challenge Camelyon17.

Brief introduction of MNIST dataset:

The MNIST dataset is a large database of handwritten digits and each image has one label from 0 to 9. It has a training set of 60,000 images and a test set of 10,000 images. The size of each one is 28 x 28 and greyscale.

Figure 2: MNIST sample

Brief problem statement of Multiple Instance Learning:

Each instance xi in one bag has a label yi. We define the label of the bag as:

Y = 1, if there exists a yi such that yi ==1

Y = 0, if for every yi such that yi == 0

The pipeline of applying MIL on MNIST dataset:

As an overview, we will first pre-train the ResNet model on the instance-labeled dataset and then feed the bag-labeled dataset in the model and extract features. Finally, we apply MIL on them.

Step 1: Split the original MNIST dataset to the bag-labeled set for proper MIL training and the instance-labeled set.

Creating MIL-MNIST toy dataset:

The MNIST dataset is a standard benchmark for classification tasks. To make it a MIL problem, we need to first construct the MIL-MNIST dataset by grouping several digits (instances) into a bag.

In our MNIST example, we will assign the bag label as “1” if one instance has a label “1”; if all the instance labels are “0–9” other than “1”, then we assign the bag label as “0”.

As the graph shown below, the bag filled with red has a bag label “1” and the bag filled with blue has a bag label “0”.

Figure 3: bags and instance labels

We randomly put each image into a bag, where each bag contains 3 to 7 instances. In order to save memory, we use the index to represent the image (as shown below).

Figure 4: Code to generate MIL-MNIST data

To generate the bag label:

Figure 5: Code to generate bag label

Step 2: Pre-train on 2 parts of MNIST dataset

1. Construct a 2D convolutional neural networks with kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)

2. Train for 5 epoch with batch size 256

3. Save model

Step 3: Load pre-trained model and extract features from the last layer

Split the rest of the data to train, validation, and test set Get features for training, validation and test set Get bag_indices and bag_labels Map bag_indices with features based on indices and create bag_features

In order to get rid of the last layer:

Extract features:

The following code shows how we get the bag indices and bag features from data generation function:

Figure 6: Code to generate bag features

The bag indices, bag labels, and bag features are as shown below:

Figure 7: Bag index with image indices

Figure 8: Bag labels

Figure 9: Bag features

Step 4: Train the model for MIL on bag_features and bag_labels and evaluate on the test set

Since each bag has a different number of instances, we need to pad tensors to the same size before we throw them in the model.

Model for Multiple Instance Learning:

The algorithm performs three steps. Either of them can be fixed functions or optimizable functions (neural networks):

Transforms instances to an embedding of low-dimensionality. (fixed) Passes the embedding through a permutation-invariant aggregation function. (optimizable) Transforms into the bag probability. (optimizable)

Figure 9: MIL diagram on the MIL-MNIST toy dataset

In general, the workflow is like the following:

Figure 10: MIL diagram on pathology slides (see reference #5)

We did step 1 as fixed for simplicity. For step 2, although we can still use fixed functions such as max or mean, in order to enable parameter optimizations that can be learned end-to-end by backpropagation, we use Neural Networks as the aggregation function. For step 3, we also want to use backpropagation to be able to make parameters optimizable.

Linear layer and LeakyReLu

Note: we set n = 7*512, where 7 is the number of instances in one bag and 512 is the size of each feature.

2. Aggregation function: AttensionSoftmax

3. The neural network with LeakyReLu as the activation function in the middle, dropout, and sigmoid as the final activation function:

4. Optimizer: SGD

5. Loss function: BCELoss

6. Accuracy: ~0.99

Conclusion:

We got around 0.99 accuracy on the MNIST dataset using MIL, which is a satisfying result. If we want, we can use even more complex aggregation functions as our middle transformation and build more complicated NN models for the final transformation to the bag level. The result also shows that MIL is a good tool to save the labeling effort and leverage weakly labeled data.

What’s next:

Apply MIL on Camelyon17 dataset.

Jupyter notebook demo link:

https://github.com/lsheng23/Practicum/blob/master/MIL_MNIST/end_to_end_mnist_MIL.ipynb

Acknowledgments ：

Special thanks to Dmytro Lituiev, a Sr Machine Learning Scientist at UCSF, and Andrew Shaw, the WAMRI Staff Research Director at USF, for proofreading and guidance.

Reference:

Ilse, Maximilian, et al. “Attention-Based Deep Multiple Instance Learning.” ArXiv.org, 28 June 2018, https://arxiv.org/abs/1802.04712.

Carbonneau, Marc-André, et al. “Multiple Instance Learning: A Survey of Problem Characteristics and Applications.” ArXiv.org, 11 Dec. 2016, https://arxiv.org/abs/1612.03365.

LeCun, Yann, et al. “THE MNIST DATABASE.” MNIST Handwritten Digit Database, Yann LeCun, Corinna Cortes, and Chris Burges, http://yann.lecun.com/exdb/mnist/.

https://www.researchgate.net/figure/Example-images-from-the-MNIST-dataset_fig1_306056875

Dmytro S Lituiev, Sung Jik Cha, Andrew Bishara, Jae Ho Sohn, Eugenia Rutenberg, Dexter Hadley, Zoltan Laszik. “A Deep Learning Model for Identifying Rejection in Transplant Kidney Biopsies”. Poster presentation at USCAP 2020, Los Angeles