[Paper Review] Attention-based Deep Multiple Instance Learning

Maximilian Ilse, Jakub M. Tomczak+1 more|arXiv (Cornell University)|Feb 13, 2018

Image Retrieval and Classification Techniques36 references671 citations

TL;DR

Introduces a neural network–based, permutation-invariant MIL framework with a trainable attention pooling operator; achieves competitive results and provides interpretable instance-level importance (ROIs) across several datasets.

ABSTRACT

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.

Motivation & Objective

• Reformulate MIL as learning a Bernoulli bag-label distribution parameterized by neural nets.
• Develop a permutation-invariant aggregation operator that is trainable (attention-based).
• Provide interpretable instance contributions to the bag label via attention weights.
• Enable end-to-end training of instance transformations, pooling, and bag-level prediction using neural networks.

Proposed method

• Model the bag probability as a symmetric function S(X)=g( sum_x f(x) ).
• Transform each instance via a neural network f_ψ to low-dimensional embeddings h_k.
• Aggregate embeddings with a differentiable, trainable attention-based pooling z = sum_k a_k h_k where a_k are learned attention weights.
• Use a gating mechanism (tanh and sigmoid) to enhance expressiveness of attention weights (a_k).
• Train end-to-end by maximizing the log-likelihood of the Bernoulli bag label Y given bag X.
• Demonstrate interpretability by showing that attention weights highlight key instances/ROIs in images.

Experimental results

Research questions

• RQ1 Can a neural attention-based MIL pooling achieve competitive bag-level accuracy on standard MIL benchmarks?
• RQ2 Does the proposed pooling provide interpretable instance-level contributions (key instances/ROIs) for decision justification?
• RQ3 How does embedding-based MIL with attention compare to instance-based MIL pooling (mean/max) across diverse datasets?
• RQ4 Is the approach effective in small-sample medical imaging settings where pixel/patch annotations are weak or scarce?

Key findings

• The attention-based deep MIL approach achieves performance on par with the best classical MIL methods on benchmark datasets and outperforms other methods on MNIST-based MIL and two histopathology datasets.
• The attention weights enable identification of key instances, providing interpretable ROIs in medical imaging tasks.
• Embedding-based models generally outperform instance-based models, and gated attention improves performance over plain attention on at least some datasets.
• Mean pooling performs worse than max pooling in the MNIST-bags experiments, while the gated-attention variant shows robustness across datasets.