S PARSE M IXTURE - OF -E XPERTS ARE D OMAIN G ENERALIZABLE L EARNERS
Anonymous authors
Paper under double-blind review
A BSTRACT
Human visual perception can easily generalize to out-of-distributed visual data, which is far beyond the capability of modern machine learning models.
Domain generalization (DG) aims to close this gap, with existing DG methods mainly focusing on the loss function design.
In this paper, we propose to explore an orthogonal direction, i.e., the design of the backbone architecture.
It is motivated by an empirical ﬁnding that transformer-based models trained with empirical risk minimization (ERM) outperform CNN-based models employing state-ofthe-art (SOTA) DG algorithms on multiple DG datasets.
We develop a formal framework to characterize a network’s robustness to distribution shifts by studying its architecture’s alignment with the correlations in the dataset.
This analysis guides us to propose a novel DG model built upon vision transformers, namely
Generalizable Mixture-of-Experts (GMoE) .
Extensive experiments on DomainBed demonstrate that GMoE trained with ERM outperforms SOTA DG baselines by a large margin.
Moreover, GMoE is complementary to existing DG methods and its performance is substantially improved when trained with DG algorithms.
I NTRODUCTION 1. M OTIVATIONS
Generalizing to out-of-distribution (OOD) data is an innate ability for human vision, but highly challenging for machine learning models (Recht et al., 2019;
Geirhos et al., 2021; Ma et al., 2022).
Domain generalization (DG) is one approach to address this problem, which encourages models to be resilient under various distribution shifts such as background, lighting, texture, shape, and geographic/demographic attributes.
From the perspective of representation learning, there are several paradigms towards this goal, including domain alignment (Ganin et al., 2016; Hoffman et al., 2018), invariant causality prediction (Arjovsky et al., 2019; Krueger et al., 2021), meta-learning
(Bui et al., 2021; Zhang et al., 2021c), ensemble learning (Mancini et al., 2018; Cha et al., 2021b), and feature disentanglement (Wang et al., 2021; Zhang et al., 2021b).
The most popular approach to implementing these ideas is to design a speciﬁc loss function.
For example, DANN (Ganin et al., 2016) aligns domain distributions by adversarial losses.
Invariant causal prediction can be enforced by a penalty of gradient norm (Arjovsky et al., 2019) or variance of training risks (Krueger et al., 2021).
Meta-learning and domain-speciﬁc loss functions (Bui et al., 2021; Zhang et al., 2021c) have also been employed to enhance the performance.
Recent studies have shown that these approaches improve ERM and achieve promising results on large-scale DG datasets (Wiles et al., 2021).
Meanwhile, in various computer vision tasks, the innovations in backbone architectures play a pivotal role in performance boost and have attracted much attention (He et al., 2016;
Hu et al., 2018; Liu et al., 2021).
Inspired by these empirical successes, we conjecture that backbone architecture design would be promising for DG .
To verify this intuition, we evaluate a transformer-based model and compare it with CNN-based architectures of equivalent computational overhead, as shown in Fig. 1(a).
To our surprise, a vanilla ViT-S/16 (Dosovitskiy et al., 2021) trained with empirical risk minimization (ERM) outperforms ResNet-50 trained with SOTA DG algorithms (Cha et al., 2021b; Rame et al., 2021; Shi et al., 2021) on DomainNet, OfﬁceHome and VLCS datasets, despite the fact that both architectures have a similar number of parameters and enjoy close performance on in-distribution domains.
We theoretically validate this effect based on the algorithmic alignment framework (Xu et al., 2020a; Li et al., 2021).
We ﬁrst prove that a network trained with the ERM loss function is more robust to distribution shifts if its architecture is more similar to the invariant correlation, where the similarity is formally measured by the alignment value deﬁned in Xu et al.
(2020a).
On the contrary, a network is less robust if its architecture aligns with the spurious correlation.
We then investigate the alignment between backbone architectures (i.e., convolutions and attentions) and the correlations in these datasets, which explains the superior performance of ViT-based methods.
To further improve the performance, our analysis indicates that we should exploit properties of invariant correlations in vision tasks and design network architectures to align with these properties.
This requires an investigation that sits at the intersection of domain generalization and classic computer vision.
In domain generalization, it is widely believed that the data are composed of some sets of attributes and distribution shifts of data are distribution shifts of these attributes (Wiles et al., 2021).
The latent factorization model of these attributes is almost identical to the generative model of visual attributes in classic computer vision (Ferrari & Zisserman, 2007).
To capture these diverse attributes, we propose a Generalizable Mixture-of-Experts (GMoE), which is built upon sparse mixture-of-experts (sparse MoEs) (Shazeer et al., 2017) and vision transformer (Dosovitskiy et al., 2021).
The sparse MoEs were originally proposed as key enablers for extremely large, but efﬁcient models (Fedus et al., 2022).
By theoretical and empirical evidence, we demonstrate that MoEs are experts for processing visual attributes, leading to a better alignment with invariant correlations.
Based on our analysis, we modify the architecture of sparse MoEs to enhance their performance in DG.
Extensive experiments demonstrate that GMoE achieves superior domain generalization performance both with and without DG algorithms.
1. C ONTRIBUTIONS
In this paper, we formally investigate the impact of the backbone architecture on DG and propose to develop effective DG methods by backbone architecture design.
Speciﬁcally, our main contributions are summarized as follows:
A Novel View of DG: In contrast to previous works, this paper initiates an exploration of the backbone architecture in DG.
Based on algorithmic alignment (Xu et al., 2020a), we prove that a network is more robust to distribution shifts if its architecture aligns with the invariant correlation, whereas less robust if its architecture aligns with spurious correlation.
The theorems are veriﬁed on synthetic and real datasets.
A Novel Model for DG: Based on our theoretical analysis, we propose Generalizable Mixture-ofExperts (GMoE) and prove that it enjoys a better alignment than vision transformers.
GMoE is built upon sparse mixture-of-experts (Shazeer et al., 2017) and vision transformer (Dosovitskiy et al., 2021), with a theory-guided performance enhancement for DG.
Excellent Performance: We validate GMoE’s performance on all 8 large-scale datasets of DomainBed.
Remarkably, GMoE trained with ERM achieves SOTA performance on 7 datasets in the train-validation setting and on 8 datasets in the leave-one-domain-out setting.
Furthermore, the GMoE trained with DG algorithms achieves better performance than GMoE trained with ERM.
P RELIMINARIES 2. N OTATIONS
Throughout this paper, a , a , A stand for a scalar, a column vector, a matrix, respectively.
O ( · ) and ω ( · ) are asymptotic notations.
We denote the training dataset, training distribution, test dataset, and test distribution as E tr , D tr , E te , and D te , respectively.
2. A TTRIBUTE F ACTORIZATION
The attribute factorization (Wiles et al., 2021) is a realistic generative model under distribution shifts.
Consider a joint distribution of the input x and corresponding attributes a 1 , · · · , a K (denoted as a 1: K ) with a i ∈
A i , where A i is a ﬁnite set.
The label can depend on one or multiple attributes.
Denote the latent factor as z , the data generation process is given by
[Equation]
The distribution shift arises if different marginal distributions of the attributes are given but they share the same conditional generative process.
Speciﬁcally, we have p train ( a 1: K ) (cid:54) =
p test ( a 1: K ) , but the generative model in equation 1 is shared across the distributions, i.e., we have p test ( a 1: K , x ) =
p test ( a 1: K ) (cid:82) p ( x | z ) p ( z | a 1: K )
d z and similarly for p train .
The above description is abstract and we will illustrate with an example.
Example 1.
(DSPRITES (Matthey et al., 2017)) Consider A 1 = { red, blue } and A 2 = { ellipse, square } .
The target task is a shape classiﬁcation task, where the label depends on attribute a 2 .
In the training dataset, 90% ellipses are red and 50% squares are blue, while in the test dataset all the attributes are distributed uniformly.
As the majority of ellipses are red, the classiﬁer will use color as a shortcut in the training dataset, which is so-called geometric skews (Nagarajan et al., 2020).
However, this shortcut does not exist in the test dataset and the network fails to generalize.
In classic computer vision, the attributes are named visual attributes and they follow a similar data generation process (Ferrari & Zisserman, 2007).
We shall discuss them in detail in Section 4.2.
2. A LGORITHMIC A LIGNMENT
We ﬁrst introduce algorithmic alignment, which characterizes the easiness of IID reasoning tasks by measuring the similarity between the backbone architecture and target function.
The alignment is formally deﬁned as the following.
Deﬁnition 1.
(Alignment; (Xu et al., 2020a))
Let N denote a neural network with n modules {N i } ni =1 and assume that a target function for learning y = g ( x ) can be decomposed into n functions f 1 , · · · , f n .
The network N aligns with the target function if replacing N i with f
i , it outputs the same value as algorithm g .
The alignment value between N and f is deﬁned as
[Equation]
where M ( f i , N i , (cid:15), δ ) denotes the sample complexity measure for N i to learn f i with (cid:15) precision at failure probability δ under a learning algorithm when the training distribution is the same as the test distribution.
In Deﬁnition 1, the original task is to learn f , which is a challenging problem.
Intuitively, if we could ﬁnd a backbone architecture that is suitable for this task, it helps to break the original task into simpler sub-tasks, i.e., to learn f 1 , · · · , f n instead.
Under the assumptions of algorithmic alignment (Xu et al., 2020a), f can be learned optimally if the sub-task f 1 , · · · , f n can be learned optimally.
Thus, a good alignment makes the target task easier to learn, and thus improves IID generalization, which is given in Theorem 3 in Appendix B.1.
In Section 3, we extend this framework to the DG setting.
O N THE I MPORTANCE OF N EURAL A RCHITECTURE FOR D OMAIN G ENERALIZATION
In this section, we investigate the impact of the backbone architecture on DG, from a motivating example to a formal framework.
3. A M OTIVATING E XAMPLE : CNN S VERSUS V ISION T RANSFORMERS
We adopt DomainBed (Gulrajani & Lopez-Paz, 2021) as the benchmark, which implements SOTA DG algorithms with ResNet50 as the backbone.
We test the performance of ViT trained with ERM on this benchmark, without applying any DG method.
The results are shown in Fig. 1(a).
To our surprise,
ViT trained with ERM already outperforms CNNs with SOTA DG algorithms on several datasets, which indicates that the selection of the backbone architecture is potentially more important than the loss function in DG.
In the remaining of this article, we will obtain a theoretical understanding of this phenomenon and improve ViT for DG by modifying its architecture. [Figure] t
[Figure]
3. U NDERSTANDING FROM A T HEORETICAL P ERSPECTIVE
The above experiment leads to an intriguing question: how does the backbone architecture impact the network’s performance in DG?
In this subsection, we endeavor to answer this question by extending the algorithmic alignment framework (Xu et al., 2020a) to the DG setting.
To have a tractable analysis for nonlinear function approximation, we ﬁrst make an assumption on the distribution shift.
Assumption 1.
Denote N 1 as the ﬁrst module of the network (including one or multiple layers) of the network.
Let p train , N 1 ( s ) and p test , N 1 ( s ) denote the probability density functions of features after N 1 .
Assume that the support of the training feature distribution covers that of the test feature distribution, i.e., max s p test , N 1 ( s ) p train , N 1 ( s ) ≤ C , where C is a constant independent of the number of training samples.
Remark 1.
(Interpretations of Assumption 1)
This condition is practical in DG, especially when we have a pretrained model for disentanglement (e.g., on DomainBed (Gulrajani & Lopez-Paz, 2021)).
In Example 1, the training distribution and test distribution have the same support.
In DomainNet, although the elephants in quickdraw are visually different from the elephants in other domains, the quickdraw picture’s attributes/features (e.g., big ears and long noise) are covered in the training domains.
From a technical perspective, it is impossible for networks trained with gradient descent to approximate a wide range of nonlinear functions in the out-of-support regime (Xu et al., 2020b).
Thus, this condition is necessary if we do not impose strong constraints on the target functions.
We deﬁne several key concepts in DG.
The target function is an invariant correlation across the training and test datasets.
For simplicity, we assume that the labels are noise-free.
Assumption 2.
(Invariant correlation) Assume there exists a function g c such that for training data, we have g c ( N 1 ( x ))
=
y, ∀ x ∈ E tr , and for test data, we have P D te
[ (cid:107) g c
( N 1 ( x )) − y (cid:107) ≤ (cid:15) ] > 1 − δ .
We then introduce the spurious correlation (Wiles et al., 2021), i.e., some attributes are correlated with labels in the training dataset but not in the test dataset.
The spurious correlation exists only if the training distribution differs from the test distribution and this distinguishes DG from classic
PAC-learning settings (Xu et al., 2020a).
Assumption 3.
(Spurious correlation) Assume there exists a function g s such that for training data g s ( N 1 ( x ))
=
y, ∀ x ∈ E tr , and for test data, we have P D te
[ (cid:107) g s ( N 1 ( x ))
− y (cid:107) > ω
( (cid:15) )] > 1 − δ .
The next theorem extends algorithmic alignment from IID generalization (Theorem 3) to DG.
Theorem 1.
(Impact of Backbone Architecture in Domain Generalization)
Denote N (cid:48)
=
{N 2 , · · · , N n } .
Assuming we train the neural network with ERM, and Assumption 1, 2, 3 hold, we have the following statements:
If Alignment ( N (cid:48) , g c , (cid:15), δ ) ≤ |E tr | , we have P D te [ (cid:107)N ( x ) − y (cid:107) ≤ O ( (cid:15) )] > 1 − O ( δ ) ; 2.
If Alignment ( N (cid:48) , g s , (cid:15), δ ) ≤ |E tr | , we have P D te [ (cid:107)N ( x ) − y (cid:107) > ω ( (cid:15) )] > 1 − O ( δ ) .
Remark 2.
(Interpretations of Theorem 1)
By choosing a sufﬁciently small (cid:15) , only one of
Alignment ( N (cid:48) , g c , (cid:15), δ ) ≤ |E tr | and Alignment ( N (cid:48) , g s , (cid:15), δ ) ≤ |E tr | holds.
Thus, Theorem 1 shows that the networks aligned with invariant correlations are more robust to distribution shifts.
In Appendix B.2, we build a synthetic dataset that satisﬁes all the assumptions.
The experimental results exactly match Theorem 1.
In practical datasets, the labels may have colored noise, which depends on the spurious correlation.
Under such circumstances, the network should rely on multiple correlations to ﬁt the label well and the correlation that best aligns with the network will have the major impact on its performance.
Please refer to Appendix B.1 for the proof.
ViT with ERM versus
CNNs with DG algorithms We now use Theorem 1 to explain the experiments in the last subsection.
The ﬁrst condition of Theorem 1 shows that if the neural architecture aligns with the invariant correlation, ERM is sufﬁcient to achieve a good performance.
In some domains of OfﬁceHome or DomainNet, the shape attribute has an invariant correlation with the label, illustrated in Fig. 1(b).
On the contrary, a spurious correlation exists between the attribute texture and the label.
According to the analysis in Park & Kim (2022), multi-head attentions (MHA) are low-pass ﬁlters with a shape bias while convolutions are high-pass ﬁlters with a texture bias.
As a result, a ViT simply trained with ERM can outperform CNNs trained with SOTA DG algorithms.
To improve ViT’s performance, Theorem 1 suggests that we should exploit the properties of invariant correlations.
In image recognition, objects are described by functional parts (e.g., visual attributes), with words associated with them (Zhou et al., 2014).
The conﬁguration of the objects has a large degree of freedom, resulting in different shapes among one category.
Therefore, functional parts are more fundamental than shape in image recognition and we will develop backbone architectures to capture them in the next section.
G ENERALIZABLE M IXTURE - OF -E XPERTS FOR D OMAIN G ENERALIZATION
In this section, we propose Generalizable Mixture-of-Experts (GMoE) for domain generalization, supported by effective neural architecture design and theoretical analysis.
4. M IXTURE - OF -E XPERTS L AYER
In this subsection, we introduce the mixture-of-experts (MoE) layer, which is an essential component of GMoE.
One ViT layer is composed of an MHA and an FFN.
In the MoE layer, the FFN is replaced by mixture-of-experts and each expert is implemented by an FFN (Shazeer et al., 2017).
Denoting the output of the MHA as x , the output of the MoE layer with N experts is given by
[Equation]
where W is the learnable parameter for the gate, W 1 FFN i and W 2 FFN i are learnable parameters for the i -th expert, φ ( · ) is a nonlinear activation function, and TOP k ( · ) operation is a one-hot embedding that sets all other elements in the output vector as zero except for the elements with the largest k values where k is a hyperparameter.
Given x in as the input of the MoE layer, the update is given by
[Equation]
where f MHA is the MHA layer, LN represents layer normalization, and x out is the output of the MoE layer.
4. V ISUAL A TTRIBUTES , C ONDITIONAL S TATEMENTS , AND S PARSE M O E S
In real world image data, the label depends on multiple attributes.
Capturing diverse visual attributes is especially important for DG.
For example, the deﬁnition of an elephant in the Oxford dictionary is “a very large animal with thick grey skin, large ears, two curved outer teeth called tusks, and a long nose called a trunk”.
The deﬁnition involves three shape attributes (i.e., large ears, curved outer teeth, and a long nose) and one texture attribute (i.e., thick grey skin).
In the IID ImageNet task, using the most discriminative attribute, i.e., the thick grey skin (Geirhos et al., 2018), is sufﬁcient to achieve high accuracy.
However, in DomainNet, elephants no longer have grey skins while the long nose and big ears are preserved and the network relying on grey skins will fail to generalize.
Deﬁne intervals
I i
⊂ R , i = 1 , · · · , M
Deﬁne functions h i , , i = 1 , · · · , M +switch h 1 ( x ) do case I i do apply h i +1 to x
To efﬁciently capture the visual attributes and combine them for
DG, conditional statements (e.g., IF/ELSE in programming) are needed.
First, to obtain these diverse attributes, different ﬁlters should be applied to different regions of the image.
For example, the operation for curved outer teeth is that if the patches belong to the teeth, we apply a ﬁlter to obtain its shape.
Second, for image recognition, the network should leverage conditional statements to integrate multiple visual attributes for learning the deﬁnition of this class.
For example, the recognition of an elephant should check if the key attributes are placed in the proper position.
The general conditional statements are shown in Algorithm 1.
In literature, the MoE layer is considered as an effective approach to implement conditional computations (Shazeer et al., 2017; Riquelme et al., 2021).
We formalize this intuition in the next theorem.
Theorem 2.
An MoE module in equation 3 with N experts and k = 1 aligns with the conditional statements in Algorithm 1 with
Alignment =
[Equation]
if N < M,
[Equation]
where M ( · , · , · , · ) is deﬁned in Deﬁnition 1, and M ∗ P is the optimal objective value of the following optimization problem:
[Equation]
where 1 I j is the indicator function on interval I j .
Remark 3.
(Interpretations of Theorem 2)
In algorithmic alignment, the network better aligns with the algorithm if the alignment value in equation 2 is lower.
The alignment value between MoE and conditional statements depends on the product of N +
and a sample complexity term.
When we increase the number of experts N , the alignment value ﬁrst decreases as multiple experts decompose the original conditional statements into several simpler tasks.
As we further increase N , the alignment value increases because of the factor N + 1 in the product.
Therefore, the MoE aligns better with conditional statements than with the original FFN (i.e., N = 1 ).
In addition, to minimize equation 5, similar conditional statements should be grouped together.
By experiments in Section 5.4 and
Appendix E.1, we ﬁnd that sparse MoE layers are indeed experts for visual attributes, and similar visual attributes are handled by one expert.
Please refer to Appendix B.3 for the proof.
4. A DAPTING M O E TO D OMAIN G ENERALIZATION
In literature, there are several variants of MoE architectures, e.g., Riquelme et al.
; Fedus et al.
and we should identify one for DG.
By algorithmic alignment, in order to achieve a better generalization, the architecture of sparse MoEs should be designed to effectively handle visual attributes.
In the following, we discuss our architecture design for this purpose.
Routing scheme Linear routers
(i .e. , equation 3) are often adopted in MoEs for vision tasks (Riquelme et al., 2021) while recent studies in NLP show that the cosine router achieves better performance in cross-lingual language tasks (Chi et al., 2022).
For the cosine router, given input x ∈ R d , the embedding W x ∈ R d e is ﬁrst projected onto a hypersphere, followed by multiplying a learned embedding E ∈ R d e × N .
Speciﬁcally, the expression for the gate is given by
[Equation]
[Figure]
where τ is a hyper-parameter.
In the view of image processing, E can be interpreted as the cookbook for visual attributes (Ferrari & Zisserman, 2007;
Zhou et al., 2014) and the dot product between E and
W x with (cid:96) 2 normalization is a matched ﬁlter.
We opine that the linear router would face difﬁculty in DG.
For example, the elephant image (and its all patches) in the Clipart domain is likely more similar to other images in the Clipart domain than in other domains.
The issue can be alleviated with a codebook for visual attributes and matched ﬁlters for detecting them.
Please refer to Appendix D. for the ablation study.
Number of MoE layers Every-two and last-two are two commonly adopted placement methods in existing MoE studies (Riquelme et al., 2021;
Lepikhin et al., 2021).
Speciﬁcally, every-two refers to replacing the even layer’s FFN with MoE, and last-two refers to placing MoE at the last two even layers.
For IID generalization, every-two often outperforms last-two (Riquelme et al., 2021).
We argue that last-two is more suitable for DG as the conditional sentences for processing visual attributes are high-level.
From experiments, we empirically ﬁnd that last-two achieves better performance than every-two with fewer computations.
Please refer to Appendix C.1 for more discussions and
Appendix D.5 for the ablation study.
The overall backbone architecture of GMoE is shown in Fig.
To train diverse experts, we adopt the perturbation trick and load balance loss as in Riquelme et al.
Due to space limitation, we leave them in Appendix C.4.
E XPERIMENTAL R ESULTS
In this section, we evaluate the performance of GMoE on large-scale DG datasets and present model analysis to understand GMoE.
5. D OMAIN B ED R ESULTS
In this subsection, we evaluate GMoE on DomainBed (Gulrajani & Lopez-Paz, 2021) with 8 benchmark datasets: PACS, VLCS, OfﬁceHome, TerraIncognita, DomainNet, SVIRO, Wilds-Camelyon and
Wilds-FMOW.
Detailed information on datasets and evaluation protocols are provided in Appendix
D.1.
We present results in Table 1 with train-validation selection , which include baseline methods and recent SOTA DG algorithms and GMoE trained with ERM.
The results demonstrate that GMoE without DG algorithms already outperforms counterparts on almost all the datasets.
Meanwhile,
GMoE has excellent performance in leave-one-domain-out criterion , and we leave the results in
Appendix D.3 due to space limit.
In the lower part of Table 1, we test our methods on three largescale datasets: SVIRO, Wilds-Camelyon, and Wilds-FMOW.
The three datasets capture real-world distribution shifts across a diverse range of domains.
We adopt the data preprocessing and domain
[Table]
split in DomainBed.
As there is no previous study conducting experiments on these datasets with
DomainBed criterion
[Table]
we only report the results of our methods
[Table]
which reveal that GMoE outperforms the other two baselines.
5. GM O E WITH DG A LGORITHMS
GMoE’s generalization ability comes from its internal backbone architecture, which is orthogonal to existing DG algorithms.
This implies that the SOTA DG algorithms can be applied to improve the
GMoE’s performance.
To validate this idea, we apply two SOTA DG algorithms to GMoE, including one modifying loss functions approaches (Fish) and one adopting model ensemble (Swad).
The results in Table 2 demonstrate that adopting GMoE instead of ResNet-50 brings signiﬁcant accuracy promotion to these DG algorithms.
5. S INGLE -S OURCE D OMAIN G ENERALIZATION R ESULTS
[Table]
In this subsection, we create a challenging task, singlesource domain generalization , to focus on generalization ability of backbone architecture.
Speciﬁcally, we train the model only on data from one domain, and then test the model on multiple domains to validate its performance across all domains.
This is a challenging task as we cannot rely on multiple domains to identify invariant correlations, and popular DG algorithms cannot be applied.
We compare several models mentioned in above analysis ( e.g., ResNet, ViT, GMoE) with different scale of parameters, including their ﬂoat-point-operations per second (ﬂops), IID and OOD accuracy.
From the results in Table 3, we see that GMoE’s OOD generalization gain over ResNet or ViT is much larger than that in the IID setting, which shows that GMoE is suitable for challenging domain generalization.
Due to space limitation, we leave experiments with other training domains in Appendix D.4.
5. M ODEL A NALYSIS
In this subsection, we present diagnostic datasets to study the connections between MoE layer and the visual attributes.
Diagnostic datasets: CUB-DG We create CUB-DG from the original Caltech-UCSD Birds (CUB) dataset (Wah et al., 2011).
We stylize the original images into another three domains, Candy,
Mosaic and Udnie.
The examples of CUB-DG are presented in Fig. 3(a).
We evaluate GMoE and other DG algorithms that address domain-invariant representation (e.g., DANN).
The results
[Table]
are
Appendix image.
Those additional information enable us to measure the correlation between visual attributes and the router’s expert selection
[Table]
Visual Attributes & Experts Correlation We choose GMoE-S/16 with 6 experts in each MoE layer.
After training the model on CUB-DG, we perform forward passing with training images and save the routers top-1 selection.
Since the MoE model routes patches instead of images and CUB-DG has provided the location of visual attributes, we ﬁrst match a visual attribute with its nearestpatches, and then correlate visual attribute and its current 9 patches’ experts.
In Fig. 3(b), we show the 2D histogram correlation between the selected experts and attributes.
Expert Selection To further understand the expert selections of the entire image, we record the router’s selection for each patch (in GMoE-S/16, we process an image into 16 × 16 patches), and then visualize the top-1 selections for each patch in Fig. 3(c).
In this image, we draw different visual attributes of the bird with large circles in different colors.
We see consistent relationship with Fig.
3(b).
For example, we see (1) experts 0 and 2 are consistently routed with patches that are in the background area; (2) expert 3 mainly focuses on the back and tail areas; (3) expert 4 mainly watches the breast and belly areas; (4) expert 5 mainly focuses on head and crown .
More examples are given in Appendix E.1.
[Figure]
C ONCLUSIONS
This paper is an initial step in exploring the impact of the backbone architecture in domain generalization.
We proved that a network is more robust to distribution shifts if its architecture aligns well with the invariant correlation, which is veriﬁed on synthetic and real datasets.
Based on our theoretical analysis, we proposed GMoE and demonstrated its superior performance on DomainBed.
As for future directions, it is interesting to develop novel backbone architectures for DG based on algorithmic alignment and classic computer vision.
C ONTENTS 1 Introduction
Motivations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Contributions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
2 Preliminaries 2.
Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Attribute Factorization . .
Algorithmic Alignment . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Generalization 3.
A Motivating Example: CNNs versus Vision Transformers . . . . . . . . . . . . .
Understanding from a Theoretical Perspective . . . . . . . . . . . . . . . . . . . .
4 Generalizable Mixture-of-Experts for Domain Generalization 4.
Mixture-of-Experts Layer . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Visual Attributes, Conditional Statements, and Sparse MoEs . . . . . . . . . . . .
Adapting MoE to Domain Generalization . . . . . . . . . . . . . . . . . . . . . .
5 Experimental Results 5.
DomainBed Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
GMoE with DG Algorithms . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Single-Source Domain Generalization Results . . . . . . . . . . . . . . . . . . . .
Model Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
6 Conclusions
A Related Works
A.1 Domain Generalization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Vision Transformers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.3 Sparse Mixture-of-Experts . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.4 Out-of-Distribution Generalization Theory . .
A.5 Algorithmic Alignment . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B Theorems and Proofs
B.1 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Synthetic Experiments for Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . .
B.2.
Noiseless Version . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2.
Noisy Version . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 Technical Lemmas
C Implementations
C.1 Detailed Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Model Initialization . . .
C.3 Perturbation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.4 Auxiliary loss . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Experimental Results
D.1 DomainBed Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.2 Training Details & Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . .
D.3 Leave-one-domain-out Results . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.4 Single-source Domain Generalization Results . . . . . . . . . . . . . . . . . . . .
D.
Ablation Study on Model Design . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.6 Ablation Study on ResNet with Strong Data Augmentation . . . . . . . . . . . . .
D.7 Computational cost comparison . . .
E Model Analysis
E.1 CUB-DG Results & Visualizations .
A R ELATED W ORKS A. D OMAIN G ENERALIZATION
Domain Generalization (DG) aims to maintain the good performance of machine learning models even in the domains that are different from the training (source) domain.
The following are the categories of mainstream domain generalization research.
(1) Invariant Learning : Aligning domain distributions and ﬁnding invariance across domains has been often studied with empirical results and theoretical proofs (Ganin et al., 2016; Zhao et al., 2019a).
Speciﬁcally, researchers have explicitly sought aligning feature distributions based on the maximum mean discrepancy (MMD)
(Li et al., 2018b), second order correlation (Sun
& Saenko
, 2016), moment matching
(Peng et al
etc.
Besides aligning feature
distributions, Arjovsky et al
(Arjovsky et al
proposed IRM to learn an ideal invariant classiﬁer on top of the representation space, which
has inspired many follow-up works (Krueger et al., 2021;
Ahuja et al., 2021; Zhou et al., 2022).
These studies implement invariant learning via loss function designs.
In this paper, we show that if we have a suitable backbone architecture, ERM is sufﬁcient to ﬁnd the invariant correlations, which indicates that the backbone architecture is as important as loss function design
(2) Ensemble Learning and Meta-learning : Apart from learning invariant features and correlations across the domain, model-speciﬁc or domain-speciﬁc information helps improve the performance on DG datasets.
The ensemble learning methods combine different models to exploit model-speciﬁc information (Mancini et al., 2018;
Segù et al., 2020; Arpit et al., 2021; Li et al., 2022).
In addition, SWAD (Cha et al., 2021b) inhibits models from being overﬁt to local sharp minima by averaging model weights below a validation loss threshold.
Meta-learning-based approaches (Li et al., 2018a; 2019; Dou et al., 2019; Zhang et al., 2021c; Bui et al., 2021) leverage domain-speciﬁc feature to improve the accuracy on each domain.
(3) Data Manipulation : Diverse training data are helpful for improving generalization, and researchers have proposed different manipulation/augmentation techniques (Nazari & Kovashka, 2020; Riemer et al., 2019), domain randomization (Yue et al., 2019; Zakharov et al., 2019).
Furthermore, a
line of works (Qiao et al., 2020; Liu et al., 2018; Zhao et al., 2019b) have exploited generating data samples to enhance the model generalization ability.
(4) Neural Architecture Search : Recently, neural architecture search-based methods were employed for DG.
In Zhang et al.
(2021a), the lottery ticket hypothesis was extended to DG settings: a biased full network contains an unbiased subnetwork that can achieve better OOD performance.
The network pruning method was further adopted to identify such subnetworks.
Differentiable architecture search was employed in Bai et al.
(2021) to ﬁnd a good CNN architecture for DG.
These works are independent of backbone architecture design.
One unique advantage of the proposed method is that it is orthogonal to the above approaches, meaning that the performance of GMoE might be enhanced when combined with these approaches.
A. V ISION T RANSFORMERS
Originated from the machine translation tasks, Transformer (Vaswani et al., 2017) has recently received great attention in computer vision (Dosovitskiy et al., 2021; Liu et al., 2021;
Wang et al., 2022) for their unprecedented performance in image recognition, semantic segmentation, and visual question answering.
In addition to strong performance, some recent works empirically demonstrate the robustness of the ViTs over CNNs in terms of adversarial noise (Qin et al., 2021) and distribution shifts (Zhang et al., 2022).
Nevertheless, the theoretical understandings remain elusive and this paper theoretically validates these effects.
The analysis also motivates us to employ mixture-of-experts to enhance the performance of ViT models.
A. S PARSE M IXTURE - OF -E XPERTS
Mixture-of-Experts models, or MoEs, make use of the outputs from several sub-models (experts) through an input-dependent routing mechanism for stronger model performance (Jacobs et al., 1991;
Jordan & Jacobs, 1994).
This training paradigm has led to the development of a plethora of methods for a wide ranges of applications (Hu et al., 1997; Tani & Nolﬁ, 1999).
However, integrating MoEs and big models will inevitably introduce even larger model sizes and longer inference time.
Sparse
MoEs (Shazeer et al., 2017) were proposed with their routers to select only a few experts so that the inference time is on-par with the standalone counterpart.
They were considered as promising ways to scaling up vision models (Riquelme et al., 2021).
Existing works conjectured that the effectiveness of MoE comes from conditional computations and this is theoretically veriﬁed in Theorem 2.
In addition, we show that MoE is a powerful architecture for DG.
A. O UT - OF -D ISTRIBUTION G ENERALIZATION T HEORY
The classic PAC learning (Valiant, 1984) is only applicable to IID settings and recently there have been a line of works to investigate the theory of OOD generalization, including structured causal models (SCM)
(Arjovsky et al., 2019; Ahuja et al., 2021;
Zhou et al., 2022) and extrapolation theory (Xu et al., 2020b; Ziyin et al., 2020; Ye et al., 2021).
The SCM-based approaches often assume that the invariant correlation is linear (Arjovsky et al., 2019; Ahuja et al., 2021;
Zhou et al., 2022).
The extrapolation theory focused on the situation where the supports of training distribution and test distribution are different (Xu et al., 2020b;
Ziyin et al., 2020; Ye et al., 2021), where the networks have very limited ability for nonlinear function approximation (Xu et al., 2020b).
Different from existing works, this paper studies the situation where the invariant correlation is an arbitrary nonlinear function, and one of our main focuses is to justify why these assumptions hold in DG of vision tasks.
A. A LGORITHMIC A LIGNMENT
The algorithmic alignment framework (Xu et al., 2020a) was originally proposed to understand the effectiveness of specialized network structures in reasoning tasks.
Speciﬁcally, it characterizes the generalization bound of reasoning tasks by studying the alignment between the reasoning algorithm and the computation graph of the neural network.
The underlying intuition is that if they align well, the neural network only needs to learn simple functions, which leads to better sample efﬁciency.
This framework was further extended in Li et al.
(2021) to investigate the impact of backbone architectures on the robustness of noisy labels.
This paper adopts algorithmic alignment to provide a novel perspective of DG and develops a new architecture based on it.
B T HEOREMS AND P ROOFS B. P ROOF OF T HEOREM
We ﬁrst state the algorithmic alignment theorem in the IID setting and then extend it to the DG setting.
Theorem 3.
(Alignment improves IID generalization; (Xu et al., 2020a))
Fix (cid:15) and δ .
Given a target function g and a neural network N , suppose { x i } Mi =1 are i.i.d.
samples drawn from some distribution
D , and let y i
= g ( x i ) .
Assumptions: (a) Algorithm stability.
Let A be a learning algorithm for N i ’
s. Suppose f = A ( { x i , y i } ) and
ˆ f = A ( { ˆ x i , y i } ) .
For any x , (cid:107) f ( x ) − ˆ f ( x ) (cid:107) ≤ L 0 · max i (cid:107) x
i − ˆ x i (cid:107) , where x i is the i -th coordinate of x .
(b) Sequential learning.
We train N i sequentially: N 1 has input samples { x (1) i , f 1 ( x (1) i ) } N i =1 , with x (1) i obtained from the training dataset.
For j > 1 , the inputs ˆ x ( j ) i for N j are the outputs from the previous modules, but labels are generated by the correct functions f j − 1 , · · · , f 1 on ˆ x (1) i .
(c)
Lipschitzness.
The learned functions ˆ f j satisfy (cid:107) ˆ f j ( x ) − ˆ f j (ˆ x ) (cid:107) ≤ L 1 (cid:107) x − ˆ x (cid:107) for some L 1 .
Under assumptions (a)(b)(c), Alignment ( N , g, (cid:15), δ ) ≤ M implies there is a learning algorithm
A such that
[Equation]
where N is the network generated by A on the training data { x i , y i } Mi =1 .
Remark 4.
(Explanations of Assumptions in Theorem 3)
The ﬁrst and third assumptions are common assumptions in machine learning and are practical.
The second assumption is impractical as we do not have auxiliary labels.
However, the same pattern is observed for end-to-end learning in experiments (Xu et al., 2020a; Li et al., 2021).
We now prove Theorem 3.
Theorem 4.
(Impact of Backbone Architecture in DG) Let N (cid:48)
=
{N 2 , · · · , N n } .
Assuming we train the neural network with ERM, and Assumption 1, 2, 3, we have the following statements:
If Alignment ( N (cid:48) , g c , (cid:15), δ ) ≤ |E tr | , we have P D te [ (cid:107)N ( x ) − y (cid:107) ≤ O ( (cid:15) )] > 1 − O ( δ ) ; 2.
If Alignment ( N (cid:48) , g s , (cid:15), δ ) ≤ |E tr | , we have P D te [ (cid:107)N ( x ) − y (cid:107) > ω ( (cid:15) )] > 1 − O ( δ ) .
Proof.
The proof is mainly based on Theorem 3.
We tackle the distribution shift in x with Lemmaand analyze the distribution shift in y with Assumption 2, 3.
First condition:
From Theorem 3,
[Equation]
By Lemma 1, we have
[Equation]
Combing with the second condition of Assumption 2, we have
[Equation]
where the ﬁrst inequality follows O ( (cid:15) )
+ O ( (cid:15) )
=
O ( (cid:15) ) , and the second inequality follows the second condition of Assumption 2 and equation 6.
Second condition:
Following the proof of equation 6, we obtain
[Equation]
Combing with the second condition of Assumption 3, we have
[Equation]
where the ﬁrst inequality follows ω ( δ ) − O ( δ )
=
O ( δ ) , and the second inequality follows the second condition of Assumption 3 and equation 7.
This ﬁnishes the proof for Theorem 1.
B. S YNTHETIC E XPERIMENTS FOR T HEOREM
In this subsection, we set up synthetic experiments to illustrate Theorem 1.
B.2. N OISELESS V ERSION
We ﬁrst set up experiments following the assumptions in Theorem 1.
Training dataset generation: Consider data pairs ( x , y ) ∈ ( R K ) P × { 1 , · · · , K } in the training and validation datasets generated as follows.
• Generate label y ∈ { 1 , · · · , K } uniformly.
• Generate K mutually orthogonal feature vectors { c k } such that c k ∈ R K and (cid:107)
c k (cid:107) 2 = 1 .
Generate x as a collection of P patches: x = ( x (1) , · · · , x ( P ) ) ∈ ( R d ) P – Pixel-level feature.
For the ﬁrst patch, the y -th pixel is set as 1 and other pixels drawn from N (0 , 1) .
– Patch-level feature.
Uniformly select one and only one patch given by c y .
– Random noise.
The rest patches are Gaussian noise drawn independently from N (0 , I K ) .
The generation of the validation dataset follows that of the training dataset.
Test dataset generation: Two OOD test datasets are given.
The ﬁrst dataset differs from the training dataset in the pixel-level feature.
Speciﬁcally, we set the ( y + 1) mod K -th pixel of the ﬁrst patch as 1 and other pixels in the ﬁrst patch follow N (0 , 1) , i.e., making the pixel-level features spurious.
The second dataset is distinguished from the training dataset such that the patch-level feature is c ( y +1) mod K , i.e., making patch-level features spurious.
Backbone architectures: We adopt MLPs and fully convolutional networks (FCNs) as MLPs align with pixel-level features while FCNs align with patch-level features.
It is clear that both MLPs and
FCNs can ﬁt the label only from pixel-level features or patch-level features.
Parameters in experiments: We set P = 10 , K = 4 .
We generate 100 , 000 samples for the training dataset and 2 , 000 test samples for the validation and test datasets.
We use a three-layer MLP with hidden size { 40 , 100 , 100 , 4 } and a two-layer FCN with 20 ﬁlters.
The experimental results are shown in Fig.
The MLP completely fails on the second dataset while the FCN completely fails on the ﬁrst dataset, which is consistent with Theorem 1.
B.2. N OISY V ERSION
In the previous experiment, both pixel-level features and patch-level features can perfectly ﬁt the label.
We give a noisy version as follows:
Pixel-level feature.
Given a probability value p 1 .
With probability p 1 , the y -th pixel of the ﬁrst patch is set as 1 and other pixels are drawn from N (0 , 1) .
With probability 1 − p 1 , the y (cid:48) -th
( y (cid:48)
(cid:54) =
y ) pixel of the ﬁrst patch is set as 1 and other
[Figure]
cc u
[Figure]
a c y
Validation
[Figure]
Test
[Figure] [Figure]
Training Epoches
A cc u r a c y
Validation Test
Test (b) FCN’s performance.
[Figure]
A cc u r a c y
Validation
[Figure]
Patch-level feature.
Given a probability value p 2 .
With probability p 2 , uniformly select one and only one patch given by c y .
With probability 1 − p 2 , uniformly select one and only one patch given by c y (cid:48)
( y (cid:48)
(cid:54) =
y ).
The experiments with different p 1 and p 2 are shown in Fig.
When the signal-to-noise ratios of pixel-level feature and patch-level feature are similar, the MLP learns the pixel-level feature, which corresponds to its alignment.
It shows that the network prefers learning the correlation that aligns with its architecture, even when exploiting other correlations may lead to performance gains.
When the difference between p 1 and p 2 is large, the network will still learn its aligned correlation at the early stage (with above 90% accuracy on Test1), but may change when the number of epochs is large.
This suggests that early stopping is beneﬁcial to DG as suggested in Cha et al.
(2021b)
under the condition that the backbone aligns to some invariant correlations.
B. P ROOF OF T HEOREM
Theorem 5.
An MoE module in equation 3 with N experts and k = 1 aligns the conditional sentences in Algorithm 1 with
Alignment =
[Equation]
if N < M
[Equation]
where M ( · , · , · , · ) is deﬁned in Deﬁnition 1, and M ∗ P is the optimal objective value of the following optimization problem:
[Equation]
subject to
[Equation]
where 1 I i is the indicator function on interval
I i .
Proof.
For N ≥ M , we assign one expert for each function.
For N < M , similar functions should be assigned to one expert to minimize the sample complexity.
Case N ≥ M : We deﬁne the mapping H ( x ) =
[1 I 1 , · · · , 1 I M ] ∈ R M , where 1 I i is the indicator function on interval
I i .
To show the alignment between MoE and conditional sentences, we deﬁne the module functions f 1 , · · · , f N and neural network modules N 1 , · · · , N N as
[Equation]
Replacing N i with f i , the MoE network simulates Algorithm 1.
Thus, for N > M , MoE algorithmically aligns conditional sentences with
[Equation]
where M ( · , · , · , · ) is deﬁned in Deﬁnition 1.
Case N < M : We deﬁne the mapping H 0 ( x ) =
[1 ∪ j ∈I 1 I j , · · · , 1 ∪ j ∈I N I j ] ∈ R N , where I i is the solution to the optimization problem in equation 5.
To show the alignment between MoE and conditional sentences, we deﬁne the module functions f 1 , · · · , f N and neural network modules
N 1 , · · · , N N as
[Equation]
Replacing N i with f i , the MoE network simulates Algorithm 1.
Thus, for N < M , MoE algorithmically aligns with conditional sentences with
[Equation]
where M ∗ P is the optimal objective value to equation 5 and M ( · , · , · , · ) is deﬁned in Deﬁnition 1.
B. T ECHNICAL L EMMAS
Lemma 1.
Under Assumption 1, for a given functions f ( · ) and an interval A , P D tr [ f ( N 1 ( x )) ∈
A ] < δ implies P D te [ f ( N 1 ( x ))
∈ A ] ≤ Cδ .
Proof.
We denote the inverse mapping as f − 1 ( A ) =
{ s | f ( s ) ∈ A } .
Note that
P D te [ f ( N 1 ( x ))
∈ A ] = (cid:90)
[Equation]
This ﬁnishes the proof for Lemma 1.
Remark 5.
Assumption 1 is similar to µ -exploration in batch reinforcement learning (RL), which aims at controlling the distribution shift in RL (Xie & Jiang, 2021).
Lemma 1 suggests that the generalization with the same training and test support resembles the IID generalization.
Unfortunately, this is the largest distribution shift we can assume as the impossibility of nonlinear out-of-support generalization is proved in Xu et al.
(2020b).
C I MPLEMENTATIONS C. D ETAILED A RCHITECTURE
For image input X , we ﬁrst process the images as a sequence of equal-sized patches x i ∈
X using a 1-layer convolutional neural network with layer normalization.
Then those patches are added with positional embeddings, and the patch embeddings (tokens) are ready to be processed by later L blocks.
To make a fair comparison with other reported methods on DomainBed, we choose the ViT-S/ which has similar parameters and run-time memory cost with ResNet-50 as the basic architecture for our main model GMoE-S/16.
The chosen ViT-S/16 model has an input patch size of 16 × 16 , 6 heads in multi-head attention layers, and 12 transformer blocks.
We consider two-layer conﬁgurations, every-two and last-two , that are widely adopted in Sparse MoE design (Riquelme et al., 2021).
The detailed architectures are illustrated in Figure 6.
Besides, we also consider evaluating GMoE’s generalization performance with larger scale models ( e.g., ViT-Base).
V-MoE (Riquelme et al., 2021) trains 32 experts model on ImageNet-21K dataset.
However, considering the number of training data, we need to adopt a smaller number of experts.
More experts require larger datasets to ensure that every expert is adequately trained since each expert only sees a small portion of the dataset.
Domain generalization datasets are usually 1-2 orders of magnitude smaller than ImageNet-21K.
Based on this fact, each GMoE block contains 6 experts.
For each image patch, the cosine router selects the TOP-2 index out of 6 and routes the patch to the corresponding 2 experts.
We put the ablation study results on layer conﬁguration and larger size backbone model on Sec.
D.5.
C. M ODEL I NITIALIZATION
To make a fair comparison with other algorithms from DomainBed (mostly with ResNet-50 pretrained on ImageNet-1K), we also utilize the pretrained ViT models on ImageNet-1K from DeiT to initialize our GMoE.
Our practice avoids the generalization performance improvement of GMoE coming from a stronger pretrained model ( e.g., directly from V-MoE pretrained on ImageNet-21K).
Meanwhile, there is no released pretrained MoE model on ImageNet-1K.
In Table 4, we see ViT-S/16 and GMoE-S/ achieve excellent generalization performance even with fewer parameters and smaller architecture.
In the GMoE block, we initialize the multiple experts with pretrained ViT block’s FFN in the same position ( e.g., block-8 and 10) and multiply it by E copies ( E is the number of experts).
For cosine routers, we randomly initialize them in each GMoE block, to make them evenly choose experts at ﬁrst.
[Figure]
Images
Patch Embedding 12 B l o c k s
Transformer Block
GMoE Block (
[Figure]
Classification Layer
GMoE Block
Transformer Block Transformer Block
GMoE Block
Transformer Block Transformer Block
Images
Patch Embedding
[Figure]
[Table]
P ERTURBATION
Another speciﬁc design for the router is adding Gaussian noise to increase stability.
Hard routing brings sparsity beneﬁts while also causing discontinuities in routing.
In a TOP-2 selection routing, a small perturbation would cause order changes if the largest output and second-largest output are close.
In domain generalization, data come from diverse domains and we need to learn a model to capture the invariant semantic meaning of an object.
Existing work (Chen et al., 2022) theoretically claims that adding noise to the gating function provides a smooth transition of the gating function during training.
Empirically, we ﬁnd adding noise with a standard deviation of 1 N would empirically improve performance and stabilize training dynamics.
C. A UXILIARY LOSS
It has been observed that the gating network tends to converge to a self-reinforcing imbalanced state, where a few experts are more frequently activated than others.
Following Riquelme et al.
Shazeer et al.
(2017), we deﬁne an importance loss L imp to encourage balanced usage of experts.
The importance of expert e is deﬁned as the normalization of the gating function’s output in a batch of images:
[Equation]
And the loss L imp is deﬁned by the squared coefﬁcient of variation of imp ( X ) :
[Equation]
While importance loss seeks that all experts have on average similar output routing weights, it is still likely some experts would not be routed (see example in Table 5).
[Table]
To address this issue, besides balancing across patches, we also need to encourage balanced assignment across experts.
Ideally, we expect a balanced assignment via a load loss L load .
However, the assignment is a discrete value; then we expect a proxy to make it differentiable in back-propagation.
Following the idea in Riquelme et al.
; Shazeer et al.
; Mustafa et al.
(2022), for each token
x and each expert e i , we could compute the probability of e i once being selected in TOP k after SoftMax operation, and then kept if re-sampling again only the noise for expert e i .
In other words, we are computing the probability that having expert e i still being among the TOP k while re-sampling again only the noise of the expert.
This probability is formally deﬁned as in Eq.
10, where η k the K -th largest entry after SoftMax operation and Φ is the cumulative distribution function of a Gaussian distribution.
[Equation]
Then the expert load could be deﬁned as load e ( X ) =
(cid:80) x
∈ X p e ( x ) and the whole experts load is deﬁned as load ( X ) = { load e ( X ) } Ee =1 .The
load loss L load is deﬁned by
[Equation]
With the above adaptations for DG tasks, GMoE still works with a simple overall loss:
[Equation]
with a hyperparameter λ > 0 to balance the main task loss and the auxiliary loss, in which usually we view equal importance of the load and importance loss.
D E XPERIMENTAL R ESULTS
D.
D OMAIN B ED D ETAILS
Benchmark Datasets Following previous DG studies, we mainly evaluate our proposed method and report empirical results on DomainBed (Gulrajani & Lopez-Paz, 2021) with 8 benchmark datasets,
PACS, VLCS, OfﬁceHome, TerraIncognita, DomainNet, SVIRO, Wilds-Camelyon and Wilds-FMOW.
Besides DomainBed, we also conduct experiments on the self-created CUB-DG dataset for model analysis.
In the following, we will provide the details of different datasets.
Dataset Details In Table 6, we provide statistics for the 8 datasets in DomainBed, as well as our self-created dataset CUB-DG (e.g., number of domains, categories and images).
[Table] [Table]
In detail, the 9 multi-domain image classiﬁcation datasets are comprised of:
PACS (Li et al., 2017) comprises four domains d ∈ {art, cartoons, photos, sketches}.
This dataset contains 9, 991 examples of dimension (3, 224, 224) and 7 classes.
VLCS
(Fang et al., 2013) comprises photographic domains d ∈ {Caltech101, LabelMe, SUN09, VOC2007}.
This dataset contains 10, 729 examples of dimension (3, 224, 224) and 5 classes.
Ofﬁce-Home (Venkateswara et al., 2017) includes domains d ∈ {art, clipart, product, real}.
This dataset contains 15, 588 examples of dimension (3, 224, 224) and 65 classes.
TerraIncognita (Beery et al., 2018) contains photographs of wild animals taken by camera traps at locations d ∈ { L100 , L38 , L43 , L46 } .
Our version of this dataset contains 24,examples of dimension (3, 224, 224) and 10 classes.
DomainNet (Peng et al., 2019) has 6 domains d ∈ {clipart, infograph, painting, quickdraw, real, sketch}.
This dataset contains 586, 575 examples of size (3, 224, 224) and 345 classes.
SVIRO (Cruz et al., 2020) is a Synthetic dataset for Vehicle Interior Rear seat Occupancy across 10 different vehicles, the 10 domains.
The domains are d ∈ { aclass , escape , hilux , i3 , lexus , tesla , tiguan , tucson , x5 , zoe } .
This dataset for image classiﬁcation contains 56,000 image examples of size (3, 224, 224) andclasses.
Wilds-Camelyon (Koh et al., 2021) dataset contains histopathological image slides collected and processed by different hospitals and curated by Wilds benchmark.
It contains 455, examples of dimension (3, 224, 224) and 2 classes from 5 hospitals.
Wilds-FMOW (Koh et al., 2021) dataset is a variant of the functional map of the world dataset (Christie et al., 2018) and contains satellite images of 62 buildings or land classes across 6 regions.
The image example is of size (3, 224, 224).
[Table]
CUB-DG dataset is built from the original Caltech-UCSD Birds (CUB) dataset (Wah et al., 2011).
We stylize the original images into three domains, Candy, Mosaic, and Undie, while keeping the original images as the Natural domain.
Overall, CUB-DG has 4 domains d ∈ {Natural, Candy, Mosaic, Udnie}.
It contains 47,152 bird examples and are categorized into 200 classes.
Each image has detailed annotations: 1 category label, 15 positional attributes ( e.g., left wing, back), 312 binary attributes ( e.g., has_bill_shape=dagger ), andbounding box information.
We treat the 15 positional attributes as visual attributes of a bird image, and thus we could evaluate the correlation between visual attributes and expert selections.
Evaluation protocols We evaluate our proposed method with two protocols on DomainBed (Gulrajani & Lopez-Paz, 2021) benchmark.
For train-validation selection , we split each training domain into training and validation subsets.
Then, we pool the validation subsets of each training domain to create an overall validation set.
Finally, we choose the model maximizing the accuracy on the overall validation set, and report the ﬁnal accuracy on one leave-out test domain.
For leave-one-domain-out validation selection , we train the model on all training domains and one domain out as the validation set.
We choose the model maximizing the accuracy on the leave-out validation domain, and report accuracy on another leave-out test domain.
We should emphasize that the leave-one-domain-out validation setting means we choose two domains as leave-out domains.
This term is used in conformity in literature (Cha et al., 2022; 2021b).
Standard error bars We train the model three times with random seeds on weight initializations.
The mean and standard error of these repetitions are reported.
D. T RAINING D ETAILS & H YPERPARAMETERS
In this section, we report the training details of our experiments.
We optimize models using Adam optimizer (Kingma & Ba, 2015) with slightly different parameters on different datasets (see Table 8).
The training and inference batch size is set to 32 for each domain.
Since Gulrajani & Lopez-Paz (2021) only evaluates baselines for 5K iterations on DomainNet, while 5K iterations are less than 2 epochs on DomainNet.
This results in insufﬁcient training on this dataset and it can not fully reveal models’ performance with different algorithms.
Subsequent state-of-the-art works (Cha et al., 2021b)
(Cha et al., 2022) re-evaluated ERM and proposed their methods on 15, iterations of DomainNet.
To compare with those SOTA counterparts, we report the results of 15K iterations on DomainNet in
Table 1.
However, we also train GMoE for 5K iterations to make a fair comparison with both types of algorithms (see Table 9).
For the rest of the datasets on DomainBed, we train GMoE for 5K iterations following the original setting.
D. L EAVE - ONE - DOMAIN - OUT R ESULTS
We have demonstrated the experimental results in Table 1 with train-validation selection.
In this section, we will further provide the results of GMoE on leave-one-domain-out selection, which is a more challenging setting.
In this setting, we should train GMoE on 5 datasets with in total ofindividual experiments and summarize them as average performance.
Due to the huge amount of experiments, recent methods rarely evaluate their methods on such setting, except for those provided
[Table]
by DomainBed.
In the Table
[Table]
we can ﬁnd that GMoE outperforms previous methods in alldatasets
[Table]
-of-domain accuracies () with leave-one-domain-out selection criterion.
[Table]
D. S INGLE - SOURCE D OMAIN G ENERALIZATION R ESULTS
In this section
[Table]
we will demonstrate the single source domain generalization results on DomainNet’s 6 domains.
We select different training domains in turn and take the remaining 5 domains as test domains.
In Table [Table] we show the OOD results when models are trained on only one domain and test on 5 test domains, as well as the IID results on the training domain’s validation set.
We specify
IID and OOD
[Table]
ResNet 16.
ViT-S/ 22.
GMoE-S/ 22.
Real
Clipart
Info
Paint
Quick
Real
Sketch
ResNet 39.
ResNet 43.
ViT-S/ 46.
GMoE-S/ 47.
Sketch
Clipart
Info
Paint
Quick
Real
Sketch
ResNet 48. ResNet 48.
ViT-S/ 52.
GMoE-S/ 52.
D. A BLATION S TUDY ON M ODEL D ESIGN
In order to validate discussions in Sec. 4.3 and Sec C.
[Table]
we conduct experiments with different layer conﬁgurations, as well as adopting larger backbone ( e.g., ViT-Base) model.
From the results in
Table
[Table]
we see that the last two conﬁguration largely exceeds every two conﬁguration, thus verifying our idea that visual attributes exist in relatively high-level signals.
This result also demonstrates that the cosine router and larger model lead to better generalization performance.
[Table] [Table]
[Table] [Table]
D. A BLATION S TUDY ON R ES N ET WITH S TRONG D ATA A UGMENTATION
In this part
[Table]
we will go into greater depth about the pre-trained model’s effect on DG performance
[Table]
As stated in Sec.
C.
[Table]
GMoE is initialized from the same architecture ViT models, where pre-trained weights are from Touvron et al.
In addition, the pre-training on ViT and ResNet models may adopt different training recipes.
And it is generally believed that models with stronger data augmentation during pre-training could have better performance on downstream tasks.
To verify this difference in detail, we consider comparing GMoE-S/16’s pre-trained model, ViT-S/16, with a stronger pre-trained ResNet-50 V2 model 1 with signiﬁcant training tricks data augmentations.
The pre-training details on ViT-S/16 and ResNet V2 are listed in Table 13.
Due to the space limit, we use abbreviations for some terms.
In detail, they are:
• LR Opt. stands for learning rate optimization.
• TrivialAug. stands for trivial augmentation (Müller & Hutter, 2021).
• Ep.
stands for pre-training epochs.
• Rand Er. stands for random erasing (Zhong et al., 2017).
• Label Sm. stands for label smoothing (Szegedy et al., 2016).
• FixRes Mt. stands for FixRes mitigations (Touvron et al., 2019).
• WDT stands for weight decay tuning.
• IRT stands for inference resize tuning (Touvron et al., 2019).
The results in Table 14 demonstrate that GMoE and ViT could still remain higher performance than the strong pre-trained model ResNet V2 due to the beneﬁts of the backbone architecture.
D. C OMPUTATIONAL COST COMPARISON
Having shown the effects of GMoE’s performance on different domain generalization benchmarks, we now conduct efﬁciency analysis with respect to iteration time and run-time memory during training and inference.
Since algorithms developed from DomainBed mainly adopt the same architectures ( e.g., ResNet50, Linear Classiﬁer), traditional model complexity measures such as ﬂops and model parameters can not truly reﬂect the difference in efﬁciency.
By evaluating the above two metrics, different algorithms with more complex loss design or gradient constraints will cause larger overheads with respect to iteration time or run-time memory.
From the results in Table 15, we observe that
GMoE achieves relatively low run-time memory and training step time among other competitors. [Table] [Table] Run
[Table]
time Memory (GB) ↓
[Table]
E M ODEL A NALYSIS E. CUB-DG R ESULTS & V ISUALIZATIONS
Generalization across Image Stylization To evaluate the model performance in generalizing to image stylization, we conduct experiments following DomainBed setting on CUB-DG.
We compare
GMoE with other three invariant learning DG algorithms, (i.e., DANN (Ganin et al., 2016), Fish (Shi et al., 2021) and Fishr (Rame et al., 2021)).
Among them, DANN and Fish are widely discussed and cited papers, and Fishr is considered the best algorithm of its kind so far.
In this experiment, all three invariant learning methods, including ERM, adopt ResNet-50 as the backbone model.
GMoE does not have any speciﬁc loss design and constraint.
The only difference is the model architecture.
In Table 16, we see that GMoE, the new backbone model for DG, is signiﬁcantly more effective than those methods speciﬁcally designed to learn invariance with certain losses and gradient constraints.
[Table]
Experts Selection Following the analysis in Sec.
5.4, we will provide more visualization results.
In Figure 8-11, we provide the experts selection results for the same image on four domains (with four types of stylization).
We see that, in most cases, the same visual attributes in different domains are handled by the same experts.
For example, in Figure 8, right eye marked by green dot ( ) is processed by Expert 5 across four domains.
The results reveal that experts are invariant in dealing with visual attributes across domains.
Multi-head Attention Visualization Ideally, the multi-head attention mechanism can be viewed as jointly attending to multiple places by ensembling multiple attention heads.
Each attention head would focus on its speciﬁc attention relationship between all patches and inherently collaborate with each other.
[Figure]
However, recent studies (Cordonnier et al., 2020) demonstrate that multiple-head attention (MHA) layers could learn redundant key/query projections, i.e. some heads might attend to similar features in input space.
This issue demonstrates that two heads head i and head j are computing the same key/query representations up to a unitary matrix
I ∈ R d × d such that W Q i =
W Q j I and
W K i
=
W K j I .
In this case, even though the two heads are computing identical attention scores, i.e.
W Q i II T W TK i , the concatenation [ W Q i , W Q j
] ∈ R d × 2 d can also be full rank, which indicates that some attention heads would focus on the same content but are agnostic to each other.
We opine that the mixture-of-experts (MoE) layer could alleviate redundancy and increase disentanglement to some extent.
In Figure 7, presented in a grid fashion, the MHA layer focuses on different signals (features) in different heads while MoE layer disentangles the information by handling each patch to different experts.
This mechanism extends the network sparsity and hence is able to improve the model generalizability.
In sum, the MoE layer is designed to leverage different experts to capture distinct visual attributes, as well as their correlations (via attention mechanism) to each other.
To see more about the learned representation in multi-head attention layer, we provide the visualization of last layer’s attention activations of ViT and GMoE on Figure 12 13 14 15.
[Figure] [Figure]
[Figure] [Figure]
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T G M o E V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E
V i T
G M o E