START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation

Thanks for your valuable reviews. Due to the character limit of $6000$ here, the tables are displayed in the ''glocal_tables.pdf''.

Q1: The number of source domains.

Thanks for your constructive feedback. Following previous works [9,13,66], we consider $N \geq 1$ source domains to theoretically study model generalization. Our analysis and methods apply to varying numbers of source domains $N$. Indeed, models trained on more diverse source domains (i.e., larger $N$) could exhibit better generalization performance. As shown in Eq.($5$) of the paper, $\bar{D}\_T \in \Lambda_S$ denotes the closest domain to target domain $D\_T$ within source domains $D\_S$. Increasing $N$ could bring $\bar{D}\_T$ closer to $D\_T$, reducing the second term in Eq.($5$), $d\_{\\text{To-MMD}}(D\_T, \bar{D}\_T)$. Meanwhile, a larger $N$ helps the model learn domain-invariant features more effectively, reducing the third term, $\sup\_{i, j \in [N]} d_{\\text{To-MMD}}(D_S^i, D_S^j)$. We aim to help the model learn domain-invariant information from a fixed number $N$ of source domains, rather than increasing $N$. Thus, our analysis focuses on the accumulation of feature-level domain gaps during training. In our experiments, we address the DG task with multiple source domains. Following existing DG methods [10,58,63], for a dataset with $K$ domains (e.g., $K=4$ for PACS), we train on $K-1$ domains and use the remaining domain for evaluation. We will include these details in the revised version.

Q2: For useful information in tokens with high saliency.

Yes, directly replacing high-saliency tokens could harm valuable information. To address this, START modifies only the style information of salient tokens to preserve semantics (L$219-221$ of the paper). Besides, we apply START to random $50\\%$ of the samples per iteration, which could prevent excessive noise (L$235-237$ of the paper). These designs generate diverse style perturbations while protecting original semantic features, improving model generalization. We will include this discussion in the revision.

Q3: Performance for SDG.

In our experiments, we evaluated our method under multiple-source-domains generalization setting, following existing DG methods [10,58,63]. We trained on all but one domain, using the remaining domain as the target. For example, $N=3$ for PACS with $4$ domains and $N = 5$ for DomainNet with $6$ domains. We also evaluated our method under single-source-domain generalization setting. As shown in Tab.$6$ of ''global_tables.pdf'', START-M significantly improves the baseline, outperforming it by $2.12\\%$ ($71.64\\%$ vs. $69.52\\%$). These results prove that our method enhances model generalization by simulating domain shifts through salience-driven token transformation, improving performance in both multi-source and single-source DG tasks.

Q4: Discussion about performance on OfficeHome.

OfficeHome is more challenging than PACS due to its larger number of categories ($65$ vs. $7$) and samples ($15,588$ vs. $9,991$). Pure VMamba achieves $76.43\\%$ on OfficeHome, surpassing many SOTA DG methods, making further improvements difficult. Nonetheless, our method improves VMamba by $0.66\\%$ ($77.09\\% \pm 0.16\\%$ vs. $76.43\\% \pm 0.15\\%$), averaged over five runs. We tested START on larger datasets (TerraIncognita and DomainNet). As shown in Tab.$1$ of ''global_tables.pdf'', START achieves the best average performance across five datasets, excelling on TerraIncognita and DomainNet. The results confirm our method's ability to stably improve model generalization.

Q5: Effects across different stages.

Our theoretical analysis examined how domain gaps accumulate within each SSM layer. Since one layer's output serves as the next layer's input, domain-specific features from earlier stages increases domain gaps in later stages. To address the issue, we applied START to all layers to reduce domain gaps comprehensively. We also tested START in either shallow or deep layers separately. As shown in Tab.$8$ of "global_tables.pdf", using START in both shallow and deep layers simultaneously performs best, aligning with our theoretical analysis. Applying START-M or START-X randomly across layers also improves performance, though less effectively than using START-M or START-X alone. This may be because START-M and START-X target different domain-related information, leading to incomplete suppression when mixed. We appreciate your feedback and will further investigate START's impact at different stages, exploring adaptive selection strategies for START-M and START-X in future work.

Q6: Effects of Mamba and its variants.

Yes, we believe Mamba is a strong alternative to ViT and CNN in image processing. Pioneering works have shown Mamba's effectiveness [22,23], and recent studies have confirmed its success across various visual tasks [48,49,50]. Mamba's selective scanning mechanism allows it to capture both global and local token dependencies, potentially combining the advantages of CNNs and ViTs. This inspired our work, which is the first to explore the advantages of Mamba in DG from a theoretical perspective. The experiments prove the superiority of Mamba. We hope our work will encourage further research in the field.

Q7: Limitations of our method.

Thanks. The limitations of our method have been discussed in Appendix A.$4$, highlighting challenges in accurately identifying and suppressing domain-specific information. Our analysis of hidden space modeling in the Mamba architecture aims to theoretically assess its generalization capability. Our analysis and method could be transferred to other architectures. As shown in Tab.$5$ of ''global_table.pdf'', we re-designed our method for ViTs, and preliminary results demonstrate its generalizability. Thanks and we will explore extending our work to other architectures in future research.

Q1: Comparision with CNN-based or ViT-based DG methods on VMamba.

Thanks for your feedback. Indeed, the previous DG method could be transferred to VMamba. However, these CNN-based or ViT-based DG methods ignore the accumulation of domain gaps in the state space modeling process of Mamba, leading to their insufficient improvement of model generalization. To address this issue, we developed a sequence-level saliency-driven token transformation, which aims to suppress domain-specific information within the SSM layers explicitly. We compare our method with representative SOTA DG methods on the VMamba backbone, including the CNN-based methods (MixStyle and DSU) and MLP-based method (ALOFT). As shown in Tab.$4$ of the ''global_tables.pdf'', on the strong baseline VMamba, our method achieves the best information compared with previous DG methods, suppressing the CNN-based DSU by $1.06\\%$ ($91.77\\%$ vs. $90.71\\%$) and the MLP-based ALOFT by $0.88\\%$ ($91.77\\%$ vs. $90.89\\%$). The results demonstrate that our method effectively mitigates the accumulation of domain differences in the input-dependent matrixes and enhances the model learning of domain-invariant information.

Q2: Performance of our START on ViTs.

Recalling that our method is derived from the theoretical analysis that input-dependent matrices in Mamba could accumulate domain-related information during training, our START aims to improve the generalization of the Mamba architecture. Nevertheless, the core concept, adaptively perturbing salient tokens in input-dependent matrices, is also applicable to ViTs. Considering that in ViTs, the attention matrix is computed by query $Q$ and key $K$, and then multiplied by the original feature $V$ to obtain the final representation. We develop our START-M to START-ViT-M, which calculates token saliency from the input-dependent matrices (i.e., $Q \times K^T$), and START-X to START-ViT-X, which uses the activation value of representation $x$ to approximate saliency. The experiments are conducted on the representative ViT architecture, i.e., DeiT-Small, with the PACS dataset. As shown in Tab.$5$ of the ''global_tables.pdf'', 1) on the DeiT-Small baseline, our START re-designed for ViTs still can effectively improve the Baseline by a significant margin, e.g., START-ViT-M outperforms the baseline by $1.20\\%$ ($87.05\\%$ vs. $85.85\\%$). The results prove the effectiveness of our START's variants on ViTs; 2) we notice that the VMamba-T ($22M$ parameters) is a stronger baseline model than the DeiT-Small ($22M$ parameters), exceeding it by a large margin of $4.09\\%$ ($89.94\\%$ vs. $85.85\\%$). The results also reveal the advantage of Mamba architecture to learn domain-invariant token dependencies in compressed state space, and our START can further enhance the generalization ability of Mamba.

Q3: Model accuracy on ImageNet in Tab.1 and Tab.2.

Thanks for your advice. Regarding DG issues, existing works typically use ImageNet pre-trained models and train them on downstream DG standard datasets with large domain gaps (e.g., PACS, DomainNet) [1,2,3]. In Tab.$1$ and $2$ of the paper, we evaluate the model by selecting one domain from the DG dataset as the target domain (e.g., Photo) and using the remaining domains (e.g., Art, Cartoon, Sketch) as source domains for training. We compute the average model performance across different target domains within the dataset.

Since the category space changes after training, models trained on DG datasets cannot be directly evaluated on ImageNet. Therefore, for the trained models on DG datasets, we select similar categories from the ImageNet1K Validation Set for evaluation. Specifically, for models trained on PACS, we selected $6$ similar categories from the ImageNet1K Validation Set, excluding "giraffe," which has no corresponding category in ImageNet1K. The experiments were conducted on models trained on PACS with "Art" as the target domain.

As shown in the table below, the trained VMamba model achieves an improvement over the pre-trained version, indicating that training on DG datasets with distribution shifts helps the model learn semantic information. Furthermore, both our START-M and START-X methods further improve the accuracy of VMamba by $4.00\\%$ ($90.33\\%$ vs. $86.33\\%$) and $3.34\\%$ ($89.67\\%$ vs. $86.33\\%$), respectively. These results demonstrate the effectiveness of our methods in reducing the accumulation of domain-specific information and promoting the learning of domain-invariant semantic information. We appreciate your suggestion and will incorporate this discussion into the revised paper for improved clarity.

[1] Mixstyle neural networks for domain generalization and adaptation. IJCV 2024

[2] Rethinking multi-domain generalization with a general learning objective. CVPR 2024

[3] Dgmamba: Domain generalization via generalized state space model. ACM MM 2024

Q1: Computational costs including inference time.

Thanks for your valuable advice. We have provided a comparison of inference times. The batch size for evaluating inference time is set to $64$, and the inference time is averaged over $100$ experiments. Since STARR-M and START-X are only activated during training and disabled during inference, they introduce no additional inference time. As shown in Tab.$7$ of the ''global_tables.pdf'' and the table below, our methods have computational costs comparable to existing SOTA methods on CNN or ViT while achieving the highest generalization performance, demonstrating the effectiveness of our methods. Thanks and we will add the results and discussions into the revised version of the paper.

Thanks for your valuable reviews. Due to the character limit of $6000$ here, the tables are displayed in ''glocal_tables.pdf''.

Q1: Revisions of Proposition 1 and 2.

Thanks. We have modified Propositions 1 and 2 to be more clear and detailed:

Proposion 1 (Accumulation of Domain Discrepancy). Given two domains $D_S$ and $D_T$, token-level domain distance $d\_{\\text{To-MMD}}(D\_S, D\_T)$ depends on $d\_{C \\tilde{\\Delta} B x}(\\bar{x}\_{i}\^S, \\bar{x}\_{i}\^T)$ and $d\_{\\tilde{\\Delta}}(\\bar{x}\_{i}\^S, \\bar{x}\_{i}\^T)$ for the $i$-th token. For the entire recurrent process, domain-specific information encoded in $S\_\\Delta$, $S\_C$, and $S\_B$ will accumulate, thereby amplifying domain discrepancy.

Proposition 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by $S\_\\Delta$, $S\_C$, and $S\_B$ can enhance their learning of domain-invariant features, thus effectively mitigating the accumulation issue in these input-dependent matrices.

Q2: About theoretical results.

As shown in Eq.($5$) of Theorem $1$, generalization error bound of Mamba depends on $\\kappa_T = d\_{\\text{To-MMD}}(D\_T, \\bar{D}\_T)$ and $\\kappa_S = \\sup\_{i, j \in [N]} d\_{\\text{To-MMD}}(D\_S\^i, D\_S\^j)$. The smaller the two terms, the lower the upper bound of generalization error. Following [25,57,60], Gaussian kernel could estimate $\kappa_T$ and $\kappa_S$. Given $\bar{y}^S$ and $\bar{y}^T$ as averaged feature maps from $D_S$ and $D_T$, domain distance is formulated as $k(\bar{y}^S, \bar{y}^T) = \exp(-||\bar{y}^S - \bar{y}^T||^2 / \gamma)$, where $\gamma$ is the kernel parameter. We analyze the effects of input-dependent matrices on $||\bar{y}^S - \bar{y}^T||^2$, which is applicable to both $\kappa_T$ and $\kappa_S$. As proven in Propositions 1 and 2, mitigating the accumulation of domain-specific information in input-dependent matrices reduces domain gaps, thus lowering generalization error bound.

Q3: Compared with DGMamba on five DG datasets.

We report performances of DGMamba from the paper [78] as its codes have not been open-sourced. As shown in Tab.$1$ of ''global_tables.pdf'', START consistently outperforms DGMamba across all datasets, exceeding it by $1.59\\%$ ($72.23\\%$ vs. $70.64\\%$) on average. On relatively small datasets (PACS, VLCS, and OfficeHome), START achieves modest but stable improvements from DGMamba. The reason might be that performance on these tasks tends to be saturated, while our method still enhances generalization. On larger and more challenging datasets (TerraIncognita and DomainNet), our method outperforms DGMamba largely, e.g., surpassing it by $3.19\\%$ ($52.79\\%$ vs. $49.60\\%$) on DomainNet. The improvements stem from START's focus on mitigating domain gap accumulation in input-dependent matrixes of all layers, effectively enforcing the model to capture semantics and sufficiently improving its ability to handle large datasets.

Q4: Compared with other feature identification methods.

We compare the "GradCAM" and "Attention Matrix" methods. For "GradCAM," we first obtain feature gradients using backpropagation without updating, then compute token saliency and augment salient tokens at each iteration. For "Attention Matrix," since Mamba lacks explicit attention matrices, we use $\alpha$ in Eq.($3$) of the paper to calculate token saliency. As shown in Tab.$3$ of ''global_tables.pdf'', on the strong baseline, START still performs much better than these advanced methods, exceeding "GradCAM" by $0.91\\%$ ($91.77\\%$ vs. $90.86\\%$) and "Attention Matrix" by $1.00\\%$ ($91.77\\%$ vs. $90.77\\%$), owing to the ability of START to explicitly suppress domain-specific features within input-dependent matrixes.

Q5: Some Typos.

Thanks. We will re-check the paper and fix all typos in the revision.

Q6: Differences from DGMamba.

Our START differs greatly from DGMamba. 1) Motivations. DGMamba observes that hidden states could amplify domain-related information, and proposes a heuristic method to address the issue. However, it lacks a deep analysis of the phenomenon. Differently, we first theoretically delve into the generalizability of Mamba, revealing how input-dependent matrices contribute to domain gap accumulation. Based on the analysis, we developed START to enhance Mamba's generalization. 2) Goals and Methods. DGMamba aims to enforce model to focus on object tokens, perturbing object tokens while replacing context tokens. It ignores that object tokens could be misclassified as context tokens, replacing which would hinder the model from learning semantics. Inversely, START aims to suppress domain-specific information in tokens focused on by input-dependent matrixes, perturbing only styles while keeping contents unchanged. We will add this discussion to Related Work.

Q7: Ablation studies on larger datasets.

We provide ablation studies on Officehome and TerraIncognita. As shown in Tab.$2$ of the ''global_tables.pdf'', our methods perform the best among all variants, e.g., on TerraIncognita, outperforming the variant w.o. Saliency Guided by $0.97\\%$, proving the effectiveness of all modules.

Q8: For domain-invariant features in input-dependent matrices.

Thanks. START is based on the idea that Input-Dependent Matrices could learn and accumulate domain-related features, which hinders them from learning domain-invariant features. To resist this challenge, thanks to our START, the domain-invariant features are preserved and enhanced by only perturbing style information while keeping semantics unchanged (as in L$219$-L$228$ of the paper). In this way, the negative effect of perturbation on the model could be effectively mitigated. We also agree that "class-aware feature augmentation" could alleviate the issue, which would be one of the possible directions for our method in future work. The discussion will be added to the revision.

Dear authors,

Thanks for your rebuttal.

Is there a chance that Input-Dependent Matrices could also learn and accumulate domain-invariant features? Despite being harder, this is totally possible because there is no explicit regularization that forces it to learn only domain-related features. In that case, augmenting the invariant features would somewhat lead to degenerated performance.

Q1: Motivation of our method.

Thanks for your feedback. Some pioneering work have demonstrated the effectivenes of Mamba architecture on various supervised visual tasks. However, few works have studied the generalization ability of Mamba under distribution shift, especially the problem that how to construct a effective Mamba-based model for DG. It motivated us to investigate Mamba's generalization from a theoretical perspective, in which we reveal the issue that the input-dependent matrixes inevitably accumulate domain-specific information and hinder model generalization during training. Based on the theoretical analysis, we propose a saliency-driven token transformation for Mamba, with which we built a generalizable model as a competitive alternative to CNNs and ViTs for DG.

Q2: Difference from DGMamba.

In essence, our approach significantly differs from DGMamba in motivations, goals and methods.

1. Different Motivations: DGMamba was motivated by the observation that hidden states could amplify domain-related information, and proposes a heuristic method to address the issue. However, it lacks a deep analysis of the phenomenon and fails to reveal its relationship with the model's generalization ability. In contrast, our method is the first to theoretically study the accumulation of domain gaps, revealing that input-dependent matrices could learn and accumulate domain-specific features. Based on the theoretical analysis, we develop START to enhance Mamba's generalization. As shown in Tab.$5$ of the paper and the table below, we also quantitatively validate the effectiveness of our method in reducing domain gaps.

2. Different Goals and Methods: DGMamba aims to enforce the model to focus on object tokens, which is achieved by perturbing the object tokens while shuffling and replacing the context tokens. However, this strategy ignores that object tokens could be misclassified as context tokens, which is likely to damage key semantic information when some object tokens are replaced. Besides, the HSS module of DGMamba aims to mask and suppress features that are lowly activated in hidden states, ignoring that the remaining hidden states could also contain domain-related information. Differently, our method aims to suppress domain-specific information in tokens focused on by input-dependent matrixes, which perturbs style information while keeping semantic content unchanged. In this way, our START could effectively mitigate the domain gap accumulation issue in Mamba, thus sufficiently improving generalizability. Notably, DGMamba uses GradCAM for context patch identification, requiring two backpropagations per iteration. Conversely, our START uses input-dependent matrixes to calculate token saliency during forward propagation, needing only one backpropagation and thus reducing training time.

We appreciate your feedback and will include this discussion in the revised version.

Q3: Experimental comparison with DGMamba.

As shown in Tab.$1$ of the ''global_tables.pdf'' and the table below, we compared our method to DGMamba across five DG datasets. Since the code of DGMamba has not been open-source, we report the results from the DGMamba paper. The results indicate that our method can stably outperform DGMamba on all five datasets, achieving an average performance improvement of $1.59\\%$ ($72.23\\%$ vs. $70.64\\%$), thereby proving its effectiveness. Specifically, on the large-scale datasets TerraIncognita and DomainNet, our method shows substantial gains over DGMamba, e.g., START-M outperforms DGMamba by $3.56\\%$ ($58.16\\%$ vs. $54.60\\%$) on TerraIncognita and $3.19\\%$ ($52.79\\%$ vs. $49.60\\%$) on DomainNet. On OfficeHome, our results are similar to DGMamba, likely because VMamba is a strong baseline that surpasses existing SOTA DG methods, and the performance on OfficeHome is nearly saturated. More advanced methods should be studied to improve model performance on OfficeHome further, and combining our START with DGMamba could be a promising direction. We will add this discussion in the revised version.

Q4: When do START-M and START-X happen?

Yes, both START-M and START-X methods are used from the first epoch. In DG setting, existing methods are trained on DG datasets using ImageNet pre-trained models [1,2,3,4]. If the model is not pre-trained, we recommend first training the model with a warm-up period for a few epochs, and then applying our START to mitigate overfitting to the source domain. Thanks for your question. We will explore this direction in future work.

[1] Mixstyle neural networks for domain generalization and adaptation. IJCV 2024

[2] Rethinking multi-domain generalization with a general learning objective. CVPR 2024

[3] Dgmamba: Domain generalization via generalized state space model. ACM MM 2024

[4] MADG: Margin-based Adversarial Learning for Domain Generalization. NeurIPS 2023.

Q5: Symbol improvements.

Thanks for your suggestion, we will remove the overloading of $h$ in the revised version and recheck all symbols to improve the readability of the paper.