START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation
Authors: Jintao Guo
Lei Qi
Yinghuan Shi
Yang Gao
Venue: NeurIPS 2024 poster (2024)
Keywords: Domain generalization
state space models
Primary area: deep_learning_architectures

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.

START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation
Jintao Guo Lei Qi
Yinghuan Shi!* Yang Gao
Nanjjng University Southeast University
guojintaodsmail.ng nju. edu.cn, qilei@seu. edu. cn, {syh, gaoy) enoy)eeuuu.duu edu. cn
Abstract
Domain Generalization (DG) aims to enable models to generalize to unseen target domains by learning from multiple source domains. Existing DG methods primar- ily 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 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 saate-of-the-art art (SOTA) perfor- mances and offers a competitive alternative to CNNs and ViT's. 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 demon- strate that START outperforms existing SOTA DG methods with efficient linear complexity. Our code is available at https://tthbb.com///ngrngggggggp
Introduction
Deep learning models have achieved impressive progress various computer vision tasks over the past years [133]. Such huge success is mostly based on the independent and identically distributed (i.i...) assumption, i.e., the training and testing data follow the same distribution [4]. However, when evaluated test data following different distributions from the training data, these models often suffer severe performance degradation. This issue, which is known as domain shift 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 6) However, DA methods cannot guarantee the performance of the model on unknown target domains that have not been observed during training g. Since the accessibility of the target domain could not always be satisfied in real scenarios,
"Corresponding authors: Yinghuan Shi and Lei Qi.
'Jintoo Guo, Yinghuan Shi, and Yang Gao are with the National Key Laboratory for Novel Software Technology and the National Institute of Healthcare Data Science, Nanjing University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Input Sequence: [xq
SSM Layer
Output Sequence:
Baseline Model
Domain
Eisbalding
y=
Domaate Bieteace
Diserpancy Accumnalatee
Domain
Domouin
Linear
Sa(ie)
"Llomato Shifk
Linesr
C= Sc(rs)
Domain N
Linea SoliMax
3 softplus(S,
Input-Dependers Matrices: Accumulating Domain Discrepancy during the Recurrent Process.
"Mitigating Domain Diserepancy Accumulation in the Input- Depuudeeeendntttt Matrices,"
Input-dependent & (x10)
Input-dependent B
Inpat-dependent
Response sequence
Denein Distace
1.45
Reduce Gap
Reduce Gap
Donncain Distance
Reduce Gap
Demorinn Distance
Redier Gap
1.80
1.23
Li6
Basoline
START-M START-
Bawiliee START-M STARTA
Raseltee START-M START-X (Oers) (Ours)
Rawliee START-M START-- (Ours) (Oues)
(Ours)
(Ours)
(Ours) (Oun))
Figure I: 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 inpul-dependent matrices can accumulate the domain-specijn 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 222
Domain Generalization (DG) is proposed to develop domain -emnainngeneralizable model on unseen target domains by learning multiple different but related source domains 8
Most existing DG methods focus on learning domain-invariant representations across source domains, primarily via domain alignment [TO [1 meta-learning 7 12 and data augmentation [3 14. These methods heavily rely on convolutional neural networks (CNNs), which have limited receptive fields due local convolutions. Consequently, the CNN-based methods inevitably tend learn local texture information, leading overfitting to source domains and poor generalization domains 15 16). Recent works in DG have introduced Vision Transformers (ViT') as the backbone for DG, utilizing the global receptive field of the self-attention mechanism to mitigate local texture bias 17 However, the complexity of self-attention increases quadratically with input length, resulting in significant computational overhead for ViTs when modeling long sequences 22 21
target
To address this issue, some pioneers have proposed advanced state space models (SSMs) 20
represented by Mamba (20), which selectively models token dependencies in input sequences in compressed state space. The selective scan mechanism allows Mamba 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 strongly related to the input-dependent matrices within the model. These matrices accumulate and amplify domain- speciffc 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 mppu-deeennden matrices of the last network layer. As shown in Fig. we observe that for the baseline model, input-dependent matrices (Â, B, and C) are prone to learning domain-specing recurrent product of r and these matrices, domain- -ppecfic 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 specific information in input-dependent matrices during training. Building on the latest Mamba-based model [22], we develop strong baseline for DG that outperforms many SOTA DG methods, which selectively learns global dependencies among tokens with linear complexity in sequence length.
features from the input x. Since the output is calculated by the
Boken-Awward Transformation (START), which can reduce domain-

---

Moreover, based on theoretical analysis. we design the siliency-driverr token- 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 salieney- driven 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 Terralncognita.
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, [3]. 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 12 (2]. Data augmentation is also an effective method to enhance model robustness to domain shifts by generating diverse data invariants through adversarial style perturbation 3 14, , and learnable parameters [35] 36 However, CNN- based DG methods suffer from the limited receptive field of convolutions, often leading to a texture bias and overfitting to source domains [15]
generation [3] 3d
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 [8] [7| 19]. Despite their advantages, ViT-based methods 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 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 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.
State space models. Recently, state space models (SSMs) have demonstrated promising performance across various vision tasks [38] 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 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 45)45), thus achieving linear-time inference and efficient training mechanisms. Based on the success of Mamba, Vision Mamba (Vim) 2 applies Mamba to VIT architecture, combining bidirectional SSM for data- dependent global visual context modeling. Meanwhile, VMMmbaa 22 designs a cross-scan module to bridge the gap between ID array scanning and 2D plain traversing. laamba-based architectures have exhibited superior performance across various supervised vision tasks, including medical image segmentation 46 point cloud analysis 51) and remote sensing analysis (52 [3]. However, few works explore the performance of Mamba under domain shifts for DG. Although DGMamba [54] has recently introduced a pioneering Mamba- 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 domain- specific information, thereby impeding model generalization. Consequently, we propose generalized SSM-based architecture for DG, incorporating a non-maraetrrie module to selectively perturb salient tokens within input-dependent matrices, thus enhancing model generalization to unseen domains.
3

---

3 Method
3.1 Preliminary
State Space Models (SSMs). The SSM is type of linear time-invariant systems that map input sequence 24 E TR/ response sequence Yt E Rh through hidden state h4 TR.. Mathematically, this process is formulated as the subsequent linear ordinary differential equations (ODEs): h Ah, + BIg; V/- Ch2, where A RNNN is the evolution parameter, and B Rx C TR N 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 A:
Å AA, B = (AA) (c
AB,
h2 - A + Brg, V1 = Chy
where Á and B are discrete counterparts of the continuous parameters A and B. and A E R 0 is the sampling timescale for the discretization process. Although the discrete SSMs can achieve linear time complexity, they rely static parameterization, i.e., A. 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 E R
C IR L and E R LL on the input 2t E LR
B Sn(r), C= Cc = softplus(S (()) (3:)).
SB, Sc, Sb 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 [2T] The parameter matrixes can further expressed as
Á A AL). B [B1.... Br C [C1, Cz). where L is the sequence length.
Considering the initial state hy = 0, Eq. ①I can be unrolled as 121]:
y1
C B
0
0
T1
C2
C₂B₂
y2
22
y = aT,
(3)
:
Ur.
C. T Br Ct Cr Br
Br
where 04,y C. IT Ag for S < is L, characterizing S6 layer as data-dependent self-attention (56). The attention matrix determined by both the input and the parameter matrices.
3.2 Theoretically Analysis for the Generalization Ability of Mamba
Previous DG methods have primarily focused on enhancing the generalizability of CNNs or ViTs, lack- ing theoretical investigations into the Mamba model. We theoretically explore the generalization error bound of Mamba, proving that perturbing the domminnsspccjff feutures within the dependen matrices of Mamba effectively diminish the upper bound of the model's generalization risk.
Notations. Given training set of N source domains Ds [D] D............... .... DD' }. the objective of DG is to use Ds to train a model that is expected to perform well on unseen target domain Dy. Let X > y' be a hypothesis from the candidate hypothesis space H. where X' and ) denote the input space and the label space, respectively. Since Mamba learns dependencies among tokens from continuous sequences, we study its generalizabilityy the token level. Let v)()) 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 (Token-level Maximum Mean Discrepancy). Given two different distributions of Ds and Dr, let denote the number of tokens the response sequence of v/()), ),, then we define the TOken-level Maximum Mean Discrepancy (To-MMD) between (DDs) and (Dr) as:
// [d(() -b -dd
d (Ds, DT) = Σ sup sup
(4)
t=1 V.EW. 11
where V represents the hypothesis space for each token, she (D) denotes the distribution of the 1-th token for domain D. and Fu is 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 the domain D is defined as: Rp(h) (g-nn((((2)) h1 (z)) where
f: 3 x y + R-+ is convex loss-function that measures the distance between h and the true labeling
function h" Moreover, following 8 60), for multiple source domains Ds = { D3 D
the convex hull As is defined as set of mixture of source domains, i.e., As {D D(-)
Ex Ty TT Ex= Te = 1, Tu E [0, 1]}. The Dr E As is defined as the closest domain to the target domain Dr. Based on Eq. 4 the following generalization risk bound can be derived.
Theorem (Generalization Risk Bound). With the previous setting and assumptions, let D's and Dr be two sets with M samples independently drawn from D'z and 'Dr respectively. For any 8 E (0, 1) with probablity of at least 8. for all h E H. the following inequality holds:
N
(D,
S
Ry (h) + dio-MMD (Dr, (Dr,Dr) Dr) + sup (D's - 2d to, (5) MJE[M]
n=1
where A 1E T' tr tr(Kpg) + Dee
+
log(2/e)
and a is the
2M
minimum combined error of the ideal hypothesis ht* on both Ds and Dr. Let KT dTD (Dr, Dr) and Kg suPp.je[N]
(D's D'g), respectively.
The proof of Theorem 1 is provided in Appendix A.I The inequality indicates that the generalization error bound depends on KT denoting the token-level maximum distance between source and target domains, and K 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 unit ball in the Fs and using Gaussian kernel with parameter y to estimate Kr and Kg. Let F R_ and x" E R² to denote the mean embeddings of samples from Dy and Dr. where L represents the token sequence length. We explore a simplified problem in conjunction with single S6 layer, Le. y = azE with a E RRLL being the data-dependent matric. The domain distance between yº and 8² is formulated as k(yS, y"") exp(-111 1 where y is the kernel parameter. Specifically, for input-dependent matrices B, C, A, we denote softmax(S ( (-)) as (.)) Then, we analyze the impact of these input-dependent matrices on || y² 1 which is applicable to both KT and Kgg. For the i-th tokens x.. and of. we define:
C -. = Sc ) (B))) - ( (B(
(6)
(x , SA (z,)
With the recurrent property of the S6 layer in Eq. we can derive the following propositions:
Proposion 1 (Accumulation of Domain Discrepancy). Given two distinct domains Ds and Dr, the token-level domain distance dgo-Mm (Ds, Dr) depends on deABo (3. x! and d (क ) for the i-th token. For the entire recurrent process, domain- speccfc-pecific information encoded in Sa. Sc. and SB will accumulate, thereby amplifying domain discrepancy.
Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain domain-specific features in tokens focused on by Sa, Sc, and SB can enhance their learning of Homain-invartie features, thus effectively mitigating the accumulation issue in these input-dependent matrices.
Propositions 1 and 2 are proved in Appendix A.I Based on the propositions, we develop a saliency- driven 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 -ependee matrices, i.e., SA. Sc. and Sg. As presented in Tab. (6) we also empirically validate the effectiveness of our method in reducing domain discrepancy in these matrices.
3.3 Saliency-driven Token-Aware Transformation for Mamba
To boost the generalization ability of the Mamba model, leveraging the Proposion and Proposion we propose novel Saliency-driven Token-AwaRe Trunsformation parudigm (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 [13155) START incorporates
5

---

Stage
Stage
Stage
Stage
L
La
X k
X44
II
Linerr Proeetion
Pooling Prediettio
Dosen Samphea
S
Floten
VSS Block
VSS Block
VSS Block
VSS Rleck
Down
Iaput Somple
Input
Ouiput
Sequence-level Token Mask M
Original Sequence x
Augmented Sequence
START-M By Matrices
Maximum Saliency
MBx + MYOE
Saltencys
Token Replacement
== 5z(3)
Stylized
Smmll
START-X By Input
Salleacy- -hased Selection
5z(x)
Original
Style Interpolation
Av - 8= 1
sullptu=(S_(x))
Saliencys
Random x
Calculated Salieney
Figure 2: Overall Architecture of the Proposed START Framework. The core of the START framework is the Saliency-driven Token-Aware Transformation, which uses 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 Ss. SB. and Sc. 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 desgo (x) _ and d ( ( Specifically, as presented in Eq. (6). deāBa is contingent on (cx))SSa ( which is the response of the SSM to Se (c,))) (z))) could be regarded as a self atffnattetionn matrix, which implicitly offers measure of saliency for token z;. To this end, we propose START-M, which utilizes the input-dependent matrices to identify salient tokens. Concretely, given input sequence {=; H= where L denotes the sequence length, we first compute the input-dependent matrices based on Eq. (2) Then, calculate the saliency value for each token based on Eq. (6), i.e., for the i-th token ıı;
Saliencym (2x)) Sc((x))offmmxx (z;) (z.)a
Afterward, generate binary mask Ms E BBL with the element being set if the corresponding element Saliencys (zx) is in the top Ptokens percentage elements.
Mcanwhile, synthesize the style-augmented sequence, which is achieved by mixing the mean and variance of different samples. Following [33
62
we first compute the style statistics as:
(zx)= = Elz z;, 2,, (zz) / (z, - 2(x))² Then we randomly select another sample x! from the current batch, utilizing its statistics to synthesize the stylized version of z:
ñ = <p(x) + (1 ())("') ā =x) + (1 - (x'),
-= x- - m())
(8)
<~ (ett((..,,00.))
(xx) ū+
Finally, with the token-level mask Ms, we mix and ă to generate the augmented sequence Faug- where tokens with maximum saliency are style-augmented, while other tokens remain unchanged:
Paug .M + Ms) O I.
where @ element-wise multiplication. Note that when Ptoken 1, START-M degenerates channel-level augmentation, i.e., MixStyle [13). However, note that style statistics could be one kind

---

Table 1: Performance (%) comparisons with the SOTA DG methods on PACS and OfficeHome.
PACS
Office Home
Method
Params.
Avg.
Art
Clipurt
Product
Avg-
CNN: RosNet-50 50
DeepAll
(MAA222))
52.40
66.50
PCL
66
(CVP
23M
90.20
83,90
98.10
82.60
88.70
67.30
59.90
78.70
80.70
71.60
EoA
23M
90.50
83.40
98.00
82.50
88.60
69.10
59.80
79.50
81.50
72.50
EQRM
63
(Ncai1P5/22)
23M
86.50
82.10
96.60
80.80
86.50
60.50
56.00
76.10
77.40
67.50
SAGM
69
(CN991229
23M
87.40
80.20
98.00
80.80
86.60
65.40
57.00
78.00
80.00
70.10
iDAG
90.80
79.70
71.80
Drop
[60] (0CVV13)
23M
89.82
84.22
98.02
85.98
89.51
67.33
60.39
79:05
80.22
71.75
CCFP 7]
JCCV222)
23M
87.50
81.30
96.40
81.40
86.60
63.70
55.50
77.20
79.20
68.90
MADG
(NoxiiP5113) 13)
23M
87.80
82.20
97.70
78.30
86.50
54.10
78.40
$0.30
70.10
PGrad
(CLJ)22)
23M
87.60
79.10
97.40
76.30
85.10
64.70
56.00
77.40
78,90
69.30
AGEA
(CLB
79.30
71.50
GMDG
(UN091224)
23M
84,70
81.70
97.50
80.50
85.60
68.90
56.20
79.90
82.00
70.70
ViT-hased or MI P-like models
MI.P-R 8
50M
85,00
77.86
94.43
65.72
80.75
63.45
56.31
77.81
79.76
69.33
SDVT 18
(AEUX22))
22M
87.60
82.40
98.00
77.20
86.30
68.30
56.30
79.50
81.80
71.50
ResMLP-S P-S
77
(TAAMF22)
40M
85,50
78.63
97.07
72.64
83.46
62.42
51.94
75.40
77.21
66.74
V
84.22
98.38
69.55
61.51
79.34
73.38
GMoE-S
19
(CLR22)
34M
89.40
83.90
99.10
74.50
86.70
69.30
58.00
79.80
82.66
72.40
SSM hased models
MGMam
(MUM MM134)
91.30
87.00
99.00
87.30
91.20
76.20
61.80
83.90
86.10
77.00
Strong Baseline
22
22M
91.55
85.11
99.14
83.97
89.94 1142
75.06
60..4
84.71
85.45
16.43
START-M
91.77=+040
85.31
77.0998..
START-X X (Ours)
22M
92.76
87.43
99.22
87.46
91.72,104
75.48
62.06
85.24
85.47
77.07 Lear
of domain-specific feature, while other forms of domain-specific features may also exist, especially within the image backgrounds (59 (3). 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 nomain-invariant information without activating domain-related noise. Ablation study in Section also proves the effectiveness of the saliency-dri driven selection scheme.
START based on input sequences (START-X), Based on Proposition 2, recalling that A, 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 to approximate the saliency of tokens directly. Specifically, for the i-th input token I,, we directly compute its saliency value as: Saliencyx (x:) = Ty. With the saliency for each token, compute the token-level binary mask MM that of START- X and employ Eq. (9) to generate the augmented sequences. In practice, we randomly apply our START method 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, 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-drin -riven token-aware transformation for SSMs, designing two different variants for identifying and perturbing the tokens focused on by the input-dependent matrices B, C, A. 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 Vammaa as presented in Appendix A.2
4 Experiments
4.1 Experimental Setup
Datasets. We perform an extensive evaluation on five DG datasets: PACS 24
comprises 9, 991
images of 7 classes from domains: Photo, Art Painting, Cartoon, and Sketch. OfficeHome [79] includes 15, 588 images of 65 classes from four diverse domains: Artistic, Clipart, Product, and
Real-Worles exhibiting a large domain gap. VLCCS [0] contains 10, 729 images of 5 categories from
4 domains: Pascal, LabelMe, Caltech, and Sun. Terralncognita [8]] comprises photographs of wild
7

---

Table 2: Performance (%) comparisons with the SOTA DG methods on VLCS and Terralncognita.
VLCS
Terrancoggita
Method
Params:
Caltech
LabelMe
SUN
PASCAL
Avg.
L100
L38
1.4
L46
Avg.
CNN:
ResNet-50
DeepA1
65
23M
97.20
64.30
73.40
74.60
77.50
49.80
42.10
56.90
35,70
46. 16.10
PCL
(2V9:22
23M
99.00
63.60
73.80
75.60
78.00
58.70
46.30
60.00
43.60
52.10
BoA
67
22
23M
99.10
63.10
75.90
78.30
79.10
57.80
46.50
61.30
43.50
52.30
(Ncstt 221
23M
47.80
SAGM
69
23M
99.00
65.20
75.10
80.70
80.00
54,80
41.40
57.70
41.30
48.80
iDAG
20
00CV133)
23M
98.10
62,70
69.90
77.10
76,90
58, 58,70
35.10
57.50
33,00
46. 46.10
CCFP
23M
64,90
74,50
37.50
PGrad
(CCR)33)
23M
98.30
64.40
74.40
79. 99.990
79.30
51.20
43.40
60.00
41.30
49.00
AGFA
14
(CL8222
23M
99,00
64.50
75.40
78.90
79.50
61.00
46.20
60.30
42.30
52.40
(CN09224)
65.90
73.40
79.30
79.20
59.80
45.30
57.10
38.20
50.10
ViT-hased models
SIVTT
18 (M071777))
22M
96.80
64.20
76. 26.20
78:50
78,
55,90
31:70
52.20
37,40
44,30
GMoB-S (CER223)
72.30
79.50
78.00
39.20
34,00
50.70
45.60
SSM-bused models
DCMamba 54
JAC MM 241
22M
98.90
64.30
79.20
80,80
80,80
62,70
48,30
61.10
46.40
$4.60
Strong Baselino
22M
75.81
79.97
29.42=022
47.27
48.56
56. 16:464
START-M (Ours)
22M
98.80
66.98
77.18
82.33
81.32 10.3
70.13
49.98
63.02
49.49
58. 16
START-X (Ours)
22M
98.66
66.64
76.97
82.58
81.21
70.70
49.47
63.96
48.95
58.27
animals taken by 4 camera- ttrap domains, with 10 classes and a total of 24, 788 images. DomainNet [5] is large-scale with 586, 575 images, having 345 classes from 6 domains. i.e., Clipart, Infograph, Painting, Quickdraw, Real, and Sketch. Results on DomainNet are reported in Appendix A.2
Implementation details. We closely follow the implementation of VMamba [22) and use the VMamba-T, which has similar parameters with ResNet-50 (22M vs. 23M), as the backhone. The backbone is pretrained on the ImageNet [82] for all our experiments. We partition the input image into X 4 patches without further flattening the patches into a ID sequence. The network depth of the VMamba-T backbone is 4 the same ResNet-50, consisting of 2, 2, 9, and VSS layers, respectively. The embedding dimensions of blocks in the 4 stages are fixed as [96, 192, 384, 768]. Following existing DG methods [15 83 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 Pioken of augmented tokens to 0.75. We apply the reave-one-domain-ouu 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.
4.2 Main Results
Evaluation on PACS. We first compare our method with SOTA CNN-based DG methods on ResNet-50. As shown in Tab. [ the strong baseline (VMamba) achieves a promising performance, exceeding ResNet-50 by 4.44% (89.94% vs. 85.50%), which indicates its superiority for DG. Moreover, we apply our START to the strong baseline and build advanced models, which can achieve significant improvements without introducing extra parameters. Notably, START-M achieves the SOTA performance, improving baseline by 1.83% (91.77% vs. 89.94%) and yielding the latest CNN-based DG method GMDG 775 by 6.17% (91.77% vs. 85.60%). START-X can also improve the baseline significantly by 1.78% (91.72% vs. 99.94%). Compared with SOTA ViT-based methods, START-M still performs excellent, yielding GMoE-S 19 by 5.07% (91.77% vs. 86.70% with small network sizes (22M 34M). Finally, our methods beat the recent DGMamba [54] exceeding it by 0.57% (91.77% vs. 91 91.20% on average, which proves the effectiveness of our method for DG.
Evaluation OfficeHome. We evaluate the effectiveness of our method OfficeHome and present the results in Tab. T Our methods achieve significant improvements compared with CNN-based methods, e.g., START-M outperforms the SOTA method EoA [67] by 4.59% (77.09% vs. 72.50%) on ResNet-50. Based on the Strong Baseline with high performance, our method can still improve by 0.66% (77.09% vs. 76.43%). START-M precedes the best MLP-like model ViP-S 778. which learns long-range dependencies along height and weight directions, with a large improvement of 4.71% (77.09% vs. 73.38%). The results justify the superiority of START.
Evaluation on VL.CS. presented in Tab. 2, our START achieves the best performance among all competitors, surpassing the top CNN-based method SAGM [69] by 1.32% (81.32% vs. 80.00%)
8

---

Additionally, our method significantly improves upon the baseline, outperforming the latest Mamba- based method DGMamba by 0.52% (81.32% vs. 80.80%).
Evaluation on Terralncognita. As shown in Tab. 2 We observe that the VMamba baseline significantly outperforms previous methods, achieving SOTA performance of 56.16% Based on the strong baseline, our method can further achieve substantial improvement by 2.11% (58, 158.27% vs. 56.16%), proving that our method effectively suppresses domain-specific features learned by Mamba.
4.3 Ablation Study and Analytical Experiments
Ablation study. We here validate the effectiveness of each operation in START. Specifically, who Saliency Guided denotes random selection of tokens for perturbation within input sequences, while w/o Token Selection means perturbing the entire input sequences. Tab. 3 presents the results using VMamba backbone PACS. Both variants show improvements over the baseline, indicating that perturbing style information can mitigate overfitting issues. However, as discussed in Section
3.3
w/o Saliency Guided, which randomly perturbs tokens, fails to provide strong regularization. Besides, the result of w/o Token Selection is inferior to our START, suggesting that perturbing back- ground tokens could activate other forms of domain- speciffc features, potentially hindering model generalization.
Table 3: Ablation study on the PACS dataset.
Method
Art
Cartoon
Photo
Sketch
Avg.
Baseline 22
91.55
85.11
99.14
83.97
89.94
w/o. Saliency Guided
92.11
99.10
85.82
90.81
w/o. Token Selection
92.05
86.55
98.90
86.35
90.94
START-M (Ours)
93,29
87.56
99.14
87.07
91.77
START-X (Ours)
92.76
87.43
99.22
87.46
91.72.0.49
Parameter sensitivity. We explore the sensitivity of our method to the "yper-paramete Ptoken, the percentage of perturbed tokens in input sequences. As shown in Fig. [3] START-M consistently performs well across different Proken 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 Pioken is high but is less effective at lower Ptoken values. This indicates differences between attention regions of input-dependent matrices and input sequences. Both methods achieve the highest accuracy at Ptoken 0.75, which is adopted for all experiments.
63.3
41.5
Accuracy (%)
11.3
133
Ratio of Perturbed Tokens
Figure 3: Sensitivity to Ptoken-
Comparisons with other feature iden- tification methods. We provide compar- isons with the "GradCAM" and "Atten- tion 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 itera- tion. For the "Attention Matrix" method, since the Mamba architecture lacks explicit attention matrices, we instead use a in Eq. ( 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. 90.86%) and "Attention Matrix" by 1.00% 777 vs. 90.77%). It is owing to the ability of START to explicitly suppress domain-specific features within input-dependent matrixes.
Table 4: Comparison (%) with other salient feature identi- fication methods on PACS with VMamba as the backbone.
Method
Art
Cartoon
Photo
Sketch
Avg.
Baseline 22]
91.55
85.11
99.14
83.97
89
89.94
GradCam
92.56
86.99 99
98.98
84.92
90.86.9244
Attention Matrix
91.75
86.68
98.88
85.76
90.77=0.0
START-M (Ours)
93.29
87.56
99.14
87.07
91.77=0.0
START-X (Ours)
92.76
87.43
99.22
87.46
91.72=0.49
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 [5]. As shown in Tab. [5] all the aug- mentation methods bring performance improvements, indicating that increasing data diversity is beneficial for the gen- eralization ability of Mamba. Notably,
Table 5: Comparisons (%) with SOTA augmentation meth- ods on PACS with VMmmba as the backbone.
Method
Art
Cartoon
Photo
Sketch
Avg.
Baseline
91.55
85.11
99.14
83.97
89.94 51
MixStyle
92.05
86.55
98.90
86.35
90.94.015
DSU [4
92.58
85.91
98.98
85,39
90.71
ALOFT [5
93.07
86.04
99.16
85.31
90.89
START-M (Ours)
93.29
87.56
99.14
87.07
91.77=0.40
START-X (Ours)
92.76
87.43
99.22
87.46
91.72=0.99
START outperforms all the SOTA augmentation methods, i..., yiclding 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 impe-ddeendentt matrices.
Domain gaps in input-dependent ma- trices. To verify the effectiieness of our method in reducing domain gaps within input-dependent matrices, we compare domain gaps across different methods using PACS with VMamba. The ex- periments focus on the last block of VMamba, examining the output feature maps ("Feat.") and the input-dependent matrices A, B, and C' of the first SS2D. The results in Tab. align well with theoretical analysis in Section 3.2 proving that START effectively reduces domain gaps in the input-dependent matrices.
Table 6: Domain gaps within input-dependent matrices.
Method
å(L)
B (U)
C
Feat. (L)
Baseline
22
1.48
1.52
2.08
2.97
MixSyyle
1.36
1.91
DSU 14
1.38
1.28
2.18
1.59
AL.OFT
1.25
1.67
START (Ours)
1.16
0.98
1.80
1.30
START-X (Ours)
1.23
0.91
1.52
1.37
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 nomain-specifee features during training. To address the issue, we propose a generalized state space model with 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 significant margin with linear complexity and small-sized network. We hope our work inspires further research in DG and contributes valuable insights to the community.
6 Acknowledgment
This work was supported by the National Key R&D Program of China 20033DD0120700 20232D0120701), NSFC Project (62222604, 62206052), China Postdoctoral Science Foundation (2024M750424), the Fundamental Research Funds for the Central Universities (0202143801 43)0120), the State Key Laboratory Fund (ZZKT2024A14) the Postdoctoral Fellowship Program of CPSF (GZC20240252). and the Jiangsu Funding Program for Excellent Postdoctoral Talent (20247B242).
References
[1] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitchead, Alexander C Berg, Wan-Yer Lo, et al. Segment anything. In ICCV, 2023.
[2] Lihe Yang, Bingyi Kang, Zilong Huang, Xiangang Xu, Jiashi Feng, and Hengshuang Zhao. Depth anything: Unleashing the power of large-scale unlabeled data. In CVPR, 2024.
[3] Arun James Thirunavukarasu, Darren Shu Jeng Ting, Kabilan Elangovan, Laura Gutierrez, Ting Fang Tan, and Daniel Shu Wei Ting. Large language models in medicine. Nature medicine, 2023.
[4] Sinno Jialin Pan Qiang Yang. A survey on transfer learning. TKDE, 2009.
[5] Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Sacnko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV, 2019.
[6] M Jehanzeb Mirza, Jakub Micorek, Horst Possegger, and Horst Bischof. The norm must go on: Dynamic unsupervised domain adaptation normalization. In CVPR, 2022.
[7] Yang Shu, Zhangjie Cao, Chenyu Wang, Jianmin Wang, and Mingsheng Long, Open domain generalization with domain domain-aagmennteeee augmented meta-learning In CVPR, 2021.
[8] Kaiyang Zhou, Ziwei Liu, Qiao, Tao Xiang, and Chen Change Loy, Domain generalization: A survey. TPAMI, 2022.
[9] Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng. and S Yu Philip. Generalizing unseen domains: A survey domain generalization. TKDE, 2022.
10

---

[10] Fu-En Yang, Yuan-Chia Cheng. Zu- Zu-Yun Shiau, and Yu-Chiang Frank Wang. Adversarial teacher-student representation learning for domain generalization. In NeurlPS, 2021.
[11] Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of- distribution generalization. In ICMI,, 2022.
[12] Chaoqi Chen, Jiongeheng Li, Xiaoguang Han, Xiaoging Liu, and Yizhou Yu. Compound domain general- ization via meta-knowe meta-knowledde encoding. In CVPR, 2022.
[13] Kaiyang Zhou, Yongxin Yang. Yu Qiao, and Tao Xiang. Mixstyle neural networks for domain generalization and adaptation. IJCV, 2024.
[14] Xiaotong Li, Yongxing Dai, Yixiao Ge, Jun Liu, Ying Shan, and Ling-Yi Yu Duan. Uncertainty modeling for out-of-distribution generalization. ICLR, 2022.
[15] Jintao Guo, Na Wang, Lei Qi, and Yinghuan Shi. Aloft: A lightweight mip-like architecture with dynamic low-freyenncy transform for domain generalization. In CVPR, 2023.
[16] Jiawang Bai, Yuan, Sau-Tao Xia, Shuicheng Yan, Zhifeng Li, and Wei Liu. Improving vision transform- ers by revisiting high- frequeeyyyncy components. In ECCV, 2022.
[17] Zangwei Zheng, Xiangyu Yue, Kai Wang, and Yang You. Prompt vision transformer for domain general- ization. arXiv preprint arXiv:220s. 089144 2022.
[18] Maryam Sultana, Muzammal Naseer, Muhammad Haris Khan, Salman Khan, and Fahad Shahbaz Khan. Self-distilled vision transformer for domain generalization. In ACCV, 2022.
[19] Bo Li, Yifei Shen, Jingkang Yang. Yezhen Wang. Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of- expeeerss are domain generalizable learners. In ICLR, 2023.
[20] Albert Gu and Tri Dao. Mamba: Lincar-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00072. 2023.
[21] Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590. 2024.
[22] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401. 2024.
[23] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. In ICML, 2024.
[24] Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In ICCV, 2017.
[25] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza. Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 2010.
[26] Wei Zhu, Le Lu, Jing Xiao, Mei Han, Jiebo Luo, and Adam P Harrison. Localized adversarial domain generalization. In CVPR, 2022.
[27] Sudao Hc, Fuyang Chen, and Hongtian Chen. A latent representation generalizing network for domain generalization in cross-scenario monitoring. TNNLS, 2023.
[28] Fangrui Lv, Jian Liang, Shuang Li, Bin Zang, Chi Harold Liu, Ziteng Wang. and Di Liu. Causality inspired representation learning for domain generalization. In CVPR, 2022.
[29] Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts. In NewrIPS, 2022.
[30] Hanlin Zhang, Yi-Fan Zhang, Weiyang Liu, Adrian Weller, Bernhard Schölkopf, and Eric P Xing. Towards principled disentanglement for domain generalization. In CVPR, 2022.
[31] Qiucheng Wu, Yujian Liu, Handong Zhao, Ajinkya Kale, Trung Bui, Tong Yu, Zhe Lin, Yang Zhang, and Shiyu Chang. Uncovering the disentanglement capability in text-to-image diffusion models. In CVPR, 2023.
[32] Jin Chen, Zhi Gao, Xinxiao Wu, and Jiebo Luo, Meta-causal learning for single domain generalization. In CVPR, 2023.

---

[33] Qiuhao Zeng, Wei Wang, Fan Zhou, Charles Ling, and Boyu Wang. Foresee what you will learn: data augmentation for domain generalization in non-stationary environment. In AAAI, 2023.
[34] Qinwei Xu. Ruipeng Zhang. Yan Wu, Ya Zhang. Ning Liu, and Yanfeng Wang. Simde: simple domain expansion approach for single-soure source domain generalization. In CVPR, 2023.
[35] Yue Wang. Lci Qi, Yinghuan Shi, and Yang Gao. Feature-based style randomization for domain general- ization. TCSVT, 2022.
[36] Shiqi Lin, Zhizheng Zhang, Zhipeng Huang. Yan Lu, Cuiling Lan, Peng Chu, Quanzeng You, Jiang Wang. Zicheng Liu, Amey Parulkar, et al. Deep frequency filtering for domain generalization. In CVPR, 2023.
[37] Sachin Mehta and Mohammad Rastegari. Mobilevit: Light-weight, general-purpose, and mobile-friendly vision transformer. In ICLR, 2022.
[38] Yubiao Yue and Zhenzhang Li. Medmamba: Vision mamba for medical image classification. arXiv preprint arXiv:2403.0388 2024.
[39] Weibin Liao, Yinghao Zhu, Xinyuan Wang, Cehngwei Pan, Yasha Wang, and Liantao Ma. Lightm-unet: Mamba assists in lightweight unet for medical image segmentation. arXiv preprint arXiv:2403.05244 2024.
[40] Xuanhua He, Ke Cao, Keyu Yan, Rui Li, Chengjun Xie, Jie Zhang, and Man Zhou. Pan-mamba Effective pan-sharpening with state space model, arXiv preprint arXiv:2402.12192 2024.
[41] Rudolph Emil Kalman. A new approach to lincar filtering and prediction problems. 1960.
[42] Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In ICLR, 2022.
[43] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In NeurlPS, 2022.
[44] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective structured state spaces. NeurlPS, 2022.
[45] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In ICLR, 2022.
[46] Jun Ma, Feifci Li, and Bo Wang. U-mamba: Enhancing long- longerange dependency for biomedical image segmentation. arXiv preprint arXiv:240l.04722. 2024.
[47] Jiacheng Ruan and Suncheng Xiang. Vm-unet: Vision mamba unct for medical image segmentation. arXiv preprint arX::
2024.
[48] Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, and Lei Li. Mamba-unet: Unet-like pure visual mamba for medical image segmentation. arXiv preprint arXiv:240..05079, 2024.
[49] Dingkang Liang, Xin Zhou, Xinyu Wang, Xingkui Zhu, Wei Xu, Zhikang Zou, Xiaoqing Ye, and Xiang Bai. Pointmamba: A simple state space model for point cloud analysis. arXiv preprint Xixx:2402.11733 2024.
[50] Jiuming Liu, Ruiji Yu. Yian Wang, Yu Zheng, Tianchen Deng. Weicai Ye, and Hesheng Wang. Point mamba: A novel point cloud backbone based on state space model with octree-basse ordering strategy. arXiv preprint ar 24l::2400.0666 2024.
[51] Tao Zhang, Xiangtai Li, Haobo Yuan, Shunping Ji, and Shuicheng Yan. Point could mamba: Point cloud learning via state space model. arXiv preprint arXiv:2403.0076) 2024.
[52] Keyan Chen, Bowen Chen, Chenyang Liu, Wenyuan Li, Zhengxia Zou, and Zhenwei Shi. Rsmamba: Remote sensing image classification with state space model. arXiv preprint arXiv:2403. 19654, 2024.
[53] Qinfeng Zhu, Yuanzhi Cai, Yuan Fang, Yihan Yang. Cheng Chen, Fan, and Anh Nguyen. Samba: Semantic segmentation of remotely sensed images with state space model. arXiv preprint 204.0024 2024.
[54] Shaocong Long. Qianyu Zhou, Xiangtai Li, Xuequan Lu, Chenhao Ying. Yuan Luo, Lizhuang Ma, and Shuicheng Yan. Dgmamba: Domain generalization via generalized state space model. In ACM MM, 2024.
12

---

[55] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In NeurIPS, 2020.
[56] Michael Poli, Stefano Massaroli, Fric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In ICML, 2023.
[57] Sinno Jialin Pan, Ivor W Tsang, James T Kwok, and Qiang Yang. Domain adaptation via transfer component analysis. TNN, 2010.
[58] levgen Redko, Emilie Morvant, Amaury Habrard, Mare Sehban, and Younès Bennani, survey on domain adaptation theory: learning bounds and theoretical guarantees. arXiv preprint arXiv:2004. /829, 2020.
[59] Yu Ding, Lei Wang. Bin Liang, Shuming Liang, Yang Wang, and Fang Chen. Domain generalization by learning and removing domain-specific features. In NeurlPS, 2022.
[60] Jintao Guo, Lei Qi. and Yinghuan Shi. Domaindrop: Suppressing comain-sensitive channels for domain generalization. In ICCV, 2023.
[61] Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In /CML, 2015.
[62] Mchrdad Noori, Milad Cheraghalikhani, Ali Bahri, Gustavo A Vargas Hakim, David Osowicchi, Ismail Ben Ayed, and Christian Desrosiers. Tfs-vit: Token-level feature stylization for domain generalization. PR, 2024.
[63]
Rang Meng, Xianfeng Li, Weijie Chen, Shicai Yang. Jie Song, Xinchao Wang, Lei Zhang, Mingli Song, Di Xie, and Shiliang Pu. Attention diversification for domain generalization. In ECCV, 2022.
[64] Chaoqi Chen, Luyao Tang, Feng Liu, Gangming Zhao, Yue Huang, and Yizhou Yo. Mix and reason: Reasoning over semantic topology with data mixing for domain generalization. In NeurlPS, 2022.
[65] Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Deep domain-adversari image generation for domain generalisation. In AAAI, 2020.
[66] Xufeng Yao. Yang Bai, Xinyun Zhang. Yuechen Zhang, Qi Sun, Ran Chen, Ruiyu Li, and Bei Yu. Pel: Proxy-based contrastive learning for domain generalization. In CVPR, 2022.
[67] Devansh Arpit, Huan Wang, Yingho Zhou, and Caiming Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In NeurlPS, 2022.
[68] Cian Fastwood, Alexander Robey, Shashank Singh, Julius Von Kügelgen, Hamed Hassani, George J Puppas, and Bernhard Schälkopf. Probable domain generalization via quantile risk minimization. In NeurIPS, 2022.
[69] Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-awe aware gradient matching for domain generalization. In CVPR, 2023.
[70] Zenan Huang, Haobo Wang, Junbo Zhao, and Nenggan Zheng. idag: Invariant dag searching for domain gencralization. In ICCV, 2023.
[71] Chenming Li, Daoan Zhang, Wenjian Huang, and Jianguo Zhang. Cross contrasting feature perturbation for domain generalization. In ICCV, 2023.
[72] Aveen Dayal, Vimal KB, Linga Reddy Cenkeramaddi, Mohan, Abhinav Kumar, and Vineeth N Balasub- ramanian. Madg: Margin-based adversarial learning for domain generalization. In NeurlPS, 2023.
[73] Zhe Wang, Jake Grigshy, and Yanjun Qi. Pgrad: learning principal gradients for domain generalization. In ICLR, 2023.
[74] Minyoung Kim, Da Li, and Timothy Hospedales. Domain generalisation via domain adaptation: An adversarial fourier amplitude approach. ICLR, 2023.
[75] Zhaorui Tan, Xi Yang, and Kaizhu Huang. Rethinking multi-domain generalization with general lcarning objective. In CVPR, 2024.
[76] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer An all-mlp architecture for vision. In NewrIPS, 2021.
13

---

[77] Hugo Touvron, Piotr Bojanowski, Mathilde Caron. Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave. Gautier Izacard, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, al. Resmlp: Feedforward networks image classification with data-efficient training. TPAMI. 2022.
[78] Qibin Hou, Zihang Jiang. Li Yuan, Ming-Ming Cheng, Shuicheng Yan, and Jiashi Feng. Vision permutator: permutable mip-like architecture for visual recognition. TPAMI, 2022.
[79] Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. CVPR, 2017.
[80] Antonio Torralba and Alexei A Efros. Unbiased look dataset bias. CVPR, 2011.
[81] Sara Beery, Grant Van Hom, and Pietro Perona. Recognition in terra incognita. In ECCV, 2018.
[82] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Sathcesh, Scan Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge.
IJCV, 2015.
[83] Kyungmoon Lee, Sungyeon Kim, and Suha Kwak. Cross-domain ensemble distillation for domain generalization. In ECCV, 2022.
[84] Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data augmentation. In ICML, 2019.
[85] Qinwei Xu, Ruipeng Zhang. Ya Zhang, Yanfeng Wang, and Qi Tian. A fourier-based framework for domain generalization. In CVPR, 2021.
[86] Zhuoxun He, Lingxi Xie, Xin Chen, Ya Zhang. Yanfeng Wang, and Qi Tian. Data augmentation revisited: Rethinking the distribution gap between clean and augmented data, arXiv preprint arXiv: 1r09v:1909.09 1909/9999999990 2019.
[87] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In JCML, 2021.
[88] Yongming Rao, Wenliang Zhao, Zheng Zhu, Jie Zhou, and Jiwen Lu. Gfnet: Global filter networks for visual recognition. TPAMI, 2023.
189] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization, ICCV, 2017.
14

---

A Appendix supplemental material
A.1 Theoretical Proofs
Lemma 1 [58). Let === {f E Hi: 11 1} denote function class, where Hy be RKHS with its associated kernel k. Let Ci. I - L(h(x), f(x) be a convex loss-function with parameter form [h(x) [(x)] for some q O, and defined Vh, F,, obeys the triangle inequality. Let S and T be two samples of size M drawn i.i.d from Dş and Dr, respectively. Then, with probability of least - (6 E (0, for all F. the following holds:
RD, SRDs [h] + dMwD (Ds, Drs (VttDs/Vtr(KDs))
(10)
log(2)
ExnDr (vtrKVt + +2V
to,
2M
KDS and DD are kernel func- tions computed on samples from Dg and Dr, respectively. the combined error of the ideal hypothesis he Ds and Dr-
where dMMD (D,, Dr) suplifll,, S1 Is
Theorem 1 (Generalization risk bound). With the previous setting and assumptions, let D's and Dr be two sets with M samples independently drawn from D and Dr, respectively. For any ò E (0, with probablity of at least - for all h E H. the following inequality holds:
RD (h) S Σ Tn TallD (h) + dTM-MMD (DT, DT)
sup d (D', D's) + 2X + (11) [JE[N]
where Ay 1 Sylla--gg//rKKDD (r(KKg) + Ez- I (abrrrrrriii
log(2/c) 2M
and a is the
minimum combined error of the ideal hypothesis h on both Ds and Dr. Let yr dIo (Dr, Dr) and Vs = suPi,jE[N] d1o-MMD (D's D's), respectively.
Proof. We initially investigate the relationship between the MMD 58 and To-MMD distances based on Definition (as presented in Eq. ). With the feature extractor (-)), we have:
// fd(((Ds) Fd(e(Ds)- U(Dr))
GAMD(Dss Dr) = sup
|| $1
/ fall Σ sup (Bi(Ds) vi (Dr))) 1=1 N.E..
(12)
< sup ||l $1
d Ds, Dr).
Then, for pair of source domain D and Dr, the following inequality holds:
d (D', Dr) S dIo Dr) + dTo (Dr, Dr),
(13)
with which we can derive weighted sum of To-MMD between source domains and target domain:
N
n dra MMD (MD Dr) S
n dr. MMD(D'G Dr) + d MMm Dr)
(14)
dIo-MMD D'S) dIo-MMD(DT. Dr). k.je[]
sup
With the above preparations, we now derive the generalization enror bound of the Mamba model on the unseen target domain. Recalling that Lemma indicates the generalization error bound between two different distributions, generalize it to the scenario of multiple source domains:
15

---

N
RD, (h) s
D (h)
Σ
nn dT (D3 Dr)+
==
n=1
un E By (i((k)rr())))))))))))))))))))( RRa Dr IV (r((K,))) 1og(2/c)
2(
2M + a (15)
nn R (h) + d MMD MMD Dr)+ + sup (.j[N]
dID (D's D's) + 22 t
where ' = M (E Th Be (r(K + trt.DD///rrK(oa... +
log(2/c)
and a is the 2M
minimum combined error of the ideal hypothesis h² on both D5 and Dr-
Proposion (Accumulation of Domain Discrepancy). Given two distinct domains Ds and Dr, the token-level domain distance (oo.MMD(Dss DT) depends on dcā (x, xx.) and d (a x f the i-th token. For the entire recurrent process, domain-specific information encoded in Sa, Sc, and SB will accumulate, thereby amplifying domain discrepancy.
Proof. For simplification, we use x E R and a7 R to denote the sample mean embeddings for the Ds and Dr. respectively. L represents the token sequence length. To investigate the generalization error boundary of Mamba, we explore simplified problem in conjunction with a single $6 layer, Le., at, where a E RLXL is the data-dependent matric. Empirically, based Eq. (3) and Eq. (4), we estimate the token-level domain gap using Fuclidean distance of y and y
L Σ
11 1-112=
Σ
-
a (ai
(16)
Combined with Eq. 2 and Eq. (3, we represent a as a direct function of the input ;::
aijj = Sz(
exp(
softmax(S<(E)))
(oftmax(Sa(3))S8(())
(17)
km
Since the SSM layer calculates based on the continuous subsequence [ž1, 02, ... F.J. we here analyze the domain gap of the extracted features the token level, i.e., ly y Specifically, assuming that we have calculated ly
of
s < L.), from Eq. (16), derive:
i+1
1 - Mit1l - -Iws -
1
Σ
aa
a
Σ (a
jaf -af
j=1
(18)
Ia - 2 (af+1.1
o + (a
(
of
1.41,441
With the defined a in Eq. 17 and denoting softmax(Sa())) as 5-(+) for brevity, we can express:
++1
0111 - 04. =So(F(+1)
exp(
(1,)))
k=/+1
- -e((z))
exp(
Σ (a(t:)A)
Sa(Z,)5((35)
(19)
k=f+1
Sc(T++1)
=1 Sc
exp(
SA (5:4 1] as
16

---

Considering that the differences between adjacent tokens are generally small, (c( S could be approximated to 1. When the dimension of z is relatively large, then for Eq. 19 we have:
ai+1, - 04. 2 (Z(+++)) 0s.
(20)
Then, substitute Eq. 20 into Eq.
18 to derive the following formula:
10 vit1l - 1y:² -
A ( ) a. - -Sa -0,55)) (4 +(+
(21)
-1 (5A( -
+ (a (a a af+1
Finally, recalling that 1 - v = B, we can express 1 vit1 as following:
1 vit1l S 3 - Av?
(So 18 (SE Des1 15
(22)
Let doāBo 55 1441, 2 Sc(E -, (
and d (I (11 = ( S Then, the above equation reveals that for the input tokens 2541 and E'+1. the distance between their extracted features, alongside the gap of historical sequences, primarily depends on å Bz (#f+1, : and d (# #1+1). Besides, the recurrent process in Eq. (22) could also lead to the accumulation or even enhancement of domain-specifie information, i.e., if the model extracts domain-relater information from the ith token, this part of the information will be retained in the features extracted by the 1-tthhhh token. For the whole recurrent process, domain-related information encoded in S.. Sc. and SB 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 SA, Sc, and SB.
Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by Sa. Sc, and SB 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, SA. Sc, and Sg are all linear projection layers, which map the input sequence ē E R. to the data-dependent matrixes A, C. and B, respectively. Hence, we analyze the influence of tokens in z these matrixes. Taking the matric B as an example, explore the simplified problem with Sn(2) = W, where WB E IR. LxN and N denotes the dimension of the hinder state. Inspired by previous works [84] 851. assume that the input sequence F could be decomposed to comaann-ppcific features i! and fomain--ppecifffeed I, in F, the projection matric B could be denoted
hSS Then, for the i-th token
B: Sp( W Wn,
| W W x3
(23)
Previous theoretical works [84] [6) 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 a.S. its corresponding weights S would be restricted to Ó. Simultaneously, due to minimal changes in the mamain--nvvrrin feature æ/ the learning of its corresponding weights W is promoted, thus enhancing Sn learning of ionaiininvariiat 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
17

---

Table 7: Performance (%) comparisons with SOTA DG methods on the DomainNet dataset with VMamba as the backbone. The best bolded.
Method
Params
Clipart
Infograph
Painting
Quickdraw
Real
Sketch
Avg.
CNN: ResNet-50
DecpAll
(AAAT20)
23M
63.00
21.20
50.10
63.70
52.00
44.00
PCL. M0
(CVPR222)
23M
67.90
24.30
55.30
15.70
66.60
56.40
47.70
EoA
07
(2c
23M
68.30
23.10
54.50
16.30
66.90
$7.00
47.70
EQRM
68
(Naur1P5223
23M
56.10
19.60
46.30
12.90
61.10
50.30
41.00
SAGM
(CVPPR133)
23M
64.90
51.50
64.10
53.60
45.00
IDG
70
23M
67.90
24.20
55.00
16.40
66.10
56,90
47.70
Domain Drop
60 (3CCV23)
23M
62.40
21.00
50,50
13.80
64.60
52.40
44.10
CCFP
(CCV(22)
23M
66.40
22.90
54.00
16.20
64.50
56.70
46.80
PGrad
73
(CLB/23)
23M
57.00
18.20
48.40
13.00
60.90
48.80
41.00
AGFA
(CLR)233
23M
66.70
22,90
54.00
16.70
65.90
47.10
GMDG
75
(CVPR'34)
23M
63.40
22.40
51,40
13.40
64.40
52.40
44.60
ViT-based or MI.P-like models
DoPrompt (22X)r22))
86M
67.70
24.60
54,90
17.50
69.60
55,20
48.30
SDVIT 18 (ACCV222)
22M
63.40
22,90
53,70
15.00
67.40
52.60
45,80
SSM-based models
Strong baseline
22
22M
74.12
28.06
58.26
17.85
70.10
60.33
51.45
START-M (Oun)
22M
75.11
29.41
60.25
19.31
71.05
61.58
52.79
START-X (Ouss
22M
75.28
29.36
60.33
19.55
71.01
61.30
52.81
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 reelaee spurious features that are challenging to fully extract and perturb. As result, perturbing 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.
A.2 Additional Experiments
Evaluation on DomainNet. We investigate the effectiveness our method on the large-scale dataset DomainNet. As shown in Tab. 7 on the challenging benchmark, find that the strong baseline VMamba can achieve the promising performance of 51 11.45% proving the superiority of Mamba models large-scale datasets. Based the strong baseline, our START can still significantly improve the performance, exceeding the baseline by 1.35% (52.81% vs 51 51..5%). Besides, compared with CNN-base methods, method can still achieve significant improvements, e.g., outperforming the latest SOTA method AGFA (74) by 5.71% (52.81% vs 47. 17.10%). Our methods also beat the best ViT-based method DoPrompt (17]. yielding by a large margin of 4.51% (52.81% 48.30%). The results prove the effectiveness of our method to help the model learn domain intttiinvvaatt representation.
Ablation studies larger datasets. We provide the ablation studies the Officehome and Terralncognita datasets. As shown in Tab. [ our methods perform the best among all variants, e.g., on Terralncognita, 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 Masea-based architectures. To validate the effectiveness of our START on other Mamba architectures, we conducted experiments 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 initial learning rate of 6.25e 6 and weight decay of lE 8. As shown in Tab. 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.3 87.38%), and on ViM-T, START-X exceeded the baseline by 1.66% (86.7 74% vs. 85.08%). These results demonstrate the generalizability of our method across different Mamba architectures.
Effeets on the ViT architecture. Recalling that our method is derived from the theoretical analysis that dependdnn matrices in Mamba could accumulate domain-related information during training,
18

---

Table 8: Ablation studies on different components of START. The experiments are conducted on large datasets, including OfficeHome and Terralncognita, with VMmmba as the backbone.
OfficeHome
Terralncognita
Method
Art
Clipart
Product
Real
Avg.
L100
L38
L43
L46
Avg.
Baseline
22
75.06
60.48
84.71
85.45
16.43
66,39
47.27
62.42
48.56
56, 46.162041
w/o. Saliency Guided
75.12
61.06
84,91
85,42
76.63
69.49
49.10
62.70
47.92
57.30
w/o. Token Selection
75.11
61.77
84,97
85.26
76.78=000
68.97
49.19
62.87
48.74
57.44
START-M (Oues)
75.15
62.04
85,31
85.84
77.09=0.16
70.13
49.98
63.02
49.49
58. 76
START-X (Dum)
75.48
62.06
85.24
85.47
77.07=0000
70.70
49.47
63.96
48.95
78.27.005
Table Effects (%) of our START on the Vim 2 architectures. The experiments are conducted on the PACS dataset with the ViM-T and ViM-S as the backbone, respectively.
Method
Params.
Art
Cartoon
Photo
Sketch
Avg.
Vimn-T 2X]
7M
88.59
80.45
98.52
72.74
85.08
SEART-Vim-T-M. (Ours)
7M
90.01
80.90
98.25
74.22
85,85
START-Vim-T-X
7M
90.97
81,53
98.90
75.57
86.74=067
Vim-S 23
26M
90.86
80.70
99.16
78.81
87.38_000
START-Vim-SS Vim-S-M (Ours
26M
93.77
83.79
99.46
79.45
89.12=0.0
START-Vime (Owr))
26M
93.98
84.23
99.44
80.00
89.41=002
Table 10: Effects (%) of our START on the ViT 87 architecture. The experiments are conducted on the PACS dataset with the DeiT-Small the backbone.
Method
Params.
Art
Cartoon
Photo
Sketch
Avg.
DeiT-Small 87
22M
87.55
82.16
98.45
75.24
85.85
START-ViTE (0ur)
22M
88.57
83.22
98.60
77.80
87.05.0034
START-ViT-X (Ous)
22M
88.72
83.01
98.50
76.78
86. 86.75+0022
Table 11:
Performance (%) of our START under the single-source domain generalization (SDG) setting. The experiments are conducted on the PACS dataset with VMamba as the backbone.
Method
Art
Cartoon
Photo
Skerch
Avg.
Strong baseline 22
75.10
83.29
44.92
74.76
69.52+064
START-M (Ours)
78.08
85.44 85.66
45.02 44.87
78.03
71.642023
START-X (Oar)
79.40
76.24
71.54=015
our START aims to improve the generalization of the Mamba architecture. Nevertheless, the core concept, adaptively perturbing salient tokens 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 START- ViT-M, which calculates token saliency from the input- matrices (i.e., Q x KT), and START-X START-ViT-X, which uses the activation value of representation æ to approximate saliency. The experiments are conducted on the representative ViT architecture, i.e., DeiT-Small, with the PACS dataset. As shown To 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-N outperforms the bascline 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 (22. parameters) is 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-invarian 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
generalization setting. As shown Tab. TT. START-M significantly improves the baseline, outperforming it by 2. 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.
19

---

Table 12: Effectiveness (%) of our START on SOTA augmentation methods. The experiments are conducted on the PACS dataset with VMamba as the backbone.
Method
Art
Cartoon
Photo
Sketch
Avg.
Strong baseline
22
91.55
85.11
99.14
83.97
89.94-0.52
DSU 14
92.58
85.91
98.98
85,39
90.71=0.2
START-DSU-M (Oun)
92.38
88.18
99.22
86.26
91.51
START-DSU-X
92.84
88.17
99.34
86.18
91.63.:225
ALOFT 15
93.07
86.04
99.16
85,31
90.89=0224
START-ALOFT-N (Ous)
93.07
88.01
99.34
85.72
91.54=0.24
START-ALOFT- (Ous)
93.12
88.23
99.40
85.65
91.60:00.8
Table 13: Performance (%) of START in different layers of the network. The experiments conducted on the PACS dataset with VMamba as the backbone.
Method
Art
Cartoon
Photo
Sketch
Avg.
Baseline
91.55
85.11
99.14
83.97
89.94=0.12
START-M (L1&2)
92.85
87.07
98,96
85.36
91.0610.19
START-M (L384)
92.77
86.50
98.66
85.85
90.95 1::22
START-M (L1&2&3&4)
93.29
87.56
99.14
87.07
91.77=0040
START-X (L1&2)
92.46
86.33
98.96
85.82
90.92.10.0
START-X (L.3&44
92.43
85.54
99.06
86.74
90.94=2002
START-X (L1&2&3344
92.76
87.43
99,22
87.46
91.72.10.99
START-M/X
92.59
87.14
98.88
85.70
91.08
Effectiveness with other SOTA augmentation methods. In our method, we utilize a statistics- 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 Vammba baseline, e.g., 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 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 the next layer's input, domain-specific features from earlier stages increase domain gaps later stages. To address the issue, we applied START to all layers 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 simultancously performs best, aligning with theoretical analysis. Applying START-M START- randomly across layers also improves performance, though less effectively than using START-M 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 M and START-X are only activated during training and disabled during inference, they introduce additional inference time. As shown in Tab. 14 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 source domains) by 6.2 6.22% (91.77% vs. 85.50%). demonstrating the superiority our START.
Visualization explanations. To provide visual evidence of the effectiveness of our START in suppressing domain-spesn 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.
20

---

Table 14: Comparison of the computational efficiency of SOTA DG methods and our START on
PACS. The experiments are conducted with 224 x 224 image size on one NVIDIA Teska V100 GPU.
Method
Backbone
Params (M)
GFlops (G)
Time (ms)
Avg. (%)
DeepAll
65
(AAAT20)
ResNet-50
23
8.26
85.50
IDAG
70 (CCV'23)
ResNet-50
23
8.00
94
88,80
iDAG
70
(CCVV23)
ResNet-101
41
15.00
495
89.20
GMoE-S
10
(CLLR22)
DeïT-S
34
5.00
136
88.10
GMoE-B
19 (ICLR22))
DeiT-B
133
19.00
361
89.20
ViP 78 (TPAM)22))
ViP-$
25
13.84
88.27
GFNet
88
(TPAM)22)
GFNct-H-Ti
13
4.10
87.76
DGMamba
34
(ACM
VMamba-T
31
5.00
233
91.20
Strong Baseline
22
VMamba-T
22
5.68
252
89.94
START-M
VMamba-T
22
5.68
252
91.77
START-X (Ours)
VMamba-T
22
5.68
252
91.72
VMamba (Baseline)
START-X (Ours)
START-M
MMamba (Baseline)
START-X (Ours)
START-M (Ours)
Origin
(Ours)
Origin
Art (Target)
Photo
Cartoon
Sketch
Figure 4: Visualization results of our START. The experiments are conducted on the PACS dataset with the "Art" as the target domain. We visualize the attention maps of the last layer the VMamba backbone. For each sample, the first column is the original image, the second column is the attention map 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 semantic features, e.g., holistic shape structure, than the pure VMamba baseline.
the VMamba baseline tends focus on specific local patches that encode Nomain-specinn features, leading to overfitting to source domains. In contrast, our START methods effectively reduce the model's focus on domain-specifi 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 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 promising method for DG tasks.
Difference from the related work. 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 heuristic method address the issue. However, it lacks deep analysis of the phenomenon. Differently, we first theoretically delve into the generalizability of Mamba, revealing how input-dependent matrices contribute domain gap accumulation. Based on the analysis, we developed START to enhance Mamba's generalization. 2) Different Goals and Methody: 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-specifff information in tokens focused input-dependent
21

---

matrixes, perturbing only styles while keeping contents unchanged. Notably, DGMamba uses the GradCAM [89 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.
A.3 Broader Impact
Our work aims enhance model generalization and computational efficiency, enabling robust performance across diverse domains with varying distributions. By mitigating the overfitting issue on limited source domains and improving performance on unseen target domains, we believe will have a positive societal impact.
A.4 Limitations of Our Work
In the paper, we employ style perturbation to perturb domain-specifie information within salient tokens of input sequences. However, there various forms of domain-specific information in the images, which could not be completely suppressed by our method. To address this issue, a potential solution design advanced feature disentanglement methods that adaptively distinguish and perturb domain-specific information. Designing class-aware feature augmentation could also alleviate the accumulation of domain-related features. We will explore these solutions in future work.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See Appendix A.4
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well- specifficci,,,nn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the
implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
22

---

The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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 faireess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
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 proof's in Section 3.2 Appendix A.
Guidelines:
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 cross- referenced.
All assumptions should be clearly stated or referenced 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 provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification:
See in Section
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer 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 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 specific model and empirical evaluation, it may be necessary to either make 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
23

---

While NeurIPS does not require releasing code, the conference does require all submis- sions 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 reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The source code is provided alUupx///ghhbbbcoo////ggengggggg
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
public/ /ggides/ gwwdeee/ooeeuuumm....em.oomo ClodeSubmissionPP
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. wwtpppblccgguuieeeoooooooocpccpe http//uddes///// wwwoooeomiisiionne...
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why. At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: See in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them. The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
24

---

Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See in Section
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
25

---

Answer: [Yes]
Justification:
See in Appendix A.3
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal
impact 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), fairesss 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 particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 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 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).
11. Safeguards
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)?
Question:
Answer: [NA].
Justification: The paper does not release the data or models with high risk for misuse.
Guidelines:
The answer NA means that the paper poses 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 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 (8), Terralncognita [8]. 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 dataset.
The authors should state which version of the asset used and, possible, include
URL.
26

---

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, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
27

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "ZJKNwODtbG", "reviewer": "Reviewer_EPCk", "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.", "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?\n- 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.\n- 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": "The authors have addressed the limitations and, if applicable, the potential negative societal impact of their work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- This paper examines the generalization ability of SSMs (Mamba), a topic that few studies have explored. This aspect is particularly important if Vision Mamba is to be widely adopted.\n- The generalization problem of Mamba and the subsequent solution, i.e. START, is motivated and supported by theoretical analysis.\n- START utilizes the input-dependent matrices to identify salient tokens that are related to domains, which is novel. Then, START uses mix-style to augment domain-related features to force the model to learn domain-invariance features, which is tailored to enhance SSMs' generalization ability and has proven effective compared to the strong baseline in the ablation study.", "weaknesses": "- 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.\n- 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?\n- 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.\n- 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.\n- Some typos. For example, line 219, xaug should be $x_{aug}$\n\nDespite lacking some comparison, I believe this paper contains novelty and is technically sound. I will raise the point if the above questions are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable reviews. Due to the character limit of $6000$ here, the tables are displayed in ''glocal_tables.pdf''. \n\n### **Q1: Revisions of Proposition 1 and 2.**\nThanks. We have modified Propositions 1 and 2 to be more clear and detailed: \n\n**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._\n\n**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._\n\n\n### **Q2: About theoretical results.**\nAs 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)$.\nThe smaller the two terms, the lower the upper bound of generalization error.\nFollowing [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.\nWe 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$.\nAs 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.\n\n### **Q3: Compared with DGMamba on five DG datasets.**\nWe 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.\n\n### **Q4: Compared with other feature identification methods.**\nWe 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.\n\n\n### **Q5: Some Typos.**\nThanks. We will re-check the paper and fix all typos in the revision.\n\n### **Q6: Differences from DGMamba.**\nOur 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.\n\n### **Q7: Ablation studies on larger datasets.**\nWe 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. \n\n### **Q8: For domain-invariant features in input-dependent matrices.**\nThanks. 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TAZfSr1JEb", "reviewer": "Reviewer_AMLN", "summary": "This paper studies the role of MAMBA architectures on Domain Generalization benchmarks and adapts the architecture to achieve robust generalization. \nThe 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.\nTo 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. \nIn 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. \nThey have conducted experimental analysis on standard DG benchmarks and demonstrate empirical gains compared to the baselines.", "questions": "please see weaknesses", "limitations": "The authors have sufficiently addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper comprehensively studies the generalization capabilities of MAMBA (state-space) models with respect to their ability to handle distribution shifts. The theoretical analysis is somewhat straightforward, relying on the MMD measure, but it is an important contribution to the field. The experiments follow standard protocols. The proposed methods are mathematically grounded and produce empirical gains compared to baselines.", "weaknesses": "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.\nThere 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.\nIn 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.\nSecond, 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?\nSmall 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Q1: Motivation of our method.**\nThanks 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. \n\n\n### **Q2: Difference from DGMamba.**\nIn essence, our approach significantly differs from DGMamba in motivations, goals and methods.\n\n1) **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.\n\n| Method | Domain gap of $\\tilde{\\Delta} (\\downarrow)$ | Domain gap of B $(\\downarrow)$ | Domain gap of C $(\\downarrow)$ | Domain gap of Feat. $(\\downarrow)$|\n| :----: | :----: | :----: | :----: | :----: |\n| Baseline | 1.48 | 1.52 | 2.08 | 2.97 |\n| MixStyle  | 1.73 | 1.36 | 1.90 | 1.91 |\n| DSU | 1.38 | 1.28 | 2.18 | 1.59 |\n| ALOFT | 1.37 | 1.25 | 2.33 | 1.67 |\n| START-M (Ours) | **1.16** | 0.98 | 1.80 | **1.30** |\n| START-X (Ours) | 1.23 | **0.91** | **1.52** | 1.37 |\n\n2) **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.\n\nWe appreciate your feedback and will include this discussion in the revised version.\n\n### **Q3: Experimental comparison with DGMamba.**\nAs 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. \nThe 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.\nSpecifically, 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.\nOn 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.\nMore advanced methods should be studied to improve model performance on OfficeHome further, and combining our START with DGMamba could be a promising direction.\nWe will add this discussion in the revised version.\n\n| Method | VLCS | PACS | OfficeHome | TerraInc | DomainNet | Avg. |\n| :----: | :----: | :----: | :----: | :----: | :----: | :----: |\n| DGMamba | 80.80 | 91.20 | 77.00 | 54.60 | 49.60 | 70.64 |\n| **START-M (Ours)** | **81.32** |  **91.77** |  **77.09** | 58.16 | 52.79 |  **72.23** |\n| **START-X (Ours)** | 81.21 | 91.72 | 77.07 |  **58.27** |  **52.81** | 72.22 | \n\n\n\n### **Q4: When do START-M and START-X happen?**\nYes, 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]. \nIf 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.\nThanks for your question. We will explore this direction in future work.\n\n[1] Mixstyle neural networks for domain generalization and adaptation. IJCV 2024\n\n[2] Rethinking multi-domain generalization with a general learning objective. CVPR 2024\n\n[3] Dgmamba: Domain generalization via generalized state space model. ACM MM 2024\n\n[4] MADG: Margin-based Adversarial Learning for Domain Generalization. NeurIPS 2023. \n\n### **Q5: Symbol improvements.**\nThanks 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d53nGxcuZv", "reviewer": "Reviewer_qvXD", "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.\n\nTo 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.", "questions": "Please refer to the weakness.", "limitations": "This paper discusses the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This passage discusses a theoretical analysis of the Mamba model's performance in domain generalization. The analysis reveals that the Mamba model's input-dependent matrices have a tendency to accumulate domain-specific features during their recurrent processing. This accumulation could potentially limit the model’s ability to generalize effectively across different, unseen domains. Essentially, the matrices that should help the model adapt to new data instead make it more sensitive to the specific characteristics of the data it was trained on, thereby impeding its general applicability.\n\n2. This statement outlines the proposal of a new state space model (SSM)-based architecture, which incorporates a saliency-driven, token-aware transformation aimed at addressing the issues previously identified with the Mamba model in domain generalization tasks. This novel architecture is designed to serve as an effective alternative to both convolutional neural networks (CNNs) and vision transformers (ViTs). It boasts an ability to generalize excellently across diverse domains while maintaining the efficient linear complexity characteristic of state space models, making it a promising option for handling domain shifts in a computationally efficient manner.", "weaknesses": "It would be better to report computation costs like inference time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Q1: Computational costs including inference time.**\nThanks 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.\n\n| Method | Backbone | Params (M) | GFlops (G) | Time (ms) | Avg. (%) |\n| :----: | :----: | :----: | :----: | :----: | :----: |\n| DeepAll [64] | ResNet-$50 | 23 |8.26 |- |85.50|\n| iDAG [69]  | ResNet-50 |23 |8.00 |94 |88.80|\n| iDAG [69]  | ResNet-101 |41 |15.00  |495 |89.20|\n| GMoE-S [19]  | DeiT-S |34 |5.00 |136 |88.10|\n| GMoE-B [19]  | DeiT-B |133 |19.00 |361 |89.20|\n| ViP [77] | ViP-S |25 |13.84 |- |88.27|\n| GFNet [87] | GFNet-H-Ti |13 |4.10 |- |87.76|\n| DGMamba [78]  | VMamba-T |31 |5.00 |233 |91.20|\n| Strong Baseline [22] | VMamba-T |22 |5.68 |252 |89.94|\n| **START-M (Ours)** | VMamba-T |22 |5.68 |252 |91.77|\n| **START-X (Ours)** | VMamba-T |22 |5.68 |252 |91.72|", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QbUcdz9GoT", "reviewer": "Reviewer_3dvs", "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.", "questions": "See weaknesses. Another suggestion is to list model accuracy on ImageNet in Table 1 and Table 2.", "limitations": "The authors have discussed limitations in the appendix.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- This paper is early work to adapt Mamba for Domain Generalization. \n\n- The baseline is Vmamba and START can improve VMamba for Domain Generalization benchmark.", "weaknesses": "- 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?\n\n- ViT has attention maps, which are also input-dependent. Can your methods transfer to ViT? Or your method can only be applied to Mamba?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Q1: Comparision with CNN-based or ViT-based DG methods on VMamba.**\nThanks for your feedback. \nIndeed, 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.\nTo address this issue, we developed a sequence-level saliency-driven token transformation, which aims to suppress domain-specific information within the SSM layers explicitly. \nWe 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). \nAs 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\\\\%$).\nThe 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.\n\n\n| Method | Art | Cartoon | Photo | Sketch | Avg. |\n| :----: | :----: | :----: | :----: | :----: | :----: |\n| Baseline [22] | 91.55 | 85.11 | 99.14 | 83.97 | 89.94 ± 0.52 |\n| MixStyle [13] | 92.05 | 86.55 | 98.90 | 86.35 | 90.94 ± 0.18 |\n| DSU [14] | 92.58 | 85.91 | 98.98 | 85.39 | 90.71 ± 0.22 |\n| ALOFT [15] | 93.07 | 86.04 | 99.16 | 85.31 | 90.89 ± 0.24 |\n| **START-M (Ours)** | **93.29** | **87.56** | 99.14 | 87.07 | **91.77** ± 0.40|\n| **START-X (Ours)** | 92.76 | 87.43 | **99.22** | **87.46** | 91.72 ± 0.49 |\n\n\n### **Q2: Performance of our START on ViTs.**\nRecalling 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.\n\n\n| Method | Art | Cartoon | Photo | Sketch | Avg. |\n| :----: | :----: | :----: | :----: | :----: | :----: |\n| DeiT-Small | 87.55 | 82.16 | 98.45 | 75.24  | 85.85 ± 0.30   |\n| **START-ViT-M (Ours)** | 88.57 | **83.22** | **98.60** | **77.80** | **87.05** ± 0.34 |\n| **START-ViT-X (Ours)** | **88.72** | 83.01 | 98.50 | 76.78 | 86.75 ± 0.22 |\n\n\n### **Q3: Model accuracy on ImageNet in Tab.1 and Tab.2.**\nThanks 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].\nIn 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.\n\nSince 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.\n\nAs 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.\nWe appreciate your suggestion and will incorporate this discussion into the revised paper for improved clarity.\n\n| Method | Pretrained VMamba | Trained VMamba | START-M | START-X | \n| :----: | :----: | :----: | :----: | :----: |\n| ACC. | 73.33 | 86.33 | **90.33** | 89.67 |\n\n[1] Mixstyle neural networks for domain generalization and adaptation. IJCV 2024\n\n[2] Rethinking multi-domain generalization with a general learning objective. CVPR 2024\n\n[3] Dgmamba: Domain generalization via generalized state space model. ACM MM 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aUwIJuV3rQ", "reviewer": "Reviewer_4RVV", "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.", "questions": "Besides the questions I raised in Weakness, I have one more questions for the author about the necessity of this research.\n\n6.\tI 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": "No Limitation section is mentioned in the main paper.\n\nI am not sure about the generalization ability of this analysis to different architectures.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1.\tClear paper writing.\n2.\tThe analysis of SSM on computer vision tasks is important, which provide concrete analysis of domain bias in the network architecture.\n3.\tBased on the analysis, the salience is an interesting matric in indicating the selection of tokens to mitigate the performance gap across domains.", "weaknesses": "1.\tI 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.\n2.\tSelecting 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.\n3.\tFor your experiments, what is the value of N. Have you considered single-source domain generalization in your evaluation, i.e., N=1?\n4.\tHonestly, the performance gain is not that significant on Office-Home, comparing with PACS. Is there any reason for such trivial gain?\n5.\tThough 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable reviews. Due to the character limit of $6000$ here, the tables are displayed in the ''glocal_tables.pdf''. \n\n### **Q1: The number of source domains.**\nThanks for your constructive feedback.\nFollowing 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.\n\n### **Q2: For useful information in tokens with high saliency.**\nYes, 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.\n\n\n### **Q3: Performance for SDG.**\nIn 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.\n\n\n### **Q4: Discussion about performance on OfficeHome.**\nOfficeHome 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.\n\n### **Q5: Effects across different stages.**\nOur 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.\n\n\n### **Q6: Effects of Mamba and its variants.**\nYes, 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.\n\n### **Q7: Limitations of our method.**\nThanks. 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank the ACs and reviewers for their constructive comments on our paper. We are encouraged by the positive feedback, including remarks such as "the aspect is important", "this paper contains novelty and is technically sound", "the theoretical analysis is a significant contribution to the field", and "the method is novel and mathematically grounded".

In our rebuttal, we answer the reviewers' concerns point-by-point, covering: 1) explanations of the theoretical analysis, 2) differences and comparisons with related works such as DGMamba, 3) protection of semantic features in input-dependent matrices, 4) additional ablation studies, 5) implementation details, and 6) scalability of our analysis and method. 

Due to the $6,000$-character limit for responses to each reviewer, we have included the "global_tables.pdf" file, which contains all relevant tables. We sincerely hope our responses could address all concerns and look forward to further discussion during the author-reviewer discussions phase.

Meta-review: 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.
Decision: Accept (poster)