Samba: Severity-aware Recurrent Modeling for Cross-domain Medical Image Grading

Authors: qi  bi
jingjun  yi
hao  zheng
wei  ji
haolan  zhan
yawen  huang
yuexiang  li
yefeng  zheng
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-23 10:20:16

Abstract: Disease grading is a crucial task in medical image analysis. Due to the continuous progression of diseases, i.e., the variability within the same level and the similarity between adjacent stages, accurate grading is highly challenging.
Furthermore, in real-world scenarios, models trained on limited source domain datasets should also be capable of handling data from unseen target domains.
Due to the cross-domain variants, the feature distribution between source and unseen target domains can be dramatically different, leading to a substantial  decrease in model performance.
To address these challenges in cross-domain disease grading, we propose a Severity-aware Recurrent Modeling (Samba) method in this paper.
As the core objective of most staging tasks is to identify the most severe lesions, which may only occupy a small portion of the image, we propose to encode image patches in a sequential and recurrent manner.
Specifically, a state space model is tailored to store and transport the severity information by hidden states.
Moreover, to mitigate the impact of cross-domain variants, an Expectation-Maximization (EM) based state recalibration mechanism is designed to map the patch embeddings into a more compact space.
We model the feature distributions of different lesions through the Gaussian Mixture Model (GMM) and reconstruct the intermediate features based on learnable severity bases.
Extensive experiments show the proposed Samba outperforms the VMamba baseline by an average accuracy of 23.5\%, 5.6\% and 4.1\% on the cross-domain grading of fatigue fracture, breast cancer and diabetic retinopathy, respectively. 
Source code is available at \url{https://github.com/BiQiWHU/Samba}.

Main paper: 
 introduction: 
 Disease grading aims to assess the severity level of a disease or a pathological region from a medical image [46, 31, 52, 50, 6] . It is more challenging than conventional deterministic classification with distinctive categories (e.g., cat vs. dog), owing to the inherent severity ambiguity within and between levels. This ambiguity arises because the progression of a certain disease or a pathological region is a transitional, continuous and time-growing process (illustrated in Fig. 1a ). On the one hand, different medical images within a same severity level can have rather different disease or pathological developments (shown in Fig. 1b ). On the other hand, medical images among different severity levels can share similar patterns, as low-level lesions may persist throughout the disease's progression.The past decade has witnessed the rapid development of disease grading methods [30, 34, 45] owing to the deep learning techniques [25, 24, 28, 26, 57] . However, most of these methods were developed by experts from a specific clinic field (ophthalmology, gynecology, etc.). Furthermore, these models usually assume that the medical images used for training and inference are independently and identically distributed (i.i.d.). In practical clinical scenarios, a grading model trained on a number of medical images (source domain) is often required to handle images it has not encountered before (unseen target domains). Due to variations between patients, scanners, imaging parameters, clinic centers, etc., the feature distributions of the source domain and unseen target domains can be dramatically different [32, 68, 23, 8] . When the severity level of a disease is measured by the distribution of lesions, this cross-domain variance can lead to the misdetection of crucial lesions, resulting in grading errors [4, 11] . Especially when the appearance of lesions is significantly affected by the style change, it may be observed that medical images from the same domain instead of from the same grade are clustered in the feature space (illustrated in Fig. 1c ). This suggests that the model has learned features with limited generalization ability. Domain generalized disease grading learns models from only a source domain, but is expected to be applicable to unseen target domains. The key to addressing this problem is accurately identifying the lesions that have a decisive impact on grading [49, 45] . As the disease progresses, multiple lesions may coexist in the image, and the critical aspect of grading is identifying the most severe one among them. However, the most severe lesion may be localized in a small region in the image, exhibiting variable shapes, and being influenced by cross-domain style changes [4, 11] . To overcome these challenges, this paper proposes a severity-aware recurrent modeling method (Samba). Samba encodes image patches in a recurrent manner and recalibrates the state distributions based on learnable bases.In many disease grading scenarios, the decisive lesions only occupy a small portion of the total area. For instance, in retinal photographs, the affected blood vessels may only involve a small section at the distal. Similarly, in computed tomography (CT) or magnetic resonance imaging (MRI) scans, malignant tumors can also present as small lesions with a diameter less than 3 mm. These small lesions are easily influenced by style variations, which can lead to incorrect grading. Therefore, the model needs to pay sufficient attention to these detailed patches to classify them accurately. To address this issue, we treat the image patches as sequential data and encode them in a recurrent manner. This approach allows the information of decisive lesions to be stored in the hidden states and propagated to subsequent sequences. Furthermore, we adopt bidirectional encoding, enabling critical local information to influence the overall representation. More specifically, we incorporate a bidirectional Mamba [17] layer into the Samba, which supports sequence-to-sequence transformation and efficiently selects data in an input-dependent manner.The Mamba model achieves its selection mechanism by parameterizing the State Space Model (SSM) based on the input. While this selection mechanism [17, 69, 35] aids in identifying decisive lesions and propagating critical information, these input-dependent parameters are also vulnerable to the influence of image style transformations. When the feature distribution is affected by cross-domain variations, both the update of hidden states and the gating mechanism are disrupted. To resolve this problem, we utilize learnable tokens to capture the lesion representations, which are then used as bases to map the feature embeddings into a more compact space. To preserve the semantic information within this process, we further employ the Expectation-Maximum (EM) algorithm [14] initialized by these bases to estimate the lesion feature distribution for each image and reconstruct the features accordingly. We refer to this process as EM-based state recalibration in this paper.Our contributions can be summarized as follows.• We develop a Severity-aware Recurrent Modeling, dubbed as Samba, for general disease grading within-and cross-domain medical images. • We propose to encode the image patches in a recurrent manner to accurately capture the decisive lesions and transport the critical information from local to global. • An EM-based state recalibration mechanism is designed to reduce the impacts of crossdomain variants by mapping the feature embeddings into a compact space. • Extensive experiments on three cross-domain disease grading benchmarks show the effectiveness of the Samba against the baseline. 
related work: 
 Domain generalization aims to learn a model that can be well generalized to unseen target domains when only trained by the source domain, where the cross-domain feature distribution is usually not identical [65] . In the past few years, a variety of machine learning techniques (e.g., discrepancy minimization [47, 13] , knowledge ensemble [12] , uncertainty quantification [39, 53] , optimal transport [16, 60] , self-learning [51, 43] , frequency decoupling [59, 9, 10] and casual inference [37, 38] ) have been proposed. In the medical imaging community, the effort of bridging the domain gap between training data and unseen inference data is so far mainly focused on medical image segmentation [32, 68, 23, 8, 58] and classification [66, 54] . These methods usually rely on either learning shape-invariant representation or reaching pixel-wise consensus among the source domains. However, they are not especially devised to tackle the key challenge in cross-domain medical image grading, where the medical images from the same severity level instead of the same domain tend to cluster together.Medical grading has also been extensively studied. For Diabetic Retinopathy (DR) grading, many works aim to highlight the subtle local pathological regions to better discern different severity levels [30, 34, 44, 6, 45, 7] . Similarly, grading models have also been developed for pulmonary nodules [46] , fatigue fracture [31] , glioma [52] , acne vulgaris [50] , etc. However, most of the existing grading methods are task-specific and assume the training and inference medical images are i.i.d., which is far from reality. Practically, a medical grading model is supposed to show reliable inference on unseen target domains that have different feature distribution from the source domain. To the best of our knowledge, only [4] and [11] made an initial investigation on learning domain generalized DR grading.State Space Model (SSM) [27] contributes to a variety of fields such as robotics, navigation, and control theory, which is a foundational scientific model. In the past few years, SSM has been adapted in the context of deep representation learning, and has shown great success in sequence modeling [19, 20] . More advanced SSM, exemplified by Mamba [17] , not only shows stronger representation ability in long sequence modeling, but also exhibits linear scaling ability for longsequence data. Built upon this, multiple Mamba variations (e.g., Vim [35] and VMamba [69] ) have shown effectiveness in the computer vision field. However, these methods mainly focus on enhancing the context representation from the image by exploiting the long-range dependencies. Instead, how to model the cross-level severity development from the medical image by SSM remains unexplored.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: At the end of the conclusion section, the limitation of the proposed method has been discussed.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
problem definition & framework overview: 
 For a given disease grading task, assume we have a number of medical images x and the corresponding severity-level labels y from K different domains, which is denoted asEQUATIONn , y EQUATIONn , yEQUATIONHere N k denotes the number of images in domain k. For the cross-domain disease grading problem, the objective is to learn a grading model F θ : x → y using images only from a source domain D 1 , which is supposed to generalize well on other unseen target domains D 2 , • • • , D K . Following prior domain generalization works, each dataset is regarded as a domain D k , as the samples in a certain dataset are usually collected from the same clinical center by the same scanners and therefore share more similar feature distribution.The overview of the proposed method is illustrated in Fig. 2 . The input image is first encoded into patch embeddings through a stem unit with 4 × 4 convolutional kernels, where the stem unit partitions the input image into patches. The patch embeddings further pass through four encoding stages. Each Samba block involves a certain number of severity-aware recurrent layers and there are downsampling layers between two consecutive blocks. Finally, a grading head consisting of an average pooling layer and a linear layer generates the final prediction. Within the Samba block, the patch embeddings are first input to the bidirectional Mamba layers to extract the information about decisive lesions. After that, EM-based state recalibration is applied to map the lesion representation into a compact space by learnable bases.The core issue in most medical image disease grading scenarios is to identify the most severe lesion. However, due to the presence of lesions from different stages of the disease in the image, accurately capturing the most severe lesion is highly challenging. When the lesion occupies a large proportion of the image, the model only needs to extract stage-related features. In contrast, when the area of the critical lesion is small, the model needs to simultaneously locate the lesion and extract relevant features. This places higher demands on the model's ability to handle local information. To address this issue, in this paper, we propose to encode the image patches in a recurrent manner. Specifically, the state space model is used to process the sequential patch embeddings.State Space Model. Let x(t) denote a 1-D input signal. SSM maps it to the 1-D output signal y(t) by an intermediate N -dimensional latent state u(t), given by u ′ (t) = Au(t) + Bx(t), y(t) = Cu(t) + Dx(t), (1) where A ∈ R N ×N denotes the state matrix, while B ∈ R N ×1 , C ∈ R N ×1 and D ∈ R N ×1 denote the projection parameters. For deep sequential modeling, A, B, C and D are parameters that can be learned by gradient descent. The parameter D is omitted for exposition (i.e., D = 0) as Dx(t) can be regarded as a skip connection and is easy to compute [19, 20] .Discretization. The structured state space [20] and Mamba [17] discretize the above continuous system so as to be tailored for deep representation learning. There are usually two ways for discretization, namely, linear recurrence and discrete convolution. For linear recurrence, instead of a continuous function x(t), a discrete sequence (x 0 , x 1 , • • • ) is taken as input. Conceptually, we have x k = x(k∆). The state matrix A is approximated as A by the zero-order hold rule. The discrete SSM is a sequence-to-sequence map x k → y k , given byEQUATIONSelective Scan Mechanism. Prior SSM methods usually focus on the linear time-invariant scenario. Instead, the selective scan mechanism [17] , which is the core of SSM operator in Mamba, learns the dynamism of weights from the input and is more aware of the context information.The Mamba model is a suitable structure that aligns with our needs. When encoding the image patches as sequential data, once important lesion information is discovered, it can be stored in hidden states and propagated to subsequent sequences. Specifically, after sliding the image x ∈ R H×W ×3 into a variety of patches, the input is formed as a sequences of 2-D patches, each of which has a spatial position of H/4 × W/4. Then, in each Samba block, the bi-directional state space modeling module has both feedforward and backforward SSM, where the selective scan mechanism allows to handle the patches in a recurrent manner. The input patches are traversed along two different scanning paths (horizontal and vertical), and each sequence is independently processed by the SSM. Subsequently, the results are merged to construct a 2D feature map as the final output.By a bidirectional design, the severity information can be transported to each patch. The local-toglobal transportation of severity information plays a vital role in the whole process, especially in the selective mechanism. With the guidance of global severity awareness, the update of hidden states can selectively ignore information about low-level lesions, primarily preserving information about the most severe lesions. Specifically, to encode the 2D images, we follow the design of vision Mamba [69] which processes the input features in the forward and backward directions. As illustrated in Fig. 2 , the outputs are gated and added together, while there is a skip connection before input to the EM-based state recalibration module.Another core issue in cross-domain disease grading is the domain generalization ability of the model. Both the intermediate features and the input-dependent parameters in Mamba are affected by the cross-domain variance. To reduce the impact of domain shift, we aim to map the features into a more compact space by feature recalibration. Specifically, the feature distribution of background and grading-related lesions is modeled by a Gaussian Mixture Model (GMM) [42] , given byEQUATIONwhere K is the total number of the Gaussian models, f n is the feature embedding of the n-th patch in image x, µ k and Σ k denote the mean and covariance of the k-th Gaussian basis, respectively. For simplicity, we set Σ k as the identity matrix I. For easy computation, the mixing coefficients of GMM are left out and the exponential inner dot kernel is used. After that, each Gaussian basis is represented by µ k , which is called severity base in this paper. These bases are learnable parameters to capture the representation of lesions. In the recalibration process, instead of directly reconstructing the features based on the bases, we estimate the lesion distribution of each image which is initialized by the severity bases. This is to prevent the loss of useful information during the compression. Concretely, we adopt the EM algorithm to estimate the GMM of each image.Within each iteration, we first estimate z nk in the E-step, which denotes the responsibility of the k-th basis to f n . Here we have 1 ≤ n ≤ N and 1 ≤ k ≤ K. The posterior probability of f n given µ k can be formulated as p(f n |µ k ) = K(f n , µ k ) by a kernel function K. Consequently, estimating the responsibility can be re-formulated into a more general form, given byEQUATIONwhere for simplicity we directly use the exponential inner dot exp(f T µ) as the kernel function.Given the estimated Z t , the severity base likelihood maximization, functioning as the M-step, is realized by updating µ. As the bases are supposed to be aligned to the embedding space of each image, the weighted sum is used to update the bases, given byEQUATIONwhere t refers to the t-th iteration and N p denotes the number of patch embeddings.Assume that the E-step and M-step execute alternately for T c iterations and the convergence criterion has been reached [14] . The final µ Tc and Z Tc are used to recalibrate the image feature F , resulting in F . Here µ Tc = {µ Tc k } and Z Tc = {z Tc nk } refer to the Gaussian basis and the responsibilities of all the patch embeddings from a sample, respectively. This process is mathematically computed asEQUATIONThen, the recalibrated feature F is fed into the next Samba module. During this process, gradingrelated features are mapped to a more compact space, while style differences introduced by image sources are partially removed. Consequently, the critical information transportation within the Mamba model can be more stable in unseen target domains.To alleviate the potential unstable issue, moving averaging is adapted to update the bases µ 0 during the training process. After the T -th iteration, the generated µ T is first averaged over a mini-batch to get µ T . Then, the update of µ 0 with momentum α ∈ [0, 1] is given byEQUATIONConsider the source domain D 1 ∼ P ( X(1) ) and a certain unseen target domainD k ∼ P ( X(k) ),where k = 2, • • • , K. Given a hypothesis h ∈ H, according to the domain adaptation/generalization bound theory [5, 2] , the relation between the target risk ϵ D k (h) and the source risk ϵ D1 (h) can be modeled by a relation inequality, given byEQUATIONwhere d H∆H P ( X(1) ), P ( X(k) ) denotes the distribution gap between the source domain and an unseen target domain, and the right-most term refers to minimal total risk over both domains. In other words, the risk of the proposed Samba on the target domain is bounded by the source domain.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: The computation resources and details are discussed in experiment section.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
datasets & evaluation protocols: 
 Cross-domain Fatigue Fracture Grading Benchmark [31] consists of a total number of 1,785 normal X-ray images and 940 X-ray images with fatigue fracture. They are collected from two hospitals with different types of sensors, which we denote as Domain-1 and Domain-2, respectively. These fatigue fracture images were graded into four stages by three physicians according to the severity level. For simplicity, we denote the grades (including the normal grade) from level-1 to level-5. domain and the other is used as the unseen target domain. According to the severity of breast invasive ductal carcinoma, three grades, namely, rare, frequent and abundant, are annotated. For simplicity, we denote them from level-1 to level-3, respectively.Cross-domain Diabetic Retinopathy Grading Benchmark consists of six DR retinal image datasets, namely, DeepDR [33] , Messidor [1], IDRID [40] , APTOS [3] , FGADR [67] , and RLDR [49] . Following recent work [11] , the single-domain generalization protocol is adapted. Specifically, one of the above six datasets is used as the source domain, and all the rest datasets are used as unseen target domains. Following [11] , two extra large-scale datasets, DDR [29] and EyePACS [15] are used to enrich the source domain for each experiment setting. The development of DR is graded into five levels according to the severity, namely, normal, mild nonproliferative diabetic retinopathy (npdr), moderate npdr, severe npdr and pdr. For simplicity, we denote them from level-1 to level-5.Under the single-domain generalization protocol, three most common evaluation metrics for grading are used, namely, area under the curve (AUC), accuracy (ACC), and F1-score (F1).We conduct extensive ablation experiments to study the impact of iteration number T and optimization of severity basis µ on the unseen target domain. Images from the first clinical center (Domain-1) are used as the source domain, while images from the second clinical center (Domain-2) are used as the unseen target domain for testing. The vanilla Mamba [35] under the Empirical Risk Minimization (ERM) is the baseline.Iteration Number T . In the EM algorithm, the iteration number T plays an important role, because it implements the approximation by iteratively conduct the E and M step. Keeping other hyperparameters and module designs the same, we report the results when the iteration number T of the EM algorithm varies from 1 to 8. The top row of Fig. 3 shows how T impacts the AUC, ACC and F1-score on the unseen target domain. Notice that, the VMamba-ERM baseline does not have EM-based state recalibration. Therefore, the performance of Vmamba-ERM is consistent to T . A too-small T does not reach the convergence criterion, and reduces the effectiveness of feature recalibration. Therefore, a clear performance decline on all the metrics is observed. On the other hand, when T is too large, the representation ability saturates, resulting in little performance improvement, while wasting computation resources.Severity Base Update. We further study how different optimization approaches of the severity base µ k impact the generalization performance on unseen target domain. We study three different settings, namely, no update, only back propogation, and moving average (Eq. 7). The bottom row of Fig. shows the results of the above three settings under a variety of iteration number T . Using moving average to update the severity base µ k is able to improve the performance substantially. It may be explained that the proposed state recalibration is differentiable, thereby enabling the application of back-propagation to update µ 0 . However, the stability of the update cannot be guaranteed due to the EM iterations. Moving average can update µ 0 to avoid collapse.Effectiveness on Recurrent Patch Modeling. The proposed Samba realizes the recurrent patch modeling by harnessing the selective state space model. To demonstrate its effectiveness compared with other recurrent or long-context based representation learning methods, we compare the proposed Samba with vanilla VMamba [69] , Vision Transformer [48] , LSTM [22] and UR-LSTM [18] . Grade-wise Improvement Analysis. We provide a break-down analysis on the grade-wise performance of the proposed Samba and the baseline, i.e., VMamba under the empirical risk minimization (ERM). Table 3 reports the performance. The proposed Samba shows a significant performance improvement on each grade level. Especially, on level-1, level-2 and level-3, the accuracy improvement over the ERM baseline is 17.0%, 7.1% and 7.0%, respectively. Compared to VMamba-ERM baseline, the EM-based State Recalibration in Samba models the feature distribution of lesions via Gaussian Mixture Models with learnable severity bases, and re-estimates by E-M algorithm. The grading features are mapped to a more compact space, and are more stable on unseen target domains.Ablation Studies on the number of Gaussian components K. We study how the number of components K impacts the generalization ability on the unseen target domain. By default K is set to 64, and we further test the performance when K is set to 16, 32, 48 and 96, respectively. The results are reported in Table 4 . When K is set to 64, the proposed Samba achieves the best grading performance, i.e., 60.5%, 72.3% and 59.1% in terms of ACC, AUC and F1-Score, respectively.Scalability to Clinical Computational Pathology. The domain shift in the Cross-domain Breast Cancer Grading Benchmark is from a machine learning perspective, and only handles the magnification difference. However, from a clinical perspective, the computational pathology has to handle the domain shift from not only the magnification difference, but also the staining procedure. However, most existing clinical computational pathology datasets only support the classification task, i.e., separating tumor category from normal category, which is not strongly relevant to our grading task. Appendix A.4 studies the performance of the proposed Samba along with the vanilla VMamba baseline on the CAMELYON17 dataset3 for a clinical sanity check.Comparison with State-of-the-art. We compare the proposed Samba with methods from three primary categories: 1) generic domain generalization methods, including Mixup [61] , MixStyle [64] , DDAIG [63] , ATS [55] , Fishr [41] , and MDLT [56] ; 2) state-of-the-art DR grading methods, which focus on DR grading without explicitly addressing domain generalization, including GREEN [34] , CABNet [21] , Swin-Transformer [36] and MIL-ViT [7] ; 3) domain-generalized DR grading methods, including DRGen [4] and GDRNet [11] . Additionally, the vanilla Mamba [35] results under Empirical Risk Minimization (ERM) are provided as a baseline reference. By default, the results are cited directly from [11] .Table 5 presents a comparison between Samba and existing methods within the context of singledomain generalization. Notably, for DG grading tasks, the metric of accuracy and F1-score are more meaningful than AUC, as the AUC can be made artificially high due to the large amount of negative samples belonging to other stages. Therefore, we only involve ACC and F1 for comparison.The proposed Samba achieves a substantial improvement over state-of-the-art domain generalized DR grading methods. Especially, on APTOS, DeepDR, FGADR, Messidor and RLDR, it outperforms the second-best in terms of ACC and F1 by 5.2% and 2.2%, 27.2% and 5.7%, 60.7% and 31.3%, 6.7% and 1.3%, 28.3% and 4.7%, respectively. The significant improvement on FGDAR may be explained that it has a different severity-level sample distribution than other datasets. The samples without DR (level-1) only occupy only 5.5% among all the training samples, which are far less than others (e.g., level-1 samples occupy 49.3% in APTOS). Therefore, existing methods may overfit other severity levels and underfit level-1. In contrast, the selective scan mechanism of the Vmamba-ERM and Samba is robust to this severity distribution shift. The EM state re-calibration in Samba makes the feature space more compact, and improves the generalization.Additionally, Samba shows a marked improvement over the baseline VMamba model under Empirical Risk Minimization. Especially, on APTOS, DeepDR, FGADR, IDRID, Messidor and RLDR, it outperforms the VMamba-ERM baseline in terms of ACC and F1 by 1.3% and 1.7%, 2.2% and 2.1%, 3.0% and 1.6%, 3.7% and 2.6%, 7.3% and 2.7%, 7.4% and 3.4%, respectively. These results demonstrate its effectiveness in handling domain gap.Understanding Recurrent Patch Modeling. We validate if the proposed recurrent patch modeling can store and transport the lesion information. An intuitive way is to inspect the correlation between the patch embeddings before and after the recurrent patch modeling. Therefore, we extract the patch embeddings before and after the fourth block. We compute the correlation matrix between the patch-wise embeddings and visualize the results in Fig. 4 . After processed by the Recurrent Patch Modeling module, more regions in the correlation matrix have higher responses. Specifically, after passing through certain high-response positions, the relevant information is transmitted to the subsequent patches in the forward direction. The high-response patches have grade-related lesions and the information is transported in the recurrent process.t-SNE Visualization. To assess the generalization capacity of the proposed Samba, we analyze the feature distribution across the source and unseen target domains using t-SNE visualization in Fig. 5 , which compares the t-SNE plots of the ERM baseline (left) and the proposed Samba (right).The results indicate that, after the EM-based state feature calibration, the proposed Samba enables feature embeddings from different domains to achieve a more uniform distribution, thereby reducing the domain gap. This improved uniformity in the feature space suggests that Samba can enhance generalization, contributing to better performance on unseen domains. 
conclusion: 
 In this paper, we aimed to tackle a practical but challenging task, learning domain generalized medical image grading. We mainly focused on two issues: the identification of decisive lesions and the impact caused by inter-domain differences. The proposed severity-aware recurrent modeling adopts a state space model to store and transport the severity information from local to global. To further mitigate the impact of cross-domain variants, an EM-based state recalibration was designed to map the patch embeddings into a compact space. The proposed method can be used in a variety of disease grading scenarios, providing an effective tool for automatic medical image analysis.Limitation Discussion & Broader Societal Impact. The feature distribution of lesions is modeled by the Gaussian mixture model and estimated by the Expectation-Maximization algorithm. However, when the training source domain has severe class imbalance, the estimated probability distribution by the proposed Samba may not necessarily reflect the domain-agnostic lesion distribution. Nevertheless, the proposed method can be combined with other techniques specifically designed for addressing class imbalance. The proposed method advances the versatility of automatic disease diagnosis, which benefits the human well beings. We do not envision negative societal impact. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: At the end of the conclusion section, the broader impacts have been discussed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This work does not have such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The asserts used in this paper are all public available for academic researches. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: A.1 Severity Base NormalizationDuring the iteration steps t = 1, • • • , T , the severity base µ(1),t k may not deviate too much from each other, which otherwise can lead to collapse when back propagation. We study the scenarios when no normalization, L-1 normalization and L-2 normalization are used on these severity basis. Fig. 6 shows the results of the above three settings when under a variety of iteration number T on the unseen target domain. L-2 normalization achieves the best performance on all the metrics, especially when T becomes larger. Fig. 4 in the main text only visualizes the correlation between the patch embeddings before and after the recurrent patch modeling when trained on the APTOS dataset. Here we further demonstrate the results from FGADR, IDRID, Messidor and RLDR. They are visualized in Fig. 7a, b, c and d , respectively. On all these datasets, we can observe a common pattern. After processed by the Recurrent Patch Modeling module, more cells in the correlation matrix have higher response. Usually, a handful of the patches inside the image have grade-related lesions. After the processing of our module, the information of these grade-related lesions is transported to other patches. It allows the model to perceive a more global-wise representation. Consequently, more patches that contain the grade-related lesion information are activated, and more cells are highly responded in the correlation matrix.We model the relation between patch embedding from SSM and severity level by drawing inspiration from the class activation map (CAM) mechanism [62] . Specifically, we take the patch embeddings from the last Samba block as input to generate the per-level severity activation patterns. Then, the activated severity patterns are displayed on the original images. We use FGADR as the unseen target domain. The results are shown in Fig. 8 , where the activated patches are highlighted in blue boxes.From the first to the fifth row, the samples from level-1 to level-5 are provided accordingly. From the first to the fifth column, the patch activation maps from level-1 to level-5 are displayed. Notice that, as level-1 refers to the normal scenario, each sample has activations on level-1, meaning some patches are normal.It is important to note that the domain shift in the Breast Cancer Grading Benchmark is technically from a machine learning perspective, and only handles the domain shift from the magnification difference. However, from a clinical perspective, the computational pathology has to handle the domain shift not only from the magnification difference, but also from the staining procedure. Therefore, it is beneficial to test if the hypothesis works in a real-world computational pathology scenario. However, a bottleneck is that, most clinical computational pathology dataset so far conducts the classification task, i.e., separating tumor category from normal category. Therefore, in this The correlation matrix of each patch embedding before and after processed by the recurrent patch modeling in the forward direction, denoted as 'Before' and 'After', respectively. The higher correlation, the more red a cell is. subsection, the proposed Samba along with the vanilla VMamba baseline are benchmarked on the CAMELYON17 dataset 4 for the cross-domain computational pathology classification task.The first experiment is the impact of the number of components K in GMM, where ACC, AUC and F1 are used as evaluation metric. The results are reported in Table 6 . By default K is set to 64, and we further test the performance when K is set to 16, 32, 48 and 96, respectively. When K is set to 64, the proposed Samba achieves the best grading performance. This observation is consistent to the performance on Cross-domain Breast Cancer Grading Benchmark, where a number of 64 Gaussians achieves the optimal performance.The second experiment is to analyze the trade off between the classification performance and the baseline model. Both the VMamba-ERM and the proposed Samba are tested, where only accuracy is used as the evaluation metric. The results are reported in Table 7 . The trend is the same as the trend on the Breast Cancer Grading Benchmark. Using Samba on each type of the VMamba backbone shows a clear performance improvement on unseen domains.Fig. 9 and Fig. 10 show some attention maps of the Samba on unseen retinal images. The proposed Samba is able to model the recurrent relation among patches. Therefore, the activation regions can general cover the lesions and are more robust to the domain shift.Level-2 Level-3 Level-4 Level-5Level-1Level-2Level-3Level-4Level-5Figure 8 : Per-severity activation map of the proposed Samba. From the first to fifth column are the activated patches from level-1 to level-5, highlighted in blue boxes. From the first to the fifth row are the samples with an annotation from level-1 to level-5. FGADR is the unseen target domain. Justification: This paper propose a Severity-aware Recurrent Modeling (Samba) for medical image grading problems on unseen target domains.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: The authors have read the code of ethics. The experiments are all on public datasets without obeying the code of ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).https://github.com/YANRUI121/Breast-cancer-gradinghttps://camelyon17.grand-challenge.org/https://camelyon17.grand-challenge.org/ 


Decision: The paper was reviewed by four experts. Although there was some disagreement about the novelties and clarity of the paper before the rebuttal and discussion stages, the authors were able to address most of the concerns. All reviewers recommend acceptance of the paper.  The authors are advised to address the requested clarifications in the final version of the paper.

Reviews: review: paper_summary: paper_summary: The authors propose a model which they call Samba, for Severity-aware recurrent modelling, which is a method designed for cross-domain medical image grading. They introduce several challenges in medical grading, namely the difficulty models encounter in generalising to unseen domains, as well as the existence of ambiguity in lesion severity grading. The model is comprised of two main blocks: recurrent bidirectional Vision Mamba layers and a Expectation-Maximisation State Recalibration (EMSB) module which consists of learnable tokens which capture lesion representations, and which are then used as bases to map to more compact feature embeddings. These are then used to initialise an Expectation-Maximisation (EM) algorithm and the lesion feature distribution is estimated using Gaussian Mixture Models (GMMs) for each image. The Mamba layers treat the image patches as sequential data, with the rational being the relevant information will be propagated through the hidden states. The EM module models the feature distribution of lesions using the GMMs to try and eliminate domain shift. Samba is applied to three benchmarks, where images are separated between a source domain and a target domain: Diabetic Retinopathy fundus images (DR), Fatigue Fracture X-rays and Breast Cancer histopathology images. The model is trained on the source domain and applied to the target domains. Ablation is carried out, comparing Samba to different baseline and SOTA models, as well as looking at iteration number and severity base update methods. The results obtained show Samba generally outperforms existing methods to some degree. The authors also provide a theoretical analysis on the generalisation risk bound in the Appendix.
strength_weakness: strength_weakness: In my opinion the paper's main weaknesses are the following:

1 - Unclear explanation of the EM-based State Recalibration module

I think the introduction to the EM-based State Recalibration module in this paper is too brief, perhaps assuming reader familiarity, and does not provide intuition as to why using GMMs + EM makes sense in this context or is more suitable to this task than other potential approaches. The implementation details of how the the GMM \& EM is integrated into the overall Samba architecture are not clearly or fully presented. The initialisation strategy for the EM is unclear to me: the paper mentions using learnable severity bases to initialise EM, but doesn't explain how these bases are learned or why this initialisation is most beneficial. It also doesn't discuss the convergence criteria used for the EM and although it shows ablation into how many iterations are best, this comes without proper introduction in the background or methods sections. The same goes for the number of Gaussian basis used: you show ablation on this, but don't introduce this in background or methods sections. Finally, although some equations are presented outlining the E- and M-steps, variables are not always defined and the choice of kernel function seems entirely ad-hoc and unjustified. In particular

You define your GMM as:

$$p(f_n) = \sum_{k=1}^K z_{nk} \mathcal{N}(f_n|\mu_k, \Sigma_k),$$ 

but $z_{nk}$ typically represents responsibilities which are estimated in the E-step, not given as part of the model definition, where $\pi_k$ would represent the mixing coefficients of the GMM. Then you derive the E-step as:

$$z_{nk} = \frac{\mathcal{K}(f_n, \mu_k)}{\sum_{i=1}^K \mathcal{K}(f_n, \mu_i)}$$

whereas a standard formulation would be 

$$z_{nk} = \frac{\pi_k \mathcal{N}(f_n|\mu_k, \Sigma_k)}{\sum_{i=1}^K \pi_i \mathcal{N}(f_n|\mu_i, \Sigma_i)}.$$

However, you don't justify the use of the arbitrary kernel function $\mathcal{K}$ (exponential inner dot $\mathrm{exp}(f^T\mu)$ over standard Gaussian probabilities. Could you explain this in more detail? Then you don't explicitly define what $Z^t$ represents, nor what the relationship is with $z_nk$. Likewise with $F$, you do not explicitly define what it represents, but I will assume you mean the feature matrix of the input image. How does $F$ relate to equation (5)? How does this derive into equation (6)? The recalibration step $\tilde F$ is not well explained in terms of matrix operations or dimensions. This lack of clarity makes it difficult to interpret the method in terms of a well-established optimisation algorithm, which is an important issue as explaining the transformations applied to the feature representations is crucial to understanding Samba. Overall, while the general idea of using GMM-EM for feature recalibration is interesting, the mathematical formulation presented in the paper has inconsistencies and deviations from standard GMM-EM that I don't think are well-justified and which I would like to see more fully explained and motivated. 

2 - Unclear structure of results and lack of clear baseline comparisons

While the paper compares three datasets to various baseline and baseline methods, each dataset is treated separately, so it's not always clear the results shown extend to the other datasets, if the comparisons are fair or how the baselines were implemented. I feel like the main results are the ones presented for the Diabetic Retinopathy dataset, yet these are presented last. I think the results section would read better and make more sense if these were presented first. However, one of my main issues here is that you only mention your main baseline comparison method on p.7, but you compare against in Figure 3 without even mentioning it. This should be properly introduced and you should explain why VMamba under empirical risk minimisation (VMamba-ERM) is a good baseline model against which to compare. You also need to explain how Samba differs from VMamba-ERM and what advantages these differences bring to the task that the other doesn't provide. Finally, you don't provide ablation on Samba itself, showing the individual contribution of the Mamba layers vs EMBS layers. With regards to the theoretical analysis on the generalisation risk bound, it feels disconnected to the rest of the paper and not particularly relevant to the stated aims. Is there some way to tie it more or explain why its important? Finally, I think an Appendix section on how the datasets were preprocessed and showing the hyperparameters employed for all the other models presented in the results sections would enrich this paper. There are other things which I find unclear, which I have put below in Questions.
questions: questions: Line 44 - I don't fully understand this sentence. Maybe you could rephrase?

Line 54 - at the distal of what? 

Line 123 - what's a stem unit?

Figure 2 - I think you need to describe the Samba blocks in more detail - here it looks like the bi-directional SSM and EM recalibration occurs in the downsampling block. Is this a mistake? What'a VSS block? 
What's the difference between the training and inference arrows in the figure?

Figure 3 - Iteration Number T - What models are you comparing yours against? What is Vmamba-ERM and what is SOSS?? Why is Vmamba showing a constant trend across iteration number T? What are the different columns? 

Line 235 - why does the number of iterations play an important role in the EM algorithm? You should introduce this beforehand and the ablation study should answer this question.

Line 237 - do you mean Figure 3 instead of Figure 6 here? - Figure 6 shows the effect of regularisation techniques and is situated in the appendix. 

Line 244 - why are these these methods appropriate to test different optimisation techniques here? Please explain this. 

Line 252 - these are percentage points, not percentage of improvement. 

Table 1- Why is Table 1 showing three F1 scores highlighted in brighter green? and two blocks of yellow and light green? Is this a mistake? Also what do these colours represent? 

Table 3 - OK, why do you think Samba is doing better there?

Line 283 - percent points, not percents...

Table 4 - why do you think Samba and Vmamba-ERM are doing so well on FGADR compared to other methods?

A.3 - the text needs a rewrite to make sense I think. Please check this. 

Figure 8/9 - Do you have any ground truths available to compare the attention maps shown in Figure 8 and 9?
limitations: limitations: The authors point out some limitations of their work with regards to class imbalance. I have pointed out other perceived limitations in the Weakness section above.
paper_summary: paper_summary: The paper introduces a new method for disease grading for both within and cross-domain medical images. Three different imaging modalities have been used for experimentation including Retinal, X-ray, and H&E images. In terms of methodological novelty, the authors propose to encode image patches in a recurrent manner to capture informative lesions. Further, they use an EM-based recalibration method to reduce the cross-domain variance by compacting the feature space. Overall, their experiments shows that the method proposed is superior to the baselines.
strength_weakness: strength_weakness: There are a few major issues with the experimental design:
1. The Cross-domain Breast Cancer Grading Benchmark is not a well-defined cross-domain problem in computational pathology. Firstly, this is not a widely used dataset in the field. Secondly, the images come from the same center yet with different scanning magnifications. It is not common to use images from two different magnifications as the "domains". The cross-domain problem in computational pathology is when the images are actually from two different centers and with two different staining protocols (images might be scanned in varying magnifications). An example of such a problem can be found here: https://camelyon16.grand-challenge.org/Data/

2. In Table 4, the authors have reported the benchmark from [8]. It has not been mentioned whether the proposed model was trained on the exact same seed, device, and cross-validation folds. If any of these are different, the comparison is not fair. Instead, I'd suggest the authors benchmark these with the same exact setting and report the results.

3. For highly imbalanced datasets, especially medical data, it is proper to report Balanced Accuracy to show the performance of the model on the rare classes. Both ACC and AUC fail to represent the rare classes. This is important for comparison as results in Table 1 show that F1 is significantly lower than ACC and AUC. Therefore, Balanced accuracy should be reported for all the benchmarks.

4. There are quite a few modules in different parts of the model that have not been studied in an ablation study. These modules are the bach norm, linear layer, etc. in the EM-based recalibration. Adding the ablation study for these modules can enhance the work :)
questions: questions: I am curious if the authors can elaborate more on this conclusion made in page 15: "After processed by the Recurrent Patch Modeling module, more regions in the correlation The high-response patches have grade-related lesions and the information is transported in the recurrent process matrix have higher response."
limitations: limitations: It has been justified properly.
paper_summary: paper_summary: The authors introduce a new method named Severity-aware Recurrent Modeling (Samba) for disease grading in medical imaging. Specifically, they propose encoding image patches in a recurrent manner to accurately capture decisive lesions and transmit critical information from local to global contexts. Additionally, an Expectation-Maximization (EM) based state recalibration mechanism is designed to map feature embeddings into a compact space, thereby reducing the impacts of cross-domain variations.
strength_weakness: strength_weakness: 1.	The significance of using the Mamba architecture in cross-domain tasks is uncertain; in fact, restricting the Mamba module may lead to excessive specialization.
2.	The ablation experiments regarding the specific structure of Samba are not sufficiently comprehensive.
3.	The open-source code is not explicitly provided.
questions: questions: 1.	Among the compared methods mentioned in Table 4, some ViT-based and ResNet-50-based methods are not specifically designed for domain adaptation tasks. Is this comparison fair?
2.	The article states: "The Mamba model is a suitable structure that aligns with our needs. Guided by global severity awareness, the update of hidden states can selectively ignore information about low-level lesions, primarily preserving information about the most severe lesions." Can you explain how SSM selects severe lesion information? A theoretical justification should be provided, as this is a point of concern.
3.	In the comparative experiments with other SOTA methods, are all the source domains supplemented with the two additional large-scale datasets DDR [24] and EyePACS [13]?
limitations: limitations: The authors have addressed the limitations of the current work.
paper_summary: paper_summary: Accurate disease grading in medical image analysis is challenging due to the variability within disease levels and the similarity between adjacent stages. Additionally, models must handle data from unseen target domains, where differences in feature distribution can significantly reduce performance. To address these issues, this paper proposes the Severity-aware Recurrent Modeling (Samba) method, which encodes image patches sequentially to capture severity information and employs an Expectation-Maximization based recalibration mechanism to handle cross-domain variations. The method also uses a Gaussian Mixture Model to model feature distributions and reconstructs intermediate features using learnable severity bases.
strength_weakness: strength_weakness: - Some parts need more detailed explanation. Although space is limited, certain sections assume a high level of pre-knowledge. For example, more explanation on Samba would be useful.
- For me it is not really clear how the patches are used in a recurrent manner and what their sequence should be. Can you please elaborate more on the recurrent part. 
- There is no study demonstrating that the model specifically attends to patches related to severity. Is there a way to illustrate this?
questions: questions: - Adding specific numbers and metrics to the abstract would be beneficial. Currently, lines 18-19 are quite vague.
- On page 7, the authors likely intend to refer to Figure 3 for the results, rather than Figure 6.
- In Figure 5, the T-SNE plot is difficult to interpret due to small icons. Could the datasets be represented by icons and the severity by color for better qualitative evaluation?
limitations: limitations: The authors addressed the limitations of the work.

Comments: title: Re: Official Comment by Reviewer LrSF
Re: Response to Author's Rebuttal
Response to Author's Rebuttal
Re: Official Comment by Reviewer b4sx
Reply to Q2: Cross-scale or Cross-domain?
Reply to Q1: Effectiveness on a common computational pathology dataset.
Response to Reviewer hFq3
comment: We are glad that your concerns have been properly addressed.

We will significantly polish and improve our work per you suggestions, particularly on prioritizing C17 over BCGR for a professional standard of the pathology perspective. 

Finally, thanks again for your time and effort in helping us improve our work.
I appreciate the detailed response and the effort of the authors to provide further insights. 
I understand that DG on cross-scale data is a valid problem from an ML point of view. Yet, my main concern is rooted in the choice of the dataset, where among many well-studied datasets present, the BCGR dataset had been chosen. Still, I am not convinced it is a good representative unless an extensive benchmark with different methods is provided. 

However, by including C17, and providing the ablation studies and the benchmarking results on that, my initial concern was addressed. With that, I strongly suggest the author use the C17 as the main pathology representative dataset and use the BCGR results as a secondary dataset. Also, a clarification section needs to be added to the work to explain that BCGR is a cross-scale generalization and is not multi-domain as it is known in pathology. This will ensure the reliability of the result for a pathology-related audience. Given the new set of results, I would raise my rating to a weak accept :)
We are glad to see your concerns have been addressed.

We will improve our paper carefully according to your valuable suggestions.
Thanks to the authors for the time and effort put into the rebuttal. I have read their response, as well as the other reviews and their associated comments. 

I agree with the authors that testing across multiple magnification (i.e. scales), although not the standard approach, can constitute cross-domain adaptation if it can be shown there exists a clear difference in the data distribution across domain. They have shown this to be the case with their UMAP embeddings and changes in magnification and resolution across scanners resulting in models underperforming is a know phenomenon in computational pathology. Additionally, they also test on X-ray fatigue fractures and Diabetic Retinopathy datasets. 

I appreciate the authors promise to update the structure of the paper: modifying/correcting pipeline figure, expanding on background and motivation, clarifying model explanation and mathematical formulation in 3.3, moving 3.4 to Appendix, introducing baseline models clearly, introducing first results with Figure 4 (DR dataset), correcting Table 1 and expanding on model ablation as shown in Table R1. Given the above reasons and the substantial work required to put these into effect, I am maintaining my original rating.
We are glad that you are satisfied with the revision and find the concerns resolved.

We will improve our paper carefully per your valuable suggestions.
Thanks to the authors for actively responding to the raised concerns and promising revisions in the updated version. However, considering the innovativeness of Samba (based on Mamba for various tasks), the score suitable for a Borderline Accept will not be altered.
**Q2**: Cross-scale problem or Cross-domain problem?

**R**: 
First of all, we would like to thank the reviewer for the valuable feedback on the *cross-scale* perspective.

However, we humbly suggest the reviewer that there might be some misunderstanding on *cross-domain*. 

From a machine learning perspective, where the venue of NeurIPS focuses more on, as long as the source and unseen target domains are *not independent and identically distributed* / *not the same* [a,b,c], the domain gap exists and the domain generalization techniques can be applied.
The change of scale, caused by the magnification, also poses distribution shift between the source and target domain, can also be treated as a type of domain generalization, according to this machine learning definition.

[a] Zhou, Kaiyang, et al. "Domain generalization: A survey." IEEE Transactions on Pattern Analysis and Machine Intelligence 45.4 (2022): 4396-4415.

[b] Zhou, Kaiyang, et al. "Domain Generalization with MixStyle." International Conference on Learning Representations. 2021.

[c] Wang, Jindong, et al. "Generalizing to unseen domains: A survey on domain generalization." IEEE transactions on knowledge and data engineering 35.8 (2022): 8052-8072.

Therefore, we humbly suggest that conceptually from a machine learning perspective, the Cross-domain Breast Cancer Grading Benchmark is applicable to validate the generalization capacity of a model.
Besides, as we have already provided in the 1-pg PDF file, the distribution shift exists between the Doman-1 and Domain-2 of this dataset. 
This clearly-existed domain gap, along with the definition of domain generalization, means the Cross-domain Breast Cancer Grading Benchmark is rational to benchmark the domain shift of grading problems.
In real-world scenarios, some hospitals may use magnification factors that are not presented in the training set. Hence, testing the cross-magnification generalization ability by the Cross-domain Breast Cancer Grading Benchmark still has the practical significance.

**Summary**:
As we mentioned at the beginning, we respect and value the reviewer’s perspective from computational pathology. 
We are glad and open if the reviewer feel it would be better to treat the results on cross-domain breast cancer grading are a discussion on the scale generalization or replace them by the results on CAMELYON17.

We hope we can reach a consensus, and look forward to your feedback and suggestions.
**Q1**: Effectiveness on a common computational pathology dataset.

**R**: We respect and value the reviewer’s perspective from computational pathology.
We adopted CAMELYON17 to conduct the two experiments in this paper where Cross-domain Breast Cancer Grading Benchmark has been used for validation.
Same as the earlier rebuttal, five individual domains marked by the staining protocols and sites in CAMELYON17 are used for the cross-domain experiments. We use Domain-1 as the source domain, and reports its performance on the rest four unseen target domains, i.e., from Domain-2 to Domain-5.

The first experiment (Table 2 in the main text) is the impact of the number of components K in GMM, where ACC, AUC and F1 are used as evaluation metric. The performance on CAMELYON17 is attached as follows.

Table 2: Impact of the number of components K in GMM on unseen target domain performance. Experiments are conducted on the CAMELYON17. Domain-1 is used as source
domain. Metrics presented in percentage (%).

| | | Domain-2 | | | Domain-3 | | | Domain-4 | | | Domain-5 | | 
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
| K value | ACC | AUC | F1 | ACC | AUC | F1 | ACC | AUC | F1 | ACC | AUC | F1 | 
| 16 | 81.80 | 92.09 | 79.96 | 79.45 | 88.96 | 76.60 | 83.92 | 95.38 | 82.10 | 75.84 | 81.99 | 74.16 | 
| 32 | 82.68 | 93.81 | 80.95 | 80.39 | 91.28 | 77.97 | 84.38 | 96.20 | 82.96 | 76.53 | 82.46 | 74.92 |
| 48 | 84.06 | 94.25 | 81.83 | 81.84 | 92.60 | 78.49 | 85.87 | 97.16 | 83.87 | 77.69 | 83.85 | 75.04 | 
| 64 | **84.59** | **95.67** | **83.10** | **82.48** | **93.32** | **79.85** | **86.50** | **97.82** | **85.13** | **78.64** | **84.70** | **75.32** |
| 96 | 84.27 | 95.48 | 82.53 | 82.06 | 92.90 | 79.44 | 86.16 | 97.45 | 84.80 | 78.15 | 84.28 | 74.86 |
| 128 | 83.65 | 94.70 | 81.97 | 81.50 | 92.41 | 78.36 | 85.47 | 96.90 | 84.62 | 77.38 | 83.66 | 74.14 |
||||||||||||||

Default K is set to 64, and we further test the performance when K is set to 16, 32, 48 and 96, respectively. 
The results are reported in Table 2. When K is set to 64, the proposed Samba achieves the best grading performance.
This observation is consistent to the performance on Cross-domain Breast Cancer Grading Benchmark, where a number of 64 Gaussians achieves the optimal performance.

The second experiment (Table 3 in the main text) is to analyze the computational cost and performance trade between the VMamba-ERM and the proposed Samba, where only accuracy is used as the evaluation metric. The performance on CAMELYON17 is attached as follows.

Table 3: Computational cost comparison between VMamba-ERM and the proposed Samba. Experiments are conducted CAMELYON17. Domain-1 is used as source domain. Metrics in percentage (%).

|Method | Backbone | GFLOPS | Para. | Domain-2 | Domain-3 | Domain-4 | Domain-5 | avg. |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|
| EMR | VMamba-T | 3.7 | 32.7 | 70.08 | 67.29 | 72.96 | 63.16 | 68.37 |
| Samba | VMamba-T | 5.5 | 32.7 | **78.74** | **76.15** | **80.06** | **71.05** | **76.50** | 
| EMR | VMamba-S | 7.9 | 63.4 | 72.86 | 69.50 | 75.08 | 65.72 | 70.79 |
| Samba | VMamba-S | 11.3 | 63.4 | **81.01** | **78.96** | **82.75** | **73.88** | **79.15** |
| EMR | VMamba-B | 14.0 | 112.4 | 76.23 | 74.17 | 79.53 | 69.87 | 74.95 |
| Samba | VMamba-B | 19.6 | 112.4 | **84.59** | **82.48** | **86.50** | **78.64** | **83.05** |
||||||||||

The same trend is observed on this dataset, where using Samba on each type of the VMamba backbone shows a clear performance improvement on unseen domains.

We sincerely wait for the reviewer’s feedback about the above experiments, which function the same as the Cross-domain Breast Cancer Grading Benchmark in the submission.
If they could meet the reviewer’s standard on computational pathology, we are happy and open to incorporate it into our work, subtyping the computational pathology part.
I would like to thank the authors for their significant effort and for providing more details in the rebuttal phase.

However, I believe the Breast Cancer Grading Benchmark is not a representative choice for claiming that the proposed approach is a good solution for a computational pathology task. Multiple well-studied and standard datasets in the field, such as TCGA sub-datasets, could have been chosen. These datasets can be used for grading (to match the rest of the paper) or subtyping tasks based on the proposed research. Camelyon17 is also a viable choice if the authors decide to work on subtyping.

Another important factor is that several ablation studies have been conducted on the Breast Cancer Grading Benchmark, which is still questionable in terms of reliability and generalizability of results. This dataset should ideally be replaced with another dataset, as mentioned above.

I also appreciate the authors' effort in providing details in the PDF file. It is important to note that there are stain variations within the same center samples due to the staining procedure, which accounts for in-domain data variance. Additionally, differences in magnification are not technically a domain-shift problem; rather, they are known as a cross-scale problem [1][2], which is not aligned with the paper's topic and the rest of the experiments. Thus, to fairly test the hypothesis in this field, the study should cite and compare proper literature on a standard dataset.

I have also considered the paper that released the Breast Cancer Grading Benchmark [3] dataset. However, within that paper, the authors do not consider their dataset as a cross-domain dataset and instead build a cross-scale model.

Based on the above rationale, I am not convinced to change my rating.

[1] Sikaroudi, M., Ghojogh, B., Karray, F., Crowley, M., and Tizhoosh, H.R., 2021, April. Magnification generalization for histopathology image embedding. In 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI) (pp. 1864-1868). IEEE.

[2] Chhipa, P.C., Upadhyay, R., Pihlgren, G.G., Saini, R., Uchida, S., and Liwicki, M., 2023. Magnification prior: a self-supervised method for learning representations on breast cancer histopathological images. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 2717-2727).

[3] Yan R, Ren F, Li J, Rao X, Lv Z, Zheng C, Zhang F. Nuclei-Guided Network for Breast Cancer Grading in HE-Stained Pathological Images. Sensors. 2022; 22(11):4061.
Thanks for your swift response. We are glad to see your questions resolved.

We will improve our paper carefully according to your valuable suggestions.
I can certainly appreciate the significant effort put into the rebuttal and I remain as a weak accept.