Retrieval-Augmented Multiple Instance Learning
Authors: Yufei CUI
Ziquan Liu
Yixin CHEN
Yuchen Lu
Xinyue Yu
Xue Liu
Tei-Wei Kuo
Miguel R. D. Rodrigues
Chun Jason Xue
Antoni B. Chan
Venue: NeurIPS 2023 poster (2023)
Keywords: Multiple Instance Learning
Whole Slide Imaging
Nearest Neighbor Retrieval
TL;DR: Augmenting multiple instance learning with optimal Transport-based retrieval achieves state-of-the-art performance in both in-domain and out-of-domain whole slide image classification.

Abstract: 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. *Code can be found at \url{https://github.com/ralphc1212/ram-mil*.

Retrieval-Augmented Multiple Instance Learning
Yufei Cuil Ziquan Liu?, Yixin Chen², Yuchen Lu", Xinyue Yu', Xue Liu', Tei-Wei Kuo Miguel R. D. Rodrigues", Chun Jason Xue", Antoni B. Chan
Mila, McGill University University College London "City University of Hong Kong Mila, Université de Montréal National Taiwan University MBBUAA
Abstract
Multiple Instance Learning (MIL) is a crucial weakly supervised learning method applied across various domains, e.g., medical diagnosis based 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 out-of-domain test set, exemplified by WSIs sourced from a novel hospital. To address this challenge, this paper introduces the 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 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 statc-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 l2 distance, and allows for visualization for human experts.
MIL (RAM-MIL) framework, which
Introduction
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, notable drawback of single instance learning lies in its reliance substantial volume of labeled data [2], which poses practical challenges as a result the high expense and time-consuming nature of fine-grained data labeling. Consequently, multiple instance learning (MIL) [3] has gained increasing popularity, given 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 MII. on cancer diagnosis [13, 12] 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 distinct regions, each may employing varying imaging
'E-mail: yufeicui92 @ gmail.com, yufei.cui @ mila.quehec Code can be found at httpp://ee //tttttbt./tt/////////// com//ralphc1212/ram-niii
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

OT
Wholi Stide Awnges
Pai: vise OT Compatzioon
D.
Computaioon
Founve
Attentios Network
X
Naighbor
Pre trained Attention Network
Altenton
Attectiin Helgfiiiig
Big Prediction Netwodk
Bag Prodiciim Network
Cost mustres C
Transpport
Bog
Classficttion
OT disuicc day L).
Cariplicttin
Figure I: 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, Optipal-Transport 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 s feature is then merged with its retrieved bag, which used the input feature to 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.
This paper first unveils that the performance degradation indeed happens for state-of-the-art MIL models (Fig. 2) and proposes a retrieval-basee MIL solution for improving both in-domain and out-of-domain performance. We prove theorem that demonstrates the reduced input dimensions leads to improved MIL performance. Inspired by the con- nection between feature merging learning and intrinsic di- mension [16], propose the
0.9
0.8
0.6
CLAM
TransMILL
BayesMIL
RAM MILL (Oun)
MILL (RAM-MIL) algorithm to learn low intrinsic dimension feature space and enhance the model generalization, es- pecially for out-of-domain data. In the retrieval process, the optimal transport (OT) the distance metric between two bags because our main theorem is proved with the generalized OT measure. In our setting, the in- stances in bag and their associated attention values in attention-bea based MIL used form discrete distribu- tion 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 retrieval dataset, employing merge function such as convex combination, to generate feature repre- sentations. Subsequently, 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 in the out-of-domain adaptation setting. Our empirical study confirms that RAM-MII. 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 l2 and OT-variant distances demonstrated by our ablation study. addition, the OT-based retrieval result can be visualized using the transportation matrix, which makes the retrieval process interpretable. Our contributions are three-fold:
In-dumain
Out-oo-dumain
Figure 2: The comparisons of accu- racy MIL methods under in-domain and out-of-domain (OOD) settings. The out-of-domain is tested on Camelyon [12] with models trained on Camelyon 16 [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.
dataset
1. This work is among the first to investigate the -ft---n--mmin 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 MII. 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-MII. outperforms state-of-the-art MII. methods in both in-domain and out-of-domain adaptation scenarios. Additionally, the transportation matrix of OT distance provides 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 J, E IR X E X' and K is the number of
MII,, the input is a bag of instances, X = (..,...., Fx), ) ay instances, which varies for different bags. There is bag-level label Y, which is positive if 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 {yi...... }. which are unknown during training. There are N such bag-label pairs constituting the dataset D = {Xn: Yn J==z- The objective of MIL to learn an optimal function for predicting the bag- lagelevel label with the bag of instances as input. To this end, the MIL model should be able to aggregate the information of instances {æx | to make the final decision. A well-adopted aggregation method is the embedding- based approach which maps X to a bag-level representation = E IR and use 2 to predict Y. [3] extends the embedding based aggregation approach by leveraging the attention mechanism, namely attention- based deep MII. (ABMIL). First, a transformation g(-) computes a low-dimensional embedding h4 g(xx)) E ie for each instance Zg. The attention module aggregates the set of embeddings (hx) ) into a bag level embedding x = Ek=1 h = Σk=n (*g(<x)) = A(X), where the attention aj for the k-th instance is computed via a softmax function,
exp (xp(mm m. (tanh( h. @sigmaidd V;" Ag,1] Σx (mp (tanh(V, h.sigiioi V. h,)))
where m, V1, V/2 are the network parameters. The bag embedding = is then mapped to the logits u with feed forward layer with parameter W for the bag-level classification, u = WT z.
Assuming there is ground-truth set scoring function S(X): X IR 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 eтror of a Lipschitz continuous set scoring function is upper bounded by 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(.) is a Lipschitz continuous function with respect to Wasserstein distance Wp (1 & p < oo) and the Lipschitz constant is L. The bag of instances X is sampled from a probability distribution μlæ²) with distribution dimension dy (intrinsic dimension). For any invertible map o. X + Rd,, 3 function a and 7/, such that for any X X,
SS(X) - 7 (xxex(((() I DE EXII S
K K.
(2)
The theorem has two implications for the ABMIL problem. First, using the Wasserstein metric in the input space, we establish this dimension-depennent error bound for the approximation, indicating the good property of Wasserstein metric in the theoretical sense. Second, this dimension-deeeendem 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 effective 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 W1. We use the 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 rigorous and meaningful metric to quantify the similarity between two distributions. takes into account not only the individual differences of instances but also their global arrangement within bag, hence,
OT as the approximation of OT distance in our
3

---

Algorithm Retrieval-AugMMented Multiple Instance Learning (RAM-MIL) Algorithm
1:
Pre-train an MIL model on D. = [X,, Yn Jm
2: Extract the sets of instance representations, i.e. H. [Hnnme, and H. { J and
compute their attentions, i.e., {a,, { and {ām m==1
N.N.
J 3: Extract the sets of bag representations [==] and (#nJn=1 with pre-trained model.
4: Input: (Xn: Yn; Hn; Zn, NN= (Xm; Ym, Hm, Rm, am JmJ m=
5: for from 1 to No do
6: = arg min Hm EH, dor(Hu: um), where dor (Myn mm) is solved by (4)
7: Store 2, = T (Z 2m-), where m = getIndex(u")
8: end for
9: Train a logistic regression bag classifier on {2n; Yn Y
providing holistic comparison. Note that although Theorem is proved for the input space X, it is trivial to prove a similar result for the feature space H.l (see the supplemental).
Building upon the benefits brought by OT, we propose 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- 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 retrievial and classification process more interpeeablllepplannhht.
3 Methodology
This section presents the RAM-MIL algorithm based on OT distance and its application in the out-of-domain adaptation.
3.1
Retrieval-Angmeented 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 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 D,- = (Xn, Yn J and the unlabeled retrieval set represented as D, {X ) where N. and N, signify the number of bags in the two respective sets. For each instance Ix in a bag X, we extract an instance representation h g(æx),, where g()) is neural net for feature extraction, such as ResNet50 [19]. Recall that an MIL model maps the bag of instance representations to the logits u for loss computation, expressed
W z. The intermediate representation z A(X), computed by K averaging instance features and attention weights z = A(X) Σk 0g hg = Ei= == ang(ars) is regarded as the bag representation.
The objective of the retrieval scheme is to select the bag in the retrieval set that is nearest 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 obtained, its bag representation 2" is then merged with the original bag representation z. The merged feature, denoted as & = T(2, "²), is used for bag classification. where π (-. :) 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 ((hmk kK J.N - Here K
retrieval
operates on the bags of instance representations H, {Hm MM represents the number of instances, Hn indicates the bag of instance representations and hink an instance representation derived via h = g(@mm)) For each bag in D.. we similarly extract bags of instance representations as Ho.
For each bag H Hog we find its nearest neighbor H, using the following OT problem to calculate
distance between the two bags H and H E Hr. Let μ ΣK azôn, and v Σ Σ äzōh, be

---

discrete distributions supported on H = {ha]k=r and H = [he]k=n where 6, = 8(h (4) is the translated Dirac delta function. Here, a and ūj are the attention scores from the MII. model for instance in bag H and instance j in bag H, respectively, forming the probability vectors for the bags. The OT problem between these two distributions can be formulated as follows:
K
doT(H,
min TET(a,
c(h TT
In this equation, T denotes the transport plan matrix where each element Tiy specifies the amount of "mass" to be transported from h; to h,. The function ch,,, h2) is cost function that quantifies the cost of transporting unit of mass from h; hj- A common choice of c(h,; h2j) is the squared l2 distance between the instance features, i.e., c(h,, h2) = l|h: 1 T(a, à) represents the set of all matrices T that satisfy the marginals T 1 a and T1 k ā. Here, 1, and 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]:
K
R
doT(U,))
min ΣΣ i=1j=1
che h2) + B
Σ Ti log Ty
TET (a,a)
s.t., T
s kk a, a T,T
In the context of the retrieval-based MIL problem, the nearest neighbor bag for given bag H would be the one H in the retrieval set that results in the smallest OT cost with H:
arg min dor (p,, Pm)
EH,,H,
In the retrieval process, we use MIL model pre-trained on DD to extract the attention vectors a and à, which serve as mass probability distributions for calculating the OT, in addition to extracting bag representations and 2. 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 source domain to train 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 H. to represent the source domain and H4- 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 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 'ss distribution. We use the two NN estimator [22] for computing the intrinsic dimension for the bag representations, either the original features for MII. or merged features for retrieval models. Figure 3 shows the intrinsic dimension of CLAM [4], Manifold Mixup [16], OT based
5

---

retrieval and different variants of retrieval methods. The l2 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 ap- proximation performance of a ground-truth set scoring function is theoretically constrained by
the intrinsic dimension of the distribution, when
CLAM
Mixup
L2 Retrival
Appix--TT-mi OTimin
Appror- OT svyg
Appos- T-max
Hausterff
RAM-MIL
measured under the Wasserstein distance. RAM-
MII. 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.
In-domain
Out-of-domain
Figure 3: The intrinsic dimension of bag repre- sentations for different MIL methods and retrieval models. CLAM uses the original feature and the rest use merged feature from retrieval. RAM-MIL the most effective method at reducing intrinsic dimension of feature space.
4 Related works
Iitention-based multiple instance learning ABMIL [3] first introduced the attention mech-
anism for MIL to an interpretation of importance for instances. Building on this, DSMIL [9] used contrastive learning for feature extraction and established global connections between instance atten- tions. TransMIL [5] took this a step further and developed a correlated MIL, employing multi-head self-attention and spatial information encoding for thorough global correlation. CLAM [4] extends ABMIL to the case of multiple classes and builds an integrated toolbox for visualizing the uncertainty. DTFD-MIL [23] proposes a two-stage feature distillation MIL framework for enhancing the perfor- mance. Bayes- MIL [11] studies the fundamental interpretability problem in the attention based MIL framework and proposed to address with probabilistic method. In our RAM-MIL, the attention weight is utilized as measure of the probability density, indicating the amount of "mass" being moved from the source bag to the target bag. Greater attention weights RAM-MIL indicate more significant relationships between instances in the bags being compared.
Domain adaptation for medical imaging There have been large volume of research studying the domain adaptation in medical imaging [24, 25, 26]. However, most of the preceding studies on whole slide images domain generalization have primarily focused on instance-level classification [27, 28, 29, 30, 31. 32, 33, 34, 35, 36, 37, 38, 39], rather than tackling the more complex challenge of weakly-supervisee bag-level classification (MIL). In contrast, our approach seeks address this more complex problem of unsupervised domain adaptation for MIL. We review these methods in details and make extensive comparisons in the supplemental.
Retrieval method for whole slide images Recent works also study how to efficiently retrieve relative WSI from database. Yottieel [40] builds search engine for indexing WSIs at scale. SISH [41] uses tree structure for fast search of WSI followed by an uncertainty-basee ranking algorithm for retrieval. In experiment, we integrate the open-sourced SISH as search engine into our RAM-MIL pipeline for comparisons (See supplemental). HHOT [42] proposes to use optimal transport (OT) as a distance measure to compare different WSIs, or different WSI datasets. Our paper, focusing on classification tasks. studies the principle why OT is suitable for WSI classification and propose retrieval-bai based classification process. We compare RAM-MIL and HHOT in our experiments.
5 Experiments
The proposed methodology is evaluated on whole slide image (WSI) datasets (Camelyon16 [43, 13]. Camelyon17 [12], TCGA-NSCLC, CPTAC-UCEC and CPTAC-LSCC) and general MIL datasets (See results in supplemental), The evaluation is divided into in-domain and UDA settings. For the in-domain setting, the retrieval set H, the union of training and validation set of Ho. For the UDA setting, the retrieval set is a held-out dataset with no labels involved in training or inference.

---

Table 1: Results on CAMELYONI6 and CAMELYONI7 for in-domain classification and unsuper- vised domain adaptation, for comparisons with related methods.
In-Domain (CAM16)
ofooof-Daain (CAMI7)
AUC
Accuracy
AUC
Accuracy
ABMIL
0.9010+0.026
0.8750++0.020
0.7287
0.7190+0.049
DSMIL
0.8944+00051
0.8682++00000
CLAM
0.9177=+0044
0.8650++0.060
0.7613+00054
0.7214+0.044
TransMIL
0.930711.0024
0.8837100001
0.569711.118
0.6451
Bayes-MIL
0.9432:10.049
0.8875-10.052
0.783910.004
0.7435=0.058
HHOT-ENN (k=1)
0.7007=0.035
0.7318+0.051
0.7173=00061 0.6939+00.066
HHOT-KNN (k=3)
0.7018+00026
0.7544=0.055
0.7263+0.048
0.7523+0.076
Mixup Retn
0.92601100051
0.802510.005
Mixup Retrio
0.9271=0.048
0.8850
0.7658+0.052
0.759410.044
Mixup Retro
0.9271
0.8825+0.048
0.7641
0.7593=0.058
12 Retn
0.0087
0.880
12 Retno
0.9241
0.848
0.7627+0.055
0.7353
l2 Retro
0.009
0.8950+00005
0.7038+00.00
0.7493+10.050
RAM-MIL Retr,
0,9451=100036
0.8025
RAM-MIL. Retro
0.9065++0.022
0.950
0.7974+0.054 0.7681
0.7433+0.003
RAM-MII. Retro
0.9419+0.008
0.9175+0.051
0.7795+0.021
We compare our OT-baaed retrieval with Euclidean distance retrieval, Hausdorff distance retrieval and Manifold MixUp as baselines. The Euclidean (l2 norm) distance di, (--)) and Hausdorff' distance based retrieval are implemented in the following ways:
Direct bag representation retrieval. As a baseline method, we propose an approach to augment bag features by directly retrieving the nearest neighbor. The setup follows the initial steps of Algo. I, from line 1 to line 3. Instead of computing the instance correlations, we directly com- pute the la distance between bag representations, formulated as di (Hn, Hm) = || - &m 12. For the n-th bag representation, we select the nearest neighbor from the retrieval set: Hm arg min HH. di2 (H,, Hm). The merged bag representation En = π (Zzn, km+) is subsequently employed for bag classification, where m² = getIndex (i Approximate OT retrieval. Recall that prior to applying the Sinkhorn's 's algorithm to solve OT, it is necessary to compute the cost matrix C = (ohh,,,)le each element is the squared 12 distance between instance representations. To highlight the significance of solving the OT problem, we introduce three simple approximations of OT by directly employing the minimum, average and maximum values of the cost matrix. The dCmin Hn, Hm) = (oin,)l(hh,, hy)lis focuses on the closest and most similar instances between bags. The dCavg Hm)) = avcrage,jlchh, hy)ls offers a holistic measure of bag similarity but might lose some finer details. The demax (,,, Hm) maxi [c(h,, hi)]is focuses the most dissimilar instances between the two bags. These three measures can be considered as OT approximations that use information from the squared l2 cost matrix.
||h, - 12. where
Hausdorff distance retrieval. The symmetric Hausdorff distance is employed for retrieval, computed as du HH, Hm) = max (H,, Hm,, di (Hm. Hn)). Here, (H,, Hm) = max H., m
e Fr.. ||; h is the directed Hausdorff distance. The Hausdorff distance is appropriate retrieval metric when assessing the maximum dissimilarity between two sets, emphasizing the instances that differ the most between the sets. In comparison, OT provides a more comprehensive and versatile choice for retrieval tasks, given its ability to consider all possible matchings between elements in the sets and finds an optimal one.
5.1 Classification of Whole Slide Image
Experimental setup: In this section, we validate the methodology within both in-domain and out-of-domain contexts. In both contexts, the training set is from Camelyon16. In the in-domain setting, we evaluate our models using the Camelyon16 dataset, where the combination of training and validation sets is treated as D,, which is denoted as Retry. For the out-of-domain context, we use Camelyon 16 as D.. As for D,, we consider two configuratioss: 1) retrieval from the training and validation sets of both Camelyon16 and Camelyon17, and 2) retrieval solely from the training
7

---

Table 2: Results on CPTAC-UCEC in domain classification, domain adaptation from CPTAC-LSCC to CPTAC-UCEC.
In Domain
Out-of-Domain
AUC
Accuracy
AUC
Accuracy
CLAM RAM-MIL
0,9500 0.9667
0.9285 0.9382
0.4986 0.6056
0.6521 0.6526
Table 3: Results on TCGA-NSCLC (subtyping).
Table 4: Tumor stage classification results on CAMELYONI7.
Method
AUC
Accuracy
CLAM
0,9420
0.8640
Method
AUC
Accuracy
Scaling VIT
0,9516
0.8821
CLAM
0,7803
0.60
TransMIL
0.9603
0.8835
Bayes-MIl MIL
0.8070
0.64
Bayes-MIl. 0.9451
0.8965
RAM-MIL
0.8139
0.65
RAM-MIL
0.9457
0.8988
and validation sets of Camelyon17, 7, referred to as Retrio and Retro, respectively. In the Camelyon 17 challenge, each slide labeled as either isolated tumour cells (ITC), micro-met, macro-met, or negative. ITC are strictly defined as single tumor cells or clusters smaller than 0.2 mm or less than 200 cells [44]. Following the principles of MII, we treated ITC as the positive class. More results about treating ITC the negative class are shown in the supplemental. To ensure robustness and avoid the influence of outliers, each experiment is executed 10 times on randomly partitioned train/ crinn/vlliatiinn/esst sets.
The baseline model we use for comparison CLAM, an extension of attention-based Multiple Instance Learning (MIL) that incorporates an instance clustering loss. In our case, the pre-trained model CLAM model for binary classification. We also make comparisons with other models such as ABMIL, DSMIL, TransMIL. Mayes-MIL, HHOT, and Mixup. For the ablation study, we add our proposed l2 retrieval method with direct bag representation retrieval. For both the Mixup and 12 retrieval models, we ensure that the source of the augmented data aligns with our in-domain and out-of-domain settings. For the OT-based pathology work. HHOT focuses on comparing whole slide images using uniform weight-bases OT and employs KNN discriminative method. Note that Mixup and our proposed methods use CLAM as the MIL method. Furthermore. Mixup shares the same merging function with retrieval method, which is simple addition convex combination, in the reported experiments. Our method does not need the ground-trus -tuth label information for retrieval set while Mixup uses that for the in-domain experiments. More results about the merging function is shown in the supplemental.
When we compute the RAM-MIL, we often deal with very large WSIs that can contain over 100,000 instances. This makes the calculation of OT time-consuning. To speed things up, we first perform pre-selection of potentially important instances. Instead of using all the instances in bag H.. pick the top 20% of instances based on their attention scores and normalized the attention scores. We provide more detailed comparisons in percentage of instances used, in the supplemental. In this way, we are focusing on the most important instances and making the whole process more efficient.
In-domain comparisons: Table presents the in-domain results. Simply using a uniform weight- based Optimal Transport (OT) in HHOT- KNNN fails to deliver satisfactory performance. Mixup. le retrieval, and RAM-MIL all enhance the classification performance for in-domain data beyond the CLAM baseline. The straightforward l2 retrieval- based MIL. outpaces Mixup, demonstrating that selecting the nearest neighbor using basic metric offers superior augmentation compared to random data selection. Using OT as a distance metric, the retrieval-based MIL method surpasses all the other compared methods for the in-domain dataset. Notably, the Retno configuration. which employs the most extensive data set for retricval (CAMELYONI6 and CAMELYONI7), offers the most significant boost in in-domain classification accuracy. Table 2 shows RAM-MIL consistently presents higher AUC and accuracy on CPTAC-UCEC. We also explore the performance of subtypes classification in Table 3, showing that RAM-MIL outperforms existing methods in accuracy, while providing AUC higher than the CLAM baseline. Compared with the results on tumor stage classification, our method enhances the model performance as shown in Table 4.
Out-of-domain comparisons: The right-hand side of Table showcases the results for the out-of- domain setting. Notably, all retrieval methods, including ours, enhance the classification performance
8

---

Table 5: Ablation study evaluated on CAMELYONI6 and CAMELYON17.
In- In-Domain (CAM16)
Out-of-Domand (CAM17)
AUC
Accuracy
AUC
Accuracy
Approx-OT-min Retr
0.930
0.8800++0.033
Approx-OT-min Retro
0.8601+00000
0.8000+0.062
0.7640+0.058
0.7353++0.043
Approx-OT-min Retto
0.9196110.03
0.8750:100056
0.768010.006
0.727411.006
Approx- Opprox-OT-ave Retri
0.9227=00.08
0.8600+0.036
Approx- av- Retrio
0.8865=00.03
0.8425+0.061
0.7625+0.060
0.7174+0.041
Approx-OT-ave Retro
0.9195==0.063
0.8750=00056
0.7641=0.054
0.7374+:0.061
Approx-OT-mas Retn
0.9174.10.049
0.8025-10.044
Approx-OT-max Retro
0.9182.10.047
0.8650110.048
0.757610.000
0.7214110.062
Approx- OT-max Retro
0.9195=00063
188750=0.056
0.7681
0.7414+00.088
12 Retn
0.908110.004
0.8800-L0.055
12 Retro
0,9241+00048
0.8950=10.048
0.762
0.7353+10.051
12 Retro
0.9398++0.045
0.8950++0.055
0.700
0.7493++0.050
Hausdorff Retn
0.9273:1.005
0.8850=0.004 049
Hausdorfi Retro
0.9226+00.05
0.8725+00060
0.7651
0.7434+0.051
Hausdorff Retro
0.9322++0.046
0.8950+0.048
0.059
0.7513++0.051
RAM-MIL Retn
0.9451:1.0036
0.8925=0.0000
RAM-MIL Retro
0.9365110.052
0.9200-0.050
0.797410.054
0.7433110.073
RAM-MIL Retro
0.941910,048
0.9175110005
0.768110.058
0.7795
for out-of-domain data, without accessing its labels. This suggests that retrieval benefits from the transfer of representative information from the source to the target domain via the merging of bag features. The l2 retrieval, while trailing Mixup in terms of accuracy, surpasses it in terms of AUC. Our OT-based retrieval achieves the highest performance across both AUC and accuracy metrics. Table 2 presents the results on domain adaptation from CPTAC-UCEC to CPTAC-LSCC. RAM-MIL demonstrates a notable improvement of 1065 in AUC compared to CLAM. The underlying reason for the good performance might be it could consider all possible matchings between elements in the sets and finds an optimal one.
A blation study Table presents the ablation study using different retrieval methods, which shows that even l2 retrieval, which operates without considering instance-level distances, surpasses the performance of the three OT approximation methods in both in-domain and out-of- settings. Among these, only the Approx-OT-min when used for in-domain data retrieval, achieves perfor- mance comparable to that of l2 retrieval. This underscores the fact that relying on simple statistics of the cost matrix cannot provide a robust and consistent approximation of the optimal transport. This outcome reinforces the necessity of solving the complete optimal transport problem, using Sinkhorn's ss algorithm, to attain reliable results.
5.2 Dimensionality Reduction
We next show that retrieval-basee methods can effectively reduce the intrinsic dimension. Figure shows the intrinsic dimension for various methods. The proposed RAM-MIL based optimal transport delivers the lowest intrinsic dimension, among all methods, both retrieval and vanilla models. Interestingly, Mixup seems to increase the intrinsic dimension, suggesting that random selection and merging of bag representations do not necessarily lead to a more compact representation. This could explain why Mixup only offers marginal improvements.
Figure 3 shows that while the CLAM model has an intrinsic dimension between 5-7, RAM-MIL achieves dimension between 2-4. Based on these observations, we choose dimension 5 to guide manual dimensionality reduction. The objectives of this exercise are two-fold: firstly, to validate that intrinsic dimension can accurately explain performance, and secondly, to illustrate that simple methods like SVD are not necessarily beneficial for dimensionality reduction of bag representations.
In our experiment, we employ SVD to reduce each bag representation of CLAM and each merged bag representation of RAM-MIL to a 5-dimensional vector. Then, we use these vectors and the bag labels to train single logistic regression classifier. The results indicate that RAM-MIL significantly outperforms CLAM classification accuracy. This experiment shows that a simple dimensionality reduction method like SVD does not significantly improve the performance of the CLAM baseline.

---

Table 6: Classification on CAMELYONI6 16 with bag representations of reduced dimensions (5-dim).
In Ho-Doominn (CAMI6)
Out-of- of-Domain (CAM17)
AUC
Accuracy
AUC
Accuracy
CLAM RAM-MII.
0.9216-0.057 0.927
0.6625+0.177 0.8475+0.157
0.7486++0.04D 0.7509+0.074
0.4991=0.1 0.7053-0.113
However, RAM-MII, which organizes the latent space by merging nearest neighbor features and then applies SVD, can achieve a commendable performance.
5.3 Interpretability with OT for MIL.
We show that by using optimal transport retrieval, novel method for highlighting the important instances could be derived. For pair of source bag H and target bag H, solving the optimal transport problem generates the transport matrix T = [Ti]kj. Suppose that we are interested in the instance labels of the source bag, and suppose that the instance labels of the target bag known. The positive instances the source bag are simply the instances transported to the positive instances in the target bag, ie., 0 =arg max, Tiy- Here, is the predicted instance label for the source bag, and and is the true instance label in the target bag. Figure 4 shows this method could provide good coverage the positive area, without model training or accessing the bag label. In contrast to OT, retrieval-basee methods utilizing other distance metrics do not possess this specific property.
Ground-truth
RAM-MIL
Ground-truth
RAM-MIL
Precision: 0.7530
Precision: 0.7222
FROC: 0.4964
FROC: 0.7303
Figure 4: The visualization of predicted positive instances (red regions) in whole slide images for RAM-MIL. FROC is defined as average sensitivity (recall) at 6 predefined false positive (FP)
rates: 1/4, 1/2, 1, 2, 4 and FPs. For reference, the averaged precision and FROC in Bayes-MIL. [11] (one of SOTA method designed for instance-level performance) are 0.8107 and 0.4919.
6 Conclusion
In this work, investigate the out-of-domain performance of multiple instance learning model, which is crucial for the ML-assisted medical diagnosis with the whole slide images. An optimal transport based retrieval method is proposed based on the understanding that there exists a negative correlation between MIL performance and input data dimension. The proposed RAM-MIL framework works for both in-domain and out-of-domain, meanwhile outperforms of-t--ofthhe-rr MIL methods. We hypothesize that OT-based retrieval could reduce the intrinsic dimension therefore improves the MIL performance and validate the hypothesis by computing the intrinsic dimension of bag representations. The transportation matrix of OT distance provides tool to visualize the correspondence between limitation the
original and retrieval instances and makes our retrieval method interpretable. efficiency of optimal transport solver is constrained by the number of instances within a bag. Although we have mitigated this problem by selecting partial instances based on high attention values, we anticipate that a more principled method will be sought in future works.
Acknowledgements
This work was supported by grant from the Research Grants Council of the Hong Kong Special Administrative Region, China (Project No. CityU 11215820), and a grant from Guangzhou Bingli Technology Co., Ltd.
10

---

References
[1] Alex Krizhevsky, Πla Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012.
[2] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248 245-255. leee, 2009.
[3] Maximilian Ilse, Jakub Tomezak, and Max Welling. Attention-based deep multiple instance learning. In International conference on machine learning, pages 2127-2136. PMLR, 2018.
[4] Ming Y Lu, Drew FK Williamson, Tiffany Y Chen, Richard J Chen, Matteo Barbieri, and Faisal Mahmood. Data-efficient and weakly supervised computational pathology on whole-slide images. Nature biomedical engineering, 5(6):555-570, 2021.
[5] Zhuchen Shao, Hao Bian, Yang Chen, Yifeng Wang, Jian Zhang, Xiangyang Ji, et al. Transmil: Transformer based correlated multiple instance learning for whole slide image classification. Advances in neural information processing systems, 34:2136-2144 2021.
[6] Chao Tu, Yu Zhang, and Zhenyuan Ning, Dual-curriclumm contrastive multi-instance learning for cancer prognosis analysis with whole slide images. Advances in Neural Information Processing Systems, 01224444224999 2022.
[7] Linhao Qu, Manning Wang, Zhijian Song, et al. Bi-directional weakly supervised knowledge distillation for whole slide image classification. Advances Neural Information Processing Systems, 35:13388-55111 153811, 2022.
[8] Syed Ashar Javed, Dinkar Juyal, Harshith Padigela, Amaro Taylor- Taylor-Weiner, Limin Yu, and Aa- ditya Prakash. Additive mil: Intrinsically interpretable multiple instance learning for pathology. Advances in Neural Information Processing Systems, 35:20689-207777 2022.
[9] Bin Li, Yin Li, and Kevin W Eliceiri. Dual-stream multiple instance learning network for whole slide image classification with self-supervised contrastive learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 14318-14328, 2021.
[10] Jinxi Xiang and Jun Zhang. Exploring low-rank property in multiple instance learning for whole slide image classification. In The Eleventh International Conference on Learning Representations, 2023.
[11] Yufei CUI, Ziquan Liu, Xiangyu Liu, Xue Liu, Cong Wang, Tei-Wei Kuo, Chun Jason Xue, and Antoni B. Chan. Bayes-MIL: A new probabilistic perspective on attention- based multiple instance learning for whole slide images. In The Eleventh International Conference Learning Representations, 2023.
[12] Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al. From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 17 challenge. IEEE transactions on medical imaging, 38(2):550-560, 2018
[13] Babak Ehteshami Bejnordi, Mitko Veta, Paul Johannes Van Diest, Bram Van Ginneken, Nico Karssemeijer, Geert Litjens, Jeroen AWM Van Der Laak, Mcyke Hermsen, Quirine F Manson, Maschenka Balkenhol, et Diagnostic assessment of deep learning algorithms for detection of lymph node metastases in women with breast cancer. Jama, 318(22):2199-22) 2017.
[14] Joaquin Cuindeero-Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008
[15] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga Richard Lanas Phillips, Irena Gao, et al. Wilds: benchmark of in-the-wild distribution shifts. International Conference on Machine Learning, pages 5637-5664. PMLR, 2021.

---

[16] Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najañi Ioannis Mitliagkas, David Lopez- Paz, and Yoshua Bengio. Manifold mixup: Better representations by interpolating hidden states. In International conference on machine learning, pages
PMLR, 2019.
[17] Jia Wan, Ziquan Liu, and Antoni B Chan. A generalized loss function for crowd counting and localization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1974-1983, 2021.
[18] Ziquan Liu, Yi Xu, Yuanhong Xu, Qi Qian, Hao Li, Xiangyang Ji, Antoni Chan, and Rong Jin. Improved fine-tuning by better leveraging pre-training data. Advances in Neural Information Processing Systems, 35:32568- 32581, 2022.
[19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[20] Marco Cuturi.
Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[21] Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation.
In International conference on machine learning, pages 1180-1189. PMLR, 2015.
[22] Elena Facco. Maria d' Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. Scientific reports, 7(1):12140, 2017.
[23] Hongrun Zhang, Yanda Meng. Yitian Zhao, Yihong Qiao, Xiaoyun Yang, Sarah Coupland, and Yalin Zheng. Difd-mil: Double-tier feature distillation multiple instance learning for histopathology whole slide image classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18802-18812 2022.
[24] Gabriele Campanella, Matthew G Hanna, Luke Geneslaw, Allen Miraflor, Vitor Werneck Krauss Silva, Klaus J Busam, Edi Brogi, Victor E Reuter, David S Klimstra, and Thomas J Fuchs. Clinical- grade computational pathology using weakly supervised deep learning whole slide images. Nature medicine, 25(8): 101-)1309999 2019.
[25] Pierre Courtiol, Charles Maussion, Matahi Moarii, Elodie Pronier, Samuel Pilcer, Meriem Sefta, Pierre Manceron, Sylvain Toldo, Mikhail Zaslavskiy, Nolwenn Le Stang, et al. Deep learning-based based classification of mesothelioma improves prediction of patient outcome. Nature medicine, 25(10):1519-1525 2019.
[26] Milad Sikaroudi, Maryam Hosseini, Ricardo Gonzalez, Shahryar Rahnamayan, and HR Tizhoosh. Generalization of vision pre-trained models for histopathology. Scientific reports, 13(1):6065, 2023.
[27] Marc Aubreville, Nikolas Stathonikos, Christof A Bertram, Robert Klopfleisch, Natalie Ter Ho- eve, Francesco Ciompi, Frauke Wilm, Christian Marzahl, Taryn A Donovan, Andreas Maier, et al. Mitosis domain generalization in histopathology images- theges-theaae midog challenge. Medical Image Analysis, 84:102699, 2023.
[28] Jiawei Yang, Hanbo Chen, Jiangpeng Yan, Xiaoyu Chen, and Jianhua Yao. Towards better understanding and better generalization of low-shot classification in histology images with contrastive learning. In International Conference on Learning Representations.
[29] Jian Ren, IIker Hacihaliloglu, Eric A Singer, David J Foran, and Xin Qi. Unsupervised domain adaptation for classification of histopathology whole-slide images. Frontiers in bioengineering and biotechnology. 7:102, 2019.
[30] Jian Ren, Ilker Hacihaliloglu, Eric A Singer, David Foran, and Xin Qi. Adversarial domain adaptation for classification of prostate histopathology whole-slide images. In Medicul Image Computing and Computer Assisted Intervention 2018: 21 International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part II 11, pages 201-20% Springer, 2018.
12

---

[31] Hyungjoo Cho. Sungbin Lim, Gunho Choi, and Hyunseok Min. Neural stain-style transfer learning using gan for histopathological images. arXiv preprint arXiv:1710.08543, 2017.
[32] M Tarek Shaban, Christoph Baur, Nassir Navab. and Shadi Albarqouni. Staingan: Stain style transfer for digital histological images. In 2019 leee 16th international symposium on biomedical imaging (Isbi 2019), pages 953-956. IEEE, 2019.
[33] Farhad G Zanjani, Svitlana Zinger, Babak E Bejnordi, Jeroen AWM van der Laak, al. Histopathology stain-color normalization using deep generative models. In Medical Imaging with Deep Learning, 2022.
[34] Maxime W Lafarge, Josien PW Pluim, Koen AJ Eppenhof, and Mitko Veta. Learning domain- invariant representations of histological images. Frontiers in medicine, 6:162, 2019.
[35] Ozan Ciga, Jianan Chen, and Anne Martel. Multi-layer domain adaptation for deep convolutional networks. In Domain Adaptation and Representation Transfer and Medical Image Learning with Less Labels and Imperfect Data: First MICCAI Workshop, DART 2019, and First International Workshop, MIL3ID 2019, Shenzhen, Held Conjunction with MICCAI 2019, Shenzhen, China, October 13 and 17, 2019, Proceedings I, pages 20- 27-27.05 Springer, 2019.
[36] Jian Ren, Jianchao Yang. Ning Xu, and David J Foran. Factorized adversarial networks for 2018.
unsupervised domain adaptation. arXiv preprint
[37] Nicolas Brieu, Armin Meier, Ansh Kapil, Ralf Schoenmeyer, Christos Gavriel, Peter D Caie, and Günter Schmidt. Domain adaptation- based augmentation for weakly supervised nuclei detection. arXiv preprint arXiv: /907. 0900999004400 04681, 2019.
[38] Ansh Kapil, Tobias Wiestler, Simon Lanzmich, Abraham Silva, Keith Steele, Marlon Rebelatto, Guenter Schmidt, and Nicolas Brieu. Dasgan-joint domain adaptation and segmentation for the analysis of epithelial regions in histopathology pd-11 images. arXiv preprint 2019.
[39] Michael Gadermayr, Laxmi Gupta, Vitus Appel, Peter Boor, Barbara M Klinkhammer, and Dorit Merhof. Generative adversarial networks for facilitating stain-independent supervised and unsupervised segmentation: study kidney histology. IEEE transactions medical imaging, 38(10):22332302, 2019.
[40] Shivam Kalra, Hamid R Tizhoosh, Charles Choi, Sultaan Shah, Phedias Diamandis, Clinton JV Campbell, and Liron Pantanowitz. Yottixel-an image search engine for large archives of histopathology whole slide images. Medical Image Analysis, 65:101757, 2020.
[41] Chengkuan Chen, Ming Lu, Drew FK Williamson, Tiffany Chen, Andrew Schaumberg, and Faisal Mahmood. Fast and scalable search of whole-slide images via self-supervised deep learning. Nature Biomedical Engineering, 6(12):1420-1433 2022.
[42] Anna Yeaton, Rahul G Krishnan, Rebecca Mieloszyk, David Alvarez-Melis, and Grace Huynh. Hierarchical optimal transport for comparing histopathology datasets. arXiv preprint arXiv:2204.08322 2022.
[43] Geert Litjens, Peter Bandi, Babak Ehteshami Bejnordi, Oscar Geessink, Maschenka Balkenhol, Peter Bult, Altuna Halilovie, Meyke Hermsen, Rob van de Loo, Rob Vogels, et al. 1399 h&c-stained sentinel lymph node sections of breast cancer patients: the camelyon dataset. GigaScience, 7(6):ii0065, 2018
[44] Neill Patani and Kefah Mokbel. Clinical significance of sentinel lymph node isolated tumour cells in breast cancer. Breast cancer research and treatment, 127:325-334, 2011.
[45] Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance. arXiv preprint arXiv:1707.0008) 2017.
[46] Thomas G Dietterich, Richard H Lathrop, and Tomás Lozano-Pérez. Solving the multiple instance problem with axis axis-parallel rectangles. Artificial intelligence, 391-))311-11 1997.
13

---

[47] Stuart Andrews, Ioannis Tsochantaridis, and Thomas Hofmann. Support vector machines for multiple-instance learning. Advances neural information processing systems, 15. 2002.
[48] Yuchen Yang, Amir Akbarnejad, Nilanjan Ray, and Gilbert Bigras. Double adversarial domain adaptation for whole-siideimmeeedsssiicaaao In Medical Imaging with Deep Learning, 2021.
14

---

A Proof
Proof. Since the bag of instances X is sampled from the probability distribution fa(z), we have the upper bound for Wasserstein distance between X and μ [45],
(X, [4] & K "
(6)
where K' is the number of samples. Next we define the function (2) as a small perturbation function (x) = + 8 and let X {a(z) z E } Using the triangle inequality, we have
2 X X)) ] E[W, |+) E[W, (,), S 2K d C,
(7)
where C is a constant as result of the perturbation. As the function S(-) is Lipschitz continuous, we have
[S(X) s(x) I. (X, X)] S O K ).).
(8)
Similar TransMIL [5], let X - R be any invertible map, where its inverse mapping is expressed as $ R + X' Then we have:
$ 1 (xxex((o(x)) a x())) = S( ( x - S())
(9)
Let So As [S(X) S(X) ( K )) have
( - ((xxex(((x)) : x x s (OLKK((+)))))
(10)
In this proof, the transformation ((-) A(-). This proof could be easily extended to the represen- tations H by assuming a probability measure over the instance representations h and replacing X with H. In this case, the transformation Φ (H) k== aghs.
Table 7: Results on general MIL datasets. Experiments were run 5 times and the average classification accuracy (# a standard error of a mean) is reported.
Method
MUSKI
MUSK2
FOX
TIGER
ELEPHANT
Attention
0.892-t0.000
0.858++0. 106
0.615=0.096
0.839+00.054
0.868+00.054
0.900=00.088
0.8634:00094
0.603=0.068
0.845+0.006
0.857+:0.064
CLAM
0.900:10.138
0.860110.128
0.610=0.125
0.805
0,860110000 0.820+1.000
RAM-MIL
0.911
0.870:10.142
0.645-0.117
0.879.0000%
B General MIL dataset
Table 7 presents the performance of RAM-MIL general MIL datasets [46, 47], offering parison with baseline methods. The results indicate that OT-based retrieval generally enhances the classification performance. The sole exception is observed with the TIGER dataset, where both CLAM and RAM-MIL are outperformed. This discrepancy might be attributed to CLAM,, as RAM-MIL uses CLAM pretrained model for attention weights and bag representation extraction. Nonetheless, RAM-MIL still improves over its CLAM baseline on TIGER. Note that our primary focus lies on the more challenging WSI datasets, hence our models not extensively optimized for general datasets. The data in these general datasets are typically of lower dimensionality and present less challenging conditions. Therefore, any potential underperformance in these contexts should not detract from the strength of our models in handling the WSI data.
C Experiment Details of WSI Classification
We present the experimental details, ablation studies and analysis step-by-step
MIL Pre-training. For the backbone MIL model we use the the same parameter setup CLAM. The model parameters are updated via the Adam optimizer with an 1.2 weight decay of le-5 and learning rate of 2c-4. Each result is obtained with 10-fold splits of training/validdaiionnetingg sets.
15

---

Table 8: Ablation study for the percentage of instances used on CAMELYONI6 and CAMELYON17.
In- Domaii (CAM16)
Out-of- Domain (CAM17)
AUC
Accuracy
AUC
Accuracy
10%
attention Retr
0.9440+0.037
0.8975+0.052
10%
attention Retno
0.9365+-0.02 0.9414110.046
0.9200=0.050
0.7974+0.054
0.7433+-0.073
10% attention Retro
0.856
0.7775:100050
0.7392-10.063
20%
attention Retri
0.9451=0.003
0.8925=0.050
20% attention Retroo
0.9341=0.051
0.8925=0.053
0.7651+0.056
0.77144:00030
20% attention Retro
0.9419::00088
0.9175=0.051
0.7681+0.058
0.7795=00022
Neighbor Selection. After pre-training the MIL model, we obtain the slide-level feature and the attention scores predicted by the network. As computing the optimal transport distance based on all instances is time-consuming, we approximate the distance with a part of samples a bag.
dor(H, = min
ΣΣ
ch, hh,,,,)T. hy)TTTTTTT. (T,))T++.. + >'T
Tip log Ti
(11)
TET (a,
s.t., 1 = =
where and â are the new attention vector obtained by selecting top 17% from and ā. In other words, we approximate a bag with 17% of instances with the highest attention values generated by pretrained MIL model. As shown in Table 8, we set 17 10 and r = 20 for the ablation study. In this experiment, we use Regularization term of 0.5 and Max number of iterations 1000.
It is observed from Table 8, improving the percentage of data improves the performance Retr, and Retro- On Retrio. the performance is saturated when using only 10% of all patches. Differentiating in-domain and out-of-domain data for retrieval could be easily accomplished by representing bag by a few amount of instances.
Table 9: Ablation study for different merge functions on CAMELYON16 and CAMELYONI7.
In Do-ain (CAM16)
Out-of of-Domain (CAM17)
AUC
Accuracy
AUC
Accuracy
Mergeadd((-featt) Retn
0.9409-100038
0.9000:L0.049
Merge (2-feats) Retho
0.9341=0.051
0.8925100053
0.765110.056
0.7714-0.030
Merge (2-feats) Retto
0.9414=0.046
0.897510.056
0.777510.050
0.7392-0.063
Mergeaaa (3-feats) Retn
0.9383=0.050
0.9175+-0.051
Mergeadd(3-feana Retno
0.9313-0.0044
0.9000+0.045
0.7641
0.7553-+0.04
Merge..da (3-feats) Retro
0.9391=1.0051
0.917511000
0.764410.059
0.7754-10022
Mergeconver (2-feats) Retri
0.946
0.802
Mergeconvez f-aats)) Retro
0.9365=0.052
0.9200=0.050
0.797410.054
0.7433-10.073
Mergeconves (2-fcats) Retro
0.9419=#.004
0.9175=0.051
0.7681
0.7795=0.021
Mergeconver (3-feats) Retry
0.909
0.897510.052
Mergeconvea (3-feats) Retno
0.943510.003 0.897510.052
0.7652110.52
0.7014
Mergeconves (3-feats) Retro
0.9417=0.048
0.9050=00.00
0.7690+10.06
0.7755=0.021
Merge Function. Table 9 presents the results using different merge functions. To generate the bag representations, we employ two merge functions: 1) simple addition, referred to as Mergeadai 2) convex combination, referred to as Mergeconver- Additionally, '2-feats' and '3-feats" refer to bag representation that are merged with nearest neighbor or nearest neighbors, respectively. For convex combination, '2-feats" uses coefficients of 0.6 and 0.4, while "3-feats' uses coefficients of 0.6, 0.2 and 0.2, where the greatest coefficient corresponds to the original representation. This experiment is done with y) 10.
derived from Table 9 that using nearest neighbor and convex combination presents the best performance. Using 2 nearest neighbors and addition presents the similar results.
16

---

Classification Training, Finally, we train single logistic regression classifier using the merged representation. The Adam optimizer is used to update the model parameters, with a L2 weight decay of le-4 and learning rate of 2e-4. The models are trained for minimum of 40 epochs and up to a maximum of 200 epochs if the early stopping criterion is not met. This criterion involves monitoring the validation loss each epoch and if it has not decreased from the previous low for over 15 consecutive epochs, early stopping is used.
D Patch-level results on tumor slides of CAMELYON16.
Table 10: Patch-level results tumor slides of CAMELYON16.
P-Prec.(t)
FROC(1)
DSMIL
0.1030
0.4443
CLAM
0.6068
0.4792
TransMIL
0.1726
0.4797
Bayes-MIL
0.8107
0,4919
RAM-MIL
0.6114
0.5281
The Pumoo-Precision is calculated by the precision of classifying the tumor patches. The Patch- Precision is calculated by averaging the precision of classifying both normal and tumor patches. The Patch-FROC is defined as the average sensitivity (recall) at 6 predefined false positive rates: 1/4, 1/2, 2, 4 and 8 FPs per WSI.
In Table 17, RAM-MIL presents the second best precision and the best FROC the patch-level segmentation. This indicates that using transport matrix for interpreting the patch-level classification achieves the best overall performance in the trade-off of false positive rate and recall. By contrast, Bayes-MIL could only obtain a high precision, which reduces the number of false alarm. However, for the application of medical WSL, reducing false negative (better recall and FROC) is supposed to be more important as classifying a positive instance to be negative is unacceptable in the application of medical pregnosis diagnosis.
E Results on CAMELYON17.
general setting is to treat ITC as the negative class in the CAMELYONI7 dataset.
Table 11: Results on CAMELYON16 and CAMELYONI7 for in-domain classification and unsuper- vised domain adaptation, for comparisons with related methods.
In- In-Domain (CAM16)
Out-of-Domall (CAM17)
AUC
Accuracy
AUC
Accuracy
CLAM
0.9177=00004
0.8650+:0000
0.8399+00.059
0.7796+00099
Mixup Retn
0.9260-1.005
0.882510.051
Mixup Retro
0.9271
0.8850++0.045
0.841840.054
0.8016+0.059
Mixup Retro
0.9271
0.8825+00008
0.8397+0.057
0.8095#0.063
12 Retn
0.9481
0.8800-L0.055
12 Retno
0.9041
0.8950-10.048
0.8403-10.049
0.7914+0.052
l2 Retro
0.9398--0.045
0.8950+00.05
0.8412-40.055
0.799
RAM- MIL Retn
0.945
0.002
RAM-MII. Retro
0.9365+00055 0.9200-00.050
0.8455-00005
0.823
RAM-MII. Retro
0.9419+-0.04
0.9175+00.05
0.8413+00005
0.8457+0.027
17

---

Table 12: Ablation study evaluated on CAMELYON16 and CAMELYONI7.
In-Domain (CAM16)
Out-of-Domaie (CAM17)
AUC
Accuracy
AUC
Accuracy
Approx-OT-min Retn 0.930
0.8800++0.033
Approx-OT-min Retro
0.8601+00000
0.8000+0.062
0.8380+0.056
0.7935+00.065
Approx-OT-min Retto
0.9196110.033
0.8750:100056
0.839710.056
0.7915110.048
Approx- Approx-OT-ave Retri
0.9227=00038
0.8600+0.030
Approx- -a- Retrio
0.8865=00.03
0.8425+0.061
0.8401#0.056
0.7902+:0.051
Approx-OT-ave Retro
0.9195==0.063
0.8750=0.056
0.8412=#0003
0.7823+:0.063
Approx-OT-mas Retn
0.9174.10.049
0.8625110.044
Approx-OT-max Retro
0.9182.10.047
0.8650110.048
0.840610.006
0.785410.0052
Approx- OT-max Retro
0.9195=00063
0.8750=00056
0.8412=#00063
0.7925=+0.058
12 Retn
0.9081100004
0.8800-10.055
12 Retro
0,9241100048
0.8950=10.048
0.840
0.7914+0.0052
12 Retro
0.9398+00.04
0.8950++0.055
0.851
0.709
Hausdorff Reth
0.9273+1.005
0.8850=0.004 049
Hausdorfi Retro
0.9226+00005
0.8725+00000
0.050
0.7932+00003
Hausdorff Retro
0.9322++0.046
0.8950+0.048
0.869
0.8000++0.061
RAM-ME MIL Retr
0.8925=0.0500 0.945
RAM-MIL Retro
0.9065110.052
0.000
1.8475100055
0.8236:10.051
RAM-MIL Retro
0,941910.048
0.9175-10.051
0.801
0.845710.027
Table 13: Ablation study for the percentage of instances used on CAMELYONI6 and CAMELYON17.
In-Domain (CAM16)
Out-of-Domain (CAM17)
AUC
Accuracy
AUC
Accuracy
10% attention Retry
0.944011.003
0.0055
10%
attention Retno
0.906510.052 0.9414=0.046
0.920010.0500
0.8075:10.055
0.8236:10.051
10% attention Retro
0.8975=0.056
0.8466:10.047
0.8037=0.062
20%
attention Retn
0.9451--.0336
0.002
20% attention Retho
0.9341+-0.051
0.8525=00053
0.8473+-0.022
0.8216+00030
20% attention Retro
0.9419+0.048
0.9175+0.051
0.8413+-0.033
0.8457+0.027
Table 14: Ablation study for different merge functions on CAMELYON16 and CAMELYON17.
In- nn-Domain (CAM16)
out-of-Domain (CAM17)
AUC
Accuracy
AUC
Accuracy
(ergeada(2-feats) Retr
0.9409=00.038
0.9000=00.49
Merge Mergeadd(2-fcats) Retrjo
0.9341=0.051
0.8925+:00033
0.8473+00.52
0.8216=0.030
Merge (2-fcats) Retro
0.9414=00006
0.8975+0.056
0.8390+0.056
0.7695=0.056
Merge (3 (3-feats) Rein
0.9083_10.050
0.9175100001
(r-geadd(3-fcate Retro
0.9313=0.044
0.9000+00.04
0.8432+00055
0.8176=00.00
Mergo (3- (cats))) Retro
0.9391=0.055
0.9175=0.045
0.8402+0.056
0.8457=0.020
Mergeconsee (2-feats) Rein
0.9451=10.036
0.802
Mergeconvez (2-feats) Retro
0.9065
0.0000-100050
0.8475-10.055
0.8236-10.051
Mergecmeees (2-feats) Retro
0.9419=0.048
0.917540.051
0.8413+00053
0.8457=00027
Mergeconver (e-feat) Retn
0.9398-10.043
0.897510.052
Mergeconvez (eeeeee/3fffeee (3-feats) Retrio
0.9435+0.033 0.8975+00002
0.8473++0.049
0.7755-00067
Mergeconvez (3-feats) Retro
0.9417=0.004
0.9050--0000
0.8474+00007
0.8097=00007
18

---

F
Comparisons with domain adaptation methods
There have been large volume of research studying the domain adaptation in patch-level models, categorized as staining transfer and domain adversarial learning.
Patch-level staining transfer: Cho et al. [31], Shaban et al. [32] and Zanjani et al. [33] use generative adversarial networks (GAN) to learn the staining difference implicitly. Images are transferred to the target domain, by the trained generators in GAN. A simple neural network is then fit on the generated
images for discriminative tasks.
Patch-level domain adversarial learning: Lafarge et al. [34] and Ciga et al. [35] apply gradient reversal method for adapting the model for discriminative tasks in medical images. Ren et [36] uses Siamese network in domain adversarial learning. Brieu et al. [37], Kapil et al. [38] and Gadermayr et al. [39] for transfer images between domains for segmentation tasks. These methods are not directly comparable with our methods, as they require patch-level labels during training, which is not realistic real- scenario. In experiment, we adapt Shaban et al. [32] with least modification, to compare with the retrieval based multiple instance learning. We run Shaban et al. [32] in patch-level in the Camelyon 17 domain adaptation setting. For obtaining the WSI features, we use averaged pooling. For fairness, the same classifier as RAM-MIL is used for Shaban et al. [32]. Here are the results:
Table 15: Comparisons with patch-level domain adaptation methods
AUC(1)
Accuracy(t)
Shaban al. [32]
0.6458
0.6889
RAM-MIL
0.8209
0.7667
Slide-level methods: Yang et al. [48] propose a slide-level domain adaptation method that involves local and global adversarial loss to train a weakly supervised learning model. RAM-MII. outperforms this method in the domain adaptation evaluation Camelyon17.
Table 16: Comparisons with slide-level domain adaptation methods
AUC(1)
Accuracy(1)
Yang [48]
0.5867
0.6786
RAM-MIL
0.7974
0.7795
Other methods: Previous methods in slide-level classification [24] and methods in other fields [25] also consider the out-of-domain testing performance. However, they did not provide method for domain adaptation, while RAM-MIL has. We compare RAM-MIL with several most advanced baselines in slide-level classification in our experiments.
G Comparisons with other retrieval methods
For the information retrieval works (Yottixel [40] and SISH [41]) of whole slide image, will add them into the related works section. However, these two methods are not directly comparable, as they focus on the retrieval problem while our goal is to use retrieval to enhance the weakly supervised classification (MIL). An executable way is to replace the OT-retrieval module in RAM-MIL with these retrieval methods. As there is no available implementation of Yottixel, and SISH is more advanced retrieval method, we integrate SISH retrieval into the RAM-MIL process. The experiment Camelyon17 (including hospitals), with 2 hospitals the in-domain dataset and the rest. being the out-of-domain dataset. Here are the experimental results:
Table 17: Comparisons with other retrieval methods
AUC(1)
Accuracy(t)
CLAM-SISH
0.7647
0.7444
RAM-MII.
0.8209
0.7667
19

---

H More results on Dimensionality Reduction.
Directly reducing the normal dimension the estimated intrinsic dimension brute-force method to improve the performance. To validate if this argument holds, we use masked auto-encoder (MAE) with output dimension to (the estimated intrinsic dimension) to extract the features from patches. To ensure good quality of extracted features, we fine-tune the MAE with all patches from Camelyon16 for 60 epochs. The performance of CLAM fitted on the extracted features is at the second row of the Table 18:
Table 18: Additional results evaluated on Camelyon16 (in-domain) and Camelyon17 (out-of-domain)
CAMELYON16
CAMELYON17
AUC
Accuracy
AUC
Accuracy
CLAM-SVD-SD CLAM-MAE-SD
0,9216
0.6625
0.8183
0.4729
0.8285
0.7368
0,7982
0.6987
RAM-MIL-SVD-SD
0,9273
0.8475
0.8199
0.7355
20


Reviews:
[{"review_id": "uSSqcFnV8z", "reviewer": "Reviewer_5E1m", "summary": "This paper introduces the Retrieval-AugMented MIL (RAM-MIL) framework, which addresses the challenge of performance deterioration in MIL models when tested on an out-of-domain test set. The proposed framework achieves state-of-the-art performance in both in-domain and out-of-domain scenarios. The RAM-MIL framework integrates Optimal Transport as the distance metric for nearest neighbor retrieval and uses the transportation matrix derived from OT to render the retrieval results interpretable at the instance level. The paper also proves a theorem that demonstrates the reduced input dimensions lead to improved MIL performance and proposes the RAM-MIL algorithm to learn a low intrinsic dimension feature space and enhance the model generalization, especially for out-of-domain data. The proposed methodology is evaluated on whole slide image (WSI) datasets and general MIL datasets.", "questions": "Please refer to Weaknesses.", "limitations": "Please refer to Weaknesses. I'm looking forward to the author's reply, and after my questions are answered, I will reconsider my rating.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This article relates the input data dimension to the performance of MIL and provides a novel theoretical explanation. This perspective is very innovative.\n\nThe article addresses the unsupervised domain adaptation problem by retrieving instances with similarities from the unlabeled domain and merging them with instances from the labeled domain. This approach impresses me deeply.", "weaknesses": "1.The author mentioned that they are the first to explore the problem of domain adaptation in pathological image domains. However, there have been some existing works in this area, including domain adaptation for whole slide images. For example:\n\n[1] Collaborative Unsupervised Domain Adaptation for Medical Image Diagnosis\n\n[2] Double adversarial domain adaptation for whole-slide-image classification\n\n[3] Adversarial Domain Adaptation for Classification of Prostate Histopathology Whole-Slide Images.\n\n[4] Unsupervised Domain Adaptation for Classification of Histopathology Whole-Slide Images\n\n2.(1) As mentioned in Comment 1, there have been numerous studies in this field. However, this article lacks a detailed review of domain adaptation issues in pathological images, including whole slide images. It also fails to provide a comprehensive explanation of the differences between the proposed method in this article and those previous studies.\n\n(2) Additionally, the experiments in this article only compare against baseline and self-variants, without comparing with the currently most advanced methods available.\n\n3.The author mentioned, \"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.\"\n\n(1) What does the term \"input data dimension\" refer to? Is it the feature dimension of the input instances? Does \"MIL performance\" refer to the performance of unsupervised domain adaptation or the performance of MIL itself?\n\n(2) As far as I know, most current MIL methods first use a pretrained self-supervised network to map instances from the image domain to the feature domain, and the dimension of these features can be defined by the pretrained self-supervised network. Therefore, can't the input data dimension issue raised in this article be simply solved by modifying and defining it through a pretrained network?\n\n(3) In conclusion, I haven't understood the relationship between the \"input data dimension\" and \"MIL performance\" in this article. Why does reducing the input data dimension improve performance? To what extent should it be reduced at least? How does this article's method determine the optimal input data dimension?\n\n4.Is the Merge operation performed by weighted averaging? How is the Merge operation specifically completed?\n\nI'm looking forward to the author's reply, and after my questions are answered, I will reconsider my rating.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. Here are the responses to weaknesses (W).\n\n**Response to W1 and W2.1:** \n\nThanks for mentioning the related works. We will add discussions about [1, 2, 3, 4] into the related works and [2] will be added to comparisons in the evaluation.\n\nNote that [1, 2, 3, 4] are **NOT** attention-based multiple Instance Learning model for whole slide image classification. [1, 3, 4] are domain adaptation methods for the instance (patch)-level classifier rather than multiple instance learning model. The instance-level classifiers neglect the global correlation between patches and require instance-level labels for training the classifier. The most related work is [2], where patch features of a whole slide image are aggregated using the Fisher vector encoding and global pooling. The fundamental difference between [2] (also noted as Yang et al in other responses) and ours is that, our method is based on attention-based multiple instance learning where the patch-level features are aggragated by learnable attentions. \n\nNevertheless, we add an experiment for comparing with [2] which will be added to the manuscript. As the datasets (CCI dataset and Warwick HER2 challenge) used in [2] are not accessible now, we evaluate [2] on with our experimental setting, using Camelyon16 as the in-domain dataset and Camelyon17 as the out of domain dataset. Accuracy and AUC of [2] are **58.67** and **67.86**, outperformed by RAM-MIL (Acc: **77.95**, AUC: **79.74**).\n\n**Response to W2.2**\n> Additionally, the experiments in this article only compare against baseline and self-variants, without comparing with the currently most advanced methods available.\n\nWe did compare with most advanced methods including CLAM, TransMIL and BayesMIL. In the rebuttal, we also add experiments comparing with [2] (domain adaptation for MIL, see responses to W1 and W2.1), HHOT (OT-method for WSI, see responses to reviewer WXnS) and SISH (Nearest Neighbour Search method, see responses to reviewer WXnS).\n\n**Response to W3.1**\n\nThanks for spotting one critic motivation of our work. The sentence at the line 70, page 2 is not accurate and will be clarified in the updated version. “input data dimension” should be “the *intrinsic dimension* of input data”, which should be distinguished from the normal dimension. We assume that the data reside on a low-dimensional manifold in the input data space, so the intrinsic dimension is less than the normal dimension.  As mentioned in the paper (line 111-113, page 3), the intrinsic dimension is measured by how many dimensions are needed to represent the data without much loss of information [Facco, et al.]. The MIL performance refers to the approximation error of an MIL model.\n\nFacco, et al. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. NeurIPS 2017.\n\n**Response to W3.2**\n\nThanks for the advices. Directly reducing the normal dimension to the estimated intrinsic dimension is a brute-force method to improve the performance. To validate if this argument holds, we use a masked auto-encoder (MAE) with output dimension to be 5 (the estimated intrinsic dimension) to extract the features from patches. To ensure a good quality of extracted features, we fine-tune the MAE with all patches from Camelyon16 for 60 epochs. The performance of CLAM fitted on the extracted features is at the second row of the Table:\n\n| Model       | C16-AUC | C16-ACC | OOD-C17-AUC | OOD-C17-ACC |\n|-------------|---------|---------|-------------|-------------|\n| CLAM-SVD    | 0.9216  | 0.6625  | 0.7486      | 0.4991      |\n| CLAM-MAE    | 0.8285  | 0.7368  | 0.7325      | 0.6200      |\n| RAM-MIL-SVD | **0.9273**  | **0.8475**  | **0.7509**      | **0.7053**      |\n\nSimply reducing the output dimension to the estimated intrinsic dimension cannot provide a decent performance, which provides a justification for our RAM-MIL method. \n\n**Response to W3.3**\n\n*The theoretical results show that a lower intrinsic dimension would lead to a better performance of the multiple instance learning model. This statement is intuitively correct.* We list the intuitions as follows:\n\n1) Lower intrinsic dimension may mean that the data lies on a simpler manifold or subspace. As the complex of the subspace is lower, learning is easier.\n\n2) Lower intrinsic dimension might also imply that the data is less noisy and redundant, which could make learning easier.\n\nAlthough there exist some methods to test the intrinsic dimension of data, there is no principled method to reduce the intrinsic dimension and determine the optimal intrinsic dimension. Our paper proposes a retrieval-augmented MIL that could potentially reduce the intrinsic dimension and it is validated in the paper (Figure 3 and Table 3).\n\n**Response to W4**\n\nAs indicated in line 155, page 4, the merging function is a convex combination. We have tested different configurations of the merging function, shown in Table 6, page 14, supplemental materials. We will make this clear in the updated version."}, {"review_id": "NC3NyadBGG", "reviewer": "Reviewer_WXnS", "summary": "This work presents a method for multiple instance learning-based method for slide retrieval based on optimal transport, called RAM-MIL. The idea is as follows: 1) first train a standard ABMIL model for supervised classification, 2) pre-extract the features from ABMIL to get slide-level features, 3) compute the OT map via SInkhorn's between bag of instance representations (using the attention weights to calculate the OT) within the retrieval set, 4) merge representations followed by training a logistic classifier. RAM-MIL was trained and evaluated in CAMELYON16 and CAMELYON17, with ablation experiments conducted with other OT-based retrieval methods.\n\n", "questions": "1. How well does this method perform against other baselines such as SISH / Yottixel for slide retrieval? ABMIL with and without linear probe?\n2. There exist many histopathology datasets that can be combined to form external cohorts. The findings of this work would be strengthened if more datasets and cancer types were evaluated.", "limitations": "Limitations of this work were addressed.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- This work represents a new method for augmenting MIL with optimal transport, particularly for handling out-of-distribution data. The idea of computing optimal transport maps between instances and bags in histopathology have successfully been performed in prior literature such as HHOT (Yeaton et al. 2022 MIDL), as well as with regard to tackling domain adaptation problems such as Falahkheirkhah et al. 2023 MIDL, Domain adaptation using optimal transport for invariant learning using histopathology datasets. Distinct from prior works is the further application of OT with feature merging at the bag-level for evaluating out-of-domain performance in slide-level classification. At the same time, this work still brings several interesting ideas on how to leverage the attention weights as the probability measures for computing OT maps.\n- Beyond the main paper, the supplement is also nicely organized and includes additional ablation experiments regarding % of attention weights used and different merge functions. Further evaluation of patch-level localization performance was also performed, which is appreciated in baselines such as CAMELYON16 (commonly used for slide-level performance only).\n- Figures are illustrative and well-designed.", "weaknesses": "1. Regarding novel contribution #1: On line 67, one of the contributions state that \"this work is among the first to investigate the out-of-domain performance for MIL, which is vital issue for the application of MIL in automated medical diagnosis with WSI. Our empirical study exposes the risk of recent MIL algorithms under distribution shifts.\" This claim is not substantiated and untrue, as most clinical studies in computational pathology do consider the problem of out-of-domain performance of MIL models (for instance, Campanella et al. 2018 Nature Med, Lu et al. 2018 Nature BME, Courtiol et al. 2019 Nat Med., etc). Regarding out-of-domain performance for image retrieval, works such as Yottixel (Kalra et al. 2020, Medical Image Analysis) and SISH (Chen et al. 2022, Nature BME) have also performed quantitative evaluation on out-of-domain external cohorts.\n2. Baselines: Though methods such as CLAM and Bayes-MIL were used as comparisons, this work lacks comparisons with other important baselines for slide retrieval in histopathology such as SISH and Yottixel, and OT-based pathology works such as HHOT (most similar to this work conceptually). Other missing baselines include performance of conventional ABMIL (that was used to train RAM-MIL), and also linear probe analysis on top of the already pre-extracted ABMIL features. The reviewer would also appreciate hyper-parameters used in linear probe analysis of this work, which may also importantly influence results.\n3. Retrieval: Though the contributions of this work are posed as solving a retrieval problem, the evaluation (task, metrics) is focused on solving a popular but narrow task in histopathology (CAMELYON16 - a needle-in-a-haystack problems). Only one dataset is used for evaluation of this method, which may not be representative of broader slide-level retrieval tasks in histopathology. Lastly, one potential practical limitation is the computational complexity of using OT for retrieval, which should be commented on in the rebuttal.\n\nI am overall positive about this method, as this work brings an interesting idea regarding augmenting OT with attention weights from ABMIL in histopathology, my main concern of this work is the lack of evaluation on other datasets such as TCGA, CPTAC, NLST, GTEX, etc. - which can be readily used for evaluating out-of-domain performance and slide-retrieval in a pan-cancer setting, as well as missing baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. Here are the responses to weaknesses (W) and Questions (Q).\n\n**Response to W1:**\n\nThanks for mentioning the related works. We will add them into the related works sections.\n\nFor the claim of contribution, we meant to say, “the out-of-domain performance of Attention-based Multiple Instance Learning (ABMIL)” rather than the first work considering out-of-domain problem for whole slide image. As it might sound ambiguous, we will change the claim of contribution to a more precise one.\n\n**Response to W2 and Q1:**\n\nFor the information retrieval works (Yottixel and SISH) of whole slide image, we will add them into the related works section. However, these two methods are not directly comparable, as they focus on the retrieval problem while our goal is to use retrieval to enhance the weakly supervised classification (MIL). An executable way is to replace the OT-retrieval module in RAM-MIL with these retrieval methods. As there is no available implementation of Yottixel, and SISH is a more advanced retrieval method, we integrate SISH retrieval into the RAM-MIL process. The experiment runs on Camelyon17 (including 5 hospitals), with 2 hospitals as the in-domain dataset and the rest being the out-of-domain dataset. Here are the experimental results:\n||AUC| ACC|\n|-----------|--------|--------|\n|CLAM|0.8155|0.7222|\n|CLAM-SISH|0.7647|0.7444|\n|RAM-MIL|**0.8209**|**0.7667**|\n\nFor the OT-based pathology work, we add HHOT into the related works section. HHOT only considers comparing whole slide images with uniform weight-based OT. As they mentioned kNN as a discriminative method, we compare HHOT-kNN with RAM-MIL (C16 as in-domain, C17 as out of domain):\n||C16-AUC|C16-ACC|OOD-C17-AUC|OOD-C17-ACC|\n|----------------|---------|---------|---------|---------|\n|HHOT-kNN (k=1)|0.7007|0.7318|0.6939|0.7173|\n|HHOT-kNN (k=3)|0.7618|0.7544|0.7263|0.7523|\n|ABMIL|0.9010|0.875|0.7287|0.7189|\n|ABMIL w/ linear probe|0.9107|0.8725|0.7312|0.7077|\n|RAM-MIL|**0.9451**|**0.9200**|**0.7974**|**0.7795**|\n\nWe also include ABMIL with/without linear probe in the above table.\n\n**Response to W3 and Q2:**\n\nWe add experiments with in-domain classification on TCGA-NSCLC, in-domain tumor stage classification with Camelyon17, in-domain classification on CPTAC-LSCC, domain adaptation from CPTAC-UCEC to CPTAC-LSCC.\nHere are the results:\n\nTCGA-NSCLC (in-domain):\n\n|| AUC |ACC |\n|-------------|---------|---------|\n| CLAM        | 0.9420  | 0.8640  |\n| Scaling ViT | 0.9516  | 0.8821  |\n| TransMIL    | **0.9603**  | 0.8835  |\n| Bayes-MIL   | 0.9451  | 0.8965  |\n| RAM-MIL     | 0.9497  | **0.8988**  |\n\nCamelyon17 tumor stage classification (in-domain):\n|| AUC |ACC |\n|-------------|---------|---------|\n| CLAM        | 0.7803  | 0.60  |\n| Bayes-MIL | 0.8070 | 0.64  |\n| RAM-MIL    | **0.8132**  | **0.65**  |\n\nLSCC as in-domain dataset and UCEC as out of domain dataset:\n|| LSCC-AUC | LSCC-ACC | UCEC-AUC | UCEC-ACC |\n|-------------|---------|---------|---------|---------|\n| CLAM       | 0.9500 | 0.9285 | 0.4986  | 0.6520  |\n| RAM-MIL | **0.9667** | **0.9382** | **0.6056** | **0.6527**  |"}, {"review_id": "K05Xak0NdW", "reviewer": "Reviewer_yGH7", "summary": "The authors propose a novel MIL framework which integrates Wasserstein Distance and Optimal Transport to match instance representations across bags. This method can be further leveraged to match representations across bags from different domains, thereby enabling Cross Domain Adaption. \n\n", "questions": "1. The authors state that \"the calculation of OT quite time-consuming\", which is why they pre-selected the top20% of instances. Can the authors give training time comparisons between OT, l2, Hausdorff and Approx-OT? This will help readers understand the practical tradeoffs between l2 and OT.\n2. In the supplement, the authors provided a table of results for OT with top10% and top20% of instances. The performance differences seem minor. What would happen with top30% of instances or is this a problem-specific choice depending on the type of disease and relative coverage of disease features in the WSI images?\n3. I am not sure if I follow the authors' reasoning in section 5.2 and the corresponding Table 3. The authors state that the intrinsic dimension of CLAM is 5-7 on CAM16 and CAM17, then why would they cap merged bag representation to 5d vector? Would this not impact CLAM more than RAM-MIL and affect its performance? What if it was capped to 7d?\n", "limitations": "The authors could discuss the general applicability of RAM-MIL. Is it only to close problems such as CAM16 and CAM17? Can it help with WSI and bag labels that are different diseases?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors' RAM-MIL method achieved state-of-the-art performance in terms AUC and accuracy on CAMELYON16 and CAMELYON17, as well as MUSH1, MUSK2, FOX, and ELEPHANT. The authors provide interesting insights on the dimension reduction properties of their method. Finally, the authors demonstrate ways to interpret and visualize their method.", "weaknesses": "1. The authors might have unmasked themselves in the double-blind review process by using the same pathology image in Figure 1 as Bayes-MIL\n\nMinor Problems:\n1. The authors state Theorem 2, without clearly stating the proof is the supplement. The authors seem to state the proof of a similar result for the feature space is in the supplement\n2. Minor typos such as \"which which\" (line 308)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. Here are the responses to weaknesses (W) and Questions (Q).\n\n**Response to W1:**\n\nThanks for mentioning the potential issue. Please note that we follow the double-blind reviewing rules and the submission does not contain any identifying information. For Figure 1, some slides are borrowed from visualization of Bayes-MIL library, which are based on OpenSlide and CLAM library. As Bayes-MIL is a published work and the codes are accessible online (https://github.com/ralphc1212/bayes-mil), we believe it does not violate the anonymous rule to use the same visualization tool in our submission. We will cite Bayes-MIL in the manuscript in Figure 1.\n\n**Response to W-Minor 1:**\n\nAs indicated in Line 480-481 of Page 13, supplemental, to get the same results for features, we only need to assume a probability measure over the instance representations h and replacing all X in the proof with H. We will provide a theorem 2 and a formal proof in the supplemental.\n\n**Response to Q1:**\n\nThanks for the advices. Here are the time comparisons for calculation one pair of slides under different metrics:\n\n|| L2 | Hausdorff | Approximate-OT | Full OT |\n|-------------|---------|---------|---------|---------|\n| Time (s)        | 0.001  | 1.4  | 0.083| 0.498|\n\nNote that Approximate-OT provide a decent trade-off between computation time and performance. Next we show using a few patches (like 10%), the performance might saturate due to the most important information is stored in a few patches. Therefore, it might be a good choice to use a few patches for OT computation.\n\n**Response to Q2:**\n\nThanks for the advices. We calculate the cumulative attention probability value for 10% and 20% of patches:\n\n|| C16 | C17 | TCGA-NSCLC |\n|-------------|---------|---------|---------|\n| 10% of patches        | 0.9661  | 0.9667  | 0.9969|\n| 20% of patches    | 0.9882  | 0.9805 | 0.9999|\n\nThis indicates most useful information are already in the 10% data, which explains why the performance does not change much when increasing the data amount.\n\n**Response to Q3:**\n\nSection 5.2 is to show directly applying simple dimensionality reduction method like SVD is not useful in improving the performance. By using retrieval then SVD, the performance is significantly better. We add an experiment with capped dimensionality 7:\n\n|| C16-AUC | C16-ACC | C17-AUC | C17-ACC |\n|-------------|---------|---------|---------|---------|\n| CLAM        | 0.8821  | 0.6688  | 0.7568  | 0.5474  |\n| RAM-MIL    | **0.9312**  | **0.905** | **0.7641** | **0.7434** |\n\n**Response to limitations:**\nWe also add experiments on other datasets: in-domain classification on TCGA-NSCLC, in-domain tumor stage classification with Camelyon17, in-domain classification on CPTAC-LSCC, domain adaptation from CPTAC-LSCC to CPTAC-UCEC. \n\nHere are the results:\n\nTCGA-NSCLC (in-domain):\n\n|| AUC |ACC |\n|-------------|---------|---------|\n| CLAM        | 0.9420  | 0.8640  |\n| Scaling ViT | 0.9516  | 0.8821  |\n| TransMIL    | **0.9603**  | 0.8835  |\n| Bayes-MIL   | 0.9451  | 0.8965  |\n| RAM-MIL     | 0.9497  | **0.8988**  |\n\nCamelyon17 tumor stage classification (in-domain):\n|| AUC |ACC |\n|-------------|---------|---------|\n| CLAM        | 0.7803  | 0.60  |\n| Bayes-MIL | 0.8070 | 0.64  |\n| RAM-MIL    | **0.8132**  | **0.65**  |\n\nLSCC as in-domain dataset and UCEC as out of domain dataset:\n|| LSCC-AUC | LSCC-ACC | UCEC-AUC | UCEC-ACC |\n|-------------|---------|---------|---------|---------|\n| CLAM       | 0.9500 | 0.9285 | 0.4986  | 0.6520  |\n| RAM-MIL | **0.9667** | **0.9382** | **0.6056** | **0.6527**  |"}, {"review_id": "PDWAFwjpiT", "reviewer": "Reviewer_q4PZ", "summary": "The paper proposes a new Multiple instance learning method called RAM-MIL. The proposed method is a two-stage approach which mainly involves the following steps: 1. Train an existing MIL model on the D0 to extract feature representations of each instance 2. Retrieve nearest neighbor from retrieval set to form merged bag representation 3. Train another classifier on the merged bag representation.  The key component in the proposed method is the use of Optimal-Transport to measure bag-to-bag distance, retrieve nearest neighbor bag and form merged bag representation. The proposed approach improves out-of-domain performance while also maintaining good in-domain performance. ", "questions": "1. Did you describe how Figure 3 is obtained?\n2. Can you clarify your experiment setting for the baseline methods (DSMIL, CLAM, TransMIL, Bayes-MIL) in Table1, did they touch the retrieval set at all? \n\n", "limitations": "the efficiency of the algorithm due to solving optimal transport problem is acknowledged.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. writing is mostly clear and easy to follow\n2. related work section is comprehensive\n3. the idea of using OT for measuring bag-to-bag distance and forming merged representations seems effective.", "weaknesses": "while the proposed method improve performance over the compared baselines (Table1), it also introduces additional steps on top of these baseline (finding nearest neighbors, forming merged bag representation and then train another bag classifier), in this sense, the additional performance gain is not that surprising.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. Here are the responses to weaknesses (W), Questions (Q) and limitations (L).\n\n**Response to W1:**\n\nThanks for the comments. As retrieval (also dubbed as external memory) has achieved great successes in applications like large language models [1][2], we believe it is worth to exploit its application in other important domains, like the challenging histopathological diagnosis. Our paper first provides the theoretical understanding, methodology, algorithms, and analysis for retrieval with multiple instance learning. For the algorithm part of this topic, it is necessary to know how retrieval could be done and how the retrieved knowledge could be merged into the original model.\n\nNote that we provide multiple choices to execute our retrieval algorithm including using L2, approximate-OT, Haursdorff distance and OT. Running the approximate version (L2 or approximate-OT) of the “additional steps” introduce negligible overhead in “finding nearest neighbor” (see stats below). Merging bag representation and train a linear bag classifier causes trivial computational overhead with modern consumer-grade GPUs. **Therefore, the additional gain does not introduce much overhead. Note that TransMIL and Bayes-MIL are heavier model with larger model size and longer running time.** We will update the parameter comparisons and running time comparisons later to show our advantage.\n\n|| L2 | Hausdorff | Approximate-OT | Full OT |\n|-------------|---------|---------|---------|---------|\n| Time (s)        | 0.001  | 1.4  | 0.083| 0.498|\n\nInspired by your comments, we think it is a good future work to study how to merge the retrieved data to MIL, in an online fashion (which takes even less time). We perform a preliminary test on this with retrieval integrated into the training pipeline of RAM-MIL. Here are the preliminary experimental results, which shows a promising online approximation to the offline-RAM-MIL:\n\n|| C16-AUC | C16-ACC | C17-AUC | C17-ACC |\n|-------------|---------|---------|---------|---------|\n| RAM-MIL-Online        | 0.9263  | 0.875  | 0.7686 | 0.7453 |\n| RAM-MIL       | 0.9451|0.9200|0.7974|0.7795|\n\n[1] Borgeaud et al. Improving language models by retrieving from trillions of tokens. ICML 2022.\n[2] Gun et al. REALM: Retrieval-Augmented Language Model Pre-Training. ICML 2020.\n\n**Response to Q1:**\n\nTo compute the intrinsic dimension, we obtain the slide features of each method then feed to the estimator provided by [Facco] for computation.\nHere are the procedures for obtaining slide features for different methods:\nCLAM: we take the weighted patch features as the slide feature.\nMixup: we randomly pick slide features at the same class as a neighbour, then merge the two features with convex combination.\nThe rest: these variants merge the slide features of source and retrieval, with retrieval implemented by different methods.\n\nFacco, et al. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. NeurIPS 2017.\n\n**Response to Q2:**\nThere is no setup of “retrieval set” in the baseline method, as this paper is the first to consider retrieval in multiple instance learning. However, **the baseline methods and our method access the same data under the in-domain setting during training.** The retrieval set is only taken from the in-domain dataset (e.g., training set and validation set of Camelyon16).\n\nFor out-of-domain setting, there is a retrieval dataset providing unlabeled out-of-domain data for retrieval. For a fair comparison, we add a method [Yang et al] that also utilize the retrieval dataset and adversarial training in the experiments with Camelyon16 and Camelyon17, as suggested by reviewer 5E1m. Here are the results:\n\n|| AUC | ACC |\n|-------------|---------|---------|\n| Yang et al        | 0.6786 | 0.5867  |\n| RAM-MIL        | **0.7974**  | **0.7795** |\n\nYang et al. Double adversarial domain adaptation for whole-slide-imageclassification. MIDL 2021.\n\n**Response to L1**\n\nWe add the efficiency results. Please refer to \"response to W1\"."}, {"review_id": "HP5ltMH3xI", "reviewer": "Reviewer_7vyh", "summary": "This work proposes RAM-MIL: an approach that uses feature alignment in MIL with an optimal transport based method to retrieve nearest bag representations between the train set and target set. This approach shows consistent improvements in ID and OOD settings. Ablation studies demonstrate the relative effectiveness of different ways to compute the nearest neighbors.", "questions": "- The authors mention using top 20% of instances for Camelyon-16 and 17. How much of this choice is a function of the Camelyon datasets where signal is often present in small foci? How will this change in MIL problems like NSCLC and RCC subtyping where the signal is diffused throughout the WSI? [1] has presented some findings on OOD generalization on these 2 datasets.\n\n- The motivation behind the Dimensionality Reduction experiments is unclear. Is the idea to use the reduced dimensions to perform inference in certain time-sensitive applications? \n\n- Is there any additional computational lift for running this method in a multi-class setting?\n\n[1] - SC-MIL: Supervised Contrastive Multiple Instance Learning for Imbalanced Classification in Pathology", "limitations": "- The authors mention the limitation of the complexity of the OT approach scaling with the number of instances in the Conclusion section.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper brings forth ideas from Unsupervised Domain Adaptation and feature alignment to improve OOD performance in the MIL setting, which is a pressing concern to address the real world deployment of these models. The assumption of using attention values as probability mass values is a logical choice. \n- Extensive experimentation and ablations have been conducted to show the effectiveness of the Optimal Transport approach over simpler options like naively calculating nearest neighbors at a bag representation level.  The results show significant improvements in both OOD and ID settings.", "weaknesses": "- My main concern is  around the computational cost of inference with RAM-MIL being high due to calculating the distance measures between the train and the retrieval set. Also, in cases where the test set is dynamic, such as one more WSI gets added to the test set, all the pairwise distances between the two sets will need to be recalculated\n- The section about interpretability with OT for MIL assumes that instance labels for target bag are known, which is a strong assumption. I would suggest removing this section from the paper and focusing on the OOD generalization part.\n- Camelyon-17 has data from 5 different medical centers. It might be worth conducting an experiment where data from a single medical center is used in training, while the other centers are used in retrieval set, since these heldout centers clearly correspond to OOD sets. It is not clear if Camelyon-17 is completely OOD wrt Camelyon-16.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. Here are the responses to weaknesses (W), Questions (Q) and limitations (L).\n\n**Response to W1:**\n\nWe report the computation time of RAM-MIL using L2, Hausdorff, approximate-OT and OT for 1 pair of whole slide images as follows. The approximate-OT provides a close performance to full OT while maintaining a low computation cost. \n\n|| L2 | Hausdorff | Approximate-OT | Full OT |\n|-------------|---------|---------|---------|---------|\n| Time (s)        | 0.001  | 1.4  | 0.083| 0.498|\n\n> Also, in cases where the test set is dynamic, such as one more WSI gets added to the test set, all the pairwise distances between the two sets will need to be recalculated\n\nThis statement is not true. Note that when a new slide is added to the test set, only the pair-wise distances between this slide and the existing retrieval slides (say the number is N) should be calculated, for future usage. We will clarify this in an updated version. \n\nFurthermore, considering a more practical case of streaming data, more efficient approximation can be used for efficiently deploying the streaming inference. For example, a naïve method is to cluster the existing retrieval slides and only compute the distance with existing cluster centers (say the number is C). This will reduce the number of OT-computation by from O(N) to O(C). In this case, OT-computation time would be trivial.\n\nWe agree that the dynamic test set is an interesting research topic, but it is out of the scope of this submission. We will leave it as future work.\n\n**Response to W2:**\n\nThanks for mentioning this typo in Section 5.3. We meant to say “Suppose that we are interested in the instance labels of the target bag, and suppose that the instance labels of the source bag are known.” Switching all “target” and “source” would disambiguate the section 5.3. We will make this revision in the manuscript to make it clear.\n\n**Response to W3:**\n\nWe add an experiment that use 2 hospitals of Camelyon17 as the source dataset and 3 other hospitals as the target dataset. Here are the results:\n||AUC| ACC|\n|-----------|--------|--------|\n|CLAM|0.8155|0.7222|\n|RAM-MIL|**0.8209**|**0.7667**|\n\n**Response to Q1:**\n\nWe list the averaged accumulated attention probability values for 10% and 20% patches from Camelyon16, Camelyon17 and TCGA-NSCLC datasets. We observe that NSCLC also present similar signal distribution as Camelyon16 and Camelyon17. We agree with [1] that the imbalanced distributions of positive patches in different WSIs discussed are a crucial issue and will add discussions about [1] to the related work section. The extention of RAM-MIL to the imbalanced positive patch distribution will be an interesting future work. \n\n|| C16 | C17 | TCGA-NSCLC |\n|-------------|---------|---------|---------|\n| 10% of patches        | 0.9661  | 0.9667  | 0.9969|\n| 20% of patches    | 0.9882  | 0.9805 | 0.9999|\n\n**Response to Q2:**\n\nNote that the intrinsic dimension is a tool to analyze and validate our methodology. Distinguished from normal dimension, intrinsic dimension is measured by how many dimensions are needed to represent the data without much loss of information. A lower intrinsic dimension would make the learning easier, as the data lies on a less complex manifold or less noisy. \n\nThe idea of Section 3.3 - Reduction of Intrinsic Dimension is to show, our proposed retrieval method could reduce the intrinsic dimension of features, thus make the learning of MIL easier. That explains why RAM-MIL has better performance compared with other methods.\n\nThe idea of Section 5.2 – Dimensionality Reduction is to show, simply reducing the normal dimension by SVD can have bad performance. Retrieval could help organizing the latent space and reducing the intrinsic dimension. Thus, SVD after retrieval has a better performance than the SVD-only approach.\n\n**Response to L1:**\n\nThanks to the reviewer’s suggestions. During the rebuttal period, we found that using attention probability to select few patches is technically sound. Using 10% patches might saturate and the performance does not increase. This limitation is no longer severe based on the observation. We will test more datasets to see if this conclusion is consistent and will post the result as soon as it is available. "}]

Author rebuttal:
We thank all reviewers for the valuable feedbacks and constructive comments. We have prepared the response and revised the manuscript accordingly to address your concerns. The major concerns and the corresponding answers/explanations/clarifications are as follows:

1. Related works (5E1m, WXnS, 7vyh) and datasets (WXnS, 7vyh): we add the experiments to compare our method with Yang et al (domain adaptation for MIL, mentioned by 5E1m), HHOT (OT-method for WSI, mentioned by WXnS) and SISH (Nearest neighbor search method, mentioned by WXnS) and will include the missing citations mentioned in the related works. RAM-MIL outperforms the related methods in the experiments.
We also add experimental results on new datasets including TCGA-NSCLC, CPTAC-LSCC and CPTAC-UCEC (WXnS). New experiments with Camelyon17 multi-class tumor stage classification and domain adaptation between hospitals from Camelyon17 (7vyh) are included. Please check the one-page pdf for details.
2. Intrinsic dimension and dimensionality reduction (5E1m, yGH7, Q4PZ and 7vyh): we made additional explanations about intrinsic dimension which is misunderstood as normal dimension. The data are assumed to lie on a low-dimensional manifold with an intrinsic dimensionality, which is lower than the normal dimensionality. 
3. Efficiency of retrieval (yGH7, q4PZ, 7vyh) and percentage of patches to calculate OT (yGH7, 7vyh): we test the retrieval efficiency of different variants of RAM-MIL - L2, Hausdorff, approximate-OT, full-OT, which have trade-off in performance and efficiency. Approximate-OT provides a decent performance with a high retrieval efficiency. Users could select the distance metrics based on the need.
During the rebuttal period, we found that using a fixed attention probability to select few patches is technically sound. Using 10% patches might saturate and the performance does not increase.
4. Others: Dimension manipulation with pre-trained model (5E1m), ABMIL with or without linear probe (WXnS), online retrieval (Q4PZ) and dynamic evaluation (7vyh) are answered with new experimental results accordingly. The minor problems like typos will be updated in the manuscript.

We hope that our response can address the mentioned weaknesses and concerns, and the reviewers could raise their scores accordingly.


Meta-review: Four out of five reviewers recommend acceptance, while 2 increased their score after rebuttal. This paper augmenta MIL with NN-based retrieval to improve performance of an important application. Although during discussion some of the acceptance reviewers expressed that they think the paper is borderline with novelty weaknesses beyond the medical-imaging domain, they all still recommended acceptance.The ACs do not have strong reason to overturn the reviewers recommendation. 

The authors should incorporate the feedback and discussions in the final manuscript.
Decision: Accept (poster)