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

Authors: jintao  guo
lei  qi
yinghuan  shi
yang  gao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 11:29:43

Abstract: Domain Generalization (DG) aims to enable models to generalize to unseen target domains by learning from multiple source domains. Existing DG methods primarily rely on convolutional neural networks (CNNs), which inherently learn texture biases due to their limited receptive fields, making them prone to overfitting source domains. While some works have introduced transformer-based methods (ViTs) for DG to leverage the global receptive field, these methods incur high computational costs due to the quadratic complexity of self-attention. Recently, advanced state space models (SSMs), represented by Mamba, have shown promising results in supervised learning tasks by achieving linear complexity in sequence length during training and fast RNN-like computation during inference. Inspired by this, we investigate the generalization ability of the Mamba model under domain shifts and find that input-dependent matrices within SSMs could accumulate and amplify domain-specific features, thus hindering model generalization. To address this issue, we propose a novel SSM-based architecture with saliency-based token-aware transformation (namely START), which achieves state-of-the-art (SOTA) performances and offers a competitive alternative to CNNs and ViTs. Our START can selectively perturb and suppress domain-specific features in salient tokens within the input-dependent matrices of SSMs, thus effectively reducing the discrepancy between different domains. Extensive experiments on five benchmarks demonstrate that START outperforms existing SOTA DG methods with efficient linear complexity. Our code is available at https://github.com/lingeringlight/START.

Main paper: 
 introduction: 
 Deep learning models have achieved impressive progress in various computer vision tasks over the past years [1] [2] [3] . Such a huge success is mostly based on the independent and identically distributed (i.i.d.) assumption, i.e., the training and testing data follow the same distribution [4] . However, when evaluated on test data following different distributions from the training data, these models often suffer severe performance degradation. This issue, which is known as domain shift [4] , has greatly hindered the applications of deep learning models in the real world.To improve the generalization of the model under domain shifts, Domain Adaptation (DA) has been widely studied, which aims to transfer the knowledge learned from labeled source domains to the unlabeled or partially labeled target domain [5, 6] . However, DA methods cannot guarantee the performance of the model on unknown target domains that have not been observed during training [7, 8] . Since the accessibility of the target domain could not always be satisfied in real scenarios, Figure 1 : Analysis of the input-dependent matrices in SSMs. We investigate domain discrepancy in the input sequence x, response sequence y, and the input-dependent matrics ∆, B, and C. The results indicate that the input-dependent matrices can accumulate the domain-specific features during the recurrent process, potentially increasing domain gap. We experiment on PACS [24] with Sketch as the target domain, analyzing the representations from the last block of VMamba backbone [22] .Domain Generalization (DG) is proposed to develop a domain-generalizable model on unseen target domains by learning multiple different but related source domains [8, 9] .Most existing DG methods focus on learning domain-invariant representations across source domains, primarily via domain alignment [10, 11] , meta-learning [7, 12] , and data augmentation [13, 14] . These methods heavily rely on convolutional neural networks (CNNs), which have limited receptive fields due to local convolutions. Consequently, the CNN-based methods inevitably tend to learn local texture information, leading to overfitting to source domains and poor generalization on target domains [15, 16] . Recent works in DG have introduced Vision Transformers (ViTs) as the backbone for DG, utilizing the global receptive field of the self-attention mechanism to mitigate local texture bias [17] [18] [19] . However, the complexity of self-attention increases quadratically with input length, resulting in significant computational overhead for ViTs when modeling long sequences [20, 21] .To address this issue, some pioneers have proposed advanced state space models (SSMs) [20, 22, 23] , represented by Mamba [20] , which selectively models token dependencies in input sequences in a compressed state space. The selective scan mechanism allows Mamba to achieve linear complexity in sequence length during training and fast RNN-like computation during inference. Despite the remarkable performance of Mamba-based methods on supervised learning tasks, few existing works have analyzed the generalization ability of Mamba under domain shift. It remains an open question whether the Mamba model can achieve excellent performance for DG tasks.In this paper, we theoretically analyze the generalization error bound of the Mamba model under domain shifts. We find that the domain distance of features extracted by the model is strongly related to the input-dependent matrices within the model. These matrices accumulate and amplify domain-specific features during training, which exacerbates the overfitting issue of the model to source domains. We empirically measure the distance among source domains within the input sequence x, the response sequence y, and the input-dependent matrices of the last network layer. As shown in Fig. 1 , we observe that for the baseline model, input-dependent matrices ( ∆, B, and C) are prone to learning domain-specific features from the input x. Since the output y is calculated by the recurrent product of x and these matrices, domain-specific features are accumulated and amplified, causing the model to overfit the source domains. To address this issue, we propose a Generalized State Space Model with Saliency-driven Token-Aware Transformation (START), which can reduce domainspecific information in input-dependent matrices during training. Building on the latest Mamba-based model [22] , we develop a strong baseline for DG that outperforms many SOTA DG methods, which selectively learns global dependencies among tokens with linear complexity in sequence length.Moreover, based on theoretical analysis, we design the saliency-driven token-aware transformation method, which simulates domain shifts during training by selectively introducing style perturbations to tokens focused on by the input-dependent matrices. START constrain these matrices to learn domain-invariant features, thus mitigating the overfitting of model on source domains. Experiments on five datasets prove the effectiveness of our method. Our contributions are summarized as follows:• We conduct a theoretical investigation into the generalization ability of the Mamba model, revealing that the input-dependent matrices in Mamba can accumulate domain-specific features during the recurrent process, thus hindering the model's generalizability. • Based on theoretical analysis, we propose a novel SSM-based architecture with saliencydriven token-aware transformation as a competitive alternative to CNNs and ViTs for DG, which performs excellent generalization ability with efficient linear complexity. • For the saliency-driven token-aware transformation, we explore two variants to identify and perturb salient tokens in feature sequences, effectively reducing domain-specific information within the input-dependent matrices of Mamba. Our method achieves SOTA performances, e.g., yielding the best CNN-based method by 5.87% (58.27% vs. 52.40%) on TerraIncognita.2 Related Works Domain generalization. Traditional DG methods, primarily based on CNN backbones, can be broadly classified into three categories: domain alignment, meta-learning, and data augmentation. Motivated by the learning theory of domain adaptation [4, 25] , domain alignment methods seek to learn domain-invariant representations through adversarial learning [10, 26, 27] , causal learning [28, 29] , or feature disentanglement [30, 31] . Another popular way to address DG is meta-learning, which partitions the training data from multiple source domains into meta-train and meta-test sets to simulate domain shifts during training [7, 12, 32] . Data augmentation is also an effective method to enhance model robustness to domain shifts by generating diverse data invariants through adversarial generation [33, 34] , style perturbation [13, 14] , and learnable parameters [35, 36] . However, CNNbased DG methods suffer from the limited receptive field of convolutions, often leading to a texture bias and overfitting to source domains [15, 37] . To address this, some researchers have introduced ViT-based methods for DG, which capture global representations by leveraging long-range spatial dependencies with attention mechanisms [18, 17, 19] . Despite their advantages, ViT-based methods are computationally intensive due to the quadratic complexity of the self-attention mechanism, limiting their practical applications [20, 21] . Inspired by the emerging Mamba model [20, 22, 23] , we explore a novel SSM-based architecture for DG that combines strong generalizability with efficient linear complexity. We theoretically analyze the generalization error bound of Mamba and design a novel saliency-driven token-aware transformation to suppress domain-specific features in the input-dependent matrices of Mamba, thereby enhancing the generalization ability of the model.Recently, state space models (SSMs) have demonstrated promising performance across various vision tasks [38] [39] [40] for their ability to effectively capture long-range dependencies while maintaining linear scalability with sequence length. Derived from the classical state space model [41] , the Structured State Space Sequence Model (S4) [42] addresses computational constraints through novel parameterizations catering to continuous-time, recurrent, and convolutional views of the state space model. Notably, Mamba [20] has emerged as a standout performer, which integrates selection mechanism and hardware-aware algorithms into previous works [43] [44] [45] , thus achieving linear-time inference and efficient training mechanisms. Based on the success of Mamba, Vision Mamba (Vim) [23] applies Mamba to ViT architecture, combining bidirectional SSM for datadependent global visual context modeling. Meanwhile, VMamba [22] designs a cross-scan module to bridge the gap between 1D array scanning and 2D plain traversing. Mamba-based architectures have exhibited superior performance across various supervised vision tasks, including medical image segmentation [46] [47] [48] , point cloud analysis [49] [50] [51] , and remote sensing analysis [52, 53] . However, few works explore the performance of Mamba under domain shifts for DG. Although DGMamba [54] has recently introduced a pioneering Mamba-based framework for DG, it lacks a deep analysis of the generalizability of Mamba. In the paper, we conduct a theoretical analysis of Mamba's generalizability, revealing that input-dependent matrices within Mamba could accumulate domainspecific information, thereby impeding model generalization. Consequently, we propose a generalized SSM-based architecture for DG, incorporating a non-parametric module to selectively perturb salient tokens within input-dependent matrices, thus enhancing model generalization to unseen domains. 
preliminary: 
 State Space Models (SSMs). The SSM is a type of linear time-invariant systems that map input sequence x t ∈ R L to response sequence y t ∈ R L through a hidden state h t ∈ R N . Mathematically, this process is formulated as the subsequent linear ordinary differential equations (ODEs): h ′ t = Ah t + Bx t , y t = Ch t , where A ∈ R N ×N is the evolution parameter, and B ∈ R N ×1 , C ∈ R N ×1 are the projection parameters. However, the differential equation is hard to solve in the deep learning setting, thus discrete SSM [55, 44] suggests discretizing the system with a time scale parameter ∆: Ā = e ∆A , B = (∆A) -1 (e ∆A -I) • ∆B,EQUATIONwhere Ā and B are discrete counterparts of the continuous parameters A and B, and ∆ ∈ R > 0 is the sampling timescale for the discretization process. Although the discrete SSMs can achieve linear time complexity, they rely on static parameterization, i.e., Ā, B, and C are time-invariant for any input, inherently limiting their ability to capture sequence context [20] . To address this issue, recently, [20] proposes Mamba, a selective SSM (S6) that effectively selects relevant context by enabling dependence of the parameters B ∈ R L×N , C ∈ R L×N , and ∆ ∈ R L×D on the input x t ∈ R L×D :EQUATIONS B , S C , S ∆ are linear projection layers and softplus(•) = log(1 + exp(•)). The input-dependent time-variant layers could enhance recurrent layers, making them more expressive and flexible in capturing complex dependencies [21] . The parameter matrixes can be further expressed asĀ = [ Ā1 , • • • , ĀL ], B = [ B1 , • • • , BL ], C = [C 1 , • • • , C L ],where L is the sequence length.Considering the initial state h 0 = 0, Eq. ( 1) can be unrolled as [21] :EQUATIONwhere α i,j = C i i k=j+1 Āk Bj , for 0 ≤ j < i ≤ L, characterizing S6 layer as a data-dependent self-attention [56] . The attention matrix α is determined by both the input and the parameter matrices.Previous DG methods have primarily focused on enhancing the generalizability of CNNs or ViTs, lacking theoretical investigations into the Mamba model. We theoretically explore the generalization error bound of Mamba, proving that perturbing the domain-specific features within the input-dependent matrices of Mamba can effectively diminish the upper bound of the model's generalization risk .D S = {D 1 S , D 2 S , • • • , D N S }, the objective of DG is to use D S to train a model that is expected to perform well on unseen target domain D T . Let h : X → Y be a hypothesis from the candidate hypothesis space H, where X and Y denote the input space and the label space, respectively. Since Mamba learns dependencies among tokens from continuous sequences, we study its generalizability at the token level. Let ψ(•) be the feature extractor of h that maps input images into feature space. Following Integral Probability Metrics [57, 58] , we define the token-level Maximum Mean Discrepancy to estimate the gap between different domains.Definition 1 (Token-level Maximum Mean Discrepancy). Given two different distributions of D S and D T , let L denote the number of tokens in the response sequence of ψ(•), then we define the TOken-level Maximum Mean Discrepancy (To-MMD) between ψ(D S ) and ψ(D T ) as:EQUATIONwhere Ψ represents the hypothesis space for each token, ψ t (D) denotes the distribution of the t-th token for domain D, and F k is a RKHS with its associated kernel k.We here investigate the generalization risk bound of the Mamba model. Theoretically, as in [59, 60] , the risk of the hypothesis h on the domain D is defined as: R D(h) = E x∼D [L(h(x)-h * (x))],where L : Y × Y → R + is a convex loss-function that measures the distance between h and the true labeling function h * . Moreover, following [8, 60] , for multiple source domains D S = {D 1 S , D 2 S , ..., D N S }, the convex hull Λ S is defined as a set of a mixture of source domains, i.e., Λ S = { D : D(•) = N n=1 π i D n s (•), N n=1 π n = 1, π n ∈ [0, 1]}.The DT ∈ Λ S is defined as the closest domain to the target domain D T . Based on Eq. ( 4), the following generalization risk bound can be derived.Theorem 1 (Generalization Risk Bound). With the previous setting and assumptions, let D i S and D T be two sets with M samples independently drawn from D n S and D T , respectively. For any δ ∈ (0, 1) with probablity of at least 1 -δ, for all h ∈ H, the following inequality holds:EQUATIONwhereλ π = 1 M ( N n=1 π n E x∼D n S [ tr(K D n S ) + E x∼D T [ tr(K D T )]) + log(2/ϵ)2M , and σ is the minimum combined error of the ideal hypothesis h * on both D S and D T . Let κ T = d To-MMD (D T , DT ) and κ S = sup i,j∈[N ] d To-MMD (D i S , D j S ), respectively. The proof of Theorem 1 is provided in Appendix A.1. The inequality indicates that the generalization error bound depends on κ T denoting the token-level maximum distance between source and target domains, and κ S measuring the maximum pairwise gap among source domains at the token level. The smaller the two terms, the lower the upper bound of generalization error. Following [25, 58, 61] , we simplify the To-MMD in Eq. ( 4) by choosing a unit ball in the F k and using Gaussian kernel with parameter γ to estimate κ T and κ S . Let xS ∈ R L and xT ∈ R L to denote the mean embeddings of samples from D S and D T , where L represents the token sequence length. We explore a simplified problem in conjunction with a single S6 layer, i.e., ȳ = αx, with α ∈ R L×L being the data-dependent matric. The domain distance between ȳS and ȳT is formulated as k(ȳ S , ȳT ) = exp(-||ȳ S -ȳT || 2 /γ), where γ is the kernel parameter. Specifically, for input-dependent matrices B, C, ∆, we denote softmax(S ∆ (•)) as S∆ (•). Then, we analyze the impact of these input-dependent matrices on ||ȳ S -ȳT || 2 , which is applicable to both κ T and κ S . For the i-th tokens xS i and xT i , we define:EQUATIONWith the recurrent property of the S6 layer in Eq. (3), we can derive the following propositions: Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by S ∆ , S C , and S B can enhance their learning of domain-invariant features, thus effectively mitigating the accumulation issue in these input-dependent matrices.Propositions 1 and 2 are proved in Appendix A.1. Based on the propositions, we develop a saliencydriven token augmentation method, which perturbs style information within the tokens that the model focuses on at the sequence level. In this way, our method enhances the extraction of domain-invariant features by the input-dependent matrices, i.e., S ∆ , S C , and S B . As presented in Tab. 6, we also empirically validate the effectiveness of our method in reducing domain discrepancy in these matrices.To boost the generalization ability of the Mamba model, leveraging the Proposion 1 and Proposion 2, we propose a novel Saliency-driven Token-AwaRe Transformation paradigm (START in short), which aims to explicitly suppress domain-related features within the input-dependent matrixes. Unlike prior methods that perturb entire feature maps at the channel level [13] [14] [15] , START incorporates a The core of the START framework is the Saliency-driven Token-Aware Transformation, which uses a saliency-driven scheme to localize tokens targeted by input-dependent matrices, subsequently perturbing domain-specific style information within these tokens. We designed two variants: START-M, which uses input-dependent matrices, and START-X, which uses input sequences to compute saliency.saliency-driven token selection scheme to perturb the prominent regions of input-dependent matrics S ∆ , S B , and S C . Based on the attention mechanism outlined in Eq. ( 3), we propose two variants to identify and perturb tokens within salient regions, including START-M that determines saliency using input-dependent matrices, and START-X computing saliency based on input sequences.START based on input-dependent matrices (START-M). As Proposition 1 reveals, for the i-th token, the token-level domain gap depends on d C ∆Bx (x S i , xT i ) and d ∆(x S i , xT i ). Specifically, as presented in Eq. ( 6), d C ∆Bx is contingent on S C (x i ) S∆ (x i )S B (x i )x i , which is the response of the SSM to x i . S C (x i ) S∆ (x i )S B (x i ) could be regarded as a self-attention matrix, which implicitly offers a measure of saliency for a token x i . To this end, we propose START-M, which utilizes the input-dependent matrices to identify salient tokens. Concretely, given an input sequence {x i } L i=1 , where L denotes the sequence length, we first compute the input-dependent matrices based on Eq. ( 2). Then, we calculate the saliency value for each token based on Eq. ( 6), i.e., for the i-th token x i :EQUATIONAfterward, we generate a binary mask M S ∈ B L with the element being set to 1 if the corresponding element Saliency M (x i ) is in the top P tokens percentage elements.Meanwhile, we synthesize the style-augmented sequence, which is achieved by mixing the mean and variance of different samples. Following [13, 62] , we first compute the style statistics as:µ(x) = 1 L L i=1 x i , σ(x) = 1 L L i=1 (x i -µ(x)) 2 .Then we randomly select another sample x ′ from the current batch, utilizing its statistics to synthesize the stylized version of x:EQUATIONFinally, with the token-level mask M S , we mix x and x to generate the augmented sequence x aug , where tokens with maximum saliency are style-augmented, while other tokens remain unchanged:EQUATIONwhere ⊙ is element-wise multiplication. Note that when P token = 1, START-M degenerates into a channel-level augmentation, i.e., MixStyle [13] . However, note that style statistics could be one kind of domain-specific feature, while other forms of domain-specific features may also exist, especially within the image backgrounds [59, 63] . As a result, directly perturbing the style information of tokens on the background might activate other forms of domain-related noise, which could still disrupt the model generalization [64] . To address this issue, our START-M proposes to selectively perturb tokens with the highest saliency, which are typically associated with foregrounds, thus enhancing the model learning of domain-invariant information without activating domain-related noise. Ablation study in Section 4.3 also proves the effectiveness of the saliency-driven selection scheme.START based on input sequences (START-X). Based on Proposition 2, recalling that ∆, B, and C are all input-dependent matrices (as in Eq. ( 2)), we design a simplified variant, namely START-X, which involves using the activation values of x to approximate the saliency of tokens directly. Specifically, for the i-th input token x i , we directly compute its saliency value as: Saliency X (x i ) = x i . With the saliency for each token, we compute the token-level binary mask M S as that of START-X and employ Eq. ( 9) to generate the augmented sequences. In practice, we randomly apply our START method to 50% of the samples in each batch, leaving the remaining samples unperturbed during each training iteration. Our START method is disabled during inference.In summary, we theoretically investigate the generalization error boundary of Mamba at the token level, highlighting that suppressing domain-related information within input-dependent matrices can effectively reduce the generalization error boundary of the model. Based on the theoretical analysis, we propose the first saliency-driven token-aware transformation for SSMs, designing two different variants for identifying and perturbing the tokens focused on by the input-dependent matrices B, C, ∆.In this way, our method can effectively enhance the reliance of the input-dependent matrices on domain-invariant features and narrow the distance between source and target domains. Notably, our START introduces no additional parameters or inference time, only involving a few matrix operations during training, thus achieving similar linear complexity to VMamba as presented in Appendix A.2. Following existing DG methods [15, 83] , we train the model for 50 epochs using AdamW optimizer and cosine decay schedule, with a batch size of 64, the initial learning rate as 5e -4, and the momentum of 0.9. For all experiments, we the ratio P token of augmented tokens to 0.75. We apply the leave-one-domain-out protocol for all benchmarks, where one domain is used for testing, and the remaining domains are employed for training. We select the last-epoch model and report the average accuracy over five runs. All the experiments are run on 4 NVIDIA Teska V100 GPUs.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: The paper has stated the full set of assumptions and provided complete proofs in Section 3.2 and Appendix A.1. 
main results: 
 Evaluation on PACS. We first compare our method with SOTA CNN-based DG methods on ResNet-50. As shown in Tab. Parameter sensitivity. We explore the sensitivity of our method to the hyper-parameter P token , the percentage of perturbed tokens in input sequences. As shown in Fig. 3 , START-M consistently performs well across different P token values, demonstrating its effectiveness in perturbing domain-specific information in salient tokens. We notice that START-X, which uses token activation to approximate saliency, performs similarly to START-M when P token is high but is less effective at lower P token values. This indicates differences between attention regions of input-dependent matrices and input sequences. Both methods achieve the highest accuracy at P token = 0.75, which is adopted for all experiments. Comparisons with other feature identification methods. We provide comparisons with the "GradCAM" and "Attention Matrix" methods. For the "Grad-CAM" method, we first obtain feature gradients using backpropagation without updating, then compute token saliency and augment salient tokens at each iteration. For the "Attention Matrix" method, since the Mamba architecture lacks explicit attention matrices, we instead use α in Eq. ( 3) to calculate token saliency. As shown in Tab. 4, 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%). It is owing to the ability of START to explicitly suppress domain-specific features within input-dependent matrixes. Comparisons with other augmentation methods. We here compare our method with SOTA DG augmentation methods on the VMamba backbone, including MixStyle [13] , DSU [14] , and ALOFT [15] . As shown in Tab. 5, all the augmentation methods bring performance improvements, indicating that increasing data diversity is beneficial for the generalization ability of Mamba. Notably, START outperforms all the SOTA augmentation methods, i.e., yielding a significant margin of 1.06%(91.77% vs. 91.72%) from DSU. The results prove the effectiveness of our methods in perturbing domain-specific features within input-dependent matrices. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: See in Section 4.1. 
conclusions: 
 In this paper, inspired by the success of Mamba in supervised tasks, we theoretically study the generalizability of Mamba and find that the input-dependent matrices in Mamba could accumulate and amplify domain-specific features during training. To address the issue, we propose a generalized state space model with a saliency-driven token-aware transformation for DG, which can selectively augment domain-specific features within salient tokens focused on by the input-dependent matrices, thus helping the model learn domain-invariant features. Our method outperforms SOTA CNN-based and ViT-based methods by a significant margin with linear complexity and a small-sized network. We hope our work inspires further research in DG and contributes valuable insights to the community.A Appendix / supplemental material A.1 Theoretical Proofs Lemma 1 [58] .Let F = {f ∈ H k : ||f || H k ≤ 1} denote a function class, where H k be a RKHS with its associated kernel k. Let L h,f : x → L[h(x), f (x)] be a convex loss-function with a parameter form |h(x) -f (x)| q for some q > 0, and defined ∀h, f ∈ F, L obeys the triangle inequality. Let S and T be two samples of size M drawn i.i.d from D S and D T , respectively. Then, with probability of at least 1 -δ (δ ∈ (0, 1)) for all h ∈ F, the following holds:EQUATIONwhere Theorem 1 (Generalization risk bound). With the previous setting and assumptions, let D i S and D T be two sets with M samples independently drawn from D n S and D T , respectively. For any δ ∈ (0, 1) with probablity of at least 1 -δ, for all h ∈ H, the following inequality holds:d MMD (D S , D T ) = sup ||f || F k ≤1 f d(h(D S ) -h(D T )) , K D SEQUATIONwhere S , D j S ), respectively. Proof. We initially investigate the relationship between the MMD [58] and To-MMD distances based on Definition 1 (as presented in Eq. ( 4)). With the feature extractor ψ(•), we have:λ π = 1 M ( N n=1 π n E x∼D n S [ tr(K D n S ) + E x∼D T [ tr(K D T )]) + log(2/ϵ) 2M ,EQUATIONThen, for a pair of source domain D n S and D T , the following inequality holds:EQUATIONwith which we can derive the weighted sum of To-MMD between source domains and target domain: With the above preparations, we now derive the generalization error bound of the Mamba model on the unseen target domain. Recalling that Lemma 1 indicates the generalization error bound between two different distributions, we generalize it to the scenario of multiple source domains:EQUATIONwhereλ π = 1 M ( N n=1 π n E x∼D n S [ tr(K D n S ) + E x∼D T [ tr(K D T )]) + log(2/ϵ) 2Mand σ is the minimum combined error of the ideal hypothesis h * on both D S and D T . Proof. For simplification, we use xS ∈ R L and xT ∈ R L to denote the sample mean embeddings for the D S and D T , respectively. L represents the token sequence length. To investigate the generalization error boundary of Mamba, we explore a simplified problem in conjunction with a single S6 layer, i.e., ȳ = αx, where α ∈ R L×L is the data-dependent matric. Empirically, based on Eq. (3) and Eq. ( 4), we estimate the token-level domain gap using Euclidean distance of ȳS and ȳS ,EQUATIONCombined with Eq. ( 2) and Eq. ( 3), we represent α as a direct function of the input x:EQUATIONSince the SSM layer calculates y i based on the continuous subsequence[x 1 , x2 , • • • , xi ],we here analyze the domain gap of the extracted features at the token level, i.e., |y S i -y T i |. Specifically, assuming that we have calculated |y S i -y T i | = β (1 ≤ i < L), from Eq. ( 16), we can derive:EQUATIONWith the defined α in Eq. ( 17) and denoting softmax(S ∆ (•)) as S∆ (•) for brevity, we can express:EQUATIONConsidering that the differences between adjacent tokens are generally small, S C (x i+1 )/S C (x i ) could be approximated to 1. When the dimension of x is relatively large, then for Eq. ( 19), we have:EQUATIONThen, we substitute Eq. ( 20) into Eq. ( 18) to derive the following formula:EQUATIONFinally, recalling that |y S i -y T i | = β, we can express |y S i+1 -y T i+1 | as following:|y S i+1 -y T i+1 | =| I + S∆ (x S i+1 )A β + S∆ (x S i+1 ) -S∆ (x T i+1 ) Ay T i + S C (x S i+1 ) S∆ (x S i+1 )S B (x S i+1 )x S i+1 -S C (x T i+1 ) S∆ (x T i+1 )S B (x T i+1 )x T i+1 | (22) Let d C ∆B x(x S i+1 , xT i+1 ) = S C (x S i+1 ) S∆ (x S i+1 )S B (x S i+1 )x S i+1 -S C (x T i+1 ) S∆ (x T i+1 )S B (x T i+1 )x T i+1 , and d ∆(x S i+1 , xT i+1 ) = S∆ (x S i+1 ) -S∆ (x T i+1). Then, the above equation reveals that for the input tokens xS i+1 and xT i+1 , the distance between their extracted features, alongside the gap of historical sequences, primarily depends on d C ∆Bx (x S i+1 , xT i+1 ) and d ∆(x S i+1 , xT i+1 ). Besides, the recurrent process in Eq. ( 22) could also lead to the accumulation or even enhancement of domain-specific information, i.e., if the model extracts domain-related information from the ith token, this part of the information will be retained in the features extracted by the (i + 1)-th token. For the whole recurrent process, domain-related information encoded in S ∆ , S C , and S B will be accumulated and amplified, which will increase the discrepancy in the features extracted by the model for different domains, thus damaging its generalization ability. Therefore, to reduce the domain gap, it is imperative to suppress domain-specific information learned by S ∆ , S C , and S B .Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by S ∆ , S C , and S B can enhance their learning of domain-invariant features, thus effectively mitigating the accumulation issue in these input-dependent matrices.Proof. Recalling that in the Mamba mode, S ∆ , S C , and S B are all linear projection layers, which map the input sequence x ∈ R L to the data-dependent matrixes ∆, C, and B, respectively. Hence, we analyze the influence of tokens in x on these matrixes. Taking the matric B as an example, we explore the simplified problem with S B (x) = W B x, where W B ∈ R L×N and N denotes the dimension of the hinder state. Inspired by previous works [84, 85] , we assume that the input sequence x could be decomposed to domain-specific features xI and domain-specific features xS . Then, for the i-th token xi in x, the projection matric B could be denoted as:EQUATIONPrevious theoretical works [84, 86] have demonstrated that when a subset of features is perturbed, its variance would be increased, and the model would be regularized to decrease the weights associated with these features to minimize the prediction loss. Therefore, by perturbing the domain-specific features x S i , its corresponding weights W S Bi would be restricted to 0. Simultaneously, due to minimal changes in the domain-invariant feature x I i , the learning of its corresponding weights W I Bi is promoted, thus enhancing S B learning of domain-invariant features. Furthermore, given the presence of foreground and background in images [54] , different tokens contain varied information. Foreground tokens primarily encode domain-invariant semantic features alongside some domain-specific details. Perturbing the domain-related features in these tokens can effectively enhance the model's learning of domain-invariant features. Conversely, background tokens encompass diverse domain-related spurious features that are challenging to fully extract and perturb. As a result, perturbing a subset of domain-related features in these tokens may inadvertently activate other forms of spurious noise, thereby hindering model generalization. To address this issue, leveraging the hidden attention mechanism of Mamba (as depicted in Eq. ( 3)), where tokens with high saliency are more likely to belong to the foreground, we propose to perturb domain-specific information solely in the tokens focused on by the input-dependent matrices. 
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: The paper does not release the data or models with a high risk for misuse.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 paper uses five public datasets, including PACS [24] , OfficeHome [79] , VLCS [80] , TerraIncognita [81] , and DomainNet [5] . See the licenses in the original pepers.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. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Evaluation on DomainNet. We investigate the effectiveness of our method on the large-scale dataset DomainNet. As shown in Tab. 7, on the challenging benchmark, we find that the strong baseline VMamba can achieve the promising performance of 51.45%, proving the superiority of Mamba models on large-scale datasets. Based on the strong baseline, our START can still significantly improve the performance, exceeding the baseline by 1.35% (52.81% vs 51.45%). Besides, compared with CNN-based methods, our method can still achieve significant improvements, e.g., outperforming the latest SOTA method AGFA [74] by 5.71% (52.81% vs 47.10%). Our methods also beat the best ViT-based method DoPrompt [17] , yielding it by a large margin of 4.51% (52.81% vs 48.30%). The results prove the effectiveness of our method to help the model learn domain-invariant representation.Ablation studies on larger datasets. We provide the ablation studies on the Officehome and TerraIncognita datasets. As shown in Tab. 8, our methods perform the best among all variants, e.g., on TerraIncognita, our START-X outperforms the variant "w.o. Saliency Guided" by 0.97% (58.27% vs. 57.30%) and the variant "w.o. Token Selection" by 0.83% (58.27% vs. 57.44%). The results demonstrate the effectiveness of all modules in our START methods.Effects on other Mamba-based architectures. To validate the effectiveness of our START on other Mamba architectures, we conducted experiments on the recent ViM [23] , using the ViM-T and ViM-S models with different network sizes. The experiments were performed on the PACS dataset with an initial learning rate of 6.25e -6 and a weight decay of 1e -8. As shown in Tab. 9, our method consistently improves performance across the models with different scales, e.g., on ViM-S, START-X outperformed the baseline by 2.03% (89.41% vs. 87.38%), and on ViM-T, START-X exceeded the baseline by 1.66% (86.74% vs. 85.08%). These results demonstrate the generalizability of our method across different Mamba architectures.Effects on the ViT architecture. 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 uses 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 × 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.10, 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.Evaluation on single-source domain generalization tasks. We here evaluate our method under the single-source-domain generalization setting. As shown in Tab. 11, 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. Effectiveness with other SOTA augmentation methods. In our method, we utilize a statistics-based style augmentation method to perturb domain-specific information within tokens. We also explore other SOTA DG augmentation methods for comparison, including the DSU [14] that models the distribution of statistics across different samples and resample new statistics from the distribution, and the ALOFT [15] that diversifies the low-frequency spectrum in the frequency domain. These methods primarily perturb style information at the channel level. As shown in Tab. 12, our method significantly improves the performance of the SOTA augmentation methods on the VMamba baseline, e.g., our START-DSU-M achieves a significant improvement over DSU by 0.8% (91.63% vs. 90.71%), exceeding the baseline by 1.69% (91.63% vs. 89.94%). The above results prove that selective perturbation of domain-specific features in salient tokens is crucial for enhancing the generalization capability of Mamba models.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 increase domain gaps in later stages. To address the issue, we applied START to all layers to reduce domain gaps comprehensively. We also tested START separately in either shallow or deep layers. As shown in Tab. 13, 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.Computational efficiency. To evaluate the computational efficiency of our proposed START, we conduct experiments on the PACS dataset and compare our method with existing CNN-based and ViT-based methods. Specifically, we compare the number of parameters, floating point operations per second (FLOPs), the inference times, and the generalization performance of each method. 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. 14, our method has significantly fewer FLOPs than ResNet-50 (5.68 vs. 8.26) while outperforming the DeepAll on ResNet-50 (the baseline that directly trains the model on source domains) by 6.22% (91.77% vs. 85.50%), demonstrating the superiority of our START.Visualization explanations. To provide visual evidence of the effectiveness of our START in suppressing domain-specific features, we use GradCAM [89] to generate attention maps of the last state space layer for both the baseline (pure VMamba) and our START models. As illustrated in Fig. 4 , The experiments are conducted on the PACS dataset with the "Art" as the target domain. We visualize the attention maps of the last layer in the VMamba backbone. For each sample, the first column is the original image, the second column is the attention map of the baseline (i.e., VMamba), and the third and last columns are the attention maps of our START-X and START-M, respectively. Our methods help the model learn more domain-invariant semantic features, e.g., holistic shape structure, than the pure VMamba baseline.the VMamba baseline tends to focus on specific local patches that encode domain-specific features, leading to overfitting to source domains. In contrast, our START methods effectively reduce the model's focus on domain-specific features, enabling it to capture generalizable global dependencies of tokens. For instance, in the case of the person image in the Art domain, the baseline focuses on multiple local regions in both the foreground and the background, making the model sensitive to domain shifts and likely to misclassify samples. Conversely, our START models mainly focus on the foreground, specifically the whole face of the person. The results prove the effectiveness of START in learning comprehensive domain-invariant features, making it a promising method for DG tasks.Difference from the related work. In essence, our method significantly differs from DGMamba [54] in their motivations, goals and methods. 1) Different 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) Different Goals and Methods: DGMamba aims to enforce the 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 input-dependent • 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.• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • 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). 


Decision: The paper receives all positive ratings after rebuttal, where three reviewers upgraded the rating. Overall, the authors have provided more technical explanations (e.g., theoretical clarification, comparisons with the DGMamba method) and experimental results (e.g., computational cost, results with ViT, comparisons with other feature identification methods) in the rebuttal. The AC carefully read the paper, reviews, and replies, and agrees with the reviewers' feedback. Hence the AC recommends the acceptance decision, based on the merit where the paper designs a Mamba framework for improving domain generalization with theoretical analysis and extensive experimental results to verify the proposed method. Also, it would be highly recommended that the authors include suggested feedback in the final version and release the code to the public for fostering the progress in this field.

Reviews: review: paper_summary: paper_summary: This paper aims to enhance the generalization ability of state space models (SSMs), i.e. Mamba. This paper first provides a theoretical investigation on the generalization ability of the Mamba model under domain shifts and finds that input-dependent matrices within SSMs could accumulate and amplify domain-specific features. Then this paper proposes saliency-based token-aware transformation (START) to selectively augment the domain-related features in salient tokens within the input-dependent matrices of SSMs. The authors also conduct extensive experiments to demonstrate better performance.
strength_weakness: strength_weakness: - Regarding the writing, I would suggest revising Proposions 1 and 2 to be more concise. Rather than using words like "it's essential for $S_\Delta, S_C, S_B$ to xxx" or "$S_\Delta, S_C, S_B$ plays a play crucial role in this context", maybe elaborate exactly how $S_\Delta, S_C, S_B$  affects the generalization ability of SSMs or how it relates to Domain Discrepancy Accumulation in the Proposions, so readers don't have to go to Appendix for details.
- The main theoretical results are a bit confusing and isolated. I don't see how Theorem 1 is related to Equation 6. Specifically, how $\Vert \bar y^S - \bar y^T \Vert^2$ relates to RHS of Equation 5. Maybe a new generalization bound tailored specifically to SSMs can be derived?
- Performance gain over previous DGMamba is marginal on three out of the four datasets. For instance, START only surpasses DGMamba by 0.09\%, 0.5\%, and 0.52\% on PACS, OfficeHome, and VLCS while the standard deviation is even larger than the gain (e.g. 0.4 on PACS and 0.16 on OfficeHome), making these improvements unconvincing.
- In essence, START identifies a salient mask, which is later used to augment the feature. However, START only compares with random and full masking. Some more advanced salient feature identification methods or Interpretability methods, e.g. gradcam or salient mask w/ attention matrix, should be compared to further demonstrate the effectiveness of START.
- Some typos. For example, line 219, xaug should be $x_{aug}$

Despite lacking some comparison, I believe this paper contains novelty and is technically sound. I will raise the point if the above questions are addressed.
questions: questions: - Can you discuss the DGMamba in related work, given this is the only work that previously explores the DG setting of Vision Mamba and demonstrates strong performance?
- I also suggest using larger datasets to run the ablation study, for instance, Office-Home, TerraIncognita or DomainNet. PACS is a relatively small dataset on which most methods achieve high enough results. Thus, improvements on PACS are generally marginal and not as convincing as the other datasets.
- START builds upon the idea that Input-Dependent Matrices focus on domain-related features. However, it's inevitable that it also focuses on some domain-invariance features. Does this pose a negative effect on the model? Or using, for example, class-aware feature augmentation be better?
limitations: limitations: The authors have addressed the limitations and, if applicable, the potential negative societal impact of their work.
paper_summary: paper_summary: This paper studies the role of MAMBA architectures on Domain Generalization benchmarks and adapts the architecture to achieve robust generalization. 
The motivation to use MAMBA-style architectures is their linear complexity. The authors theoretically analyze conventional MAMBA for DG and make an important finding that internal hidden states in MAMBA capture domain-specific characteristics, thus hurting their generalization capabilities.
To address this, they propose two techniques grounded in theory. Both of them involve suppressing the domain-specific information from salient tokens. These tokens are obtained by looking at the "attention matrix" of the last layer and determining the top K tokens which have high attention values. To "debias" these, the first method mixes these tokens (derived through a mask) with style computed from another image. This way, while the content is preserved, the style is augmented, thus making the network learn domain-invariant features. 
In the second method, they calculate the salient tokens by directly relying on the activation values and employ the same procedure as above to produce a new style image. 
They have conducted experimental analysis on standard DG benchmarks and demonstrate empirical gains compared to the baselines.
strength_weakness: strength_weakness: The main weakness is a rather weak motivation for using MAMBA for DG. It almost feels like a low-hanging fruit contribution without proper motivation for the same.
There is another very similar work on using MAMBA for DG called DGMAMBA that was recently published. I commend the authors for including that paper in their experimental comparison. But upon reading that paper, there are very similar ideas being proposed, especially the observation that hidden states of MAMBA carry domain-specific information.
In the author response phase, I would like the authors to clearly differentiate how the proposed approach is different from DG MAMBA. Furthermore, DGMAMBA and the proposed approach produce very similar numbers on at least two datasets.
Second, does the salient token perturbation of both M and X happen from epoch 1? Or is it after a few epochs of training? Since the method relies on ImageNet pre-trained models, maybe from epoch 1, the salient tokens can be meaningfully predicted. But what if the models are not pre-trained? How does the method change?
Small improvement suggestion: It would be nice to remove the overloading of $h$ - in section 3.1 it represents hidden states and in section 3.2 it denotes hypothesis.
questions: questions: please see weaknesses
limitations: limitations: The authors have sufficiently addressed the limitations.
paper_summary: paper_summary: Advancements in state space models (SSMs), particularly a model called Mamba, have demonstrated efficient performance in supervised learning, offering linear complexity in training and rapid computation during inference, similar to RNNs. This paper explores the potential of the Mamba model for DG and identifies a challenge: input-dependent matrices in SSMs can exacerbate domain-specific features, thus impairing generalization.

To overcome this limitation, the authors introduce a novel SSM-based architecture featuring a saliency-based token-aware transformation, termed START. This architecture effectively reduces domain discrepancies by selectively altering and suppressing domain-specific features within the salient tokens of input-dependent matrices. START achieves state-of-the-art performance in DG, presenting a viable and competitive alternative to both CNNs and ViTs.
strength_weakness: strength_weakness: It would be better to report computation costs like inference time.
questions: questions: Please refer to the weakness.
limitations: limitations: This paper discusses the limitations.
paper_summary: paper_summary: This paper targets adapting Mamba for Domain Generalization. The authors find that find the input-dependent matrices in SSMs could accumulate and amplify domain-specific features, so they hinders model generalization. To address this paper, the authors selectively perturb and suppress domain-specific features in salient tokens within the input-dependent matrices of SSMs. Experiments show that their performance can improve Vmamba on Domain Generalization benchmark.
strength_weakness: strength_weakness: - The only fair baseline is Vmamba (Strong Baseline denoted in the paper). Can other DG methods on CNN or ViT directly transfer to VMamba? If they can, how about their performance? Is your method still the best?

- ViT has attention maps, which are also input-dependent. Can your methods transfer to ViT? Or your method can only be applied to Mamba?
questions: questions: See weaknesses. Another suggestion is to list model accuracy on ImageNet in Table 1 and Table 2.
limitations: limitations: The authors have discussed limitations in the appendix.
paper_summary: paper_summary: This paper considers the domain generalization problem and analyze generalization risk across different domains based on the State Space Model (SSM). At the beginning, the advance of SSM over ViT and CNN is introduced. Specifically, the generalization risk bound is formulated by the token-level maximum mean discrepancy. Then, a detailed theoretical analysis is conducted to indicate the accumulation of domain bias on instance-dependent features across different stages. Then, the saliency is roughly estimated by the variables within SSM and the tokens with high saliency is replaced with instance-independent features to help suppress the domain-related information. Under this merit, there are two variants are designed for token selection (START-X, START-M). Detailed experimental study is also provided.
strength_weakness: strength_weakness: 1.	I am generally curious about the condition & assumption of this analysis. Specifically, I am not sure how the number of training sources will impact analysis and the performance of derived approach. Intuitively speaking, a model trained on more diverse source domains can perform better on the unseen domains. However, in your Eq. (5), it does not directly indicate that the performances by large N and small N will be different. Necessary explanation should be made.
2.	Selecting and replacing the tokens with high saliency can help mitigate the domain-related information. However, some useful information can also be hurt inevitably. I am wondering how the approaches can be used to deal with this trade-off.
3.	For your experiments, what is the value of N. Have you considered single-source domain generalization in your evaluation, i.e., N=1?
4.	Honestly, the performance gain is not that significant on Office-Home, comparing with PACS. Is there any reason for such trivial gain?
5.	Though theoretical evidence is additionally provided to show the accumulation of domain bias in the instance-dependent features, how your theory can be used to show the relationship between degree of bias by instance-dependent metrics across different stages. To make the approaches derived from your analysis more effective, should a shallower stage or deeper stage be strongly constrained. Further, whether START-M and START-X can be interactively used in different layers/stages?
questions: questions: Besides the questions I raised in Weakness, I have one more questions for the author about the necessity of this research.

6.	I am generally interested in whether Mamba and its variants are better structures in image processing, comparing with ViT and CNN. If not, what the importance of your paper in analyzing the property of Mamba in domain generalization?
limitations: limitations: No Limitation section is mentioned in the main paper.

I am not sure about the generalization ability of this analysis to different architectures.

Comments: title: Response to Author's rebuttal
Resonse to the comment
Discussion about Q8: Domain-invariant features in input-dependent matrices.
Discussion regarding Q8
comment: I would appreciate the authors' effort in this rebuttal and I think my questions have been well-addressed. I am generally positive towards this paper submission.
Thanks for the rebuttal. After reading the rebuttal, I decided to improve my rating to 5.
Thank you for your insightful discussion. While I am still not convinced by the motivation, I see a strong value in the theoretical analysis and impressive empirical performance. Owing to these, I am increasing my score. 

I thank the authors for especially running experiments in short notice. I saw the rebuttal pdf and they have provided a lot more additional analysis which needs to be acknowledged.
Dear authors,
Thanks for the comment. I don't have further questions. Considering the supplemented results, I will raise my score.
Dear Reviewer: 

Thanks for your constructive comment.
Yes. As presented in L$582-586$ of Appendix and Eq.(22), _the domain distance between the tokens $\\bar{x}\_{i+1}^S$ and $\\bar{x}\_{i+1}^T$ depends on both the features extracted from these tokens and the domain distances accumulated in the historical sequence._
Therefore, the Input-Dependent matrix could also learn and accumulate domain-invariant features, contributing to Mamba’s strong generalization ability, as proved in our experiments (Tab.1 and Tab.2 of the paper). 

However, a major challenge in DG is the overfitting issue to the limited source domains [1,2,3,4], which makes the model easily mistake domain-specific features for domain-invariant features. 
Despite Mamba's strong learning capacity, it still suffers from the overfitting issue on DG tasks, inevitably accumulating domain-specific information during training. 
As presented in Tab.5 of the paper or the table below, the strong baseline, _i.e._, VMamba directly trained on source domains, still exhibits notable domain gaps. 
Thus, it is crucial to enforce the model to learn domain-invariant features while suppressing domain-specific ones.

| Method | Baseline | MixStyle | DSU | ALOFT | START-M (Ours) | START-X (Ours) |
| :----: | :----: | :----: | :----: | :----: | :----: | :----: |
Domain Gap $\downarrow$ | 2.97 | 1.91 | 1.59 | 1.67 | **1.30** | 1.37 |

Nevertheless, It is difficult to directly distinguish and extract domain-invariant features, while mistakenly augmenting these features would somewhat degrade performance. 
To address the challenge, rather than roughly augmenting all features in salient tokens, our method perturbs only a part of the features, _i.e._, the style statistics that are more likely to contain domain-specific information [5,6,7,8], to reduce the risk of augmenting domain-invariant features. 
In this way, the model is compelled to mine category-related semantics from the unchanged features, leading to more accurate learning and accumulation of domain-invariant information.

Thanks for your insightful comment. We look forward to further discussion if you have any additional questions.


[1] Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. TPAMI, 2022.

[2] Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and S Yu Philip. Generalizing to unseen domains: A survey on domain generalization. TKDE, 2022.

[3] Jintao Guo, Lei Qi, and Yinghuan Shi. Domaindrop: Suppressing domain-sensitive channels for domain generalization. In ICCV, 2023.

[4] Yu Ding, Lei Wang, Bin Liang, Shuming Liang, Yang Wang, and Fang Chen. Domain generalization by learning and removing domain-specific features. In NeurIPS, 2022.

[5] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Mixstyle neural networks for domain generalization and adaptation. IJCV, 2024.

[6] Chenming Li, Daoan Zhang, Wenjian Huang, and Jianguo Zhang. Cross contrasting feature perturbation for domain generalization. In ICCV, 2023.

[7] Sangrok Lee, Jongseong Bae, Ha Young Kim. Decompose, Adjust, Compose: Effective Normalization by Playing with Frequency for Domain Generalization. CVPR 2023

[8]  Xiaotong Li, Yongxing Dai, Yixiao Ge, Jun Liu, Ying Shan, and Ling-Yu Duan. Uncertainty modeling for out-of-distribution generalization. ICLR, 2022.
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.