Source-linked AI summary
Attention-based Deep Multiple Instance Learning
Maximilian Ilse, Jakub M. Tomczak, Max Welling
TL;DR
MIL seeks to predict a bag-level label from multiple instances while identifying which instances drive that label. The paper uses neural networks to model a Bernoulli bag-label distribution and introduces permutation-invariant attention pooling. The approach matches leading MIL methods on benchmarks, outperforms others on MNIST-based and histopathology datasets, and provides interpretable instance-based regions of interest.
Problem
MIL must predict a single bag label from multiple instances while addressing the difficult and clinically important task of identifying key instances.
Method
The method models the bag label with a Bernoulli distribution, neural-network transformations, and a trainable permutation-invariant weighted average corresponding to attention.
Results
The approach is on a par with the best classical MIL methods on benchmark datasets and outperforms other methods on a MNIST-based problem and two histopathology datasets.
Takeaways & Limitations
Attention weights provide insight into each instance's contribution and can identify regions of interest in image datasets.
Takeaways & Limitations
The paper focuses on binary MIL and leaves multiclass MIL, repulsion points, and dependencies among instances for future research.
Abstract
from arXiv · showhide
Multiple instance learning (MIL) is a variation of supervised learning where a single class label is assigned to a bag of instances. In this paper, we state the MIL problem as learning the Bernoulli distribution of the bag label where the bag label probability is fully parameterized by neural networks. Furthermore, we propose a neural network-based permutation-invariant aggregation operator that corresponds to the attention mechanism. Notably, an application of the proposed attention-based operator provides insight into the contribution of each instance to the bag label. We show empirically that our approach achieves comparable performance to the best MIL methods on benchmark MIL datasets and it outperforms other methods on a MNIST-based MIL dataset and two real-life histopathology datasets without sacrificing interpretability.
1. Introduction
Multiple instance learning assigns one label to a bag of instances, creating a need to predict bag-level categories while identifying influential instances. The paper proposes an interpretable, flexible neural attention approach and reports strong performance across benchmark and applied datasets.
- MIL assigns a single class label to a bag of instances, often representing weakly annotated data in applications such as medical imaging.Medical images may receive only benign/malignant labels or roughly specified regions of interest.
- The task includes predicting a bag label and discovering key instances that trigger the label, which is particularly useful for medical interpretation.Existing methods differ in interpretability, while instance-level accuracy and agreement among methods remain concerns.
- The proposed model parameterizes MIL with neural networks and replaces max or mean pooling with a trainable weighted average corresponding to attention.Its three-stage structure transforms instances, aggregates them permutation-invariantly, and maps the result to a bag probability.
- Attention weights identify key instances that can be used to highlight possible regions of interest.The paper reports empirical evidence that the model indicates key instances in image datasets.
- The approach matches the best classical MIL methods on benchmark datasets and outperforms other methods on a MNIST-based problem and two histopathology datasets.The method is presented as improving flexibility and interpretability without sacrificing predictive performance.
2. Methodology
The method models bag labels probabilistically with neural networks while enforcing permutation invariance, and uses attention-based pooling to learn flexible, interpretable instance aggregation.
- MIL represents each example as an unordered, variable-sized bag of instances with one observed binary bag label and unobserved instance labels.
- A valid MIL model must be permutation-invariant because instance ordering and dependencies within a bag are unavailable.
- Optimizing the maximum over instance labels can cause vanishing gradients and is restricted to instance-level classifiers.
- The bag label is modeled with a Bernoulli distribution parameterized by θ(X), trained by optimizing log-likelihood.
- The general architecture transforms instances with f, aggregates them using symmetric pooling σ, and transforms the result with g.
- Neural networks parameterize instance transformations and optionally bag-level transformations, enabling end-to-end differentiable training when pooling is differentiable.
- Attention mechanism: The proposed attention pooling computes a learned weighted average of embeddings with weights summing to 1, preserving invariance to bag size.
- Gated attention mechanism: Gated attention combines tanh and sigmoid nonlinearities to introduce a learnable nonlinearity that may improve complex relation modeling.
3. Related work
MIL research combines bag-level learning with neural feature extraction and attention-based pooling. In medical imaging, these approaches address weak annotation and can support both diagnosis and region-of-interest identification.
- MIL pooling: Mean and max pooling are common MIL operators, but they are non-trainable and may limit applicability.The proposed alternative is a fully trainable pooling operator that adapts to new instances.
- MIL with neural networks: Neural networks can provide feature extraction within MIL, extending classical settings that rely on precomputed instance features.This is particularly relevant for image and text analysis, where additional feature extraction is needed.
- MIL and attention: The proposed two-layer neural attention operator outperforms commonly used MIL pooling operators.Earlier attention-based MIL variants used auxiliary regression, a one-layer network, or dot-product attention.
- MIL for medical imaging: Medical imaging uses MIL to process smaller image patches as bags because whole-image processing is infeasible and pixel-level annotations are difficult to obtain.The approach is applied to histopathology, where it can provide diagnosis and indicate regions of interest.
4. Experiments
The experiments evaluate attention-based neural MIL across benchmark, MNIST-based, and histopathology datasets using established comparisons and repeated evaluation. The approach is competitive on classical benchmarks, strongest in several small-sample settings, and supports instance-level interpretation through attention weights.
- Experimental setup: The study evaluates Attention and Gated-Attention on five classical MIL datasets, MNIST-BAGS, and two real-life histopathology datasets.Comparisons include neural MIL variants, common pooling layers, and an SVM-based MIL model where feasible.
- Classical MIL datasets: Classical benchmark results show that both proposed attention models are comparable with the best-performing classical MIL methods.Table 1 reports averages over five runs with a standard error of the mean.
- MNIST-bags: The attention-based deep MIL approach performs much better than other methods in the small-sample MNIST-BAGS regime.It remains significantly higher in AUC with 50–150 bags for approximately 10 instances per bag and 50–100 bags for approximately 50 instances per bag.
- MNIST-bags: For MNIST-BAGS, the proposed method outperforms the SVM-based approach except at large effective sample sizes, while model differences narrow with sufficient training data.Embedding-based models outperform instance-based models, and mean pooling is generally worse than max pooling but can eventually converge to the best value.
- Interpretability: Attention weights identify relevant instances: all nines are highlighted in a positive MNIST bag, and histopathology heatmaps substantially match epithelial ground truth.The heatmaps are formed by weighting image patches with their corresponding attention weights despite image-level training annotations.
- Histopathology datasets: On breast and colon cancer datasets, the proposed approach achieves high predictive performance and outperforms other methods, with attention variants differing by dataset.Gated attention performs better than plain attention on BREAST CANCER, while they behave similarly on COLON CANCER.
5. Conclusion
The paper presents a flexible, interpretable neural-network MIL approach with attention-based pooling. Across five MIL datasets, one image corpus, and two histopathology datasets, it performs competitively while providing ROI-based decision interpretations, but the study remains limited to binary MIL.
- The proposed MIL approach uses neural networks and trainable attention pooling to model permutation-invariant bag score functions.The method is described as flexible and interpretable, with attention-based pooling.
- Across five MIL datasets, one image corpus, and two real-life histopathology datasets, the method is on a par with the best methods or performs best across evaluation metrics.
- Attention provides interpretations by presenting regions of interest, an important capability for practical applications.
- The study focuses on binary MIL and leaves multi-class MIL, repulsion points, and dependencies among instances for future research.
6. Appendix
The appendix documents the deep MIL architectures, optimization procedures, datasets, and additional experimental results. It also provides examples comparing attention-based and instance-based heatmaps for histopathology images.
- Figure 6 contrasts instance-based, embedding-based, and attention-based deep MIL approaches.Instance scores are shown in red and bag-vector representations in blue.
- The appendix lists architectures and optimization details for classical MIL datasets, MNIST-bags, and histopathology datasets.It also describes the five benchmark datasets and the MNIST-bags and histopathology model configurations.
- Additional MNIST-bags results report test AUC for bags averaging 10, 50, and 100 instances.The corresponding tables vary the number of training bags.
- Figures 7–9 show attention weights for negative bags and positive bags containing one or multiple instances of the digit 9.Positive instances are highlighted with red rectangles.
- Histopathology examples compare ground-truth epithelial patches with attention and Instance+max heatmaps after rescaling their weights or scores.The figures use H&E images, cell-centered patches, positive patches, and model-derived heatmaps.