Non-I.I.D. Multi-Instance Learning for Predicting Instance and Bag Labels with Variational Auto-Encoder

1 Introduction

Multi-instance learning (MIL) [Dietterich et al., 1997] was originally proposed for drug activity prediction where the task is to predict whether a molecule is suitable for binding to a receptor. Since each molecule may take many distinct lowenergy conformations and scientists only know its suitability for binding at the molecule level, MIL is proposed to model the molecules as bags and the conformations of the molecules as instances, where only the bag labels are provided but the instance labels are unknown to the learner.

Since its inception, MIL has been studied extensively in various applications where the tasks have inherently structured representations or the fine-grained instance labels are expensive to obtain. For example, MIL has been studied in text categorization [Andrews et al., 2002; Zhou et al., 2009; Ji et al., 2020] where articles are represented by bags of sentences with only article-level labels, and in image classifications [Chen et al., 2006; Ilse et al., 2018; Skrede et al., 2020; Yao et al., 2020] where the images are divided into bags of patches with only image-level labels.

As many real-world objects are inherently structured, an important advantage of multi-instance learning is that by representing objects as bags of instances it can convey more information than using a flat single-instance representation. Since instances within a bag correspond to parts of an object, they share structural and contextual information inherited from the bag and are unlikely to be independent. To see this, let us consider the example from drug activity prediction where instances within a bag represent low-energy conformations of the same molecule. The instances are evidently not independently and identically distributed (i.i.d.) since all conformations share the same structure of single bonds and only differ in how the single bonds rotate.

Unfortunately, despite the fact that neither the instances within a bag are independently and identically distributed, nor should they be treated as such [Zhou et al., 2009; Zhang and Zhou, 2014], most of the existing MIL algorithms approach the problem by either explicitly assume all instances to be i.i.d. and directly predict the instance labels, or focus solely on predicting the bag labels by transforming the bags into a single-vector embedded space which buries instancelevel information and prohibits instance label prediction.

In this paper, we propose the Multi-Instance Variational Autoencoder (MIVAE) algorithm which explicitly models the instances within a bag as non-i.i.d. using a generative model consisted of a shared bag-level latent factor and instancelevel latent factors specific to each of the instances (Figure 1). By using the bag-level factor to capture the structural and contextual dependencies among the instances and using the instance-level factors to capture the instance-specific variations, MIVAE excels at the tasks of both instance label and bag label prediction. On the one hand, integrating the shared bag-level factor and the individualized instance-level factors encapsulate sufficient information for predicting the bag labels. On the other hand, since the shared structural and contextual dependencies are captured by the bag-level factor and the instance-level factors only capture the variations of the individual instances, the instance-level factors promotes better prediction of the instance labels .

The contributions of this work are three-fold: (i) We propose the Multi-instance Variational Autoencoder (MIVAE) model, to the best of our knowledge, the first neural networkbased MIL algorithm that simultaneously models the instances as non-i.i.d. (ii) We extend the powerful variational autoencoder framework to MIL and instantiate an implementation of MIVAE for simultaneously inferencing the bag-level and instance-level latent factors and predicting both instance labels and bag labels. (iii) We empirically demonstrate the effectiveness of MIVAE over classic and neural network-based MIL algorithms using a suite of MIL benchmark datasets and end-to-end medical image diagnosis datasets for both instance label and bag label prediction.

The rest of this paper is organized as follows. We discuss related work in Section 2 and present the proposed MIVAE framework in Section 3. In Section 4 we report the experimental results. Finally, we conclude the paper in Section 5.

2 RelatedWork

Existing MIL algorithms can be roughly divided into two categories depending whether the algorithm operates at the level of instances or the level of bags [Amores, 2013]. The first group of algorithms tackle the MIL problem at the instance-level by training a instance classifier which separates positive instances in the positive bags from the negative ones [Andrews et al., 2002; Kim and la Torre, 2010; Kandemir and Hamprecht, 2014]. The second group of methods operates at the bag-level, either by embedding the bags into a single-vector and solve a single-instance learning problem [Chen et al., 2006; Wei et al., 2017], or designing multiinstance kernels to measure similarities between the bags [Zhou et al., 2009; Xu et al., 2019].

A few MIL algorithms model instances within the same bag as not identically and independently distributed. [Zhou et al., 2009] models the relationships among the instances with an ε-graph capturing the instance similarities. [Zhang and Zhou, 2014] addresses distribution change by considering covariate shift at both instance-level and bag-level.

Recently, several neural network-based approaches have been investigated for MIL. [Wang et al., 2018] extends neural networks to MIL by using permutation-invariant pooling operations on the learned instance embeddings. [Ilse et al., 2018; Shi et al., 2020] propose to use attention in the pooling layer such that the attention weights can be interpreted as how much the instances contribute to the bag label. To the best of our knowledge, all of the existing neural network-based MIL algorithms assume the instances to be i.i.d. and the proposed MIVAE is the first deep generative model for MIL.

Some algorithms has approached MIL with generative models which directly predicting the instance labels by either assuming the Gaussian Process [Kim and la Torre, 2010; Haussmann et al., 2017] or the Dirichlet Process [Kandemir and Hamprecht, 2014]. However, all of them explicitly assume the instances within a bag to be i.i.d. In this work we show that such assumption not only does not agree with the real-world scenarios, but also hurts empirical performance for both bag and instance level label prediction.

Perhaps the work most related to ours is the proposal [Doran and Ray, 2016], where they model the multi-instance bags as distributions over instances. However, their contribution is only applicable to classical MIL tasks with pre-computed features, while MIVAE also excels at end-to-end learning from weakly-labeled data with neural networks.

3 Method

Let X = Rd denote the space of the instances and Y = {0, 1} denote the space of the label. The learner is given a dataset of m bags multi-instance bags B = {B1, · · · ,Bi, · · · ,Bm}, where each bag Bi = {xi1, · · · ,xij, · · · ,xxini } is a set containing an arbitrary number of ni instances with xij ∈ X. During training, each bag Bi is provided with a bag label yi ∈ Y; however, the label of the instances are unknown. The goal of multi-instance learning as two-fold, to predict both the bag label and the instance labels of unseen bags. When the context is clear, we drop the subscript i fromxij, Bi and yi for the conciseness of the notations.

At the heart of this work lies the claim that the instances within the same bag are not independently and identically distributed and the i.i.d. assumption is only applicable at the bag-level. On the contrary, instances within the same bag share structural and contextual dependencies inherited from the bag and are only independent to each other when the dependencies have been considered. Apart from drug activity prediction, the dependencies of instances exists widely in MIL applications. For example, in text classification instances are paragraphs/sentences of an article and obviously share the styles of the author and contexts of the topic [Ji et al., 2020]. In medical image diagnosis instances are patches of an organ from a patient which naturally share the pathology features and prognoses of the patient [Skrede et al., 2020].

In order to capture both the shared dependencies and the individual variations of the instances, we propose the MIVAE model as depicted in Figure 1. Specifically, it assumes that for each bag the observed instances are generated from two type of latent factors: a bag-level latent factor zB shared by all instances, and ni instance-level latent factors zIj , j ∈ {1, · · · , ni} specific to each of the instances. While the shared bag-level factor zB controls the structural and contextual dependencies among the instances, the instance-level factors zIj are responsible for the variations of each instances. When conditioning on the shared bag-level factor zB , the instances becomes independent to each other.

At the level of bags, MIL is supervised and we assume the bag-level factor zB to be dependent on y. However, since both the instance labels and how they relate to the bag labels are unknown, e.g., instances may not have explicit labels and the bag labels may not be determined by a single instance, assuming fixed generative process may restrict the expressiveness of the model. We address the relationship between the instance-level factors and the bag label during inference.

3.1 Multi-Instance Variational Autoencoder

In practice we only observe the bags and bag labels, our goal is thus to infer the unobservable posterior distributions for both the bag-level and the instance-level latent factors pθ(zB |B), pθ(zIj |xj), and the posterior distribution of the observable pθ(B|zB , zI1 , · · · , zIj ) where θ denotes the generative parameters. Since exact inference of the posterior is intractable due to the fact that both the marginal likelihood and the posterior lack analytical solution, we employ variational autoencoder [Kingma and Welling, 2014] parameterized by neural networks for efficient approximate inference.

As depicted in Figure 2(a), we utilize two separatelyorganized encoding networks qφB (zB |B) and qφI (zIj |xj), where φB and φI denote the parameters, to serve as variational posteriors of the bag-level and instance-level latent factors. The encoded latent factors are then feed into a single decoder pθ(xj|zB , zIj ) for reconstructing the bag of instances. Following standard VAE design, the prior distributions of the latent factors p(zB ) and p(zIj ) are chosen as multivariate Gaussians. Specifically, the prior distributions of the factors and the generative model are factorized as:

where p(xjk|zB , zIj ) is the suitable distribution for the k-th feature of the instance, i.e., Gaussians for continuous features or Bernoulli for binary features; fy is a function parameterized by neural network using the bag label as input; DzB and DzI are the parameters that determine the dimensionality of the bag-level and instance-level latent factors.

Comparing to a standard VAE, the inference model of MIVAE needs to overcome two challenges specific to MIL as depicted in Figure 2(b). Firstly, apart from inferring the instance-level factors zI specific to each instance, MIVAE needs to infer a common bag-level factor zB shared by ni individual instances xj, j = 1, · · · , ni; however, in standard VAE the latent variables are specific to each of the inputs. Secondly, without any instance labels, MIVAE needs to infer the instance labels from the instance-level factors zIj .

Given instances xj, the instance-level factors zIj can be inferred directly with a standard encoder qφI (zIj |xj) parameterized by neural networks. For the shared bag-level factor zB , a straightforward approach would be directly encoding from the concatenation of all instances within the bag; however, this approach unfortunately violates the permutationinvariant property of the instances, i.e., the encoded bag-level factor may dependent on the ordering of the instances.

To avoid this problem, we design MIVAE to first encode the intermediate bag-level factors ˆzBj using qφB (ˆzBj |xj) for each instance, and then utilize a permutation-invariant function for encouraging the intermediate instance-specific factors into a shared bag-level factor. Specifically, the variational approximations for the instance-level factors and the intermediate bag-level factors are defined as

for xj with j = 1, · · · , ni, where f ψ φI , f ψ φB and f π φI , f π φB are the means and variances for the Gaussian distributions parameterized by neural networks. Then, the shared bag-level factor is calculated as the mean of the intermediate factors as

In other words, the encoding neural networks f ψ φB and f π φB first infer the means and variances of the posterior distributions for the intermediate factors from the instances xj, then the shared bag-level factor is calculated as the average of all the intermediate bag-level factors. It is worth noting that the variance of zB then become the average of all the intermediate factor variances. In this way, as the decoder reconstructs each of the instances using the shared bag-level factor zB and individual instance-level factor zIj , all of the intermediate bag-level factors are encouraged to converge to the same value during optimization. Therefore, zB will become a shared bag-level latent factor common to all xj, while each zIj captures as the instance-specific variations.

The generative and inference models of MIVAE can be learned by maximizing the marginal likelihood over the bags

where the second term is the Kullback-Leibler (KL) divergence from the variational approximation to the true posterior. Because the KL-divergence is non-negative, the first term serves as a lower bound of the marginal likelihood which is also called the evidence lower bound (ELBO). Maximizing the marginal likelihood is equivalent to maximizing the ELBO, which can be re-written as

A missing piece of the puzzle is how to predict the instance labels from the inferred instance-level factors zI and predict the bag labels from both bag-level and instance-level factors. Instead of assuming fixed generative distributions which may restrict the applicability of the model, we employ an auxiliary classifier qωI (y|zB , zI1 , · · · , zIni ) attached to the evidence lower bound for end-to-end prediction of both instance labels and bag labels. Using the auxiliary classifier allows MIVAE to have the flexibility of utilizing any permutation-invariant multi-instance pooling techniques, i.e., max-pooling, LSE-pooling, and attention pooling [Ilse et al., 2018]. In order to demonstrate that disentangling the dependencies among instances into zB and using only the instancelevel factors zIj promotes better instance label prediction performance, we use the straightforward max-pooling operation which is defined as

where fωI is a neural network with a sigmoid activation to map the instance-level factors to probabilities. Then, the auxiliary for predicting the bag label can be expressed as

where fωB is also a parameterized neural network that outputs probabilities. Finally, the MIVAE objective can be expressed as the sum of the ELBO and the auxiliary classifier,

where α is a weighting parameter of the auxiliary objective. Using the re-parameterization trick [Kingma and Welling, 2014], the objective of MIVAE can be efficiently optimized using stochastic gradient descent.

For predicting the instance labels of unseen bags, we use the trained instance-level encoder qφI (zIj |xj) to infer the mean of the posteriors for each xj ∈ B, and use the intermediate output fωI (zInj ) from the auxiliary classifier. For predicting the bag labels, we use the bag-level encoder qφB (zB |B) to infer the mean of the bag-level latent factor zB , and directly use the output of the auxiliary classifier qω(y|zB , zI1 , · · · , zIni ) for prediction.

4 Experiments

We empirically evaluate MIVAE against state-of-the-art MIL algorithms for both bag label prediction and instance label prediction tasks. Firstly, we compare against SVM-based algorithms including mi-SVM [Andrews et al., 2002], miFV [Wei et al., 2017], mi-Graph [Zhou et al., 2009], and KISVM [Li et al., 2009]. Secondly, we compare against neural network-based algorithms including mi-Net and MI-Net [Wang et al., 2018], and AttentionMIL [Ilse et al., 2018]. Thirdly, we compare against generative MIL algorithms for predicting instance labels, including Gaussian Process MIL (GPMIL) [Kim and la Torre, 2010], Dirichlet Process MIL (DPMIL) [Kandemir and Hamprecht, 2014], and Variational Gaussian Process MIL (VGPMIL) [Haussmann et al., 2017].

We implement MIVAE using PyTorch. Unless otherwise stated, the reported results are obtained from averaging the test results of 10 times repeated 10-fold cross validation as in previous MIL literature. Training for MIVAE are conducted for 100 epochs and parameters are tuned using validation loss evaluate with 10% of the training bags. Experiments are conducted using a PC with a single Nvidia GTX 1080Ti GPU.

4.1 Bag Classification: MIL Benchmarks

We first study the performance of MIVAE for the bag label predictions on five MIL benchmark datasets including the drug activity prediction datasets Musk1 and Musk2, and three content-based image retrieval tasks Fox, Tiger, and Elephant. Following previous works, we report the average accuracy and the standard deviation obtained from repeating 10-fold cross validation for 10 times.

Since these benchmark datasets only contains a small number of bags (approximately 100 bags for the Musk datasets and 200 bags for the image retrieval datasets) and the precomputed feature vectors, we do not expect MIVAE to significantly outperform existing methods since its deep generative model may not be well suited for these type of tasks. However, from Table 1 we can see that the performance of MIVAE is competitive to the compared baselines. Specifically, MIVAE achieves the highest average accuracy on Musk1 and Tiger datasets, and performs competitively on the others. The results demonstrate that the model of MIVAE is sucessful for handling the uncertainties of multi-instance learning.

4.2 Instance Classification: 20 NewsGroups

To study the performance for predicting instance labels, we evaluate MIVAE against baselines using the 20 Newsgroup dataset [Zhou et al., 2009]. It contains a collection of 20 datasets where each of them is consisted of 100 bags. Each bag contains approximately 40 instances of articles from 20 different topics where each instance represents one article described by the top 200 TF-IDF features. A bag is considered to be positive if at least one of its instances belongs to a specific topic. The labels of the bags are distributed evenly; however, in each positive bag only approximately 3% of the instances are positive. Because of the proportions of positive instances are small, the dataset has been widely used as a standard benchmark for instance label prediction performances.

Following previous work, we report results obtained from repeating 10-fold cross validation for 10 times using the train/test splits provided in [Zhou et al., 2009]. Because the ground truth instance labels are highly imbalanced, we use Area Under the Precision-Recall Curve (AUC-PR) instead of the area under the receiver operating characteristic curve. Furthermore, in order to demonstrate that the advantage of MIVAE lies in its model instead of parameter tuning, the parameters of MIVAE are selected using only the first of the 20 Newsgroup datasets (alt.atheism), and the same parameters are applied across all the rest of the datasets.

As shown by the Table 2, without tuning parameters for each specific dataset, MIVAE outperforms the compared baselines on 15 out of 20 datasets for predicting instance labels despite the small number of training bags. Since all compared algorithms assume the instances to be i.i.d., these results confirm the advantage of MIVAE for modeling the dependencies among the instances.

4.3 Histopathology Images: Colon Cancer

Computer-aided histopathology using hematoxylin and eosin (H&E) stained whole-slide images is an important task for MIL since traditional supervised approaches require pixellevel annotations, which is extremely labor-intensive for pathologists [Skrede et al., 2020]. Therefore, a label efficient multi-instance approach using only image-level labels is essential. For this purpose, we evaluate MIVAE against state-of-the-art end-to-end neural network-based MIL algorithms on both bag and instance label prediction using H&E histopathology images from the colon cancer dataset provided in [Sirinukunwattana et al., 2016].

The dataset contains 100 H&E images where each image is from a patient's tissue of either normal or malignant regions. For each image bag, the instances are generated as patches 27 × 27 pixels using markings of major nuclei for each cell. A total amount of 22,444 nuclei instances are provided with ground truth instance class labels, i.e. epithelial, inflammatory, fibroblast, and miscellaneous. The bag is considered to be positive if it contains at least one epithelial patch.

For fairness of comparison, the training images are color normalized and augmented with the same random horizontal flipping/vertical flipping, and random rotation procedures. Furthermore, they also use the same convolutional structures for feature extraction (please refer to the supplementary for details). When evaluating bag label prediction, we use accuracy since the bag labels are evenly distributed. Because the instances label are imbalanced, we use AUC-PR for quantitative evaluating instance label prediction performance as in the last section. The reported results are the average performances from 5 times repeated 10-fold cross-validation.

As shown in the results of Table 3, MIVAE achieves both the highest classification accuracy for predicting bag labels and the best AUC-PR scores for predicting instance labels, indicating that MIVAE outperforms state-of-the-art neural network-based MIL algorithms at both tasks.

In Figure 3, we use heatmaps to visually compare the instance label prediction performances. For each image, the heatmaps are obtained by multiplying the instances prediction scores of each algorithm (normalized to the range of [0, 1]) to the pixel values of the image patches. From Figure 3 (b)-(e), it is evident that the heatmaps generated by MIVAE are the closest to the ground truths, while mi-Net exhibits lower recalls and the predictions of AttentionMIL are not as confident as the predictions of MIVAE.

It is worth noting that both mi-Net and AttentionMIL assume the instances to be independently distributed. Moreover, mi-Net also uses max-pooling when linking instance prediction scores to bag label. Therefore, contrasting the heatmaps of MIVAE and mi-Net paints a clear picture for the benefits of modelling the instance dependencies: mi-Net only selects a small portion of positive patches, while MIVAE is able to identify most of them because the shared dependencies are excluded from instance label predictions. Furthermore, comparing MIVAE with AttentionMIL shows that by using the instance-level factors, max-pooling provides more confident predictions than the attention weights. For more examples, please kindly refer to the supplementary materials.

To sum up, the experiments results demonstrates that the proposed MIVAE algorithm is effective at both bag label and instance label prediction tasks. Moreover, the instance prediction scores of MIVAE can be used for explaining the region of interests which is valuable for weakly-labeled tasks where fine-grained labels are costly to obtain.

5 Conclusion

In this paper, we propose MIVAE, a generative multi-instance learning algorithm which models the instances within a bag as non-i.i.d.. Experiments demonstrate that MIVAE is effective for both instance-level and bag-level label prediction. An interesting future direction would be to investigate how other MIL pooling techniques affect the performance of MIVAE.