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

Q1: 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.

R: Thanks for your valuable feedback. We will provide more explanations of the proposed Samba, such as:

1. After proceeded by BSSM, the feature embedding of each patch embedding $\boldsymbol{f}_n$ is firstly projected into the latent space by a linear layer to get the target state embedding $\boldsymbol{x}_n$. The severity base $\boldsymbol{\mu}_n$ is modeled as a mixture of Gaussian from $\boldsymbol{x}_n$. Then, the E-M algorithm iteratively approximates $\boldsymbol{\mu}_n$ to $\boldsymbol{x}_n$. After convergence, the embedding is fed into the rest module.
2. Each Gaussian is initialized by the Kaiming initialization.
3. The E-M empirically implements 3 iterations, according to the observation in Fig.3.
4. The number of Gaussian kernels $K$ in L171 will be defined when first introducing.
5. The specialization of Samba for cross-domain tasks, especially, the feature distribution shift between the source domain and unseen target domains affects not only the intermediate features but also the selective scan mechanism of Mamba, which poses a clear performance drop of Mamba on unseen target domains. To address issue, the proposed method not only introduces both forward and backward directions to primarily preserve information about the most severe lesions, but also introduces a EM-based State Recalibration mechanism to compact the feature space so that the feature distribution is less sensitive to the domain shift.

Q2: 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?

R: Thanks for your valuable feedback. After sliding the image $\mathbf{I} \in \mathbb{R}^{H \times W \times 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 \times 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. Specifically, the 2-D selective scan mechanism in both components is directly inherited from [15]. 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.

We will enrich these details accordingly.

Q3: There is no study demonstrating that the model specifically attends to patches related to severity. Is there a way to illustrate this?

R: Thanks for your valuable suggestion. We model the relation between patch embedding from SSM and severity level by drawing inspiration from the class activation map (CAM) mechanism [a]. We take the patch embeddings from the last Samba block as input, so as 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.~R1, 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 map from level-1 to level-5 is generated by the aforementioned methods. Notice that, as level-1 refers to the normal scenario, each sample has activations on level-1, meaning some patches are normal.

[a] Zhou, Bolei, et al. "Learning deep features for discriminative localization." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

Q4: Adding specific numbers and metrics to the abstract would be beneficial. Currently, lines 18-19 are quite vague.

R: Thanks for your valuable suggestion. We will accordingly add the following in the introduction: Extensive experiments show that 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.

Q5: On page 7, the authors likely intend to refer to Figure 3 for the results, rather than Figure 6.

R: Sorry for the typo. We will correct it from Fig.6 to Fig.3.

Q6: In Fig.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?

R: Thanks for your valuable suggestion, and we have modified it accordingly. Please refer to Fig.~R3 in the attached PDF file for reference.

Finally, should you have further suggestions and comments, we are glad to incorporate during the discussion stage.

Q1: Significance of using Mamba in cross-domain tasks; restricting Mamba module may lead to excessive specialization.

R: Thanks for your valuable comments, so that we could have a chance to clarify the generalization and universality of the proposed Samba. Specially,

(1) The feature distribution shift between the source domain and unseen target domains affects not only the intermediate features but also the selective scan mechanism of Mamba, which poses a clear performance drop of Mamba on unseen target domains. To address issue, the proposed method not only introduces both forward and backward directions to primarily preserve information about the most severe lesions, but also introduces a EM-based State Recalibration mechanism to compact the feature space so that the feature distribution is less sensitive to the domain shift.

(2) From the experimental side, the proposed Samba shows a significant unseen domain performance improvement than the VMamba baseline on three different modalities, varying from fundus images, X-ray images, and pathological images, respectively. Moreover, as discussed in the paper, the recurrent encoding of image patches also contributes to cross-domain disease grading. In Tab.4, even the VMamba baseline outperforms CNN- and ViT-based methods.

We will enrich these discussions accordingly.

Q2: Ablation studies regarding the specific structure of Samba.

R: We provide an ablation study on the specific structure of Samba. Specifically, two components, namely, Bi-directional State Space Modeling (BSSM) and EM-based State Recalibration (ESR), are incorporated into the VMamba baseline. The experiments are conducted on the DG Fatigue Fracture Grading Benchmark, and the results are reported in Table R1 in the attached PDF. It is observed that BSSM contributes to an ACC, AUC and F1 improvement of 5.2%, 1.7% and 4.9%, respectively. ESR contributes to an ACC, AUC and F1 improvement of 18.3%, 9.4% and 12.2%, respectively.

Q3: The open-source code is not explicitly provided.

R: Thanks for your valuable suggestion. Owing to the company regulation, the source code will be made available after publication. We respectively ask for the reviewer's understanding on this restriction.

Q4: 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?

R: Thanks for your value comments, so that we could have a chance to clarify the compared methods. We first compare six methods for domain generalization tasks, namely, [36,50,51,53,55,56]. Then, we compare two DR grading methods under the domain generalization setting, namely [4,8]. All these eight methods are plenty for comparison under the context of domain generalization. The rest four CNN or ViT based methods [6,19,29,31] do not have domain generalization property, and are involved just for boarder comparison. Besides, these methods are implemented under the emperical risk minimization (ERM) setting, which is the same as the VMamba baseline. This demonstrates the effectiveness of Mamba structure in the cross-domain grading tasks.

Q5: Can you explain how SSM selects severe lesion information?

R: Thanks for your valuable suggestion.

We model the relation between patch embedding from SSM and severity level by drawing inspiration from the class activation map (CAM) mechanism [a]. Specifically, give an image $\boldsymbol{x}$, assume that $\boldsymbol{\tilde{f}}_k(i)$ denote the activation of unit $k$ from the last Samba block at the patch $i$. Then, for unit $k$, after implementing a global average pooling (GAP), the feature embedding $F^k$ is $\sum_i f_k(i)$.

Thus, for a certain severity level $c$, the input to the softmax $S_c$ is $\sum_k w_k^c F_k$, where $w_k^c$ is the weight corresponding to severity level $c$ for unit $k$.
Here $w_k^c$ indicates the importance of $F_k$ for severity level $c$. Finally, the output of the softmax for severity level $c$, $P_c$ is computed as

$$P_c = \frac{{\rm exp} (S_c)}{\sum_c {\rm exp} (S_c)}.$$

By plugging $F_k = P_i f_k(i)$ into the severity level score $S_c$, we obtain

$$S_c = \sum_k w_k^c \sum_i f_k(i) = \sum_i \sum_k w_k^c f_k(i).$$

Then, the severity activation map $M_c$ for severity level $c$ is defined as follows, where the activation of each patch $i$ is computed as

$$M_c(i) = \sum_k w_k^c f_k(i).$$

We take the patch embeddings from the last Samba block as input, so as to visualize 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.~R1 in the attached 1-pg PDF, 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 map from level-1 to level-5 is generated by the aforementioned methods. Notice that, as level-1 refers to the normal scenario, each sample has activations on level-1, meaning some patches are normal.

[a] Zhou, Bolei, et al. "Learning deep features for discriminative localization." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

Q6: 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]?

R: Yes, as all the performance of the state-of-the-art methods reported in [8] is supplemented with the two additional large-scale datasets DDR [24] and EyePACS [13], for fair evaluation, we also supplement both datasets when training. We will clarify it in the main text accordingly.

Finally, should you have further suggestions and comments, we are glad to incorporate during the discussion stage.

Q1: Is Cross-domain Breast Cancer Grading Benchmark proper? 1) not widely used. 2) It is not common to use images from two different magnifications as the "domains". 3) Evaluation on CAMELYON17, from different centers and different staining.

R: Thanks for your valuable comment, so that we could have a chance to clarify some important aspects.

1. Following your suggestion, we take the five individual domains marked by the staining protocols in CAMELYON17 for further experiments. We use Domain-1 as the source domain, and reports its performance on the rest four unseen target domains, denoted from Domain-2 to Domain-5. The results are reported as follows. It is observed that, the proposed Samba still shows a significant performance improvement on all the four unseen target domains compared with the baseline VMamba-ERM. The results will be included when revision.

Table R2: Generalization performance comparison between Samba and VMamba-ERM baseline.

2. We also inspect in our dataset, if the samples vary in staining and if they have domain gap from a rigorous machine learning perspective. As shown in Fig.R2 a (attached in the 1-pg PDF), the samples from Domain-1 ($\times$20) and Domain-2 ($\times$40) are not only different in magnifications, but also varied in staining. Besides, we use t-SNE visualization to inspect the samples' feature distribution from Domain-1 and Domain-2, displayed in Fig. R2 b (attached in the 1-pg PDF), respectively. It can be observed that, samples from Domain-1 (marked in X) and Domain-2 (marked in O) have clear distance in the feature space. Besides, samples of the same severity level but from different domains are not clustered together. It indicates that the two domains in this dataset have the so-called domain gap and are suitable to benchmark the generalization capablity.

3. In real-world scenarios, some hospitals may use magnification factors that are not present in the training set. So testing the cross-magnification generalization ability by the Cross-domain Breast Cancer Grading Benchmarkalso still has practical significance.

Q2: Are evaluation in Table 4 same and fair?

R: We would like to clarify that, the proposed Samba is evaluated under all the default settings of [8] for fair evaluation. Specially, the batch size is 16, the training terminates after 100 epochs, the initial learning rate is $10^{-3}$ with a weight decay of $5\times10^{−4}$ and a momentum value of 0.9. We will explicitly mention these details.

Q3: 1) Proper to report Balanced Accuracy, e.g., Table 1. 2) Balanced accuracy for all the benchmarks.

R:

1. We adapt the balance accuracy metric (denoted as BACC), which is the mean of Sensitivity and Specificity, on DG Fatigue Fracture Grading benchmark (Table.R3) and DG Breast Cancer Grading Benchmark (Table.R4). The results are reported as follows. It can be seen that, the proposed Samba still outperforms the rest methods in term of the balance accuracy, indicating its effectiveness on rare classes.

Table R3: Effectiveness of the proposed Samba on recurrent patch modeling under BACC metric. Domain-1 and Domain-2 in the Fatigue Fracture Grading Benchmark are used as the source and unseen target domain.

Table R4: Effectiveness of the proposed Samba than baseline under BACC metric. Experiments on DG Breast Cancer Grading Benchmark.

2. We would like to raise the reviewer's attention that, despite that the severity level is highly imbalanced in grading problems especially in DR grading, ACC, F1 and AUC metrics are both acknowledged as the commonly-used evaluation metrics [4,6,8,19,29,36,50,51,53,54,55]. Therefore, we directly adapt these metrics and existing evaluation protocols for fair evaluation in Table 4. Besides, balanced accuracy probes a model's performance very similar as F1-score does, while we already use F1-score.

Q4: Ablation on mdules.

R: We provide an ablation study on each module of the proposed Samba, which is also mentioned by Reviewer#Wyiu and b4sx. Specifically, on top of the VMamba baseline, two components, namely, Bi-directional State Space Modeling (BSSM) and EM-based State Recalibration (ESR), are added. The experiments are conducted on the DG Fatigue Fracture Grading Benchmark. The results are reported in Table~R1 in the attached 1-pg PDF file. It is observed that BSSM contributes to an ACC, AUC and F1 improvement of 5.2%, 1.7% and 4.9%, respectively. ESR contributes to an ACC, AUC and F1 improvement of 18.3%, 9.4% and 12.2%, respectively. Meanwhile, we respectively ask for the reviewer's understanding that our work focuses on innovatively introducing BSSM and ESR on top of the vanilla elements (e.g., layer, norm) in VMamba without modify. The ablation of vanilla elements may beyond the scope of this work.

Q5: Elaborate more on text in p15.

R: Please refer to the general response.

Should you have further suggestions, we are glad to address during the discussion stage.

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 (%).

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 (%).

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.

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.

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.

Q1: Unclear explanation. R:

1. After proceeded by BSSM, each patch embedding $\boldsymbol{f}_n$ is projected by a linear layer to get the state embedding $\boldsymbol{x}_n$. The severity base $\boldsymbol{\mu}_n$ is modeled as a mixture of Gaussian from $\boldsymbol{x}_n$. The E-M algorithm iteratively approximates $\boldsymbol{\mu}_n$ to $\boldsymbol{x}_n$. After convergence, the embedding is fed into the rest module.
2. As introduced in L194-196, moving average is adapted to update the $\boldsymbol{\mu}^{0}$, which is used as the initial Gaussian parameter in each EM process.
3. The E-M empirically implements 3 iterations.
4. We will define the number of Gaussian kernels $K$ in L171.
5. We will mention that the mixing coefficients of GMM are left out for simplicity and easy computation and the exponential inner dot kernel is used.
6. In the $t$-th iteration, $\mathbf{Z}^t$: the responsibilities of all the patch embeddings from a sample, where $\mathbf{Z}^t=\{z_{nk}^t\}$.
7. Typo. L185, remove $\mathbf{F}$.
8. In Eq.6, $\mathbf{\mu}$: the Gaussian basis of all the patch embeddings from a sample, where $\boldsymbol{\mu}=\{\mu_{k}\}$. $\mathbf{\widetilde{F}}$: the entire image feature from all the patch embeddings.

Q2: Structure. R:

1. To clarify, as this paper handles cross-domain medical grading, each dataset we use has more than one domains. The evaluation of domain generalization is evaluated inside each dataset, not across different datasets. Tab.4: pls refer to general response.

2. We will put Table 4 and the subsection first.

3. We will metion before Fig.3. Especially, The comparison is made between the proposed Samba and the vanilla VMamba as baseline. The VMamba baseline is trained under the empirical risk minimization (ERM), which we denote as Vmamba-ERM. Notice that, training under ERM is a common baseline for domain generalization. On top of Vmaba-ERM, the advancement of Samba is two modules EMSR and BSSM.

4. Per-component ablation: in general response.

5. We will put the theoretical analysis to the supplementary, making room for details of the method and ablation study.

6. Details of data pre-processing and hyper-parameter settings on each dataset will be in Appendix.

Q3: rephase L44. R: Domain generalized disease grading aims to learn a model that can be well generalized to unseen target domains, when the model is only trained on the source domain data. Practically, the feature distribution between the source and unseen target domains usually varies.

Q4: L54. R: at the distal of the blood vessels.

Q5: Stem unit? R: A stem unit partitions the input image into patches.

Q6: Typo in Fig.2. R:

1. Both modules should occur in the Samba block. Will correct.

2. Typo. 'VSS block'->'Samba block'.

3. Training and inference input images from different domains, marked by red and green arrows, respectively. We will distinguish the source and target domains (located at the upper-left of Fig.2) using different colors.

Q7: Fig.3. compare what? What's Vmamba-ERM&SOSS? Why Vmamba is constant across T? Different columns? R:

1. The comparison is made between the proposed Samba and the vanilla VMamba as baseline, which is trained under the empirical risk minimization (ERM) and denoted as Vmamba-ERM.

2. Typo. 'SOSS' -> 'Samba'.

3. VMamba-ERM does not have EM-based State Recalibration. EMSR is parameterized by T. Therefore, the performance of Vmamba-ERM is consistent to T.

4. From column 1 to 3, it reports the AUC, ACC and F1 metric.

Q8: Why iteration number? R: E-M algorithm implements the approximation by iteratively conduct the E and M step [12]. A small iteration number does not reach the convergence criterion, which results in an poor approximation. After a too-large iteration number, the approximation may already reach the optimal, adding unnecessary computational cost and training time. We will introduce & clarify.

Q9: Line 237. R: Typo. Fig.6->Fig.3.

Q10: L244-explain why proper? R: The process of the proposed state recalibration is differentiable, thereby enabling the application of back-propagation to update $\boldsymbol{\mu}^0$. However, the stability of the update cannot be guaranteed due to the EM iterations. Therefore, we adopt moving average to update $\boldsymbol{\mu}^0$ to avoid collapse.

Q11&14: L252&283 - typos on percentage points. R: will correct.

Q12: Tab.1 R: Brighter green, yellow and light green intend to highlight the best, second and third top performance. We will correct.

Q13: Tab.3, why Samba better? R: Compared to VMamba-ERM baseline, the EM-based State Recalibration in Samba models the feature distribution of lesions via Gaussian Mixture Model with learnable severity bases, and re-estimates by E-M algorithm. The grading-related features are mapped to a more compact space. It can be more stable in unseen target domains.

Q15: Why much better on FGADR? R: FGDAR 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 is far less than others (e.g., level-1 samples occupy 49.3% in APTOS). Therefore, other methods may over-fit other severity levels and under-fit level-1. In contrast, the selective scan mechanism of the Vmamba-ERM and Samba allows the this severity distribution shift. The EM state re-calibration in Samba makes the feature space more compact, and improves the generalization.

Q17: Fig.8/9. ground truth? R: As the datasets are for grading, we respectfully ask for the reviewer's understanding, there are no available fine-grained ground truth (e.g., pixel, boxes).

Finally, should you have further questions or suggestions, we are happy to address during the discussion stage.