DGR-MIL: Exploring Diverse Global Representation in Multiple Instance Learning for Whole Slide Image Classification

Abstract

Multiple instance learning (MIL) stands as a powerful approach in weakly supervised learning, regularly employed in histological whole slide image (WSI) classification for detecting tumorous lesions. However, existing mainstream MIL methods focus on modeling correlation between instances while overlooking the inherent diversity among instances. However, few MIL methods have aimed at diversity modeling, which empirically show inferior performance but with a high computational cost. To bridge this gap, we propose a novel MIL aggregation method based on diverse global representation (DGR-MIL), by modeling diversity among instances through a set of global vectors that serve as a summary of all instances. First, we turn the instance correlation into the similarity between instance embeddings and the predefined global vectors through a cross-attention mechanism. This stems from the fact that similar instance embeddings typically would result in a higher correlation with a certain global vector. Second, we propose two mechanisms to enforce the diversity among the global vectors to be more descriptive of the entire bag: (i) positive instance alignment and (ii) a novel, efficient, and theoretically guaranteed diversification learning paradigm. Specifically, the positive instance alignment module encourages the global vectors to align with the center of positive instances (e.g., instances containing tumors in WSI). To further diversify the global representations, we propose a novel diversification learning paradigm leveraging the determinantal point process. The proposed model outperforms the state-of-the-art MIL aggregation models by a substantial margin on the CAMELYON-16 and the TCGA-lung cancer datasets. The code is available at https://github.com/ChongQingNoSubway/DGR-MIL .

1. Introduction

Histological whole slide images (WSIs) are commonly used to diagnose a variety of cancers, e.g., breast cancer, lung cancer, etc. [46]. However, the gigapixel resolution of WSIs hinders the direct translation of classic deep learning methods into WSI applications mainly due to computational intractability [4, 11, 35, 38]. Therefore, the analysis of WSIs typically starts with cropping images into small patches and then performing analysis on a per-patch basis. In addition, the absence of labor-intensive pixel/patch-level annotations poses a significant challenge for the precise localization of targets of interest (e.g., tumors in WSIs) in a fully supervised setting. As a result, Multiple Instance Learning (MIL), a weakly supervised method, is commonly employed in WSI analyses by treating an entire WSI as a bag and the cropped patches as instances.

The prevailing MIL models in analyzing WSIs have been built upon the attention-based MIL (AB-MIL) framework [28] since its introduction. However, the standard AB-MIL treats each instance independently and does not take the correlations between instances into account. Although many of its follow-ups address this challenge by a variety of means [30, 47, 58, 64], they mainly focus on modeling the correlation between instances by assigning high correlations to instances from the same category (e.g., tumor instances). However, even instances from the same category exhibit variations in phenotype, size, as well as spatial diversity marked by immune infiltration across different patients [7, 37, 66]. For example, negative instances close to the tumor boundaries typically resemble positive instances while appearing differently compared to the other negative instances [24]. As a result, instances belonging to the same category may not be assigned high correlations; similarly, instances from different categories could also receive high correlations. This spurious correlation between instances is prone to trap the MIL model by incorrectly aggregating them when making predictions. Formally, we quantify the diversity of instances between and within bags in WSIs by leveraging the rate-distortion theory [12, 15, 63], where a higher rate indicates a less compressible but more diverse collection of samples (see details of computing the diversity measure in Appendix A). As consistent with findings in pathology, we observe that both positive and negative instances in WSIs exhibit between-bag and within-bag diversity (refer to Fig. 1). Based on this fact, we argue that the diversity of instances is important in designing MIL models. Before that, clustering/prototype-based MIL methods tried to solve the diversity by utilizing attention scores as pseudo labels to provide instance-level supervision [55, 61]. This introduces a chicken-and-egg issue. The effectiveness of pseudo-labels relies on successful MIL classification pooling, which in turn depends on precise attention localization. Especially when patch representations are inferior or MIL initially guided by poor pseudo label, leading to even misleading localization and unstable optimization [32, 68]. Among them, PMIL presents an alternative method to avoiding noise attention [62], initially selecting prototypes through clustering, followed by modeling diversity via prototype and patch representation. However, the design of the multi-stage framework empirically leads to suboptimal learning outcomes, and the restricted number of prototypes, due to high computational burden, results in diminished diversity.

Fig. 1:

(a) Examples of positive instances of with-bag and between-bag diversities measured by rate-distortion theory. (b) Histogram of the diversity measure within positive bags on the CAMELYON16 dataset. (c) The between-bag distinction measures the pair-wise similarity between bags.

To this end, we propose to jointly model this diversity through a set of learnable global vectors. The learned global vectors summarize diverse instances of interest (e.g., tumors in WSIs). As a result, the diversity between instances can be implicitly modeled by computing the correlation between instance embeddings and the global vectors through a cross-attention mechanism. To enhance the ability of the global vectors to capture the most discriminative global context for WSI classification, we introduce the concept of tokenized global vectors. It is worth mentioning that the importance map for instances can be calculated based on the attention between the tokenized global vector and the embedding of each individual instance. To learn diverse global vectors, we propose two main strategies. First, we push the global vectors toward the centers of the positive bag by a positive instance alignment mechanism. Second, we propose a low-complexity and theoretically guaranteed diversity loss to enforce the orthogonality between the global vectors by utilizing the linear algebra property of the determinantal point process (DPP). In this paper, we explore the design of diverse global representation in the MIL model to model the diversity of instances in WSI. The main contributions are four-fold: (i) We introduce a new perspective on modeling the diversity of instances in WSI. (ii) We further propose a novel MIL aggregation model, termed DGR-MIL, to model diversity in MIL through a set of learnable global vectors. (iii) To learn a diverse global representation (vectors), we propose two main mechanisms: positive instance alignment and a novel diversity loss. (iv) Experimental results on two WSI benchmarks demonstrate the proposed DGR-MIL outperforms other competing MIL aggregation methods.

2. Related Work

2.1. Multiple instance learning in WSIs

MIL has been widely applied in many fields, e.g., pathology [28, 30, 47, 64], video analysis [2, 42], time series [14, 21]. In particular, the applications of the MIL in Whole Slide Image classification can be roughly summarized into two subcategories: i) instance-based MIL [22, 27, 60] and ii) bag embedding-based MIL. Instance-based methods typically require the propagation of the bag-level label to each of its instances to train the model. Consequently, the final bag-level prediction is obtained by aggregating instance-level predictions. However, empirical studies have proven its performance inferior to the embedding-based competitors because of the noisy instance-level supervision [54]. In contrast, bag-embedding-based methods start by projecting instances into feature embeddings and subsequently aggregate the information of these embeddings to obtain the bag-level prediction. Since the introduction of attention-based MIL (AB-MIL) [28], the prevailing applications of bag embedding-based MIL in WSI analysis have revolved around this framework. However, AB-MIL operates under the assumption that all instances within a bag are independent and identically distributed while failing to uncover inter-instance correlations. Therefore, numerous of its follow-up works centered around mitigating this limitation by taking advantage of non-local attention mechanism [30], transformer [47], pseudo bags [64], sparse coding [40], and low-rank constraints [58].

Most existing mainstream MIL methods have modeled correlations mainly through similarity between instances. However, they did not consider the variability of instances between and within bags. Conversely, clustering/prototype-based MIL employs attention scores for selecting prototypes [55, 61], potentially introducing noise and misleading model decisions [32,68]. Unlike attention-guided methods, PMIL [62] suggests a two-stage framework that first leverages clustering to identify reference prototypes and capture the sub-cluster representation among patch instances and prototypes. However, unrestricted optimization in prototype selection can easily lead to suboptimal outcomes, and a limited number of prototypes can result in a loss of diversity (limited by computational resources). In this paper, we explicitly model the diversity among instances in bag-embedding-based MIL through a learnable global representation. Although the proposed method falls into the category of transformer-based MILs, it differs from the previous transformer-based MILs [47,58] in two main aspects. First, we model the diversity between instances by comparing instances to the proposed global vectors via a cross-attention mechanism. Second, we propose a tokenized global vector to summarize the context information of positive instances.

2.2. Transformer

The transformer [51] has been widely applied in computer vision [9, 20, 33, 52], time series modeling [57, 67], and the natural language processing fields [18, 43, 44]. Standard transformers discover contextually relevant information by modeling the correlation between elements within a sequence through the self-attention mechanism. However, the traditional self-attention operation has quadratic time and space complexity $𝒪\left(n^2\right)$, with respect to a sequence containing $n$ elements. In the context of MIL, sequence length typically becomes quite large since one bag often approximately comprises ten thousand instances. This extremely long sequence poses significant computational intractability. Although [23,48,53] demonstrate that proper approximation of standard self-attention can reduce its quadratic complexity to linear, it still struggles to capture extremely long-term dependencies of context [6, 45, 58]. In contrast, the cross-attention mechanism [49,52], which was originally proposed to relate positions from one sequence to another, allows models to consider cross-sequence information. Inspired by this, we propose to model the diversity between and among instances through a cross-attention between instances and the proposed global vectors (see details in Section 3.1). This dramatically reduces the complexity compared to the self-attention mechanism (see Appendix C for details of model complexity) since the number of global vectors is significantly less than the sequence length.

3. Methods

The proposed DGR-MIL comprises two main parts: i) the design of the global representation in MIL pooling (Section 3.1), and ii) the strategy of learning diverse global representation (Section 3.2), where we further propose positive instance alignment and a computational-efficient diversity loss with a theoretical guarantee. The entire framework of DGR-MIL is depicted in Fig. 2.

Fig. 2:

Overview of the proposed DGR-MIL where the global vectors are used for modeling the diversity of instances. The diverse global vectors are learned through the positive instance alignment module and the diversity learning mechanism.

Preliminary.

Without loss of generality, we take binary MIL classification as an example: The objective is to predict the bag-level label $Y\in \{0,1\}$, given a bag of instances $\mathit{X}=\left\{{\mathit{x}}_1,{\mathit{x}}_2,\cdots ,{\mathit{x}}_n\right\}$, denoting a WSI with $n$ tiled patches. However, the corresponding instance-level labels ${\left\{y_i\right\}}_{i=1}^n$ are unknown in most WSI analyses due to the laboriousness of obtaining patch-level annotations. This turns the WSI classification into a weakly-supervised learning scheme according to the standard MIL formulation:

$$Y=\left\{\begin{array}{l}0,\text{iff}\sum _i{y}_i=0\\ 1,\text{otherwise.}\end{array}\right.$$ (1)

Because of the gigapixel resolution of WSIs, MIL typically cannot be performed in an end-to-end fashion [8, 34, 35] and instead necessitates a simplified learning scheme. This simplified MIL learning process comprises three main parts: i) a pre-trained feature extractor $f_{\mathit{proj}}(\cdot )$ that projects each instance into a $L$-dimensional vector, ii) a MIL pooling operator $\sigma (\cdot )$ that combines instance-level embeddings into a bag-level feature, and iii) a bag-level classifier $f_{cls}(\cdot )$ that takes the bag-level feature as input and produces the bag-level prediction as output. Mathematically, this process is given by

$$\begin{array}{c}\hat{Y}=f_{cls}\left(\sigma \left(\left\{{\tilde{\mathit{x}}}_1,\cdots ,{\tilde{\mathit{x}}}_n\right\}\right)\right),{\tilde{\mathit{x}}}_i\in {\mathbb{R}}^L\\ \text{with}\left\{{\tilde{\mathit{x}}}_1,\cdots ,{\tilde{\mathit{x}}}_n\right\}=f_{proj}\left(\left\{{\mathit{x}}_1,\cdots ,{\mathit{x}}_n\right\}\right),\end{array}$$ (2)

where $\hat{Y}$ denotes the predicted bag-level label. In the attention-based MIL (AB-MIL) [28] framework, the typical formulation for the MIL pooling operator is as follows:

$$\sigma \left({\tilde{\mathit{x}}}_i\right)=\frac{\mathrm{e}\mathrm{x}\mathrm{p}\left\{{\mathit{W}}^{\mathrm{T}}\left(\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{h}\left(\mathit{V}{\tilde{\mathit{x}}}_i\right)\right)\odot \mathrm{s}\mathrm{i}\mathrm{g}\mathrm{m}\left(\mathit{U}{\tilde{\mathit{x}}}_i\right)\right\}}{\sum _{i=1}^n\mathrm{e}\mathrm{x}\mathrm{p}\left\{{\mathit{W}}^{\mathrm{T}}\left(\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{h}\left(\mathit{V}{\tilde{\mathit{x}}}_i\right)\right)\odot \mathrm{s}\mathrm{i}\mathrm{g}\mathrm{m}\left(\mathit{U}{\tilde{\mathit{x}}}_i\right)\right\}},$$ (3)

where $\mathit{W}$, $\mathit{V}$, and $\mathit{U}$ are learnable parameters.

3.1. Global Representation in MIL Pooling

To accommodate the variability of the target lesions within and between bags, we develop a diverse global representation in the MIL pooling stage. Specifically, we define the global representation of the target (positive) instances as a set of learnable vectors given by $\mathit{G}=\left[{\mathit{g}}_1^{\mathrm{\top }},\cdots ,{\mathit{g}}_K^{\mathrm{\top }}\right]\in {\mathbb{R}}^{K\times L}$ with ${\mathit{g}}_k\in {\mathbb{R}}^L$ where $K$ is the number of global vectors. It is worth noting that a feed-forward network (FFN) is used to embed further both the input instance vectors $\tilde{\mathit{X}}={\left\{{\tilde{\mathit{x}}}_i\right\}}_{i=1}^n$ and the global vectors $\mathit{G}$ (see Fig. 2). However, we keep using $\mathit{G}\in {\mathbb{R}}^{K\times L}$ to denote global vectors for notation brevity.

Instance Correlation as Cross Attention.

The standard AB-MIL framework assumes the instances are independent and identically distributed while overlooking the correlation effect between instances. Hence, the self-attention mechanism becomes a natural choice for modeling the inter-instance correlation. However, due to the large number of instances within a bag in MIL, the quadratic time and space complexity $𝒪\left(n^2\right)$ of standard self-attention poses a significant challenge in computation. Alternatively, the previous transformer-based MIL [47] mitigates this problem by employing Nystrom-Attention [59], approximating the standard self-attention with linear complexity, which has proved effective of modeling correlation between positive and negative instances. It could be used to gather similar instances together by attention, benefiting from filtering background information. However, self-attention usage only guarantees the general separation of the positive and negative instances in a bag, which overlooks the diversity between instances and between bags.

Here, we implicitly model the diversity between instances by comparing the similarity between each instance vector and the proposed diverse global vectors. Specifically, this is achieved through a cross-attention mechanism where the global vector $\mathit{G}$ serves as queries, and a bag of instance vectors $\tilde{\mathit{X}}$ is used as key-value pairs. Formally, the $h$-th head of the proposed cross attention is given by

$$\begin{array}{c}{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_h(\mathit{G},\tilde{\mathit{X}})=\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}\left({\mathit{Q}}_h,{\mathit{K}}_h,{\mathit{V}}_h\right)\\ {\mathit{Q}}_h=\mathit{G}{\mathit{W}}_h^Q,{\mathit{K}}_h=\tilde{\mathit{X}}{\mathit{W}}_h^K,{\mathit{V}}_h=\tilde{\mathit{X}}{\mathit{W}}_h^V,\end{array}$$ (4)

where ${\mathit{W}}_h^Q$, ${\mathit{W}}_h^K$, ${\mathit{W}}_h^V\in {\mathbb{R}}^{L\times L/H}$ are learnable parameters for linear projections, where $H$ is number of heads. For the derivation purposes, we follow the traditional definition of the attention mechanism in the transformer (i.e., $\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}\left({\mathit{Q}}_h,{\mathit{K}}_h,{\mathit{V}}_h\right)=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left({\mathit{Q}}_h{\mathit{K}}_h^{\mathrm{\top }}/\sqrt{d_k}\right){\mathit{V}}_h$). The output of the yielding multi-head cross attention (MHCA) is the concatenation of the outputs from all heads through a linear projection:

$$\mathrm{M}\mathrm{H}\mathrm{C}\mathrm{A}(\mathit{G},\tilde{\mathit{X}})=\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}\left({\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_1;\cdots ;{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_H\right){\mathit{W}}^O,$$ (5)

where ${\mathit{W}}^O\in {\mathbb{R}}^{L\times L}$ is a trainable parameter. The proposed cross-attention mechanism reduces the quadratic time and space complexity $𝒪\left(n^2\right)$ in the standard self-attention mechanism to linear $𝒪\left(Kn\right)$ where $K\ll n$. In practice, we applied the Nystrom-Attention to the instance vectors and global vectors before performing the cross-attention (see Fig. 2) for two main reasons. First, applying self-attention to input instance vectors can facilitate filtering out the background. Second, applying self-attention to the global vectors can increase their discrepancies.

Tokenized Global Vector.

The vision transformer includes a class token to encode the globally discriminative representation associated with certain labels in image classification tasks. This token is typically added to the input token embedding by serving as a summary of the entire image. Building upon this inspiration, we propose to add a tokenized global vector ${\mathit{g}}_{\text{token}}$ as a summary of all the other global vectors. Now, the yielding global vectors can be denoted as $\tilde{\mathit{G}}=\left\{{\mathit{g}}_{\text{token}},{\mathit{g}}_1,\cdots ,{\mathit{g}}_K\right\}\in {\mathbb{R}}^{(K+1)\times L}$. The output of the tokenized global vectors after the cross-attention layer (Eq.(5)) is then used for bag-level classification. Following the convention in AB-MIL, the yielded importance score of each instance can be computed as

$$\sigma \left({\tilde{x}}_i\right)=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(\frac{({\mathit{g}}_{\text{token}}{\mathit{W}}_h^Q){({\tilde{\mathit{x}}}_i{\mathit{W}}_h^K)}^{\mathrm{\top }}}{\sqrt{d_k}}\right).$$ (6)

At first glance, adding the token to the global vectors instead of the input instance embedding appears counterintuitive. However, an in-depth analysis reveals its favorable properties. The proposed global vectors are learned in an unsupervised way (see details in Section 3.2), which poses a significant challenge in perfectly eliminating information from negative instances in the global vectors. This may be attributed to the similarity between positive instances and their adjacent negative instances, as tumor-adjacent regions typically exhibit high-density, quantitative expression in the spatial relationships of cells [24]. Each diverse global vector encapsulates a collection of analogous tissue features. As a result, certain global vectors emphasize certain types of positive instances. Accordingly, adding tokenized global vectors facilitates the model to capture the most discriminative global representation while suppressing the information from the negative instances (as evident in Fig. 5(b)).

Fig. 5:

Visualization of the attention map: (a) raw WSI with the ground-truth annotation, (b) the attention map computes using the tokenized global vectors, and (c-g) the attention map computes using the other ($K-1$) global vectors with $K=6$ in our experiment.

3.2. Learning Diverse Global Representation

Due to the weakly-supervised nature of MIL, how to learn the global representation of the target of interest remains an open problem. In this section, we introduce two strategies that can be used to learn a reliable and diverse global representation in MIL, respectively: i) positive instance alignment and ii) diversity learning via utilizing the linear algebra property of the DPP.

Positive Instance Alignment.

To enforce that the global representation aligns with the instances of interest (i.e., positive instances), we push the global vectors toward the positive bag centers but away from the negative bag centers. To do so, we first define the center of the positive and negative bags as ${\tilde{\mathit{x}}}_c^{\left(\mathit{pos}\right)}\in {\mathbb{R}}^L$ and ${\tilde{\mathit{x}}}_c^{\left(\mathit{neg}\right)}\in {\mathbb{R}}^L$, respectively. Similar to [25], the positive and negative centers are then updated in a momentum fashion at each training iteration:

$$\begin{gathered} {\tilde{\mathit{x}}}_c^{\left(pos\right)}=m{\tilde{\mathit{x}}}_c^{\left(pos\right)}+(1-m)\frac{1}{\left|{ℐ}_{\mathit{pos}}\right|}\sum _{i\in {ℐ}_{\mathit{pos}}}{\tilde{x}}_i \\ {\tilde{\mathit{x}}}_c^{\left(neg\right)}=m{\tilde{\mathit{x}}}_c^{\left(neg\right)}+\left(1-m\right)\frac{1}{\left|{ℐ}_{\mathit{neg}}\right|}\sum _{i\in {ℐ}_{\mathit{neg}}}{\tilde{x}}_i, \end{gathered}$$ (7)

where $m$ denotes the momentum update rate, which is set empirically to 0.4. ${ℐ}_{\mathit{pos}}$ and ${ℐ}_{\mathit{neg}}$ are the index sets of positive bags and negative bags, respectively. This indicates that the update of the positive instance center occurs only if a positive bag is fed into the network. The same strategy is applied to the negative center update (i.e., updated if and only if a negative bag is encountered). Up to now, we can formulate a set of triplet $\{\mathit{G},{\tilde{x}}_c^{\left(\mathit{pos}\right)},{\tilde{x}}_c^{\left(\mathit{neg}\right)}\}$. The triplet loss [3] is then adopted to enforce the global representation $\mathit{G}$ being close to the positive bag center while away from the negative bag center:

$${ℒ}_{tri}=\sum _{k=1}^K{\left[d_{+}\left(G_k,{\tilde{x}}_c^{\left(pos\right)}\right)-d_{-}\left(G_k,{\tilde{x}}_c^{\left(neg\right)}\right)+\mu \right]}_{+},$$ (8)

where $\mu$ is the margin parameter, and $d$ denotes the distance measure. We use cosine similarity as the distance measure.

Diversity Learning.

Although the positive instance alignment mechanism pushes the global representation to be aligned with the positive bag center, it is likely to result in a trivial solution where all the global vectors are identical. However, a diverse global representation is desired to capture the variability of positive instances. Hence, we propose our unique diversity loss inspired by DPP for data selection to maximize the diversity among global vectors and hence better summarize the instances. DPP is a well-known diversification tool [29] and is often used to select diverse subsets [10, 12, 13, 17, 50]. Inspired so, rather than use it for selection, we utilize it as a diversity measurement.

Mathematically, $𝒫$ is an L-ensemble DPP if the likelihood of an arbitrary subset $A\subseteq 𝒮$ drawn from the entire set $𝒮$ satisfies:

$${𝒫}_{\mathit{L}}\left(A\right)\propto \mathrm{d}\mathrm{e}\mathrm{t}\left({\mathit{L}}_A\right),$$ (9)

where ${\mathit{L}}_A$ denotes a submatrix of the similarity Gram matrix $\mathit{L}$ indexed by $A$. In the case of prompting diversity of global vectors $\mathit{G}=\left[{\mathit{g}}_1^{\mathrm{\top }},\cdots ,{\mathit{g}}_K^{\mathrm{\top }}\right]$, the similarity matrix is given as $\mathit{L}=\mathit{G}{\mathit{G}}^{\mathrm{\top }}\in {\mathbb{R}}^{K\times K}$, we simply set $A=ℬ=\left[K\right]$ and each global vector ${\mathit{g}}_i$, $i\in A$ is treated as a data point, and the total number of subsets can be calculated as $2^{\left|𝒮\right|}=2^K$. It is worth noting that the matrix $\mathit{L}$ is positive semi-definite.

Lemma 1.

( [29]) From a geometric perspective, the determinants in Eq.(9) can be interpreted as the squared $\left|A\right|$-dimensional volume spanned by its feature vectors:

$${𝒫}_{\mathit{L}}\left(A\right)\propto \mathrm{d}\mathrm{e}\mathrm{t}\left({\mathit{L}}_A\right)={\mathrm{V}\mathrm{o}\mathrm{l}}^2\left({\left\{{\mathit{g}}_i\right\}}_{i\in A}\right).$$ (10)

Lemma 1 immediately implies that a diverse subset is more likely to span larger volumes. This is because as the similarity between two data points (i.e., ${\mathit{L}}_{ij:i\ne j}$) increases, they will span fewer areas (see Fig. 3(a) and (b)), hence decreasing the probabilities of sets containing both of them (see Eq.(9)). Accordingly, feature vectors that are more orthogonal to each other span the largest volumes (see Fig. 3(a)), hence resulting in the most diverse subsets.

Fig. 3:

The similarity matrix for the global vectors $\mathit{G}$ learned from the CAMELYON16 dataset in two scenarios: (a) $\mathit{G}$ is orthogonal and (b) $\mathit{G}$ is non-orthogonal. To support Lemma 1 and Remark 1, we computed the area of the parallelogram corresponding to the two highly correlated global vectors. We omitted the diagonal elements in subpanel figure (b), as ${\mathit{L}}_{ii}=1,\forall i\in \left[K\right]$.

Theorem 1.

Given a set of global vectors $\mathit{G}=\left[{\mathit{g}}_1^{\mathrm{\top }},\cdots ,{\mathit{g}}_K^{\mathrm{\top }}\right]$ with $‖{\mathit{g}}_i‖=C,\forall i\in \left[K\right]$, maximizing the DPP-based diversity (i.e. $\mathrm{m}\mathrm{a}\mathrm{x}\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}\right.)$) results in orthogonal global vectors with ${\mathit{g}}_i\perp {\mathit{g}}_j,\forall i\ne j,i,j\in \left[K\right]$.

Proof.

The determinant $\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{L}\right)=\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}\right)$ is upper-bounded according to Hadamard's inequality [39]:

$$|\mathrm{d}\mathrm{e}\mathrm{t}(\mathit{L})|\stackrel{(a)}{=}\mathrm{d}\mathrm{e}\mathrm{t}(\mathit{L})\stackrel{(b)}{\le }\prod _{i=1}^K{\mathit{L}}_{ii}.$$ (11)

Condition (a) is fulfilled because the matrix $\mathit{L}$ is positive semi-definite. The equality of Condition (b) is achieved if and only if all non-diagonal entries of $\mathit{G}$ are zeros, meaning rows of the global vectors are orthogonal. The normalization constraint in Eq.(11) leads the upper bound to be the infimum, since ${\mathit{L}}_{ii}={\Vert \mathit{g}\Vert }_i^2\le C^2$ and it can be achieved if and only if the equality of Condition (b) is satisfied. This completes the proof.

According to Theorem 1, we propose a diversity loss ${ℒ}_{\mathit{div}}$ to diversify the proposed global vectors by minimizing the negative logarithm of $\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}\right)$:

$${ℒ}_{\mathit{div}}=-\mathrm{l}\mathrm{o}\mathrm{g}\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}\right),\text{s.t.}‖{\mathit{g}}_i‖=1=C.$$ (12)

Remark 1.

Theorem 1 implies that optimal diversity through minimizing our loss is theoretically achievable. This is because enforcing the constraints $‖{\mathit{g}}_i‖=1$ leads the infimum of ${ℒ}_{\mathit{div}}$ to reach zero due to $\mathrm{l}\mathrm{o}\mathrm{g}{\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}\right)}_{ii}=\mathrm{l}\mathrm{o}\mathrm{g}({‖{\mathit{g}}_i‖}^2)=0$. In contrast, the diversity loss ${ℒ}_{\mathit{div}}$ can be arbitrarily small (up to $-\mathrm{\infty }$) without the constraint $‖{\mathit{g}}_i‖=1$, which results in a unstable training.

We also add a small value $ϵ=1\times {10}^{-10}$ to prevent the logarithm of the determinant from being negative infinity (i.e. any two global vectors become collinear). The final diversity loss is given as

$${ℒ}_{\mathit{div}}=-\mathrm{l}\mathrm{o}\mathrm{g}\mathrm{d}\mathrm{e}\mathrm{t}\left(\mathit{G}{\mathit{G}}^{\mathrm{\top }}+ϵ\mathit{I}\right),$$ (13)

where $\mathit{I}$ denotes the identity matrix. It is noteworthy that the complexity to compute the loss is approximate $𝒪\left(L\right)$, which is negligible (see Appendix D).

3.3. Objective Function

The proposed MIL model is trained in an end-to-end fashion by jointly optimizing the weighted combination of cross-entropy (ce) loss that corresponds to the bag-level classification, triplet loss, and the proposed diversity loss:

$${ℒ}_{\mathit{final}}={ℒ}_{ce}+{\lambda }_{\mathit{tri}}{ℒ}_{\mathit{tri}}+{\lambda }_{\mathit{div}}{ℒ}_{\mathit{div}},$$ (14)

where ${\lambda }_{tri}$ and ${\lambda }_{tri}$ are balance parameters.

4. Experiments and Results

To validate the effectiveness of the proposed DGR-MIL, we conduct experiments on the CAMELYON16 dataset [5] and TCGA-lung cancer dataset (TCGA-NSCLC).

Dataset and Evaluation Metrics.

The two datasets are followed the experimental data partition setting in [64]. For the CAMELYON16, the training set is further divided into training and validation sets with a 9:1 ratio. We report the mean of accuracy, F1 score, and AUC with their corresponding 95% interval on the testing dataset after running five experiments. For the TCGA lung cancer dataset, we perform 4-fold cross-validation experiments, where the dataset is partitioned into training, validation, and testing sets with a patient ratio of 65:10:25. We report the mean and standard variation of accuracy, F1 score, and AUC on the testing dataset from 4-fold cross-validation.

Experiment Setup.

Three sets of instance features were extracted using different strategies to evaluate the proposed method’s adaptability across various feature embeddings. The first set provided by DTFD-MIL [64], employing OTSU’s method for patch extraction from WSIs and ResNet-50 for feature extraction, resulting in 1024-dimensional vectors per patch. For thorough validation, two additional sets of features were generated by segmenting each WSI into non-overlapping 224×224 patches using threshold filtering, resulting in 3.4 and 10.3 million patches from CAMELYON16 and TCGA lung cancer datasets [30,31,40,68], respectively. These patches were processed using ResNet-18 and Vision Transformer, pre-trained on ImageNet, to produce 512 and 768-dimensional feature vectors.

Baseline MIL Models.

We compare the proposed model to eight state-of-the-art MIL methods. These models can be roughly divided into two categories: i) AB-MIL [28] and its variants, including CLAM-SB [35], DS-MIL [30], and DTFD-MIL [64]; ii) the transformer-based methods including Trans-MIL [47] and ILRA-MIL [58]. iii) clustering/prototype-based MIL including PMIL [62].

Implementation Details.

All the models are trained using the parameter settings provided by [30, 35, 47, 58, 64]. (See Appendix B, including our method).

Additional Experiments.

We also include the experiments on using CTransPath [56] as feature extractor for CAMELYON16 dataset. Additionally, to validate the generalizability of our method on broader applications other than WSI, we conduct the experiment on MIL benchmark [1, 19]. Our method demonstrates the obvious superiority over other methods in both experiments. Please refer to Appendix F.

4.1. Experimental Results

The proposed method outperforms the other state-of-the-art MIL aggregation models by a large margin in both the CAMELYON16 and TCGA-NSCLC datasets using features extracted by three different means (see Table 1). We also show the statistical superiority of our method in Appendix E. Specifically, the proposed model outperforms the second-best models in terms of accuracy (1.7%; 1.3%), F1 score (3.1%; 1.5%), and AUC (1.1%; 1.7%) when using features extracted from ResNet-50 in CAMELYON16 and TCGA-NSCLC, respectively. A similar performance gain is observed on features extracted from ResNet-18 including accuracy (3.4%; 1.1%), F1 score (3.5%; 1.0%), and AUC (4.4%; 1.1%). We also observe an improvement in accuracy (3.4%; 1.1%), F1 score (3.5%; 1.0%), and AUC (4.4%; 1.1%) when using features extracted from the vision transformer. In general, the proposed model shows a greater performance improvement in the CAMELYON16 dataset compared to the TCGA-NSCLC dataset. This might be attributed to the fact that CAMELYON16 consists of more diverse instances than TCGA-NSCLC.

Table 1:

Main results on the CAMELYON16 dataset and TCGA-NSCLC dataset by using features extracted by different means. Our method statistically outperforms all other competitors (refer to the statistic test in Appendix E)

	CAMELYON16			TCGA-NSCLC
	Accuracy	F1	AUC	Accuracy	F1	AUC
	ResNet-50 ImageNet Pretrained
Classic AB-MIL (ICML’18)	0.845(0.839,0.851)	0.780(0.769,0.791)	0.854(0.848,0.860)	0.8690.032	0.8660.021	0.9410.028
DS-MIL (CVPR’21)	0.856(0.843,0.869)	0.815(0.797,0.832)	0.899(0.890,0.908)	0.8880.013	0.8760.011	0.9390.019
CLAM-SB (Nature Bio. Eng.’21)	0.837(0.809,0.865)	0.775(0.755,0.795)	0.871(0.856,0.885)	0.8750.041	0.8640.043	0.9440.023
CLAM-MB (Nature Bio. Eng.’21)	0.823(0.795,0.850)	0.774(0.752,0.795)	0.878(0.861,0.894)	0.8780.043	0.8740.028	0.9490.019
PMIL (MedIA’23)	0.831(0.799,0.863)	0.816(0.779,0.853)	0.845(0.813,0.876)	0.8730.010	0.8750.011	0.9330.007
Trans-MIL (NeurIPS’21)	0.858(0.848,0.868)	0.797(0.776,0.818)	0.906(0.875,0.937)	0.8830.022	0.8760.021	0.9490.013
DTFD-MIL (MaxS) (CVPR’22)	0.864(0.848,0.880)	0.814(0.802,0.826)	0.907(0.894,0.919)	0.8680.040	0.8630.029	0.9190.037
DTFD-MIL (MaxMinS) (CVPR’22)	0.899(0.887,0.912)	0.865(0.848,0.882)	0.941(0.936,0.944)	0.8940.033	0.8910.027	0.9610.021
DTFD-MIL (AFS) (CVPR’22)	0.908(0.892,0.925)	0.882(0.861,0.903)	0.946(0.941,0.951)	0.8910.033	0.8830.025	0.9510.022
ILRA-MIL (ICLR’23)	0.848(0.844,0.853)	0.826(0.823,0.829)	0.868(0.852,0.883)	0.8950.017	0.8960.017	0.9460.014
Our	0.917 (0.902,0.931)	0.913 (0.898,0.928)	0.957 (0.951,0.963)	0.908 0.015	0.911 0.018	0.963 0.008
	ResNet-18 ImageNet Pretrained
Classic AB-MIL (ICML’18)	0.805(0.772,0.837)	0.786(0.757,0.815)	0.843(0.827,0.858)	0.8740.005	0.8730.006	0.9370.001
DS-MIL (CVPR’21)	0.791(0.739,0.843)	0.776(0.712,0.840)	0.814(0.754,0.875)	0.8310.012	0.8380.008	0.8960.009
CLAM-SB (Nature Bio. Eng.’21)	0.792(0.769,0.815)	0.766(0.746,0.786)	0.811(0.777,0.845)	0.8690.010	0.8690.010	0.9310.006
CLAM-MB (Nature Bio. Eng.’21)	0.786(0.754,0.818)	0.770(0.746,0.795)	0.825(0.808,0.843)	0.8800.016	0.8800.016	0.9440.012
PMIL (MedIA’23)	0.800(0.775,0.825)	0.784(0.765,0.804)	0.829(0.807,0.851)	0.8560.006	0.8620.003	0.9330.010
Trans-MIL (NeurIPS’21)	0.839(0.822,0.856)	0.827(0.805,0.848)	0.854(0.823,0.886)	0.8770.009	0.8790.008	0.9380.014
DTFD-MIL (MaxS) (CVPR’22)	0.856(0.824,0.887)	0.792(0.742,0.842)	0.878(0.862,0.893)	0.8300.014	0.8210.020	0.8930.015
DTFD-MIL (MaxMinS) (CVPR’22)	0.833(0.807,0.858)	0.768(0.747,0.788)	0.878(0.872,0.883)	0.8530.012	0.8500.021	0.9250.013
DTFD-MIL (AFS) (CVPR’22)	0.817(0.791,0.843)	0.734(0.687,0.781)	0.868(0.841,0.896)	0.8700.007	0.8640.012	0.9350.010
ILRA-MIL (ICLR’23)	0.831(0.768,0.895)	0.819(0.768,0.871)	0.852(0.811,0.893)	0.8780.002	0.8790.001	0.9370.004
Our	0.873 (0.862,0.884)	0.862 (0.852,0.871)	0.898 (0.886,0.909)	0.891 0.029	0.890 0.021	0.955 0.023
	Vision Transformer ImageNet Pretrained
Classic AB-MIL (ICML’18)	0.851(0.837,0.865)	0.835(0.810,0.860)	0.873(0.840,0.906)	0.9040.011	0.9040.010	0.9530.013
DS-MIL (CVPR’21)	0.810(0.741,0.879)	0.806(0.742,0.869)	0.871(0.836,0.906)	0.8750.020	0.8790.016	0.9330.016
CLAM-SB (Nature Bio. Eng.’21)	0.839(0.831,0.847)	0.816(0.799,0.834)	0.864(0.841,0.887)	0.9070.008	0.9070.001	0.9540.014
CLAM-MB (Nature Bio. Eng.’21)	0.826(0.806,0.846)	0.804(0.795,0.813)	0.851(0.825,0.878)	0.9110.007	0.9110.007	0.9590.008
PMIL (MedIA’23)	0.843(0.831,0.856)	0.826(0.814,0.838)	0.843(0.820,0.867)	0.8820.009	0.8840.006	0.9400.006
Trans-MIL (NeurIPS’21)	0.862(0.841,0.883)	0.846(0.823,0.869)	0.860(0.848,0.873)	0.9090.009	0.9090.009	0.9530.006
DTFD-MIL (MaxS) (CVPR’22)	0.846(0.832,0.860)	0.767(0.746,0.787)	0.859(0.842,0.876)	0.9040.011	0.9040.010	0.9530.013
DTFD-MIL (MaxMinS) (CVPR’22)	0.839(0.826,0.851)	0.752(0.742,0.763)	0.862(0.836,0.888)	0.8950.013	0.8920.016	0.9520.011
DTFD-MIL (AFS) (CVPR’22)	0.831(0.818,0.844)	0.759(0.737,0.781)	0.880(0.864,0.897)	0.9010.005	0.9000.008	0.9590.012
ILRA-MIL (ICLR’23)	0.850(0.825,0.875)	0.838(0.812,0.865)	0.864(0.843,0.885)	0.9020.007	0.9040.007	0.9540.006
Our	0.893 (0.889,0.897)	0.882 (0.877,0.886)	0.891 (0.884,0.899)	0.926 0.008	0.925 0.008	0.969 0.004

We also observe the performance of the three sets of feature embeddings varied: the ViT feature embeddings outperform the ResNet-18 features but show inferior performance compared to the ResNet-50 features. This is mainly attributed to the fact that a greater number of positive instances is extracted by the ResNet-50 (provided by DTFD-MIL) as shown in Fig. 4(d). In contrast, a smaller portion of positive instances in the extracted patches may accompany a drop in performances [41]. This phenomenon benefits the pseudo-bag partitions in DTFD-MIL, as more positive instances within a bag are prone to result in less noisy pseudo-bag labels. This accounts for the drop in DTFD-MIL performance when applied to feature embeddings that contain a lower proportion of positive instances.

Fig. 4:

Ablation studies on (a) number of non-tokenized global vectors on both CAMELYON16 and TCGA-NSCLC datasets, (b) and (c) balance parameter ${\lambda }_{tri}$ and ${\lambda }_{div}$ on CAMELYON16 dataset, respectively. (d) Comparison in the number of positive instances per bag.

4.2. Ablation Studies

We conduct ablation studies on model design variants in the CAMELYON16 dataset with features extracted by a ResNet-50, unless specified otherwise.

Effectiveness of the Proposed Global Representation.

We ablate different components of the proposed model, i.e., the positive instance alignment module and the diversity loss. While the model without these two components serves as the baseline in Table 2. We first observe that incorporating the proposed global vectors described in Section 3.1 (without employing any of the learning strategies in Section 3.2) yielded an AUC of 0.922 and 0.928. This AUC exceeds that of most existing MIL models, except for DTFD-MIL (MaxMinS & AFS) (see Table 1 and 2). Subsequently, by including the proposed positive instance alignment module, we observe a performance gain of (2.2%, 2.8%) in accuracy, (2.3%, 2.9%) in F1 score, and (2.2%, 2.8%) in AUC. Up to now, we outperform the DTFD-MIL in terms of accuracy and F1 score (see Table 1 and 2), and achieve a similar AUC (AUC = 0.944,0.956) compare to the DTFD-MIL(AFS) (AUC = 0.946,0.951). Further incorporating the proposed diversity loss into the objective function yields a performance gain of (1.3%,0.7%) in AUC, which outperforms DTFD-MIL (AFS) by (1.1%,1.2%).

Table 2:

The ablation studies on different modules. $𝒫$: Positive instance alignment module. $𝒟$: Diversity loss.

$𝒫$	$𝒟$	CAMELYON16			TCGA-NSCLC
		Accuracy	F1	AUC	Accuracy	F1	AUC
✗	✗	0.895	0.887	0.922	0.872	0.875	0.928
✗	✓	0.906	0.900	0.938	0.896	0.896	0.952
✓	✗	0.917	0.910	0.944	0.900	0.904	0.956
✓	✓	0.917	0.913	0.957	0.908	0.911	0.963

Effectiveness of the Tokenized Global Representation.

As shown in Table 3, including the tokenized global vector ${\mathit{g}}_{\text{token}}$ yields a remarkable performance gain by improving accuracy by (1.0%, 0.5%), F1 score by (1.3%, 0.6%), and AUC by (2.2%, 0.6%). As consistent with the pathological findings that instances are diverse, we observe that different global vectors indeed corresponded to different instance representations, which can be depicted by the attention map produced by different global vectors in Fig. 5. However, we also observe that the learned global vectors still include non-tumor related representation, particularly around tumor boundaries, as positive instances around tumor boundaries have a similar appearance to surrounding negative instances (see Fig. 5.(c) and (d)). As a result, incorporating tokenized global vectors can mitigate this problem by capturing the most discriminative positive (tumor) regions (see Fig. 5.(b)).

Table 3:

The ablation studies on tokenized global representation.

$g_{\text{token}}$	CAMELYON16			TCGA-NSCLC
	Accuracy	F1	AUC	Accuracy	F1	AUC
✗	0.907	0.900	0.935	0.903	0.905	0.957
✓	0.917	0.913	0.957	0.908	0.911	0.963

Number of Global Vectors.

We find that the optimal number of global vectors $K$ in different data sets may vary due to dataset intrinsic properties. Specifically, the optimal $K$ for the CAMELYON16 and TCGA-NSCLC dataset are $K=5$ and $K=3$, respectively (Fig. 4.(a)). We observe that an overly large $K$ is likely to decrease performance as it will harden the learning task (see Fig. 4.(a)).

Loss Balance Hyperparameters.

By conducting a grid search, we find that the optimal setting of the balance parameters is ${\lambda }_{tri}=0.1$ and ${\lambda }_{div}=0.1$ (see Fig. 4.(b) and (c)). An overly small ${ℒ}_{\mathit{tri}}$ and ${ℒ}_{\mathit{div}}$ (e.g., 0.01) is likely to enforce inadequate constraints on the learned global representation by deviating it from learning meaningful information of instance of interest. While larger balance parameters (e.g., {0.5, 1.0}) distract the model from the main classification task, leading to a drop in classification performance.

5. Conclusion

Inspired by the pathological fact that instances are diverse, we propose a novel MIL model from the perspective of modeling diversity in instances through the cross-attention between instances and a set of learnable and diverse global vectors. To learn the global vectors, we propose a positive instance alignment mechanism and the DPP-driven diversity loss. Extensive experiments demonstrate that the proposed MIL model competed favorably against other existing MIL models. Importantly, our work provides an explicit way to account for the diversity in WSI. This pathology-driven approach is beneficial in capturing heterogeneity among the patient population. We also narrowed the performance gap between the diversity-drive MIL method and mainstream MIL.