Title: START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation
Venue: NeurIPS 2024
Accepted: yes
Keywords: Domain generalization

state space models
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

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.

1 Introduction

Deep learning models have achieved impressive progress in various computer vision tasks over the past years [1-3]. Such a huge success is mostly based on the independent and identically distributed ( i.i.d. ) assumption, i.e. , the training and testing data follow the same distribution [4]. However, when evaluated on test data following different distributions from the training data, these models often suffer severe performance degradation. This issue, which is known as domain shift [4], has greatly hindered the applications of deep learning models in the real world.

To improve the generalization of the model under domain shifts, Domain Adaptation (DA) has been widely studied, which aims to transfer the knowledge learned from labeled source domains to the unlabeled or partially labeled target domain [5, 6]. However, DA methods cannot guarantee the performance of the model on unknown target domains that have not been observed during training [7, 8]. Since the accessibility of the target domain could not always be satisfied in real scenarios,

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

1.4

1.5

1.6

1.7

1.8

1.9

2.0

2.1

2.2

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Domain Distance

Baseline

START-M

(Ours)

START-X

(Ours)

Response sequence y

2.97

1.30

1.37

Domain Distance

Baseline

START-M

(Ours)

START-X

(Ours)

Input-dependent

𝑪

2.08

1.80

1.52

Domain Distance

Baseline

START-M

(Ours)

START-X

(Ours)

Input-dependent

𝑩

1.52

0.98

0.91

Domain Distance

Baseline

START-M

(Ours)

START-X

(Ours)

Input-dependent

#

∆

(

×

10)

1.48

1.16

1.23

Domain 1

Domain 2

Domain N

…

…

…

…

Linear + SoftMax

Linear

#

∆

= softplus(

𝑆

∆

𝑥

"

)

𝐵 = 𝑆

#

𝑥

"

𝐶 = 𝑆

$

𝑥

"

Linear

𝒙 = [𝒙

𝟏

, 𝒙

𝟐

, … , 𝒙

𝑳

]

Embedding

𝐴

'Domain Shift'

'Domain Shift'

Domain Distance

𝒙

" ∆

𝑩

𝑪

𝒚

1.48

1.52

2.08

2.97

1.26

Discrepancy

Accumulated

Input-Dependent Matrices:

Accumulating Domain Discrepancy during the Recurrent Process.

𝒚 = [𝒚

𝟏

, 𝒚

𝟐

, … , 𝒚

𝑳

]

Output Sequence:

SSM Layer

'Mitigating Domain Discrepancy Accumulation in the Input-Dependent Matrices.'

1.0

1.1

1.2

1.3

1.4

1.5

1.6

Domain Generalization (DG) is proposed to develop a domain-generalizable model on unseen target domains by learning multiple different but related source domains [8, 9].

Most existing DG methods focus on learning domain-invariant representations across source domains, primarily via domain alignment [10, 11], meta-learning [7, 12], and data augmentation [13, 14]. These methods heavily rely on convolutional neural networks (CNNs), which have limited receptive fields due to local convolutions. Consequently, the CNN-based methods inevitably tend to learn local texture information, leading to overfitting to source domains and poor generalization on target domains[15, 16]. Recent works in DG have introduced Vision Transformers (ViTs) as the backbone for DG, utilizing the global receptive field of the self-attention mechanism to mitigate local texture bias [17-19]. However, the complexity of self-attention increases quadratically with input length, resulting in significant computational overhead for ViTs when modeling long sequences [20, 21].

To address this issue, some pioneers have proposed advanced state space models (SSMs) [20, 22, 23], represented by Mamba [20], which selectively models token dependencies in input sequences in a compressed state space. The selective scan mechanism allows Mamba to achieve linear complexity in sequence length during training and fast RNN-like computation during inference. Despite the remarkable performance of Mamba-based methods on supervised learning tasks, few existing works have analyzed the generalization ability of Mamba under domain shift. It remains an open question whether the Mamba model can achieve excellent performance for DG tasks.

In this paper, we theoretically analyze the generalization error bound of the Mamba model under domain shifts. We find that the domain distance of features extracted by the model is strongly related to the input-dependent matrices within the model. These matrices accumulate and amplify domain-specific features during training, which exacerbates the overfitting issue of the model to source domains. We empirically measure the distance among source domains within the input sequence x , the response sequence y , and the input-dependent matrices of the last network layer. As shown in Fig. 1, we observe that for the baseline model, input-dependent matrices ( ˜ ∆ , B , and C ) are prone to learning domain-specific features from the input x . Since the output y is calculated by the recurrent product of x and these matrices, domain-specific features are accumulated and amplified, causing the model to overfit the source domains. To address this issue, we propose a Generalized State Space Model with Saliency-driven Token-Aware Transformation (START), which can reduce domainspecific information in input-dependent matrices during training. Building on the latest Mamba-based model [22], we develop a strong baseline for DG that outperforms many SOTA DG methods, which selectively learns global dependencies among tokens with linear complexity in sequence length.

Moreover, based on theoretical analysis, we design the saliency-driven token-aware transformation method, which simulates domain shifts during training by selectively introducing style perturbations to tokens focused on by the input-dependent matrices. START constrain these matrices to learn domain-invariant features, thus mitigating the overfitting of model on source domains. Experiments on five datasets prove the effectiveness of our method. Our contributions are summarized as follows:

We conduct a theoretical investigation into the generalization ability of the Mamba model, revealing that the input-dependent matrices in Mamba can accumulate domain-specific features during the recurrent process, thus hindering the model's generalizability.

Based on theoretical analysis, we propose a novel SSM-based architecture with saliencydriven token-aware transformation as a competitive alternative to CNNs and ViTs for DG, which performs excellent generalization ability with efficient linear complexity.

For the saliency-driven token-aware transformation, we explore two variants to identify and perturb salient tokens in feature sequences, effectively reducing domain-specific information within the input-dependent matrices of Mamba. Our method achieves SOTA performances, e.g. , yielding the best CNN-based method by 5 . 87% ( 58 . 27% vs. 52 . 40% ) on TerraIncognita.

2 Related Works

Domain generalization. Traditional DG methods, primarily based on CNN backbones, can be broadly classified into three categories: domain alignment, meta-learning, and data augmentation. Motivated by the learning theory of domain adaptation [4, 25], domain alignment methods seek to learn domain-invariant representations through adversarial learning [10, 26, 27], causal learning [28, 29], or feature disentanglement [30, 31]. Another popular way to address DG is meta-learning, which partitions the training data from multiple source domains into meta-train and meta-test sets to simulate domain shifts during training [7, 12, 32]. Data augmentation is also an effective method to enhance model robustness to domain shifts by generating diverse data invariants through adversarial generation [33, 34], style perturbation [13, 14], and learnable parameters [35, 36]. However, CNNbased DG methods suffer from the limited receptive field of convolutions, often leading to a texture bias and overfitting to source domains [15, 37]. To address this, some researchers have introduced ViT-based methods for DG, which capture global representations by leveraging long-range spatial dependencies with attention mechanisms [18, 17, 19]. Despite their advantages, ViT-based methods are computationally intensive due to the quadratic complexity of the self-attention mechanism, limiting their practical applications [20, 21]. Inspired by the emerging Mamba model [20, 22, 23], we explore a novel SSM-based architecture for DG that combines strong generalizability with efficient linear complexity. We theoretically analyze the generalization error bound of Mamba and design a novel saliency-driven token-aware transformation to suppress domain-specific features in the input-dependent matrices of Mamba, thereby enhancing the generalization ability of the model.

State space models. Recently, state space models (SSMs) have demonstrated promising performance across various vision tasks [38-40] for their ability to effectively capture long-range dependencies while maintaining linear scalability with sequence length. Derived from the classical state space model [41], the Structured State Space Sequence Model (S4) [42] addresses computational constraints through novel parameterizations catering to continuous-time, recurrent, and convolutional views of the state space model. Notably, Mamba [20] has emerged as a standout performer, which integrates selection mechanism and hardware-aware algorithms into previous works [43-45], thus achieving linear-time inference and efficient training mechanisms. Based on the success of Mamba, Vision Mamba (Vim) [23] applies Mamba to ViT architecture, combining bidirectional SSM for datadependent global visual context modeling. Meanwhile, VMamba [22] designs a cross-scan module to bridge the gap between 1 D array scanning and 2 D plain traversing. Mamba-based architectures have exhibited superior performance across various supervised vision tasks, including medical image segmentation [46-48], point cloud analysis [49-51], and remote sensing analysis [52, 53]. However, few works explore the performance of Mamba under domain shifts for DG. Although DGMamba [54] has recently introduced a pioneering Mamba-based framework for DG, it lacks a deep analysis of the generalizability of Mamba. In the paper, we conduct a theoretical analysis of Mamba's generalizability, revealing that input-dependent matrices within Mamba could accumulate domainspecific information, thereby impeding model generalization. Consequently, we propose a generalized SSM-based architecture for DG, incorporating a non-parametric module to selectively perturb salient tokens within input-dependent matrices, thus enhancing model generalization to unseen domains.

3 Method

3.1 Preliminary

State Space Models (SSMs). The SSM is a type of linear time-invariant systems that map input sequence x t ∈ R L to response sequence y t ∈ R L through a hidden state h t ∈ R N . Mathematically, this process is formulated as the subsequent linear ordinary differential equations (ODEs): h ′ t = Ah t + Bx t , y t = Ch t , where A ∈ R N × N is the evolution parameter, and B ∈ R N × 1 , C ∈ R N × 1 are the projection parameters. However, the differential equation is hard to solve in the deep learning setting, thus discrete SSM [55, 44] suggests discretizing the system with a time scale parameter ∆ :

where ¯ A and ¯ B are discrete counterparts of the continuous parameters A and B , and ∆ ∈ R > 0 is the sampling timescale for the discretization process. Although the discrete SSMs can achieve linear time complexity, they rely on static parameterization, i.e. , ¯ 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 ∈ R L × N , C ∈ R L × N , and ∆ ∈ R L × D on the input x t ∈ R L × D :

S B , S C , S ∆ are linear projection layers and softplus( · ) = log(1 + exp( · )) . The input-dependent time-variant layers could enhance recurrent layers, making them more expressive and flexible in capturing complex dependencies [21]. The parameter matrixes can be further expressed as ¯ A = [ ¯ A 1 , · · · , ¯ A L ] , ¯ B = [ ¯ B 1 , · · · , ¯ B L ] , C = [ C 1 , · · · , C L ] , where L is the sequence length. Considering the initial state h 0 = 0 , Eq. (1) can be unrolled as [21]:

where α i,j = C i ∏ i k = j +1 ¯ A k ¯ B j , for 0 ≤ j < i ≤ L , characterizing S6 layer as a data-dependent self-attention [56]. The attention matrix α is determined by both the input and the parameter matrices.

3.2 Theoretically Analysis for the Generalization Ability of Mamba

Previous DG methods have primarily focused on enhancing the generalizability of CNNs or ViTs, lacking theoretical investigations into the Mamba model. We theoretically explore the generalization error bound of Mamba, proving that perturbing the domain-specific features within the input-dependent matrices of Mamba can effectively diminish the upper bound of the model's generalization risk .

Notations. Given a training set of N source domains D S = { D 1 S , D 2 S , · · · , D N S } , the objective of DG is to use D S to train a model that is expected to perform well on unseen target domain D T . Let h : X → Y be a hypothesis from the candidate hypothesis space H , where X and Y denote the input space and the label space, respectively. Since Mamba learns dependencies among tokens from continuous sequences, we study its generalizability at the token level. Let ψ ( · ) be the feature extractor of h that maps input images into feature space. Following Integral Probability Metrics [57, 58], we define the token-level Maximum Mean Discrepancy to estimate the gap between different domains.

Definition 1 (Token-level Maximum Mean Discrepancy). Given two different distributions of D S and D T , let L denote the number of tokens in the response sequence of ψ ( · ) , then we define the TOken-level Maximum Mean Discrepancy (To-MMD) between ψ ( D S ) and ψ ( D T ) as:

where Ψ represents the hypothesis space for each token, ψ t ( D ) denotes the distribution of the t -th token for domain D , and F k is a RKHS with its associated kernel k .

We here investigate the generalization risk bound of the Mamba model. Theoretically, as in [59, 60], the risk of the hypothesis h on the domain D is defined as: R D ( h ) = E x ∼ D [ L ( h ( x ) -h ∗ ( x ))] , where L : Y×Y → R + is a convex loss-function that measures the distance between h and the true labeling function h ∗ . Moreover, following [8, 60], for multiple source domains D S = { D 1 S , D 2 S , ..., D N S } , the convex hull Λ S is defined as a set of a mixture of source domains, i.e. , Λ S = { ¯ D : ¯ D ( · ) = ∑ N n =1 π i D n s ( · ) , ∑ N n =1 π n = 1 , π n ∈ [0 , 1] } . The ¯ D T ∈ Λ S is defined as the closest domain to the target domain D T . Based on Eq. (4), the following generalization risk bound can be derived.

Theorem 1 (Generalization Risk Bound). With the previous setting and assumptions, let D i S and D T be two sets with M samples independently drawn from D n S and D T , respectively. For any δ ∈ (0 , 1) with probablity of at least 1 -δ , for all h ∈ H , the following inequality holds:

where λ π = 1 M ( ∑ N n =1 π n E x ∼ D n S [ √ tr ( K D n S ) + E x ∼ D T [ √ tr ( K D T )]) + √ log (2 /ϵ ) 2 M , and σ is the minimum combined error of the ideal hypothesis h ∗ on both D S and D T . Let κ T = d To-MMD ( D T , ¯ D T ) and κ S = sup i,j ∈ [ N ] d To-MMD ( D i S , D j S ) , respectively.

The proof of Theorem 1 is provided in Appendix A.1. The inequality indicates that the generalization error bound depends on κ T denoting the token-level maximum distance between source and target domains, and κ S measuring the maximum pairwise gap among source domains at the token level. The smaller the two terms, the lower the upper bound of generalization error. Following [25, 58, 61], we simplify the To-MMD in Eq. (4) by choosing a unit ball in the F k and using Gaussian kernel with parameter γ to estimate κ T and κ S . Let ¯ x S ∈ R L and ¯ x T ∈ R L to denote the mean embeddings of samples from D S and D T , where L represents the token sequence length. We explore a simplified problem in conjunction with a single S 6 layer, i.e. , ¯ y = α ¯ x , with α ∈ R L × L being the data-dependent matric. The domain distance between ¯ y S and ¯ y T is formulated as k (¯ y S , ¯ y T ) = exp( -|| ¯ y S -¯ y T || 2 /γ ) , where γ is the kernel parameter. Specifically, for input-dependent matrices B,C, ∆ , we denote softmax( S ∆ ( · )) as ˜ S ∆ ( · ) . Then, we analyze the impact of these input-dependent matrices on || ¯ y S -¯ y T || 2 , which is applicable to both κ T and κ S . For the i -th tokens ¯ x S i and ¯ x T i , we define:

With the recurrent property of the S 6 layer in Eq. (3), we can derive the following propositions:

Proposion 1 (Accumulation of Domain Discrepancy). Given two distinct domains D S and D T , the token-level domain distance d To-MMD ( D S , D T ) depends on d C ˜ ∆ Bx (¯ x S i , ¯ x T i ) and d ˜ ∆ (¯ x S i , ¯ x T i ) for the i -th token. For the entire recurrent process, domain-specific information encoded in S ∆ , S C , and S B will accumulate, thereby amplifying domain discrepancy.

Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by S ∆ , S C , and S B can enhance their learning of domain-invariant features, thus effectively mitigating the accumulation issue in these input-dependent matrices.

Propositions 1 and 2 are proved in Appendix A.1. Based on the propositions, we develop a saliencydriven token augmentation method, which perturbs style information within the tokens that the model focuses on at the sequence level. In this way, our method enhances the extraction of domain-invariant features by the input-dependent matrices, i.e. , S ∆ , S C , and S B . As presented in Tab. 6, we also empirically validate the effectiveness of our method in reducing domain discrepancy in these matrices.

3.3 Saliency-driven Token-Aware Transformation for Mamba

To boost the generalization ability of the Mamba model, leveraging the Proposion 1 and Proposion 2 , we propose a novel Saliency-driven Token-AwaRe Transformation paradigm (START in short), which aims to explicitly suppress domain-related features within the input-dependent matrixes. Unlike prior methods that perturb entire feature maps at the channel level [13-15], START incorporates a

̅

Original Sequence

x

B = 𝑺

𝑩

(𝒙)

𝑪 = 𝑺

𝑪

(𝒙)

(

∆= softplus(𝑺

∆

𝒙 )

…

Saliency-based Selection

…

…

Calculated Saliency of x

Maximum Saliency

…

…

Sequence-level Token Mask

ℳ

…

…

𝜇 = 𝜆𝜇

!

+ 1 - 𝜆 𝜇

"

(

𝜎 = 𝜆𝜎

!

+ 1 - 𝜆 𝜎

"

Stylized

̅ 𝑥 =

$%&

'

3

𝜎 +

̅ 𝜇

Token Replacement

Augmented Sequence

! 𝒙

Original

x

Random

𝒙

(

Style Interpolation

7

𝑥 = ℳ⨀𝑥 + (1 - ℳ)⨀ ̅ 𝑥

𝑴

𝑪	×	 -

∆	×B	×	𝒙

START-M By Matrices

𝑿

𝒙

START-X By Input

Or

Input

Output

Flatten &

Linear Projection

Pooling & Prediction

VSS

Block

S

T

A

R

T

VSS

Block

S

T

A

R

T

VSS

Block

S

T

A

R

T

VSS

Block

S

T

A

R

T

Input Sample

H

W

Stage 1

Stage 2

Stage 3

Stage 4

×𝐿

+

×𝐿

,

×𝐿

-

×𝐿

.

saliency-driven token selection scheme to perturb the prominent regions of input-dependent matrics S ∆ , S B , and S C . Based on the attention mechanism outlined in Eq. (3), we propose two variants to identify and perturb tokens within salient regions, including START-M that determines saliency using input-dependent matrices, and START-X computing saliency based on input sequences.

START based on input-dependent matrices (START-M). As Proposition 1 reveals, for the i -th token, the token-level domain gap depends on d C ˜ ∆ Bx (¯ x S i , ¯ x T i ) and d ˜ ∆ (¯ x S i , ¯ x T i ) . Specifically, as presented in Eq. (6), d C ˜ ∆ Bx is contingent on S C ( x i ) ˜ S ∆ ( x i ) S B ( x i ) x i , which is the response of the SSM to x i . S C ( x i ) ˜ S ∆ ( x i ) S B ( x i ) could be regarded as a self-attention matrix, which implicitly offers a measure of saliency for a token x i . To this end, we propose START-M, which utilizes the input-dependent matrices to identify salient tokens. Concretely, given an input sequence { x i } L i =1 , where L denotes the sequence length, we first compute the input-dependent matrices based on Eq. (2). Then, we calculate the saliency value for each token based on Eq. (6), i.e. , for the i -th token x i :

Afterward, we generate a binary mask M S ∈ B L with the element being set to 1 if the corresponding element Saliency M ( x i ) is in the top P tokens percentage elements.

Meanwhile, we synthesize the style-augmented sequence, which is achieved by mixing the mean and variance of different samples. Following [13, 62], we first compute the style statistics as: µ ( x ) = 1 L ∑ L i =1 x i , σ ( x ) = √ 1 L ∑ L i =1 ( x i -µ ( x )) 2 . Then we randomly select another sample x ′ from the current batch, utilizing its statistics to synthesize the stylized version of x :

Finally, with the token-level mask M S , we mix x and ˜ x to generate the augmented sequence x aug , where tokens with maximum saliency are style-augmented, while other tokens remain unchanged:

where ⊙ is element-wise multiplication. Note that when P token = 1 , START-M degenerates into a channel-level augmentation, i.e. , MixStyle [13]. However, note that style statistics could be one kind

of domain-specific feature, while other forms of domain-specific features may also exist, especially within the image backgrounds [59, 63]. As a result, directly perturbing the style information of tokens on the background might activate other forms of domain-related noise, which could still disrupt the model generalization [64]. To address this issue, our START-M proposes to selectively perturb tokens with the highest saliency, which are typically associated with foregrounds, thus enhancing the model learning of domain-invariant information without activating domain-related noise. Ablation study in Section 4.3 also proves the effectiveness of the saliency-driven selection scheme.

START based on input sequences (START-X). Based on Proposition 2 , recalling that ∆ , B , and C are all input-dependent matrices (as in Eq. (2)), we design a simplified variant, namely STARTX, which involves using the activation values of x to approximate the saliency of tokens directly. Specifically, for the i -th input token x i , we directly compute its saliency value as: Saliency X ( x i ) = x i . With the saliency for each token, we compute the token-level binary mask M S as that of STARTX and employ Eq. (9) to generate the augmented sequences. In practice, we randomly apply our START method to 50% of the samples in each batch, leaving the remaining samples unperturbed during each training iteration. Our START method is disabled during inference.

In summary, we theoretically investigate the generalization error boundary of Mamba at the token level, highlighting that suppressing domain-related information within input-dependent matrices can effectively reduce the generalization error boundary of the model. Based on the theoretical analysis, we propose the first saliency-driven token-aware transformation for SSMs, designing two different variants for identifying and perturbing the tokens focused on by the input-dependent matrices B,C, ∆ . In this way, our method can effectively enhance the reliance of the input-dependent matrices on domain-invariant features and narrow the distance between source and target domains. Notably, our START introduces no additional parameters or inference time, only involving a few matrix operations during training, thus achieving similar linear complexity to VMamba as presented in Appendix A.2.

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 4 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-World, exhibiting a large domain gap. VLCS [80] contains 10 , 729 images of 5 categories from 4 domains: Pascal, LabelMe, Caltech, and Sun. TerraIncognita [81] comprises photographs of wild

animals taken by 4 camera-trap 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 ResNet50 ( 22 Mvs. 23 M), as the backbone. The backbone is pretrained on the ImageNet [82] for all our experiments. We partition the input image into 4 × 4 patches without further flattening the patches into a 1 D sequence. The network depth of the VMamba-T backbone is 4 the same as ResNet50 , consisting of 2 , 2 , 9 , and 2 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], we train the model for 50 epochs using AdamW optimizer and cosine decay schedule, with a batch size of 64 , the initial learning rate as 5 e -4 , and the momentum of 0 . 9 . For all experiments, we the ratio P token of augmented tokens to 0 . 75 . We apply the leave-one-domain-out protocol for all benchmarks, where one domain is used for testing, and the remaining domains are employed for training. We select the last-epoch model and report the average accuracy over five runs. All the experiments are run on 4 NVIDIA Teska V100 GPUs.

4.2 Main Results

Evaluation on PACS. We first compare our method with SOTA CNN-based DG methods on ResNet50 . As shown in Tab. 1, the strong baseline (VMamba) achieves a promising performance, exceeding ResNet50 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 [75] by 6 . 17% ( 91 . 77% vs. 85 . 60% ). START-X can also improve the baseline significantly by 1 . 78% ( 91 . 72% vs. 89 . 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 ( 22 Mvs. 34 M). Finally, our methods beat the recent DGMamba [54], exceeding it by 0 . 57% ( 91 . 77% vs. 91 . 20% ) on average, which proves the effectiveness of our method for DG.

Evaluation on OfficeHome. We evaluate the effectiveness of our method on OfficeHome and present the results in Tab. 1. 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 ResNet50 . Based on the Strong Baseline with high performance, our method can still improve it by 0 . 66% ( 77 . 09% vs. 76 . 43% ). START-M precedes the best MLP-like model ViP-S [78], 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 VLCS. As 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% ).

Additionally, our method significantly improves upon the baseline, outperforming the latest Mambabased method DGMamba by 0 . 52% ( 81 . 32% vs. 80 . 80% ).

Evaluation on TerraIncognita. As shown in Tab. 2, We observe that the VMamba baseline significantly outperforms previous methods, achieving a SOTA performance of 56 . 16% . Based on the strong baseline, our method can further achieve substantial improvement by 2 . 11% ( 58 . 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, w/o 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 on 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 background tokens could activate other forms of domain-specific features, potentially hindering model generalization.

Parameter sensitivity. We explore the sensitivity of our method to the hyper-parameter P token , the percentage of perturbed tokens in input sequences. As shown in Fig. 3, START-M consistently performs well across different P token values, demonstrating its effectiveness in perturbing domain-specific information in salient tokens. We notice that START-X, which uses token activation to approximate saliency, performs similarly to START-M when P token is high but is less effective at lower P token values. This indicates differences between attention regions of input-dependent matrices and input sequences. Both methods achieve the highest accuracy at P token = 0 . 75 , which is adopted for all experiments.

Comparisons with other feature identification methods. We provide comparisons with the 'GradCAM' and 'Attention Matrix' methods. For the 'GradCAM' method, we first obtain feature gradients using backpropagation without updating, then compute token saliency and augment salient tokens at each iteration. For the 'Attention Matrix' method,

0

Ratio of Perturbed T okens

89.5

90.0

90.5

91.0

91.5

92.0

Accuracy (%)

Baseline

START-X

START-M

1

0.90

0.75

0.60

0.

45

0.

30

0.

15

since the Mamba architecture lacks explicit attention matrices, we instead use α in Eq. (3) to calculate token saliency. As shown in Tab. 4, on the strong baseline, START still performs much better than these advanced methods, exceeding 'GradCAM' by 0 . 91% ( 91 . 77% vs. 90 . 86% ) and 'Attention Matrix' by 1 . 00% ( 91 . 77% vs. 90 . 77% ). It is owing to the ability of START to explicitly suppress domain-specific features within input-dependent matrixes.

Comparisons with other augmentation methods. We here compare our method with SOTA DG augmentation methods on the VMamba backbone, including MixStyle [13], DSU [14], and ALOFT [15]. As shown in Tab. 5, all the augmentation methods bring performance improvements, indicating that increasing data diversity is beneficial for the generalization ability of Mamba. Notably,

START outperforms all the SOTA augmentation methods, i.e. , yielding a significant margin of 1 . 06%

( 91 . 77% vs. 91 . 72% ) from DSU. The results prove the effectiveness of our methods in perturbing domain-specific features within input-dependent matrices.

Domain gaps in input-dependent matrices. To verify the effectiveness of our method in reducing domain gaps within input-dependent matrices, we compare domain gaps across different methods using PACS with VMamba. The experiments focus on the last block of VMamba, examining the output feature maps ('Feat.') and the input-dependent

matrices ˜ ∆ , B , and C of the first SS2D. The results in Tab. 6 align well with theoretical analysis in Section 3.2, proving that START effectively reduces domain gaps in the input-dependent matrices.

5 Conclusions

In this paper, inspired by the success of Mamba in supervised tasks, we theoretically study the generalizability of Mamba and find that the input-dependent matrices in Mamba could accumulate and amplify domain-specific features during training. To address the issue, we propose a generalized state space model with a saliency-driven token-aware transformation for DG, which can selectively augment domain-specific features within salient tokens focused on by the input-dependent matrices, thus helping the model learn domain-invariant features. Our method outperforms SOTA CNN-based and ViT-based methods by a significant margin with linear complexity and a small-sized network. We hope our work inspires further research in DG and contributes valuable insights to the community.

6 Acknowledgment

This work was supported by the National Key R&D Program of China (2023ZD0120700, 2023ZD0120701), NSFC Project (62222604, 62206052), China Postdoctoral Science Foundation (2024M750424), the Fundamental Research Funds for the Central Universities (020214380120), the State Key Laboratory Fund (ZZKT2024A14), the Postdoctoral Fellowship Program of CPSF (GZC20240252), and the Jiangsu Funding Program for Excellent Postdoctoral Talent (2024ZB242).

References

Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In ICCV , 2023.

Lihe Yang, Bingyi Kang, Zilong Huang, Xiaogang Xu, Jiashi Feng, and Hengshuang Zhao. Depth anything: Unleashing the power of large-scale unlabeled data. In CVPR , 2024.

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.

Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. TKDE , 2009.

Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV , 2019.

MJehanzeb Mirza, Jakub Micorek, Horst Possegger, and Horst Bischof. The norm must go on: Dynamic unsupervised domain adaptation by normalization. In CVPR , 2022.

Yang Shu, Zhangjie Cao, Chenyu Wang, Jianmin Wang, and Mingsheng Long. Open domain generalization with domain-augmented meta-learning. In CVPR , 2021.

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

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

Fu-En Yang, Yuan-Chia Cheng, Zu-Yun Shiau, and Yu-Chiang Frank Wang. Adversarial teacher-student representation learning for domain generalization. In NeurIPS , 2021.

Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-ofdistribution generalization. In ICML , 2022.

Chaoqi Chen, Jiongcheng Li, Xiaoguang Han, Xiaoqing Liu, and Yizhou Yu. Compound domain generalization via meta-knowledge encoding. In CVPR , 2022.

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

Xiaotong Li, Yongxing Dai, Yixiao Ge, Jun Liu, Ying Shan, and Ling-Yu Duan. Uncertainty modeling for out-of-distribution generalization. ICLR , 2022.

Jintao Guo, Na Wang, Lei Qi, and Yinghuan Shi. Aloft: A lightweight mlp-like architecture with dynamic low-frequency transform for domain generalization. In CVPR , 2023.

Jiawang Bai, Li Yuan, Shu-Tao Xia, Shuicheng Yan, Zhifeng Li, and Wei Liu. Improving vision transformers by revisiting high-frequency components. In ECCV , 2022.

Zangwei Zheng, Xiangyu Yue, Kai Wang, and Yang You. Prompt vision transformer for domain generalization. arXiv preprint arXiv:2208.08914 , 2022.

Maryam Sultana, Muzammal Naseer, Muhammad Haris Khan, Salman Khan, and Fahad Shahbaz Khan. Self-distilled vision transformer for domain generalization. In ACCV , 2022.

Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In ICLR , 2023.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 , 2023.

Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590 , 2024.

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.10166 , 2024.

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.

Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In ICCV , 2017.

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.

Wei Zhu, Le Lu, Jing Xiao, Mei Han, Jiebo Luo, and Adam P Harrison. Localized adversarial domain generalization. In CVPR , 2022.

Sudao He, Fuyang Chen, and Hongtian Chen. A latent representation generalizing network for domain generalization in cross-scenario monitoring. TNNLS , 2023.

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.

Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts. In NeurIPS , 2022.

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.

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.

Jin Chen, Zhi Gao, Xinxiao Wu, and Jiebo Luo. Meta-causal learning for single domain generalization. In CVPR , 2023.

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.

Qinwei Xu, Ruipeng Zhang, Yi-Yan Wu, Ya Zhang, Ning Liu, and Yanfeng Wang. Simde: A simple domain expansion approach for single-source domain generalization. In CVPR , 2023.

Yue Wang, Lei Qi, Yinghuan Shi, and Yang Gao. Feature-based style randomization for domain generalization. TCSVT , 2022.

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.

Sachin Mehta and Mohammad Rastegari. Mobilevit: Light-weight, general-purpose, and mobile-friendly vision transformer. In ICLR , 2022.

Yubiao Yue and Zhenzhang Li. Medmamba: Vision mamba for medical image classification. arXiv preprint arXiv:2403.03849 , 2024.

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.05246 , 2024.

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.

Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.

Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In ICLR , 2022.

Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In NeurIPS , 2022.

Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In NeurIPS , 2022.

Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In ICLR , 2022.

Jun Ma, Feifei Li, and Bo Wang. U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722 , 2024.

Jiacheng Ruan and Suncheng Xiang. Vm-unet: Vision mamba unet for medical image segmentation. arXiv preprint arXiv:2402.02491 , 2024.

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:2402.05079 , 2024.

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 arXiv:2402.10739 , 2024.

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-based ordering strategy. arXiv preprint arXiv:2403.06467 , 2024.

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.00762 , 2024.

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.

Qinfeng Zhu, Yuanzhi Cai, Yuan Fang, Yihan Yang, Cheng Chen, Lei Fan, and Anh Nguyen. Samba: Semantic segmentation of remotely sensed images with state space model. arXiv preprint arXiv:2404.01705 , 2024.

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.

Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In NeurIPS , 2020.

Michael Poli, Stefano Massaroli, Eric 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.

Sinno Jialin Pan, Ivor W Tsang, James T Kwok, and Qiang Yang. Domain adaptation via transfer component analysis. TNN , 2010.

Ievgen Redko, Emilie Morvant, Amaury Habrard, Marc Sebban, and Younès Bennani. A survey on domain adaptation theory: learning bounds and theoretical guarantees. arXiv preprint arXiv:2004.11829 , 2020.

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

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

Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In ICML , 2015.

Mehrdad Noori, Milad Cheraghalikhani, Ali Bahri, Gustavo A Vargas Hakim, David Osowiechi, Ismail Ben Ayed, and Christian Desrosiers. Tfs-vit: Token-level feature stylization for domain generalization. PR , 2024.

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.

Chaoqi Chen, Luyao Tang, Feng Liu, Gangming Zhao, Yue Huang, and Yizhou Yu. Mix and reason: Reasoning over semantic topology with data mixing for domain generalization. In NeurIPS , 2022.

Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Deep domain-adversarial image generation for domain generalisation. In AAAI , 2020.

Xufeng Yao, Yang Bai, Xinyun Zhang, Yuechen Zhang, Qi Sun, Ran Chen, Ruiyu Li, and Bei Yu. Pcl: Proxy-based contrastive learning for domain generalization. In CVPR , 2022.

Devansh Arpit, Huan Wang, Yingbo Zhou, and Caiming Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In NeurIPS , 2022.

Cian Eastwood, Alexander Robey, Shashank Singh, Julius Von Kügelgen, Hamed Hassani, George J Pappas, and Bernhard Schölkopf. Probable domain generalization via quantile risk minimization. In NeurIPS , 2022.

Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization. In CVPR , 2023.

Zenan Huang, Haobo Wang, Junbo Zhao, and Nenggan Zheng. idag: Invariant dag searching for domain generalization. In ICCV , 2023.

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

Aveen Dayal, Vimal KB, Linga Reddy Cenkeramaddi, C Mohan, Abhinav Kumar, and Vineeth N Balasubramanian. Madg: Margin-based adversarial learning for domain generalization. In NeurIPS , 2023.

Zhe Wang, Jake Grigsby, and Yanjun Qi. Pgrad: Learning principal gradients for domain generalization. In ICLR , 2023.

Minyoung Kim, Da Li, and Timothy Hospedales. Domain generalisation via domain adaptation: An adversarial fourier amplitude approach. ICLR , 2023.

Zhaorui Tan, Xi Yang, and Kaizhu Huang. Rethinking multi-domain generalization with a general learning objective. In CVPR , 2024.

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 NeurIPS , 2021.

Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Gautier Izacard, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, et al. Resmlp: Feedforward networks for image classification with data-efficient training. TPAMI , 2022.

Qibin Hou, Zihang Jiang, Li Yuan, Ming-Ming Cheng, Shuicheng Yan, and Jiashi Feng. Vision permutator: A permutable mlp-like architecture for visual recognition. TPAMI , 2022.

Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In CVPR , 2017.

Antonio Torralba and Alexei A Efros. Unbiased look at dataset bias. In CVPR , 2011.

Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In ECCV , 2018.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. IJCV , 2015.

Kyungmoon Lee, Sungyeon Kim, and Suha Kwak. Cross-domain ensemble distillation for domain generalization. In ECCV , 2022.

Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data augmentation. In ICML , 2019.

Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yanfeng Wang, and Qi Tian. A fourier-based framework for domain generalization. In CVPR , 2021.

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:1909.09148 , 2019.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML , 2021.

Yongming Rao, Wenliang Zhao, Zheng Zhu, Jie Zhou, and Jiwen Lu. Gfnet: Global filter networks for visual recognition. TPAMI , 2023.

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. In ICCV , 2017.

A Appendix / supplemental material

A.1 Theoretical Proofs

Lemma 1 [58]. Let F = { f ∈ H k : || f || H k ≤ 1 } denote a function class, where H k be a RKHS with its associated kernel k . Let L h,f : x → L [ h ( x ) , f ( x )] be a convex loss-function with a parameter form | h ( x ) -f ( x ) | q for some q > 0 , and defined ∀ h, f ∈ F , L obeys the triangle inequality. Let S and T be two samples of size M drawn i.i.d from D S and D T , respectively. Then, with probability of at least 1 -δ ( δ ∈ (0 , 1) ) for all h ∈ F , the following holds:

where d MMD ( D S , D T ) = sup || f || F k ≤ 1 ∣ ∣ ∫ fd ( h ( D S ) -h ( D T )) ∣ ∣ , K D S and K D T are kernel functions computed on samples from D S and D T , respectively. σ is the combined error of the ideal hypothesis h ∗ on D S and D T .

Theorem 1 (Generalization risk bound). With the previous setting and assumptions, let D i S and D T be two sets with M samples independently drawn from D n S and D T , respectively. For any δ ∈ (0 , 1) with probablity of at least 1 -δ , for all h ∈ H , the following inequality holds:

where λ π = 1 M ( ∑ N n =1 π n E x ∼ D n S [ √ tr ( K D n S ) + E x ∼ D T [ √ tr ( K D T )]) + √ log (2 /ϵ ) 2 M , and σ is the minimum combined error of the ideal hypothesis h ∗ on both D S and D T . Let γ T = d To-MMD ( D T , ¯ D T ) and γ S = sup i,j ∈ [ N ] d To-MMD ( D i S , D j S ) , respectively.

Proof. We initially investigate the relationship between the MMD [58] and To-MMD distances based on Definition 1 (as presented in Eq. (4)). With the feature extractor ψ ( · ) , we have:

Then, for a pair of source domain D n S and D T , the following inequality holds:

with which we can derive the weighted sum of To-MMD between source domains and target domain:

With the above preparations, we now derive the generalization error bound of the Mamba model on the unseen target domain. Recalling that Lemma 1 indicates the generalization error bound between two different distributions, we generalize it to the scenario of multiple source domains:

where λ π = 1 M ( ∑ N n =1 π n E x ∼ D n S [ √ tr ( K D n S ) + E x ∼ D T [ √ tr ( K D T )]) + √ log (2 /ϵ ) 2 M and σ is the minimum combined error of the ideal hypothesis h ∗ on both D S and D T .

Proposion 1 (Accumulation of Domain Discrepancy). Given two distinct domains D S and D T , the token-level domain distance d To-MMD ( D S , D T ) depends on d C ˜ ∆ Bx (¯ x S i , ¯ x T i ) and d ˜ ∆ (¯ x S i , ¯ x T i ) for the i -th token. For the entire recurrent process, domain-specific information encoded in S ∆ , S C , and S B will accumulate, thereby amplifying domain discrepancy.

Proof. For simplification, we use ¯ x S ∈ R L and ¯ x T ∈ R L to denote the sample mean embeddings for the D S and D T , respectively. L represents the token sequence length. To investigate the generalization error boundary of Mamba, we explore a simplified problem in conjunction with a single S 6 layer, i.e. , ¯ y = α ¯ x , where α ∈ R L × L is the data-dependent matric. Empirically, based on Eq. (3) and Eq. (4), we estimate the token-level domain gap using Euclidean distance of ¯ y S and ¯ y S ,

Combined with Eq. (2) and Eq. (3), we represent α as a direct function of the input ¯ x :

Since the SSM layer calculates y i based on the continuous subsequence [¯ x 1 , ¯ x 2 , · · · , ¯ x i ] , we here analyze the domain gap of the extracted features at the token level, i.e. , | y S i -y T i | . Specifically, assuming that we have calculated | y S i -y T i | = β ( 1 ≤ i < L ), from Eq. (16), we can derive:

With the defined α in Eq. (17) and denoting softmax( S ∆ ( · )) as ˜ S ∆ ( · ) for brevity, we can express:

Considering that the differences between adjacent tokens are generally small, S C (¯ x i +1 ) /S C (¯ x i ) could be approximated to 1 . When the dimension of ¯ x is relatively large, then for Eq. (19), we have:

Then, we substitute Eq. (20) into Eq. (18) to derive the following formula:

Finally, recalling that | y S i -y T i | = β , we can express | y S i +1 -y T i +1 | as following:

Let d C ˜ ∆ B ¯ x (¯ x S i +1 , ¯ x T i +1 ) = S C (¯ x S i +1 ) ˜ S ∆ (¯ x S i +1 ) S B (¯ x S i +1 )¯ x S i +1 -S C (¯ x T i +1 ) ˜ S ∆ (¯ x T i +1 ) S B (¯ x T i +1 )¯ x T i +1 , and d ˜ ∆ (¯ x S i +1 , ¯ x T i +1 ) = ˜ S ∆ (¯ x S i +1 ) -˜ S ∆ (¯ x T i +1 ) . Then, the above equation reveals that for the input tokens ¯ x S i +1 and ¯ x T i +1 , the distance between their extracted features, alongside the gap of historical sequences, primarily depends on d C ˜ ∆ Bx (¯ x S i +1 , ¯ x T i +1 ) and d ˜ ∆ (¯ x S i +1 , ¯ x T i +1 ) . Besides, the recurrent process in Eq. (22) could also lead to the accumulation or even enhancement of domain-specific information, i.e. , if the model extracts domain-related information from the i th token, this part of the information will be retained in the features extracted by the ( i +1) -th token. For the whole recurrent process, domain-related information encoded in S ∆ , S C , and S B will be accumulated and amplified, which will increase the discrepancy in the features extracted by the model for different domains, thus damaging its generalization ability. Therefore, to reduce the domain gap, it is imperative to suppress domain-specific information learned by S ∆ , S C , and S B .

Proposion 2 (Mitigating Domain Discrepancy Accumulation). Perturbing domain-specific features in tokens focused on by S ∆ , S C , and S B can enhance their learning of domain-invariant features, thus effectively mitigating the accumulation issue in these input-dependent matrices.

Proof. Recalling that in the Mamba mode, S ∆ , S C , and S B are all linear projection layers, which map the input sequence ¯ x ∈ R L to the data-dependent matrixes ∆ , C , and B , respectively. Hence, we analyze the influence of tokens in x on these matrixes. Taking the matric B as an example, we explore the simplified problem with S B ( x ) = W B ¯ x , where W B ∈ R L × N and N denotes the dimension of the hinder state. Inspired by previous works [84, 85], we assume that the input sequence ¯ x could be decomposed to domain-specific features ¯ x I and domain-specific features ¯ x S . Then, for the i -th token ¯ x i in ¯ x , the projection matric B could be denoted as:

Previous theoretical works [84, 86] have demonstrated that when a subset of features is perturbed, its variance would be increased, and the model would be regularized to decrease the weights associated with these features to minimize the prediction loss. Therefore, by perturbing the domain-specific features x S i , its corresponding weights W S B i would be restricted to 0 . Simultaneously, due to minimal changes in the domain-invariant feature x I i , the learning of its corresponding weights W I B i is promoted, thus enhancing S B learning of domain-invariant features.

Furthermore, given the presence of foreground and background in images [54], different tokens contain varied information. Foreground tokens primarily encode domain-invariant semantic features

alongside some domain-specific details. Perturbing the domain-related features in these tokens can effectively enhance the model's learning of domain-invariant features. Conversely, background tokens encompass diverse domain-related spurious features that are challenging to fully extract and perturb. As a result, perturbing a subset of domain-related features in these tokens may inadvertently activate other forms of spurious noise, thereby hindering model generalization. To address this issue, leveraging the hidden attention mechanism of Mamba (as depicted in Eq. (3)), where tokens with high saliency are more likely to belong to the foreground, we propose to perturb domain-specific information solely in the tokens focused on by the input-dependent matrices.

A.2 Additional Experiments

Evaluation on DomainNet. We investigate the effectiveness of our method on the large-scale dataset DomainNet. As shown in Tab. 7, on the challenging benchmark, we find that the strong baseline VMamba can achieve the promising performance of 51 . 45% , proving the superiority of Mamba models on large-scale datasets. Based on the strong baseline, our START can still significantly improve the performance, exceeding the baseline by 1 . 35% ( 52 . 81% vs 51 . 45% ). Besides, compared with CNN-based methods, our method can still achieve significant improvements, e.g. , outperforming the latest SOTA method AGFA [74] by 5 . 71% ( 52 . 81% vs 47 . 10% ). Our methods also beat the best ViT-based method DoPrompt [17], yielding it by a large margin of 4 . 51% ( 52 . 81% vs 48 . 30% ). The results prove the effectiveness of our method to help the model learn domain-invariant representation.

Ablation studies on larger datasets. We provide the ablation studies on the Officehome and TerraIncognita datasets. As shown in Tab. 8, our methods perform the best among all variants, e.g. , on TerraIncognita, our START-X outperforms the variant ' w.o. Saliency Guided' by 0 . 97% ( 58 . 27% vs. 57 . 30% ) and the variant ' w.o. Token Selection' by 0 . 83% ( 58 . 27% vs. 57 . 44% ). The results demonstrate the effectiveness of all modules in our START methods.

Effects on other Mamba-based architectures. To validate the effectiveness of our START on other Mamba architectures, we conducted experiments on the recent ViM [23], using the ViM-T and ViM-S models with different network sizes. The experiments were performed on the PACS dataset with an initial learning rate of 6 . 25 e -6 and a weight decay of 1 e -8 . As shown in Tab. 9, our method consistently improves performance across the models with different scales, e.g. , on ViM-S, START-X outperformed the baseline by 2 . 03% ( 89 . 41% vs. 87 . 38% ), and on ViM-T, START-X exceeded the baseline by 1 . 66% ( 86 . 74% vs. 85 . 08% ). These results demonstrate the generalizability of our method across different Mamba architectures.

Effects on the ViT architecture. Recalling that our method is derived from the theoretical analysis that input-dependent matrices in Mamba could accumulate domain-related information during training,

our START aims to improve the generalization of the Mamba architecture. Nevertheless, the core concept, adaptively perturbing salient tokens in input-dependent matrices, is also applicable to ViTs. Considering that in ViTs, the attention matrix uses query Q and key K , and then multiplied by the original feature V to obtain the final representation. We develop our START-M to START-ViT-M, which calculates token saliency from the input-dependent matrices ( i.e. , Q × K T ), and START-X to START-ViT-X, which uses the activation value of representation x to approximate saliency. The experiments are conducted on the representative ViT architecture, i.e. , DeiT-Small, with the PACS dataset. As shown in Tab.10, 1) on the DeiT-Small baseline, our START re-designed for ViTs still can effectively improve the Baseline by a significant margin, e.g. , START-ViT-M outperforms the baseline by 1 . 20% ( 87 . 05% vs. 85 . 85% ). The results prove the effectiveness of our START's variants on ViTs; 2) we notice that the VMamba-T ( 22 M parameters) is a stronger baseline model than the DeiT-Small ( 22 M parameters), exceeding it by a large margin of 4 . 09% ( 89 . 94% vs. 85 . 85% ). The results also reveal the advantage of Mamba architecture to learn domain-invariant token dependencies in compressed state space, and our START can further enhance the generalization ability of Mamba.

Evaluation on single-source domain generalization tasks. We here evaluate our method under the single-source-domain generalization setting. As shown in Tab. 11, START-M significantly improves the baseline, outperforming it by 2 . 12% ( 71 . 64% vs. 69 . 52% ). These results prove that our method enhances model generalization by simulating domain shifts through salience-driven token transformation, improving performance in both multi-source and single-source DG tasks.

Effectiveness with other SOTA augmentation methods. In our method, we utilize a statistics-based style augmentation method to perturb domain-specific information within tokens. We also explore other SOTA DG augmentation methods for comparison, including the DSU [14] that models the distribution of statistics across different samples and resample new statistics from the distribution, and the ALOFT [15] that diversifies the low-frequency spectrum in the frequency domain. These methods primarily perturb style information at the channel level. As shown in Tab. 12, our method significantly improves the performance of the SOTA augmentation methods on the VMamba baseline, e.g. , our START-DSU-M achieves a significant improvement over DSU by 0 . 8% ( 91 . 63% vs. 90 . 71% ), exceeding the baseline by 1 . 69% ( 91 . 63% vs. 89 . 94% ). The above results prove that selective perturbation of domain-specific features in salient tokens is crucial for enhancing the generalization capability of Mamba models.

Effects across different stages. Our theoretical analysis examined how domain gaps accumulate within each SSM layer. Since one layer's output serves as the next layer's input, domain-specific features from earlier stages increase domain gaps in later stages. To address the issue, we applied START to all layers to reduce domain gaps comprehensively. We also tested START separately in either shallow or deep layers. As shown in Tab. 13, using START in both shallow and deep layers simultaneously performs best, aligning with our theoretical analysis. Applying START-M or STARTX randomly across layers also improves performance, though less effectively than using START-M or START-X alone. This may be because START-M and START-X target different domain-related information, leading to incomplete suppression when mixed.

Computational efficiency. To evaluate the computational efficiency of our proposed START, we conduct experiments on the PACS dataset and compare our method with existing CNN-based and ViT-based methods. Specifically, we compare the number of parameters, floating point operations per second (FLOPs), the inference times, and the generalization performance of each method. The batch size for evaluating inference time is set to 64 , and the inference time is averaged over 100 experiments. Since STARR-M and START-X are only activated during training and disabled during inference, they introduce no additional inference time. As shown in Tab. 14, our method has significantly fewer FLOPs than ResNet50 (5.68 vs. 8.26) while outperforming the DeepAll on ResNet50 (the baseline that directly trains the model on source domains) by 6 . 22% ( 91 . 77% vs. 85 . 50% ), demonstrating the superiority of our START.

Visualization explanations. To provide visual evidence of the effectiveness of our START in suppressing domain-specific features, we use GradCAM [89] to generate attention maps of the last state space layer for both the baseline (pure VMamba) and our START models. As illustrated in Fig. 4,

START-M

(Ours)

START-X

(Ours)

VMamba

(Baseline)

Origin

Art (Target)

Cartoon

START-M

(Ours)

START-X

(Ours)

VMamba

(Baseline)

Origin

Photo

Sketch

PITO,

the VMamba baseline tends to focus on specific local patches that encode domain-specific features, leading to overfitting to source domains. In contrast, our START methods effectively reduce the model's focus on domain-specific features, enabling it to capture generalizable global dependencies of tokens. For instance, in the case of the person image in the Art domain, the baseline focuses on multiple local regions in both the foreground and the background, making the model sensitive to domain shifts and likely to misclassify samples. Conversely, our START models mainly focus on the foreground, specifically the whole face of the person. The results prove the effectiveness of START in learning comprehensive domain-invariant features, making it a promising method for DG tasks.

Difference from the related work. In essence, our method significantly differs from DGMamba [54] in their motivations, goals and methods. 1) Different Motivations : DGMamba observes that hidden states could amplify domain-related information, and proposes a heuristic method to address the issue. However, it lacks a deep analysis of the phenomenon. Differently, we first theoretically delve into the generalizability of Mamba, revealing how input-dependent matrices contribute to domain gap accumulation. Based on the analysis, we developed START to enhance Mamba's generalization. 2) Different Goals and Methods : DGMamba aims to enforce the model to focus on object tokens, perturbing object tokens while replacing context tokens. It ignores that object tokens could be misclassified as context tokens, replacing which would hinder the model from learning semantics. Inversely, START aims to suppress domain-specific information in tokens focused on input-dependent

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 to 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 it will have a positive societal impact.

A.4 Limitations of Our Work

In the paper, we employ style perturbation to perturb domain-specific information within salient tokens of input sequences. However, there are 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 is to 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

1. 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.

It is fine to include aspirational goals as motivation as long as it 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 a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

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 proofs in Section 3.2 and Appendix A.1.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: See in Section 4.1.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The source code is provided at https://github.com/lingeringlight/START.

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/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See in Section 4.1.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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, or 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 a 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 are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See in Section 4.1.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

Justification: The paper does not release the data or models with a high risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The paper uses five public datasets, including PACS [24], OfficeHome [79], VLCS [80], TerraIncognita [81], and DomainNet [5]. See the licenses in the original pepers.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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 not 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1.\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.\n\nWEAKNESSES:\n1.\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?\n\nQUESTIONS:\nBesides 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?\n\nLIMITATIONS:\nNo Limitation section is mentioned in the main paper.\n\nI am not sure about the generalization ability of this analysis to different architectures.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

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

STRENGTHS:
- This paper is early work to adapt Mamba for Domain Generalization. 

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

- ViT has attention maps, which are also input-dependent. Can your methods transfer to ViT? Or your method can only be applied to Mamba?

QUESTIONS:
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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-15

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAdvancements 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nIt would be better to report computation costs like inference time.\n\nQUESTIONS:\nPlease refer to the weakness.\n\nLIMITATIONS:\nThis paper discusses the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nplease see weaknesses\n\nLIMITATIONS:\nThe authors have sufficiently addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors have addressed the limitations and, if applicable, the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}