Attention based multiple instance learning for weakly supervised slide classification, in the ABMIL and CLAM style. A whole slide image is far too large to label pixel by pixel, so the practical setup is weak supervision: you only know the label of the slide as a whole, not which patches inside it carry the disease. Multiple instance learning treats the slide as a bag of patch features and learns to predict the bag label while an attention mechanism discovers which patches mattered.
This repo implements that idea on synthetic data so the behaviour is easy to verify. A bag is positive when it contains at least one signal instance, and the model never sees which instance that is. If the attention pooling works, it should both classify bags correctly and point its attention at the planted signal.
src/model.py: theAttentionMILclassifier and aGatedAttentionpooling module. The model embeds each instance, computes an attention weight per instance, normalises the weights with a softmax over the bag, and feeds the attention weighted sum to a small classifier head. It accepts either a single bag of shape(N, in_dim)or a padded batch of bags of shape(B, N, in_dim)with an optional mask so padding gets zero attention.src/data.py: synthetic bag generation. Background instances come from a standard normal. A positive bag has exactly one signal instance whose mean is shifted, which makes it a single needle inside an otherwise random bag.src/train.py: a per bag training loop, a bag accuracy evaluator, and a helper that returns the attention weights for one bag.example.py: an end to end demo that trains a model and reports accuracy and attention localisation.tests/: pytest behaviour checks.
The gated variant follows Ilse, Tomczak and Welling (ICML 2018). Two parallel linear branches act on each instance embedding, one through a tanh and one through a sigmoid gate, and their elementwise product feeds the attention score. The gate lets the network suppress instances that the tanh branch alone would respond to. Scores are turned into weights with a softmax over the instances in the bag, so the weights for one bag are non negative and sum to one. That property is what makes attention readable as a soft selection over instances.
pip install -r requirements.txt
python example.py
On one run with the seeds in example.py the model reached a held out bag
accuracy of 1.000 on 60 test bags and placed its peak attention on the planted
signal instance in 26 of 26 positive bags. These numbers come from the bundled
synthetic data and will move with different seeds, signal strength, or bag
size. The point is the qualitative behaviour: the model both classifies the bag
and localises the instance that caused the label, having only ever been trained
on the bag label.
python -m pytest tests/ -q
The tests are property and behaviour checks rather than fixed number assertions. They confirm that attention weights form a valid distribution over each bag, that the batched and masked forward path matches the single bag path exactly, that masked padding receives zero attention and does not disturb the real instances, that the synthetic data places its signal where it claims to, and finally that a trained model beats chance on bag accuracy and localises the signal instance well above the random baseline. Everything runs on CPU with tiny tensors and no downloads.
The data here is synthetic on purpose so the tests are fast and deterministic.
Swapping in real patch features is a matter of replacing make_bag_dataset
with a loader that yields one feature tensor per slide and a slide level label.
The model and training loop already accept variable length bags one at a time,
which is how a real slide pipeline would feed them.