SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yufei Cui $^{1\boxtimes}$ , Ziquan Liu $^{2}$ , Yixin Chen $^{3}$ , Yuchen Lu $^{4}$ , Xinyue Yu $^{4}$ , Xue Liu $^{1}$ , Tei-Wei Kuo $^{56}$ , Miguel R. D. Rodrigues $^{2}$ , Chun Jason Xue $^{3}$ , Antoni B. Chan $^{3}$ $^{1}$ Mila, McGill University  $^{2}$ University College London  $^{3}$ City University of Hong Kong  $^{4}$ Mila, Université de Montréal  $^{5}$ National Taiwan University  $^{6}$ MBZUAI
Multiple Instance Learning (MIL) is a crucial weakly supervised learning method applied across various domains, e.g., medical diagnosis based on whole slide images (WSIs). Recent advancements in MIL algorithms have yielded exceptional performance when the training and test data originate from the same domain, such as WSIs obtained from the same hospital. However, this paper reveals a performance deterioration of MIL models when tested on an out-of-domain test set, exemplified by WSIs sourced from a novel hospital. To address this challenge, this paper introduces the Retrieval-AugMented MIL (RAM-MIL) framework, which integrates Optimal Transport (OT) as the distance metric for nearest neighbor retrieval. The development of RAM-MIL is driven by two key insights. First, a theoretical discovery indicates that reducing the input's intrinsic dimension can minimize the approximation error in attention-based MIL. Second, previous studies highlight a link between input intrinsic dimension and the feature merging process with the retrieved data. Empirical evaluations conducted on WSI classification demonstrate that the proposed RAM-MIL framework achieves state-of-the-art performance in both in-domain scenarios, where the training and retrieval data are in the same domain, and more crucially, in out-of-domain scenarios, where the (unlabeled) retrieval data originates from a different domain. Furthermore, the use of the transportation matrix derived from OT renders the retrieval results interpretable at the instance level, in contrast to the vanilla  $l_{2}$  distance, and allows for visualization for human experts.
As the standard supervised learning paradigm, single instance learning has been the focus of machine learning research and its performance has been improved significantly since the advent of deep learning [1]. Nonetheless, a notable drawback of single instance learning lies in its reliance on a substantial volume of labeled data [2], which poses practical challenges as a result of the high expense and time-consuming nature of fine-grained data labeling. Consequently, multiple instance learning (MIL) [3] has gained increasing popularity, given its ability to be trained with limited supervision. This trend has particularly gained momentum within the domain of histopathology, where the analysis of medical images at a gigabyte scale, such as whole slide images (WSI), has emerged as a focal point in recent studies [4, 5, 6, 7, 8, 9, 10, 11].
Even though the performance of MIL on cancer diagnosis [13, 12] is remarkable, existing algorithms are only evaluated on the in-domain test set and have the risk of performance degradation when confronted with distribution shifts between the training and test sets [14, 15]. The distribution shift problem is particularly important in the context of automated medical diagnosis models, since the models are deployed across diverse hospitals in distinct regions, each may employing varying imaging
$^{1}$ E-mail: yufeicui92@gmail.com, yufei.cui@mila.quebec
Code can be found at https://github.com/ralphc1212/ram-mil.
Figure 1: The proposed RAM-MIL method. Left: An attention-based MIL model is pre-trained on the training set to produce feature representations and attention weights. Right: With pre-trained features and attention weights, Optimal-Transport (OT) is employed to compute bag-to-bag distance, based on which the nearest neighbor bag from the retrieval set is selected. A training bag's feature is then merged with its retrieved bag, which is used as the input feature to a bag classifier.
techniques. The robustness of such models in the face of distribution shift is of utmost importance within this safety-critical application, as any failure to address this challenge may result in severe consequences.
Figure 2: The comparisons of accuracy of MIL methods under in-domain and out-of-domain (OOD) settings. The out-of-domain is tested on Camelyon17 [12] with models trained on Camelyon16 [13]. MIL fails to generalize well to OOD data, while the proposed RAM-MIL improves the OOD performance as well as the in-domain performance.
This paper first unveils that the performance degradation indeed happens for state-of-the-art MIL models (Fig. 2) and proposes a retrieval-based MIL solution for improving both in-domain and out-of-domain performance. We prove a theorem that demonstrates the reduced input dimensions leads to improved MIL performance. Inspired by the connection between feature merging learning and intrinsic dimension [16], we propose the Retrieval-AugMented MIL (RAM-MIL) algorithm to learn a low intrinsic dimension feature space and enhance the model generalization, especially for out-of-domain data. In the retrieval process, we use the optimal transport (OT) as the distance metric between two bags because our main theorem is proved with the generalized OT measure. In our setting, the instances in a bag and their associated attention values in attention-based MIL are used to form a discrete distribution used to represent a bag, and OT is used to calculate the distance between the two discrete distributions of two
bags. Following retrieval, the original bags are merged with their retrieved counterparts from a retrieval dataset, employing a merge function such as convex combination, to generate feature representations. Subsequently, a classifier is trained as the MIL model. The retrieval dataset can either be the original dataset in the absence of out-of-domain data, or an unlabeled out-of-domain dataset in the out-of-domain adaptation setting. Our empirical study confirms that RAM-MIL effectively reduces the intrinsic dimensionality of the input feature space, surpassing existing MIL methods in both in-domain and out-of-domain settings. The advantage of OT distance over  $l_{2}$  and OT-variant distances is demonstrated by our ablation study. In addition, the OT-based retrieval result can be visualized using the transportation matrix, which makes the retrieval process interpretable. Our contributions are three-fold:
1. This work is among the first to investigate the out-of-domain performance for MIL, which is a vital issue for the application of MIL in automated medical diagnosis with WSIs. Our empirical study exposes the risk of recent MIL algorithms under distribution shifts.
2. Our theoretical result based on Wasserstein distance in the input space shows a negative correlation between MIL performance and the input data dimension, which motivates us
to propose the novel RAM-MIL framework based on the OT as the distance metric in the retrieval process.
3. Our experiment demonstrates that RAM-MIL outperforms state-of-the-art MIL methods in both in-domain and out-of-domain adaptation scenarios. Additionally, the transportation matrix of OT distance provides a tool to visualize the correspondence between original and retrieval instances, which makes our retrieval method interpretable.
2 Motivation of RAM-MIL
In this section, we show that why using OT could lead to a performance improvement for multiple instance learning.
We follow the standard formulation of Attention-based Multiple Instance Learning (MIL) [3, 4]. In MIL, the input is a bag of instances,  $\mathbf{X} = \{\pmb{x}_1,\dots ,\pmb{x}_K\} ,\pmb{x}_k\in \mathbb{R}^D,\mathbf{X}\in \mathcal{X}$  and  $K$  is the number of instances, which varies for different bags. There is a bag-level label  $Y$ , which is positive if at least one of the instances is positive, and negative if all instances are negative. We further assume the instances also have corresponding instance-level labels  $\{y_{1},\ldots ,y_{K}\}$ , which are unknown during training. There are  $N$  such bag-label pairs constituting the dataset  $\mathcal{D} = \{X_n,Y_n\}_{n = 1}^N$ . The objective of MIL is to learn an optimal function for predicting the bag-level label with the bag of instances as input. To this end, the MIL model should be able to aggregate the information of instances  $\{\pmb{x}_k\}_{k = 1}^K$  to make the final decision. A well-adopted aggregation method is the embedding-based approach which maps  $\mathbf{X}$  to a bag-level representation  $z\in \mathbb{R}^{d}$  and use  $z$  to predict  $Y$ . [3] extends the embedding-based aggregation approach by leveraging the attention mechanism, namely attention-based deep MIL (ABMIL). First, a transformation  $g(\cdot)$  computes a low-dimensional embedding  $h_k = g(x_k)\in \mathbb{R}^d$  for each instance  $x_{k}$ . The attention module aggregates the set of embeddings  $\{h_k\}_{k = 1}^K$  into a bag level embedding  $z = \sum_{k = 1}^{K}a_{k}h_{k} = \sum_{k = 1}^{K}a_{k}g(x_{k}) = \mathcal{A}(X)$ , where the attention  $a_{k}$  for the  $k$ -th instance is computed via a softmax function,
$$
a _ {k} = \frac {\exp \left\{\boldsymbol {m} ^ {T} \left(\tanh  \left(\boldsymbol {V} _ {1} ^ {T} \boldsymbol {h} _ {k}\right) \odot \operatorname {s i g m o i d} \left(\boldsymbol {V} _ {2} ^ {T} \boldsymbol {h} _ {k}\right)\right) \right\}}{\sum_ {j = 1} ^ {K} \exp \left\{\boldsymbol {m} ^ {T} \left(\tanh  \left(\boldsymbol {V} _ {1} ^ {T} \boldsymbol {h} _ {j}\right) \odot \operatorname {s i g m o i d} \left(\boldsymbol {V} _ {2} ^ {T} \boldsymbol {h} _ {j}\right)\right) \right\}}. \tag {1}
$$
where  $m, V_1, V_2$  are the network parameters. The bag embedding  $z$  is then mapped to the logits  $u$  with a feed forward layer with parameter  $W$  for the bag-level classification,  $u = W^T z$ .
Assuming there is a ground-truth set scoring function  $S(\mathbf{X}) : \mathcal{X} \mapsto \mathbb{R}$  that generates the label  $Y$ , our task is to approximate the function with the attention-based model. We have the following theorem demonstrating that, when employing the Wasserstein distance in the input space, the approximation error of a Lipschitz continuous set scoring function is upper bounded by a term associated with the input dimension  $d$ . The insight from the theorem offers a guidepost for the design of our methodology.
Theorem 1. Suppose  $S(\cdot)$  is a Lipschitz continuous function with respect to Wasserstein distance  $\mathcal{W}_p$  ( $1 \leq p < \infty$ ) and the Lipschitz constant is  $L$ . The bag of instances  $X$  is sampled from a probability distribution  $\mu(\boldsymbol{x})$  with distribution dimension  $d_{\mu}$  (intrinsic dimension). For any invertible map  $\Phi: \mathcal{X} \to \mathbb{R}^d$ ,  $\exists$  function  $\sigma$  and  $\gamma$ , such that for any  $X \in \mathcal{X}$ ,
$$
\left| S (\boldsymbol {X}) - \gamma \left(\Phi_ {\boldsymbol {X} \in \mathcal {X}} \left\{\sigma (\boldsymbol {x}): \boldsymbol {x} \in \boldsymbol {X} \right\}\right) \right| \leq O \left(L \cdot K ^ {- \frac {1}{d _ {\mu}}}\right). \tag {2}
$$
The theorem has two implications for the ABMIL problem. First, using the Wasserstein metric in the input space, we establish this dimension-dependent error bound for the approximation, indicating the good property of Wasserstein metric in the theoretical sense. Second, this dimension-dependent error bound motivates our RAM-MIL framework to reduce the intrinsic dimensionality of the input space for improving the performance of ABMIL, which is consistent with a previous work [16] that shows feature mixup is effective at reducing input feature dimension. This paper uses OT distance as the instantiation of the Wasserstein distance, since OT with Euclidean distance as the cost is equivalent to  $\mathcal{W}_1$ . We use the entropy-regularized OT as the approximation of OT distance in our implementation since the regularization makes the optimization process more stable and efficient and is widely used in machine learning and vision research [17, 18]. The incorporation of OT into Multiple Instance Learning (MIL) provides a new framework for understanding and assessing the similarity between different instances. Intuitively, each bag in MIL can be considered as a probability distribution of instances, where the probabilities are the attention values  $a$ . Thus, OT offers a rigorous and meaningful metric to quantify the similarity between two distributions. It takes into account not only the individual differences of instances but also their global arrangement within a bag, hence,
Algorithm 1 Retrieval-AugMented Multiple Instance Learning (RAM-MIL) Algorithm
1: Pre-train an MIL model on  $\mathcal{D}_o = \{\pmb {X}_n,Y_n\}_{n = 1}^{N_o}$    
2: Extract the sets of instance representations, i.e.,  $\mathcal{H}_o = \{\pmb {H}_n\}_{n = 1}^{N_o}$  and  $\mathcal{H}_r = \{\tilde{H}_m\}_{m = 1}^{N_r}$  , and compute their attentions, i.e.,  $\{\pmb {a}_n\}_{n = 1}^{N_o}$  and  $\{\tilde{\pmb{a}}_m\}_{m = 1}^{N_r}$    
3: Extract the sets of bag representations  $\{\pmb {z}_n\}_{n = 1}^{N_o}$  and  $\{\tilde{z}_n\}_{n = 1}^{N_r}$  with pre-trained model.   
4: Input:  $\{\pmb {X}_n,Y_n,\pmb {H}_n,\pmb {z}_n,\pmb {a}_n\}_{n = 1}^{N_o},\{\tilde{\pmb{X}}_m,\tilde{Y}_m,\tilde{\pmb{H}}_m,\tilde{\pmb{z}}_m,\tilde{\pmb{a}}_m\}_{m = 1}^{N_r}$    
5: for  $n$  from 1 to  $N_{o}$  do   
6:  $\nu^{*} = \arg \min_{\tilde{H}_{m}\in \mathcal{H}_{r}}d_{OT}(\mu_{n},\nu_{m})$  , where  $d_{OT}(\mu_n,\nu_m)$  is solved by (4)   
7: Store  $\hat{\pmb{z}}_n = \pi (\pmb {z}_n,\pmb {z}_{m^*})$  , where  $m^{*} =$  getIndex(  $\nu^{*}$    
8: end for   
9: Train a logistic regression bag classifier on  $\{\hat{z}_n,Y_n\}_{n = 1}^{N_r}$
providing a holistic comparison. Note that although Theorem 1 is proved for the input space  $\mathcal{X}$ , it is trivial to prove a similar result for the feature space  $\mathcal{H}$  (see the supplemental).
Building upon the benefits brought by OT, we propose a retrieval mechanisms to further strengthen the MIL framework. The key intuition is to leverage the rich information available across different bags. This allows for more robust and generalizable learning as each bag might not exist in isolation but is part of a broader set of data with shared and contrasting characteristics. The retrieval process, particularly when guided by OT, helps to identify and bring into context "neighbor" bags that are most relevant or similar to a given bag. This allows the model to learn more effectively by taking into account the broader context in which each bag exists. Empirically, we confirm the OT-guided retrieval could effectively reduce the intrinsic dimension of input feature space, which explains the performance of proposed methodology. Moreover, with the computed transport matrix, the correspondence between the original and selected bag can be visualized, which makes the retrieval and classification process more interpretable/explainable.
This section presents the RAM-MIL algorithm based on OT distance and its application in the out-of-domain adaptation.
3.1 Retrieval-Augmented Multiple Instance Learning based on Optimal Transport
We introduce the methodology that utilizes OT distance to address the retrieval problem for multiple instance learning (MIL). The idea of enhancing MIL using retrieval lies in augmenting a bag feature from the original dataset with its nearest neighbor bag feature from the retrieval dataset. Consider the original dataset of bags marked for model training and validation, denoted as  $\mathcal{D}_o = \{\pmb{X}_n, Y_n\}_{n=1}^{N_o}$  and the unlabeled retrieval set represented as  $\mathcal{D}_r = \{\tilde{\pmb{X}}_m\}_{m=1}^{N_r}$ , where  $N_o$  and  $N_r$  signify the number of bags in the two respective sets. For each instance  $\pmb{x}_k$  in a bag  $\pmb{X}$ , we extract an instance representation  $\pmb{h}_k = g(\pmb{x}_k)$ , where  $g(\cdot)$  is a neural net for feature extraction, such as ResNet50 [19]. Recall that an MIL model maps the bag of instance representations to the logits  $\pmb{u}$  for loss computation, expressed as  $\pmb{u} = \pmb{W}^T\pmb{z}$ . The intermediate representation  $\pmb{z} = \mathcal{A}(\pmb{X})$ , computed by averaging instance features and attention weights  $\pmb{z} = \mathcal{A}(\pmb{X}) = \sum_{k=1}^{K}a_k\pmb{h}_k = \sum_{k=1}^{K}a_kg(\pmb{x}_k)$ , is regarded as the bag representation.
The objective of the retrieval scheme is to select the bag in the retrieval set that is nearest to the input bag. To perform bag-to-bag retrieval, we use the distributional distance between two bags, which considers their constituent instances. Once the retrieved bag is obtained, its bag representation  $\tilde{z}^*$  is then merged with the original bag representation  $z$ . The merged feature, denoted as  $\hat{z} = \pi(z, \tilde{z}^*)$ , is used for bag classification, where  $\pi(\cdot, \cdot)$  is the merging function, e.g., a convex combination.
Retrieval with Optimal Transport. To leverage the probabilistic geometry between instances, instead of directly retrieving the nearest bags based on the bag representations  $\mathbf{z}$ , our retrieval operates on the bags of instance representations  $\mathcal{H}_r = \{\tilde{H}_m\}_{m=1}^{N_r} = \{\{\tilde{h}_{mk}\}_{k=1}^K\}_{m=1}^{N_r}$ . Here  $K$  represents the number of instances,  $\tilde{H}_m$  indicates the bag of instance representations and  $\tilde{h}_{mk}$  is an instance representation derived via  $\tilde{h}_{mk} = g(\tilde{x}_{mk})$ . For each bag in  $\mathcal{D}_o$ , we similarly extract bags of instance representations as  $\mathcal{H}_o$ .
For each bag  $\pmb{H} \in \mathcal{H}_o$ , we find its nearest neighbor in  $\mathcal{H}_r$  using the following OT problem to calculate distance between the two bags  $\pmb{H}$  and  $\tilde{\pmb{H}} \in \mathcal{H}_r$ . Let  $\mu = \sum_{i=1}^{K} a_i \delta_{\pmb{h}_i}$  and  $\nu = \sum_{j=1}^{\tilde{K}} \tilde{a}_j \delta_{\tilde{\pmb{h}}_j}$  be
discrete distributions supported on  $\pmb{H} = \{\pmb{h}_k\}_{k=1}^K$  and  $\tilde{\pmb{H}} = \{\tilde{\pmb{h}}_k\}_{k=1}^{\tilde{K}}$ , where  $\delta_{\mu} = \delta(\pmb{h} - \pmb{\mu})$  is the translated Dirac delta function. Here,  $a_i$  and  $\tilde{a}_j$  are the attention scores from the MIL model for instance  $i$  in bag  $\pmb{H}$  and instance  $j$  in bag  $\tilde{\pmb{H}}$ , respectively, forming the probability vectors for the bags. The OT problem between these two distributions can be formulated as follows:
$$
d _ {\mathrm {O T}} (\mu , \nu) = \min  _ {\boldsymbol {T} \in \mathcal {T} (\boldsymbol {a}, \tilde {\boldsymbol {a}})} \sum_ {i = 1} ^ {K} \sum_ {j = 1} ^ {\tilde {K}} c \left(\boldsymbol {h} _ {i}, \tilde {\boldsymbol {h}} _ {j}\right) T _ {i j} \tag {3}
$$
In this equation,  $T$  denotes the transport plan matrix where each element  $T_{ij}$  specifies the amount of "mass" to be transported from  $h_i$  to  $\tilde{h}_j$ . The function  $c(h_i, \tilde{h}_j)$  is a cost function that quantifies the cost of transporting a unit of mass from  $h_i$  to  $\tilde{h}_j$ . A common choice of  $c(h_i, \tilde{h}_j)$  is the squared  $l_2$  distance between the instance features, i.e.,  $c(h_i, \tilde{h}_j) = \| h_i - \tilde{h}_j\|_2^2$ .  $T(\mathbf{a}, \tilde{\mathbf{a}})$  represents the set of all matrices  $T$  that satisfy the marginals  $T^\top \mathbf{1}_K = \mathbf{a}$  and  $T\mathbf{1}_{\tilde{K}} = \tilde{\mathbf{a}}$ . Here,  $\mathbf{1}_n$  and  $\mathbf{1}_m$  are vectors of ones. We also impose an entropy regularization [20] term to reduce the sensitivity to outlier instances. Specifically, we solve the following OT problem with a Sinkhorn's algorithm [20]:
$$
d _ {\mathrm {O T}} (\mu , \nu) = \min  _ {\boldsymbol {T} \in \mathcal {T} (\boldsymbol {a}, \tilde {\boldsymbol {a}})} \sum_ {i = 1} ^ {K} \sum_ {j = 1} ^ {\bar {K}} c \left(\boldsymbol {h} _ {i}, \tilde {\boldsymbol {h}} _ {j}\right) T _ {i j} + \beta \cdot \sum_ {i j} T _ {i j} \log T _ {i j} \tag {4}
$$
$$
s. t., \boldsymbol {T} ^ {\top} \mathbf {1} _ {K} = \boldsymbol {a}, \boldsymbol {T} \mathbf {1} _ {\tilde {K}} = \tilde {\boldsymbol {a}}, \boldsymbol {T} \geq 0
$$
In the context of the retrieval-based MIL problem, the nearest neighbor bag for a given bag  $\pmb{H}$  would be the one  $\hat{\pmb{H}}$  in the retrieval set that results in the smallest OT cost with  $\pmb{H}$ :
$$
\nu^ {*} = \arg \min  d _ {O T} (\mu , \nu_ {m}) \tag {5}
$$
$$
\tilde {\boldsymbol {H}} _ {m} \in \mathcal {H} _ {r}
$$
In the retrieval process, we use an MIL model pre-trained on  $\mathcal{D}_o$  to extract the attention vectors  $a$  and  $\tilde{a}$ , which serve as mass probability distributions for calculating the OT, in addition to extracting bag representations  $z$  and  $\tilde{z}$ . This method effectively transforms the instance importance, as determined by the model, into a mass probability distribution. Instances with high attention scores are considered to possess larger "mass". In the context of the OT problem, such instances become more "expensive" to transport, thereby significantly influencing the computation of the transport cost. Consequently, the model is encouraged to pair similar instances carrying substantial mass, a strategy that aligns with the intuitive notion of matching similar instances together. Employing the pre-extracted attention to classify the weighted bag representation after retrieval ensures consistency with the retrieval process and enhances efficiency. The whole retrieval-enhanced multiple instance learning is shown in Algo. 1.
3.2 Unsupervised Domain Adaptation
The proposed methodology can be effectively employed in the context of unsupervised domain adaptation (UDA) [21]. UDA seeks to leverage labeled data from a source domain to train a model that can perform well on an unlabeled target domain.
Remarkably, our retrieval algorithm (as described in Algo. 1) can be seamlessly applied to the UDA task. To achieve this, we set  $\mathcal{H}_o$  to represent the source domain and  $\mathcal{H}_r$  to represent the target domain. Our approach retrieves the nearest neighbor bag representation from the target domain for each bag representation in the source domain. Subsequently, we train the final classifier using the merged representation and the source domain label.
The RAM-MIL based on OT allows leveraging the geometric structure of the instance space across domains. This promotes the alignment of the source and target domains at the instance level. By merging the bag features, the discriminative information embedded in the source domain is transferred to the target domain. As a result, the bag classifier is able to discriminate the target domain bags without needing access to their labels.
3.3 Reduction of Intrinsic Dimension
In this section, we validate our hypothesis that retrieval can effectively reduce the intrinsic dimension of each bag's distribution. We use the two NN estimator [22] for computing the intrinsic dimension for the bag representations, either the original features for MIL or merged features for retrieval models. Figure 3 shows the intrinsic dimension of CLAM [4], Manifold Mixup [16], OT based
retrieval and different variants of retrieval methods. The  $l_{2}$  distance based retrieval, approximate OT retrieval and Hausdorff distance based retrieval are variants of our retrieval method that will be elaborated in Section 5.
The figure shows that our proposed retrieval methods indeed lower the intrinsic dimensions of bag representations in both in-domain and out-of-domain settings. As shown in (2), the approximation performance of a ground-truth set scoring function is theoretically constrained by the intrinsic dimension of the distribution, when measured under the Wasserstein distance. RAMMIL delivers the lowest intrinsic dimensions among the methods compared, thus validating its superior performance in classification and aligns with our theoretical analysis (2). We offer more detailed discussion of intrinsic dimension in Section 5.2.
Figure 3: The intrinsic dimension of bag representations for different MIL methods and retrieval models. CLAM uses the original feature and the rest use merged feature from retrieval. RAM-MIL is the most effective method at reducing intrinsic dimension of feature space.