Sparse Network Inversion for Key Instance Detection in Multiple Instance Learning

MIL deals with a set of bags $\{X_{n}\,|\,n=1,\ldots,N\}$, where $N$ is the number of bags. A bag $X_{n}$ is associated with multiple instances $\{X_{n,k}\in{\rm I\!R}^{D}\,|\,k=1,\ldots,K\}$, where $K$ is the number of instances ($K$ can vary for different bags), and there are no dependency and no ordering between $K$ instances in a bag. Each instance in the bag $X_{n}$ is associated with an instance-level label $\{Y_{n,k}\in\{0,1\}\,|\,k=1,\ldots,K\}$, and the bag $X_{n}$ is associated with the bag-level label $Y_{n}=OR(Y_{n,1},\ldots,Y_{n,K})$. That is, the bag-level label is positive if at least one of the instances in the bag is a positive instance, otherwise negative. Under the MIL assumption, an attention-based deep MIL model computes a bag probability $P_{n}$ and attention scores $a_{n}$ for the $n^{\text{th}}$ bag $X_{n}$, where $a_{n}$ is a set of attention scores of $K$ instances in the $n^{\text{th}}$ bag, i.e., $a_{n}=\{a_{n,1},\ldots,a_{n,K}\}$. For the $n^{\text{th}}$ bag $X_{n}$, the bag-level prediction $\widehat{Y}_{n}$ is computed by applying a threshold to bag probability $P_{n}$. In the case of KID, whenever $\widehat{Y}_{n}=1$, for the $k^{\text{th}}$ instance $X_{n,k}$ in the $n^{\text{th}}$ bag, an instance-level prediction $\widehat{Y}_{n,k}$ is computed by applying a threshold to the normalized attention score $\widehat{a}_{n,k}$ after Min-Max normalization is applied to the attention score $a_{n,k}$.

We build the attention-based deep MIL model that uses attention-based pooling on embedding-space paradigm as shown in Fig. 1, and we train the model by optimizing negative log-likelihood function. There are two reasons we use this structure for the MIL model. First, embedding-space paradigm has better bag-level classification performance than instance-space paradigm. The second reason is that attention-based pooling allows embedding-space paradigm to detect key instances because the magnitude of the attention score corresponding to an instance in a positive bag tells how likely it is to be a key instance.

As shown in the Fig. 1, the attention-based deep MIL model transforms the instances in the bag $X_{n}$ to instance embeddings, and then the model makes a bag embedding by aggregating the instance embeddings through attention-based pooling. Finally, the model calculates a bag probability $P_{n}$ by putting the bag embedding to a bag-level classifier. For the MIL model, if we assume that $G_{n}=\{g_{n,1},\ldots,g_{n,K}\}$ is a set of the instance embeddings of the bag $X_{n}$, where $g_{n,k}$ is the instance embedding of the $k^{\text{th}}$ instance in the $n^{\text{th}}$ bag and $g_{n,k}\in$ I​R^1×N, an attention score for the $k^{\text{th}}$ instance in the $n^{\text{th}}$ bag can be written as follows [1]:

where $a_{n,k}$ is an attention score for the $k^{\text{th}}$ instance in the $n^{\text{th}}$ bag, $\text{w}\in$I​R^M×1 and $\text{V}\in$I​R^M×N are fully-connected layer’s parameters. The bag embedding for the $n^{\text{th}}$ bag based on attention-based pooling can be written as follows [1]:

where $z_{n}$ is the bag embedding for the $n^{\text{th}}$ bag. The bag embedding $z_{n}$ is put to a bag-level classifier $H\left(\cdot\right)$, and then the bag probability is calculated by using the result of bag-level classifier $H\left(\cdot\right)$. This process can be written as follows:

where $P_{n}$ is the bag probability for the $n^{\text{th}}$ bag. For the bag-level classifier $H\left(\cdot\right)$, we use a single fully-connected layer, which can learn the interactions within the bag embedding. To make the bag-level prediction $\widehat{Y}_{n}$, we apply a threshold to $P_{n}$. In the case of KID, whenever $\widehat{Y}_{n}=1$, we apply Min-Max normalization to the attention score $a_{n,k}$ of $k^{\text{th}}$ instance in the bag $n^{\text{th}}$ bag. This process can be written as follows:

where $\widehat{a}_{n,k}$ is the normalized attention score of $a_{n,k}$ and $a_{n}$ is a set of attention scores of $K$ instances in the $n^{\text{th}}$ bag, i.e., $a_{n}=\{a_{n,1},\ldots,a_{n,K}\}$. Then, we make instance-level prediction $\widehat{Y}_{n,k}$ by applying a threshold to $\widehat{a}_{n,k}$. However, since the attention-based deep MIL model has limits in KID performance, we perform KID by recomputing the instance-level prediction $\widehat{Y}_{n,k}$ of the optimized bag $\widehat{X}_{n}$ after optimizing the bag $X_{n}$ to fit the criterion of the model by applying sparse network inversion.

Neural network inversion is a method that optimizes the input to fit the criterion of the neural network model. To apply neural network inversion to an attention-based deep MIL model, we initialize the input as the values of the input image. Since we use negative log-likelihood as objective function when training the attention-based deep MIL model, we optimize the bag $X_{n}$ by fitting the same criterion of the attention-based deep MIL model. That is, whenever the bag-level prediction $\widehat{Y}_{n}$ for the bag $X_{n}$ is positive, we apply neural network inversion to the bag $X_{n}$. Then, the parts of the optimized bag $\widehat{X_{n}}$ that make the bag positive are strengthened, and the other parts are weakened. In other words, if the optimized bag $\widehat{X_{n}}$ is put to the trained MIL model again, the attention scores corresponding to the key instances become larger. Formally, the objective function of neural network inversion can be written as follows:

where $\widehat{Y}_{n}$ is bag-level prediction for the bag $X_{n}$. The reason we use $\widehat{Y}_{n}$, not $Y_{n}$, is that we cannot access the true bag-level label $Y_{n}$ during test. To update the bag $X_{n}$ by optimizing the objective function of neural network inversion, we use Stochastic Gradient Descent (SGD) optimization algorithm [16]. Each time the instances of the bag $X_{n}$ is updated, we set the range of the values of instances from $0$ to $1$ as each instance is image data with a value between $0$ and $1$.

Sparse network inversion is a kind of neural network inversion that uses a sparseness constraint. That is, sparse network inversion updates the bag $X_{n}$ by optimizing the objective function of neural network inversion with a sparseness constraint. We incorporate a sparseness constraint into neural network inversion to give regularization effect when we update the bag $X_{n}$ and to further weaken the parts of the bag $X_{n}$ that do not affect the criterion of the trained MIL model. Thus, if we put the optimized bag $\widehat{X}_{n}$ to the attention-based deep MIL model, we can detect more key instances than when using only attention-based deep MIL model or applying neural network inversion. The objective function of sparse network inversion can be written as follows:

where $l(X_{n})$ is the objective function of neural network inversion and $\lambda$ is a sparseness coefficient. Since the objective function of sparse network inversion involves a sparseness constraint for bag $X_{n}$, in order to optimize the bag $X_{n}$, we use proximal gradient method, which is an extension of typical gradient algorithm [17]. The method is used due to its simplicity and adequateness for solving data with high dimension [17, 18]. In the proximal gradient method, soft-thresholding function is used as proximal operator for a sparseness constraint. The soft-thresholding function can be written as follows [17]:

where $\lambda$ is the sparseness coefficient. We can write the soft-thresholding function compactly as follows:

where sign($x$) means the sign of $x$. Formally, when we use proximal gradient method for the bag $X_{n}$, update equation of sparse network inversion can be written as follows:

where $s_{\lambda}(\cdot)$ is the soft-thresholding function, $\eta$ is the learning rate, and $l(\cdot)$ is the objective function of neural network inversion. Since instances of the bag $X_{n}$ is image data with a value between $0$ and $1$, each time each instance of the bag $X_{n}$ is updated, we set the range of the instances in the bag from $0$ to $1$.

For all experiments, we use the same model architecture and optimizer that showed the high classification performance for each dataset as used in the precedent work [1]. To construct the attention-based deep MIL model, we apply attention-based pooling before the last layer of the deep neural network model. We call the attention-based deep MIL model Att, the method that applies neural network inversion to the attention-based deep MIL model Att+inv, and the method that applies sparse network inversion to the attention-based deep MIL model Att+sparse. In the MIL models based on instance-space paradigm, the models compute instance scores and aggregate the scores by max pooling or mean pooling. To construct the model, we applied max pooling or mean pooling after the last layer of the deep neural network model. We call these MIL models Inst+max, Inst+mean.

To do a fair evaluation, we use 10-fold cross-validation, and each experiment was performed five times independently. In the case of the MNIST-based image MIL dataset, we divided the MNIST dataset into an MNIST train dataset and an MNIST test dataset. Then, we created the MNIST-based image MIL train dataset and the MNIST-based image MIL test dataset by doing sampling from an MNIST train dataset and an MNIST test dataset independently.

To compare the bag-level classification performance, we use accuracy as a metric because the number of positive and negative bags on various MIL datasets is balanced as shown in Table 1. In the case of KID performance, since the number of positive and negative instances in positive bags on various MIL datasets are imbalanced as shown in Table 2, we use F1 measure as used in the precedent work [6, 21], as F1 measure is known to be insensitive to imbalance of the labels in data [22].

To experiment with these datasets, we used an existing model architectures: Lenet 5 model [24] for MNIST-based image MIL dataset; the proposed model [19] for two real-world histopathology datasets. When constructing the models, ReLU is used as activation function in all layers except for the two layers that is used for attention-based pooling, where tanh is used as the activation function [1]. To optimize the MIL models, Adam optimization algorithm is used [25]. For the MIL models on MNIST-based image MIL dataset, beta values of 0.9 and 0.999, learning rate of 0.0005, and weight decay of 0.0001 are used [1]. In the two histopathology datasets, beta values of 0.9 and 0.999, learning rate of 0.0001, weight decay of 0.0005 are used [1]. When optimizing the models, we used 100 epochs and we stopped optimizing models based on the lowest validation error. Also, in the case of two histopathology datasets, to prevent overfitting due to the small amount of data, we performed data augmentation by randomly rotating and flipping the data. In order to improve KID performance of the attention-based deep MIL model, when we optimize data by applying neural network inversion, we updated each data 1,000 times by using SGD optimization algorithm [16] where learning rate is 0.001, momentum coefficient is 0.9. When we applied sparse network inversion, we used proximal gradient method. For MNIST-based image MIL dataset, we used 0.0001 as learning rate and we repeated update 200 times with varying the sparseness constraint coefficient from 0.0005 to 0.003. In COLON CANCER, the learning is 0.00001 and each data is updated 200 times with varying the sparseness constraint coefficient from 0.0004 to 0.0024. In the case of BREAST CANCER, 0.0001 is used as learning rate and each data is updated 100 times with varying the sparseness constraint coefficient from 0.003 to 0.006. For KID, we set the threshold to the optimal value among [0.1, 0.2, 0.3, 0.4, 0.5] for each model and dataset.

The bag-level classification results and KID results on all datasets are in Table 3 and Fig. 2, respectively. For bag-level classification performance, the attention-based deep MIL model outperforms the other MIL models which are based on instance-space paradigm for all datasets. On the other hand, in the KID performance, the attention-based deep MIL model shows a worse performance than the other model based on instance-space paradigm for MNIST-based image MIL dataset and COLON CANCER. In the case of BREAST CANCER, although the attention-based deep MIL model has better performance than the other model based on instance-space paradigm, the absolute performance is not optimal. Thus, despite the superior bag-level classification performance of the attention-based deep MIL model, there is a problem in using an attention-based deep MIL model for KID. However, if we apply neural network inversion or sparse network inversion to data in the attention-based deep MIL model, since these methods relieve the problem that attention-based deep MIL model focuses only on few distinguishable key instances by removing the constraint that data cannot be changed, the KID performance is improved. Especially, in the case of sparse network inversion, the KID performance is significantly improved as the effects of regularization and sparseness have a positive effect on the result.

To intuitively show that our method can provide better interpretable results than the attention-based deep MIL model, we provide a visualization of the results in Fig. 3 and Fig. 4, where Fig. 3 and Fig. 4 are related to MNIST-based image MIL dataset and COLON CANCER, respectively. Each figure consists of seven images: (a) all patches of an image in the each dataset; (b) all refined patches that are the result of sparse network inversion; (c) heatmap of the attention-based deep MIL model: every patch from (a) multiplied by its corresponding Min-Max normalized attention score of every patch from (a); (d) heatmap of the sparse network inversion: every patch from (a) multiplied by its corresponding Min-Max normalized attention score of every patch from (b); (e) visualization result of the attention-based deep MIL model: every patch from (a) multiplied by its corresponding value that applies threshold to Min-Max normalized attention score of every patch from (a); (f) visualization result of sparse network inversion: every patch from (a) multiplied by its corresponding value that applies threshold to Min-Max normalized attention score of every patch from (b); (g) ground truth of the patches in instance-level. As shown in the Fig. 3-(c) and Fig. 4-(c), Min-Max normalized attention scores of the attention-based deep MIL model are skewed to few distinguishable key instances. However, if we apply our method to the patches, Min-Max normalized attention scores are scattered around more key instances than before applying sparse network inversion to the patches as shown in the Fig. 3-(d) and Fig. 4-(d). By applying threshold to the normalized attention scores of the attention-based deep MIL model and sparse network inversion, we can perform KID as shown in the Fig. 3-(e), Fig. 3-(f), Fig. 4-(e), and Fig. 4-(f). From these results, although we use only bag-level data and bag-level label during training, we can confirm that the attention-based deep MIL model finds key instances in a positive bag. However, if we just use the attention-based deep MIL model as shown in Fig. 3-(e) and Fig. 4-(e), the model finds only few key instances. On the other hand, as shown in Fig. 3-(f) and Fig. 4-(f), our method allows the model to find more key instances than before applying our method.

From these experiments on the MNIST-based image MIL dataset and two real-world histopathology datasets, we can see that our method significantly improves the KID performance while the bag-level classification performance is maintained. This means that applying sparse network inversion to the attention-based deep MIL model helps improve the KID performance of the attention-based deep MIL model.