Sparse Mixture-of-Experts are Domain Generalizable Learners
Authors: Bo Li
Yifei Shen
Jingkang Yang
Yezhen Wang
Jiawei Ren
Tong Che
Jun Zhang
Ziwei Liu
Venue: ICLR 2023 notable top 5% (2023)
Keywords: domain generalization
mixture-of-experts
algorithmic alignment
visual attributes

Abstract: 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 finding that transformer-based models trained with empirical risk minimization (ERM) outperform CNN-based models employing state-of-the-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 \emph{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.

Published as a conference paper at ICLR 2023
SPARSE MIXTURE-OF-EXPERTS ARE DOMAIN GENER- ALIZABLE LEARNERS
Bo Li Yifei Shen?* Jingkang Yang' Yezhen Wang* Jiawei Ren'
Tong Che Jun Zhang'
Ziwei Liu!
La-L Nanyang Technological University
The Hong Kong University of Science and Technology
"Mila-Queber AI Institute "Nvidia Research (libo0013, 3, z ziwei. . liv) Centu.edu.ss
ABSTRACT
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, propose explore orthogonal direction, i.e., the design of the backbone architecture. It is motivated by an empirical finding that transformer-baset models trained with empirical risk mininization (ERM) outperform CNN-based models employing state-of- the-art (SOTA) DG algorithms on multiple DG datasets. We develop formal framework to characterize a network's 's robustness to distribution shifts by studying its architecture's 's alignment with the correlations 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. Morcover, GMoE complementary to existing DG methods and its performance is substantially improved when trained with DG algorithms.
the dataset. This analysis
INTRODUCTION
1.1 MOTIVATIONS
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 background, lighting, texture, shape, and geographic/demographic attributes.
From the perspective of representation learning, there several paradigms towards this goal, including domain aligmment (Ganin et al., 2016; Hoffman et al., 2018), invariant causality predic- tion (Arjovsky et al., 2019; Krueger al., 2021), meta-learning (Bui al., 2021; Zhang et al., 2021c), ensemble learning (Mancini et al., 2018; Cha et al., 2021b), and feature disentanglement (Wang al., 2021; Zhang et al., 2021b). The most popular approach implementing these ideas is to design a specific loss function. For example, DANN (Ganin et al., 2016) aligns domain distribu- tions by adversarial losses. Invariant causal prediction can be enforced by penalty of gradient norm (Arjovsky et al., 2019) or variance of training risks (Krueger et al., 2021). Meta-learning and domain -opmaie 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 pivotal role in performance boost and have attracted much attention (He et al., 2016; Hu et al., 8; Liu et al., 2021). Additionally, it has been empirically demonstrated in Sivaprasad et al. (2021) that different CNN architectures have different performances on DG datasets. Inspired by these pionecring works, 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
"quull contribution, #Corresponding author.

---

Published as a conference paper at ICLR 2023
computational overhead, shown in Fig. I(a). To our surprise, 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., 202 Rame et al., 2021; Shi et al., 2021) on DomainNet, OfficeHome 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 first prove that 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 defined 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 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 composed of some sets of attributes and distribution shifts data are distribution shifts of these attributes (Wiles et al., 2021). The latent factorization model of these attributes is almost identical the generative model of visual attributes in classic computer vision (Ferrari & Zisserman, 2007). To capture these diverse attributes, we propose a Generalizable Mfxture-of-Expeer mixture-of-ceppert (sparse MoEs) (Shazeer et al., 2017) and vision transformer (Dosovitskiy et al., 2021). The sparse MoEs were originally proposed key enablers for extremely large, but efficient 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.
(GMoE), which is built upon sparse
1.2 CONTRIBUTIONS
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. Specifically, our main contributions are summarized as follows:
Novel View of DG: In contrast to previous works, this paper initiates formal exploration of the backbone architecture in DG. Based on algorithmic alignment (Xu et al., 2020a), we prove that network is more robust to distribution shifts its architecture aligns with the invariant correlation, whereas less robust if its architecture aligns with spurious correlation. The theorems are verified on synthetic and real datasets.
A Novel Model for DG: Based on our theoretical analysis, we propose Generalizable Mixture-of- Experts (GMoE) and prove that enjoys a better alignment than vision transformers. GMoE is built upon sparse mixture-of-experts (Shazeer 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 Do- mainBed. Remarkably, GMoE trained with ERM achieves SOTA performance on datasets in the train-validation setting and datasets in the leave-one-domain-out setting. Furthermore, the GMoE trained with DG algorithms achieves better performance than GMoE trained with ERM.
2 PRELIMINARIES
2.1 NOTATIONS
Throughout this paper, a, A stand for scalar, column vector, matrix, respectively. O(-) and w-)) are asymptotic notations. We denote the training dataset, training distribution, test dataset, and test distribution &tr, D... Erc, and Dre, respectively.
2.2
ATTRIBUTE FACTORIZATION
The attribute factorization (Wiles et al., 2021) is a realistic generative model under distribution shifts. Consider a joint distribution of the input and corresponding attributes at ... ... ak (denoted as
2

---

In Definition 1, the original task is to learn f. which is a challenging problem. Intuitively, if we could find a backbone architecture that is suitable for this task, it helps break the original task into Published as a conference paper at ICLR 2023
l:K) with a' A', where A' is finite set. The label can depend on multiple attributes. Denote the latent factor as 2, the data generation process is given by
p(z). a' p(a))z), x p(z|z), p(a" I:K . = / )dz (1)
The distribution shift arises if different marginal distributions of the attributes are given but they share the same conditional generative process. Specifically, we have Prain (a!:*) * Prea ((1:K) ), but the generative model in equation I is shared across the distributions, i.c., have Preat Prea(a!: l:K ,x) Prest(a p(x) x x))p(= ( la l: dzz and similarly for Purin. The above description is abstract and we will illustrate with example.
Example (DSPRITES (Matthey et al., 2017)) Consider A = fred, blue} and Aa fellipsee square). The target task is shape classification task, where the label depends at- tribute a². 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 red, the classifier 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.3 ALGORITHMIC ALIGNMENT
We first 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 defined as the following.
Definition (Aligument: (Xu et al., 2020a)) Let N denote a neural network with n modules {Ni and assume that a target function for learning = g(æ) can be decomposed into n The network N aligns with the target function if replacing N; with fü outputs the same value as algorithm 9. The aligument value between N and f is defined as
functions f,, In.
Alignment(N, 5, 8) == max Mi,,
(2)
where M(fisNG, 6) denotes the sample complexity measure for N₂ learn fi with precision at failure probability 8 under a learning algorithm when the training distribution is the same as the test distribution.
simpler sub-tasks, i.e., to learn fi, ....,, fa instead. Under the assumptions of algorithmic alignment (Xu al., 2020a). f can be learned optimally if the sub-task f1, fn can be learned optimally, Thus, a good alignment makes the target task easier to learn, and thus improves IID generalization. which is given Theorem 3 in Appendix B. In Section 3, we extend this framework the DG setting.
3 ON THE IMPORTANCE OF NEURAL ARCHITECTURE FOR DOMAIN GENERALIZATION
In this section, investigate the impact of the backbone architecture DG, from motivating example to a formal framework.
3.1 MOTIVATING EXAMPLE: CNNS VERSUS VISION TRANSFORMERS
We adopt DomainBed (Gulrajani & Lopez-Paz, 2021) 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. ((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 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 architecture.
3

---

Published as a conference paper at ICLR 2023
SOTA CNN based algorithms vanille
Camell Baskeebbe Apple Eigphant
80
caajos Pengen Swad
70
Accuracy
60
50
cu Jm
43
VLCS
OfficeHome
Domainfiet
Quickdraw
Cipart
Real
Shetch
Infegraph
(a) Performance comparison.
(b) Dataset visualization.
Figure 1: (a) Performance comparison of ViT-S/16 (w/21. .8M trainable parameters) with ERM and the ResNet-50 ww/2.66M backbone with SOTA DG algorithms. ERM(RN-50), Fishr(RN-50), Swad 50/Ensemble) denotes the ResNet-50 trained with ERM, Fish (Shi et al., 2021), Fishr (Rame et al., 2021), and Swad (Cha et al., 2021b), respectively. (b) Image examples from DomainNet. Each row shows one class and each column shows domain.
3.2 UNDERSTANDING FROM A THEORETICAL PERSPECTIVE
The above experiment leads to an intriguing question: how does the backbone architecture impact the network's '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 first make an assumption on the distribution shift.
Assumption 1. Denote N1 as the first module of the network (including one or multiple layers) of the network. Let NN (8) and Ni (8) denote the probability density functions of features after Ni- Assume that the support of the training feature distribution covers that of the test feature distribution,
L.., maxs Piann,1 (8] S C, where C is constant independent of the number of training samples.
Pun,Ni (#)
Remark 1. (Interpretations of Assumption 1) This condition is practical in DG, especially when we have a pretrained model for disentanglement (c.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 i attributes/features (e.g.. big ears and long noise) are covered in the training domains. From a technical perspective. is impossible for networks trained with gradient descent approximate 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 define several key concepts 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 9e such that for training data, we have ge (x)) = y, Vx E Ey,, and for test data, we have IP D. [NIl
yll S <) - 5.
We then introduce the spurious correlation (Wiles et al., 2021), i.c. some attributes are correlated with labels in the training dataset but not in the test dataset. The spurious correlation exists only the training distribution differs from the test distribution and this distinguishes DG from classic PAC-learning settings (Xu al., 2020a).
Assumption 3. (Spurious correlation) Assume there exists a function 9a such for training data 9. (N( = y, Vx E,,, and for test data, we have IP PD (N)) - yll > w(e)] > 1 8.
The next theorem extends algorithmic alignment from IID generalization (Theorem 3) to DG.
Theorem 1. (Impact of Backbone Architecture in Domain Generalization) Denote A
(N2.-. ... .Nn). Assuming we train the neural network with ERM, and Assumption 1, 2, 3 hold, we have the following statements:
1. If Aligument( (N', Sec e, 0) S (arl we have PD,, [l))())))))))))))))))())) 0(()) 0(8);
4

---

Published as a conference paper at ICLR 2023
2. If Alignment N" .9. C. 6) & |Erl we have PD. [I]) yl > w(e)) > 1 1
Remark 2. (Interpretations of Theorem 1) By choosing a sufficiently small e, only one of Alignment(A" 9c: E, o, < |Et-| and Alignment(AV, 19,, - |E,rl holds. Thus, Theorem I shows that the networks aligned with invariant correlations are more robust to distribution shifts. In Ap- pendix B.2, we build a synthetic dataset that satisfies 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 fit 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 I to explain the experi- ments in the last subsection. The first condition of Theorem shows that the neural architecture aligns with the invariant correlation, ERM is sufficient achieve good performance. In some domains of OfficeHome or DomainNet, the shape attribute has an invariant correlation with the label, illustrated in Fig. l(b). On the contrary, 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 filters with shape bias while convolutions are high-pass filters with a texture bias. As result, a VIT simply trained with ERM can outperform CNNs trained with SOTA DG algorithms.
To improve ViT's performance, Theorem I 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 configuration 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.
4 GENERALIZABLE MIXTURE-OF-EXPERT FOR DOMAIN GENERALIZATION
In this section, we propose Generalizable Mixture of Experts (GMoE) for domain generalization, supported by effective neural architecture design and theoretical analysis.
4.1 MIXTURE-OF-EXPIE
LAYER
In this subsection, we introduce the mixture-of-ene -oflure--f-expertss (MoE) layer, which is essential component of GMoE. One ViT layer is composed of an MHA and an FFN. In the MoE layer, the FFN is replaced by mixturs-of-experts and each expert is implemented by FFN (Shazeer et al., 2017). Denoting the output of the MHA as æ, the output of the MoE layer with experts is given by
N
fMOE (x) = E(t)), E(())= = Σ TOP (Softmax (((ootmaax(W)))) W of WH.,,)) x),
(3)
i=1
where Win is the learnable parameter for the gate, W FFN; and Wi,, are learnable parameters for the i-th expert, (--) is a nonlinear activation function, and TOP operation 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 a hyperparameter. Given Din as the input of the MoE layer, the update is given by
x = SMIA (LN((.).) + Zin, Cout
fME (LN(x)) +
where fMHA is the MHA layer, LN represents layer normalization, and Wour is the output of the MoE layer.
4.2 VISUAL ATTRIBUTES, CONDITIONAL STATEMENTS, AND SPARSE MoEs
In real world image data, the label depends on multiple attributes. Capturing diverse visual attributes is especially important for DG. For example, the definition 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 trunk". The definition 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, is sufficient to achieve high accuracy (Geirhos
5

---

Published as a conference paper at ICLR 2023
al., 2018). However, in DomainNet, elephants longer have grey skins while the long nose and big ears are preserved and the network relying on grey skins will fail to generalize.
The conditional statement (i.e., IF/ELSE in programming), shown in Algorithm I, is a powerful tool to efficiently capture the visual at- tributes and combine them for DG. Suppose train the network recognize the elephants on DomainNet, as illustrated the first row of Fig. ((b). For the elephants in different domains, shape and tex- ture vary significantly while the visual attributes (large ears, curved teeth, long nose) are invariant across all the domains. Equipped with conditional statements, the recognition of the elephants can be expressed as "if an animal has large ears, two curved outer teeth, and a long nose, then it is an elephant". Then the subtasks are to recognize these visual attributes, which also requires conditional statements. For example, the operation for "curved outer teeth" is that "if the patch belongs to the teeth, then we apply a shape filter to it"" In literature, the MoE layer is considered effective approach to implement conditional computations (Shazeer et al., 2017; Riquelme et al., 2021). We formalize this intuition the next theorem.
Algorithm I: Conditiona Statements
Define intervals I C R. 1, M Define functions
hi, i 1,... MM +
switch h (3) do if hog (x) E 4 then L apply h441 to
Theorem 2. An MoE module in equation 3 with N experts and k aligns with the conditional statements in Algorithm 1 with
(N + +1) max (M's M(G,
iN < M,
Alignment
(N +1) max
max (iE(1...M)
MGG.
IN M,
(4)
where M(-rs:s-) is defined in Definition I, and M' g is the optimal objective value of the following oplimization problem:
@ minimize II.Ix
max SE(I,..N)
M(fFEN. (H,,, o h1)" [hyleE,,,
(5)
subject to I I 12, 3..3.... ... M + 1},
where 11, is the indicator function interval Iş.
Remark 3. (Interpretations Theorem 2) In algorithmic aligmment, the network better aligns with the algorithm if the alignment value in equation 2 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 aligmment value first 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 + in the product. Therefore, the MoF 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. we find 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.3 ADAPTING MOE TO DOMAIN GENERALIZATION
In literature, there several variants of MoE architectures, e.g., Riquelme al. (2021); Fedus et al. (2022), and we should identify one for DG. By algorithmic alignment, in order to achieve better generalization, the architecture of sparse MoEs should be designed effectively handle visual attributes. In the following, we discuss our architecture design for this purpose.
Routing scheme Linear routers (i.c., 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 al., 2022). For the cosine router, given in- put E Rd, the embedding W æ e Rdd is first projected onto hypersphere, followed by multiplying a learned embedding E d Specifically, the expression for the gate is given by
E Wæ
G(x: TOP; Softmax
W| x
E

---

Published as a conference paper at ICLR 2023
Expert
Expert
Expert
Expert N
Cosine Router
GMME Block
Skin
Attertion
Normalized
Embedding
Embedding
Trunk
Tusks
Teeth
Figure 2: Overview architecture of GMoE. The cosine router distributes normalized image patches of different visual attributes to corresponding experts. Our analysis and experiments (in Section 4.2 and Section 5.4) demonstrate that an expert is potentially responsible for a group of similar visual attributes.
where a hyper-parameter. In the view of image processing, can be interpreted as the cookbook for visual attributes (Ferrari & Zisserman, 2007; Zhou et al., 2014) and the dot product between E and WE with l normalization a matched filter. We opine that the linear router would face difficulty in DG. For example, the elephant image (and its all patches) in the Clipart domain is likely more similar to other images the Clipart domain than in other domains. The issue can alleviated with codebook for visual attributes and matched filters for detecting them. Please refer to Appendix D.6 for the ablation study.
Number of MoE layers Every-two and last-Iwo are two commonly adopted placement methods in existing MoE studies (Riquelme et al., 2021; Lepikhin et al., 2021). Specifically, every-two refers to replacing the even layer's FFN with MoE, and last-two refers to placing MoE at the two even layers. For IID generalization, every-two often outperforms last-two (Riquelme al., 2021). We argue that last-two more suitable for DG as the conditional sentences for processing visual attributes high-level. From experiments, empirically find that last-two achieves better performance than every-two with fewer computations. Please refer to Appendix C. for more discussions and Appendix D.6 for the ablation study.
The overall backbone architecture of GMoB is shown in Fig. 2. To train diverse experts, we adopt the perturbation trick and load balance loss in Riquelme et al. (2021). Due space limitation, leave them in Appendix C.4.
5
EXPERIMENTAL RESULTS
In this section, we evaluate the performance of GMoE on large-scale DG datasets and present model analysis to understand GMoE.
5.1 DOMAINBED RESULTS
In this subsection, we evaluate GMoE on DomainBed (Gulrajani & Lopez-Paz, 2021) with 8 bench- mark datasets: PACS, VLCS, OfficeHome, TerraIncognita, DomainNet, SVIRO, Wilds-Camely and Wilds-FMOW. Detailed information on datasets and evaluation protocols are provided in Ap- pendix D.I. The experiments are averaged over 3 runs as suggested (Gulrajani Lopez-Paz, 2021).
We present results in Table with rrain-validation selection, which include baseline methods and recent 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 leave the results in Appendix D.3 due to space limit. In the lower part of Table we test our methods on three large-scale datasets: SVIRO, Silds-Camelyon, and Wilds-FMOW The three datasets capture real-world distribution shifts
7

---

Published as a conference paper at ICLR 2023
Table 1: Overall out-of-domain accuracies with train-validation selection criterion. The best result highlighted in bold. GMoE achieves the best performance PACS, VLCS, OfficeHome, and DomainNet while ranking in the third-best on Terralncognita.
Algorithm
PACS
VLCS
Officellome
TerraInc
Domain
ERM (kesNet55) (Vapnik, 1991)
85.7 + 0.5
77.4 + 0.3
67.5 - 0.5
47.2 + 0.4
41.2 + 0.2
IRM |ArXiv 20J (Arjovsky et al., 2019)
83.5 +: 0.8
78.5 0.5
64.3
2.2
47.6
0.8
33,9 2.8
DANN [JMLR
(Ganin et al., 2016)
84.6
0.3
68.6
0.4
46.4 0.8
41.8 0.2
CORAI [FCCV 16] (Sun & Saenko, 2016)
[HCCV
86.0 0.2
+..7+005
68.6 0.4
46.4 0.8
41.8 +
MMD [CVPR 18) (Li et al., 2018b)
0.2
67.1=0.1
42.2 1.4
39.4
FISH [ICLR 22] (Shi et al., 2021)
77.8
68.6=0.4
45.1+
42.7 =
SWAD [NcurIPS 21] (Cha
102
8.1
0.1
7.26=0.2
5..0
46.3
Fishr [ICML 22] (Rame ul., 2021)
0.2
0.2
68.6
47.4 1.6
41.7 +.0
MIRO |ECCV 22) (Cha et ul., 2022)
854 + 0.4
7..0
70.5
5044 1.1
44.3
(VTT-S/16) [ICLR 21] (Desovitskiy et al., 2021)
86.2
7.
72.2
42.0
47.3
GMoE-S/16 (Ours)
88.1
0.1
80.2 0.2
74.2 0.4
48.5 0.4
48.7 0.2
Algorithms
SVIRO
Wilds Camelyon
Wilds-FMOW
ERM (ResNet50) (Vapnik, 1991)
85.7 + 0.1
93.1
40.6 - 0.4
ERM (VIT-S/16) [ICLR 21] (Dosovitskiy et al., 2021)
89.6 0.0
91.1
44.8 + 0.2
GMoE-S/16 (Ours)
90.3 + 0.1
93.7 = 0,2
46.6 + 0.4
across a diverse range of domains. We adopt the data preprocessing and domain split in DomainBed. As there is no previous study conducting experiments on these datasets with DomainBed criterion, we only report the results of our methods, which reveal that GMoE outperforms the other two baselines.
5.2
GMoE WITH DG ALGORIHTMM
GMoE's generalization ability comes from its internal backbone architecture, which is orthogonal to existing DG algorithms. This implies that the DG algorithms can be applied to improve the GMoE's performance. To validate this idea, apply loss functions approaches (Fish) and one adopting model ensemble (Swad). The results in Table demonstrate that adopting GMoE instead of ResNet-50 brings significant accuracy promotion to these DG algorithms. Experiments on GMoE with more DG algorithms are in Appendix D.4.
DG algorithms GMoE, including one modifying
5.3 SINGLE-SOURCE DOMAIN GENERALIZATION RESULTS
In this subsection, we create a challenging task, single- source domain generalization, to focus generalization ability of backbone architecture. Specifically, we train the model only data from one domain, and then test the model on multiple domains to validate its perfor- mance across all domains. This is challenging task as we cannot rely on multiple domains to identify invariant correlations, and popular DG algorithms cannot be ap- plied. We compare several models mentioned in above analysis (e.g., ResNet, VIT, GMoE) with different scale of parameters, including their float-pont float-point- -opeattttttotn per
Table 2: GMoE trained with DG algo- rithms.
Algorithm
Domain
GMoE
48.7
Fish (Rame et al., 2021)
42.7
GMoE w/Fish
48.8
Swad (Cha et al., 2021a)
46.5
GMoE w/Swad
49.6
second (flops), IID and OOD accuracy. From the results in Table 3, we see that GMoB's OOD generalization gain over ResNet ViT is much larger than that in the IID setting, which shows that GMoE suitable for challenging domain generalization. Due to space limitation, we leave experiments with other training domains in Appendix D.5.
5.4 MODEL ANALYSIS
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., 1). 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 invariant representation (e.g., DANN). The results are Appendix E. which demonstrate superior performance of GMoE compared with other DG algorithms. CUB-DG datasets provides rich visual attributes (e.g. beak's 's shape, belly's color) for each
8

---

Published as a conference paper at ICLR 2023
Table 3; Single-souree DG accuracy (%). Models are trained on Paint domain and tested (1) on Paint validation set (2) the rest 5 domains' validation sets on DomainNet. The flops are reference values to compare the model's computational efficiency. IID Imp. denotes IID improvement on Paint's validation set comparing with ResNet50. OOD Imp. denotes average OOD improvement across
test domains.
MACs
Paint
Clipart
Info
Paint
Quick
Real
Sketch
IDD Imp.
OOD Imp.
4.1G
ResNet50
37.1
12.9
62.7
2.2
49.3
33.3
7.9G
ResNet101
40.5
13.1
63.4
3.1
51.2
35.4
1.1%
12.4%
4.6G
VIT-S/16
42.7
15.9
69.0
5.0
37.0
10.0%
38.6%
4.8G
GMoE-S/16
43.5
69.3
5.3
38.0
10.5%
42.3%
image. That additional information enables us to measure the correlation between visual attributes and the router's expert selection.
Visual Attributes Experts Correlation We choose GMoE-S/16 with experts 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 first match a visual attribute with its nearest patches, and then correlate the visual attribute and its current 9 patches' experts. Fig. 3(b), show the 2D histogram correlation between the selected experts and attributes. Without any supervision signal for visual attributes, the ability to correlate attributes automatically emerges during training. Specifically, each expert focus distinct visual attributes; and 2) similar attributes attended by the same expert (e.g., g., the left wing and the right wing are both attended by e4). This verifies the predictions by Theorem 2.
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, then visualize the top-1 selections for each patch in Fig. 3(c). In this image, we mark the selected expert for each patch with black number and draw different visual attributes of the bird (e.g. beak and tail types) with large circles. We see the consistent relationship between Fig. 3(b) and Fig. 3(c).
process image into 16 x 16 patches), and
In detail, we see (I) experts and are consistently routed with patches in the background area; (2) the left wing, right wing. beak and tail areas are attended by expert 3; (3) the left leg and right leg areas are attended by expert 4. More examples are given in Appendix E.
badigrunn
thooa
tai
whn
right log
Left
Natural
Candy
RighrWing
cmne
breed
beck
Mossic
Udiile
60
01
02
(a) CUB-DG examples.
(b) Visual attributes.
(c) Expert selection.
Figure 3: (a) Examples of CUB-DG datasets from four domains. (b)The y-axis corresponds to 15 + attributes (15 visual attributes + background). The z-axis corresponds to the selected expert id. (c) Finetuned GMoE's router decision of block-10. The image is from CUB-DG's natural domain.
6 CONCLUSIONS
This paper is an initial step in exploring the impact of the backbone architecture in domain gener- alization. We proved that a network is more robust to distribution shifts if its architecture aligns well with the invariant correlation, which is verified on synthetic and real datasets. Based on our theoretical analysis, proposed GMoE and demonstrated its superior performance DomainBed. As for future directions, interesting to develop novel backbone architectures for DG based on algorithmic alignment and classic computer vision.

---

Published as a conference paper at ICLR 2023
7 ETHICS STATEMENTS
Our paper suggests a new direction for Domain Generalization. We think this contribution can direct people to critically select model architecture in situations that require enhancing the model's out-of- domain generalization ability. We validated our conclusions under certain assumptions mentioned in the paper and empirically validated our findings in the image classification scenario. However, when applying our model or conclusions to other scenarios and specific usages, they still need careful validation before being actually applied. Besides, we did use any non-public data, unauthorized software, or API in our paper, there are no privacy or other related ethical concerns.
8 ACKNOWLEDGMENTS
This research/project is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: AISG2-PhDD-222--1--223 Besides, this project is also supported by NTU NAP, MOE AcRF Tier 2 Industry Alignment Fund Industry Collaboration Projects (IAF-ICP) Funding Initiative, well as cash and in-kind contribution from the industry partner(s).
and under the RIE2020
We also thank Yushi Lan, Quanzhou Li, Ziqi Huang. Siyao Li, Runqi Pan, and Wenzhen Zhu for their help in revising the manuscript. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NTU or the Singapore Government.
REFERENCES
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, loannis Mitliagkas, and Irina Rish. Invariance principle meets information bot- tleneck for out-of-distribution generalization. In Marc' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Infor- mation Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 3438 3450, 2021. URL https://proceeding nww.ips...... cc http/pppee////aa////ee Abstract html. 18, 19
Martín Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. CoRR, 2019. URL http: www.ii.ogg/aa.a..aaaa..ao.
907. 02893. 1, 8, 18, 19, 30
Devansh Arpit, Huan Wang, Yingbo Zhou, and Caiming Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. CoRR, 2021. URL https:
//arxiv. http/bbb//2111 10832. 18
Haoyue Bai, Fengwei Zhou, Lanqing Hong, Nanyang Ye, S-H Gary Chan, and Zhenguo Li. Nas-ood: Neural architecture search for out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. $320-8329 2021. 19
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In The European Conference on Computer Vision (ECCV), 2018. doi: 111100779778-3000-0277---00 httpps://e
URL
http//ii..o.....///// 10.00.110//9900 3- 030 01200- 0_28. 28
Manh-Ha Bui, Toan Tran, Anh Tuan Tran, and Dinh Phung. Exploiting domain-specifie features to enhance domain generalization. CORR, 2021. URL
09410. 1, 19
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Scunghyun Park, Yunsung Lee, and Sungrae Park. SWAD: domain generalization by seeking flat minima. In Advances in Neural Infor- mation Processing Systems, 2021a. URL https: /ttp:///oceeeeeeeeeo ngs. neurips. cc/ http///pr//////// 8a 17 60 Abst-aatt.tttt.aaa.... html.. 8, 30
h02l1hhahh/bbbr ccde4 3.3363333666
Junbum Cha. Sanghyuk Chun, Kyungjae Lee, Han-Cheu Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systemx, 2021b. 2, 4, 19, 22, 29
10

---

Published as a conference paper at ICLR 2023
Junbum Cha, Kyungjae Lee, Sungrae Park, and Sanghyuk Chun. Domain generalization by mutual- information regularization with pre-trdined models. CoRR, 2022. doi: 10448500/anXxv22003.17339 URL http::::/dooooogg////// 2t55:110..555555 48550/ aarviv 2203.10789, 8, 29, 30
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding mixture of experts in deep learning. CoRR, abs/2208.02813, 2022. doi: 1448550/aXxv.2228.000 2201.0281.. URL
https: /tw.///oi.io/////////oooo wrg/10../100.................... 48550 arxvv 2208 0.8.2.2133. 26
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, and Furu Wei. On the representation collapse of sparse mixture of experts. CoRR, abs/2204.09179. 2022, doi: 10.48550/anX1..2224..99
URL https: org/10.
48550/ 2204. 091099179 99.
Gordon A. Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world. tion, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pp. 6172-618 6180. Com- puter Vision Foundation / IEEE Computer Society, 2018. doi: 10.1109/CVPR.2018.000 URL http: httt::://ppmaaae http///paacccssscteee thecf com/ content....oooooooo ceattten___p____e___ _evp__0111111- 8 Functi functionall onal _______aa_________
2018 IEEE Conference on Computer Vision and Pattern Recogni-
__paper html. 28
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. Multi-head attention: Collaborate instead of concatenate. CoRR, 2020. URL htt.t:tt////atetee www.wwwwww.aaaaaaaaeaaeeeaoreeee hwtg//ab//// absaabas00000 16362. 34
Steve Dias Da Cruz, Oliver Wasenmüller, Hans-Peter Beise, Thomas Stifter, and Didier Stricker. SVIRO: synthetic vehicle interior rear seat occupancy dataset and benchmark. In IEEE Winter Conference on Applications of Computer Vision, WACV 2020, Snowmass Village, CO, USA, March 1-5, 2020, pp. 962-971. IEEE, 2020. doi: 10. https: / /oo. org/ //rgg//// 10. 1109 /WACV45572 2020. 9093315. 28
URL
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021. URL https:////// http///pperreeeeeeee/ee/eeeeeeet httpreewwwwweeeeeeeeeeeeeeeeeeee 19, 30
net/ wwteefottttee forum? d==Yi cbFdNTTy. 2, 8,
Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnosie learning of semantic features. Advances in Neural Information Processing Systems, 32, 2019. 19
Chen Fang, Ye Xu, and Daniel N. Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In IEEE International Conference Computer Vision, ICCV 2013, Sydney, Australia, December 1-8, 2013, pp. 1657-1664. IEEE Computer Society, 2013. doi: 10. 10.1109/ICCV2201 3.80 URL https:/ doi orrilorg/11i /0.. 10..0.11099 111111199/000 IIE////CV.... 2013.208 28
William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models in deep learning. arXiv preprint arXiv:2209.01666 2022. 2, 6
Vittorio Ferrari and Andrew Zisserman. Learning visual attributes. Advances in neural information processing systems, 20, 2007. 3.
Yaroslav Ganin, Evgeniya Ustinova. Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor S. Lempitsky. Domain-ad aabaainnadversariiin training of neural networks. J. Mach. Learn. Res., 2016. URL http:://tmm r. org http//ppprrr////e/ v17 /atesr///vv////13 html. 1, 8, 18, 30, 34
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. amagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv preprint 1xXiv:1811..222 2018. 5
Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. In Advances in Neural Information Processing Systems 34, 2021.

---

Published as a conference paper at ICLR 2023
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In 9th International net/
Conference on Learning Representations, ICLR, 2021. URL Eorum??d=110xeeDDoo
3, 4, 78, 29
Jiang Guo, Darsh J Shah, and Regina Barzilay. Multi-source domain adaptation with mixture of experts. arXiv preprint 18X9v:1809.02255 2018. 19
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016.
Judy Hoffman, Mehryar Mohri, and Ningshan Zhang. Algorithms and theory for multiple- source adaptation. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kris- ten Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Pro- cessing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 8256 8266, 2018. URL https: httt:p::/procceding .wwiinngsneueeee neurips hccppaper/2018.. 8/has// (22000766348233822 www.:aaadaaaaaaaaaaaaaaaaaaaaaae Abstracct html.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-exciaation networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018. 1
Yu Hen Hu, Surekha Palreddy, and Willis Tompkins. patient- adaptable ecg beat classifier using a mixture of experts approach. IEEE transactions on biomedical engineering, 1997. 19
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Comput., 1991. doi: 10.1162/neco. 1991.3.1.79 URL htwps:tt/////ctt //dbi. org/10 11 62 neco. 991.3.1.1
19
Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Comput., 1994. doi: 2..1162/neco.19944..2. URL https: http///i.....oooo//oo////o 1162/
neco. 1 1994.9999.... 2. 19
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: //arxiv. org/ aba/ 11t/211122.0 6980. 29
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsub- ramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran Haque, Sara M. Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 5637-5664. PMLR, 2021. URL http: http///ooceeeennee////ee mmlr. press v13 wwhhooh2lllaa html. 28
David Krueger, Ethan Caballero, Jörn-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Rémi Le Priol, and Aaron C. Courville. Out-of disttiiiittiittoon generalization via risk extrapolation (rex). In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 5815-5826, PMLR, 2021. URL http: / /tppppeeeennn.... mlr. wwwpspessssseeee (1133///Kugrre20 html. 1, 18
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In 9th International Conference Learning Representations, ICLR, 2021. URL https: openreviei net http:/ewwwee//oeteteeeeeeee www.mmiidugaea.. 7XHTmYb.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. In IEEE International Conference Computer Vision, ICCV 2017. Venice, Italy, October 22-29, 2017, 5543 IEEÉ Computer Society, 2017. doi: 10. 19.100.0.........5911
URL htw.t: /oi .org 1109 IC/V 201 2017..00 591. 28
12

---

Published as a conference paper at ICLR 2023
Da Li, Yongxin Yang, Yi-Zhe Song. and Timothy M. Hospedales. Leamning to generalize: Meta- learning for domain generalization. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, 2018a. URL https: /twp///wwww/..////////////ww aaai. org/ wctp//c///aa// indexi.exx//.// php/AAAI
AAAT18 pwwppapprre/ee//ea//pep /viee////ww/w/i/eeeeeeee 6067. 19, 30
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C. Kot. Domain generaliza- tion with adversarial feature learning. sion and Pattern Recognition, 2018b. doi: 10.1009/CVPP.2218.00555 twwccheeeveeccc com/ wooteet_cvpp__2000 201 8/html /Li __mmain________
2018 IEEE Conference on Computer Vi- URL
paper. peaapr.html......... 8, 18,
Jingling Li, Mozhi Zhang, Keyulu Xu, John Dickerson, and Jimmy Ba. How does neural network's architecture impact its robustness to noisy labels? Advances in Neural Information Processing
Systems, 34:9788-999999 -9038,9903,--0 2021. 2, 20
Yiying Li, Yongxin Yang. Wei Zhou, and Timothy Hospedales. Feature-critie networks for heteroge- neous domain generalization. In International Conference on Machine Learning, pp. 3915-3924.
PMLR, 2019. 19
Ziyue Li, Kan Ren, Xinyang Jiang, Bo Li, Haipeng Zhang, and Dongsheng Li. Domain generalization using pretrained models without fine-tuning. CoRR, 2022. doi: 10.485500axXxv.20550000000 URL https: doii www.oogg.o.o.....oo.oooo.o...... htegg/1..8550.00 48550/ arxiv. 2203. 04600.
Alexander H. Liu, Yen-Cheng Liu, Yu- Yu-Ying Ying Ych, and Yu-Chiang Frank Wang. A unified feature disentangler for multi-domain image translation and manipulation. In Advances Neural Infor- mation Processing Systems, 2018. URL https: httpttpp:pppppcpeeinggeeo neurips. cc/ popppper////// 2018 hww.//hahh/aaa 8 43887 aae 55 A
html. 19
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In 2021 IEEE/CVE International Conference on Computer Vision, ICCV, 2021. doi: 10.11099CCV44992.221.0000 URL wttph:tpp//ccc doi http://..........1..////////ee// 2ttp///88222222....220 2021. 1,
Yi Ma, Doris Tsao, and Bleung-Yeung Shum. On the principles of parsimony and self-consistency for the emergence of intelligence. Frontiers of Information Technology & Electronic Engineering, pp. 1-26, 2022.
Massimiliano Mancini, Samuel Rota Bulò, Barbara Caputo, and Elisa Ricci. Best sources forward: Domain generalization through source-specific nets. In 2018 IEEE International Conference on Image Processing, 2018. doi: 10. 109//CCP.2018.85553111 URL https: /ttp://ooooooooooooooooo ort/ 1109/ ICIP 2018 8451318. 8. 1, 18
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement 2017.
testing sprites dataset.
Samuel G Müller and Frank Hutter. Trivialaugment: Tuning-lree yet state-of-the-art data augmenta- tion. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 782,782, 2021. 31
Basil Mustafa. Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multi- modal contrastive learning with limoe: the language-image mixture experts. arXiv preprint 2rXi::2206.02770. 2022. 26
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur: Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXiv:2010.15775, 2020.
Narges Honarvar Nazari and Adriana Kovashka. Domain generalization using shape represen- tation. Proceedings of European Conference on Computer Vision ECCV 2020 Work- shops, 2020. doi: 10. 00077978-3-030-664111 978- 3- 030 66415 2_ 45. 19
URL http::::/dooiorroo 1007/
Namuk Park and Songkuk Kim. How do vision transformers work? arXiv preprint arXiv:2202.06709,
2022. 5
13

---

Published as a conference paper at ICLR 2023
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In 2019 JEEE/CVF International Conference on Computer Vision, ICCV, 2019. doi: 10.1 10.110//CCC2222900004449 URL httpt::// htttts::/do...og/2..1111 ICCV 201 9 00149. 18, 28
Fengchun Qiao, Long Zhao, and Xi Peng. Learning to learn single domain generalization. In 2020 IEEE/CVF Conference on Computer Vision and Putlern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, 2020. doi: 10. 1199CCVPR2200022200002222
URL https://ddo
org /15 1 109 CVw.42202200000 2020 01257. 19
Yao Qin, Chiyuan Zhang, Ting Chen, Balaji Lakshminarayanan, Alex Beutel, and Xuezhi Wang. Understanding and improving robustness of vision transformers through patch-based negative augmentation. CoRR, 2021. URL hwtph::tt/aatiiiieee wttp://orrgaaab//////ee
07858. 19
Nasim Rahaman, Muhammad Walced Gondal, Shruti Joshi, Peter Gehler, Yoshua Bengio, Francesco Locatello, and Bernhard Schölkopf. Dynamic inference with neural interpreters. Advances in Neural Information Processing Systems, 14410095--1999 2021. 19
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. arXiv preprint arXiv:2109.0288 2021. 2, 4, 8, 30, 34
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize imagenet? In International Conference on Machine Learning, pp. 5389-5400. PMLR. 2019. 1
Matthew Riemer, Ignacio Cases, Robert Ajemian, Miao Liu, Irina Rish, Yuhai Tu, and Gerald Tesauro. Learning to learn without forgetting by maximizing transfer and minimizing interference. In Proceedings of The 7th International Conference on Learning Representations, 2019. URL https: /tpps://openrevvv
net/ www.:/rrtoooummmooe i d= gTShAct
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, An- dré Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Advances Neurul Information Processing Systems, NeurIPS 2021, December 6-14, 2021, 2021. URL http:::s//poocceeedeeneoo neurss cc/ paper r2eee12022111 /hash/ 48237d9
-bb-taa99-3-bbbbaactaaaea html. 6, 7, 19, 25, 26
Mattia Segü, Alessio Tonioni, and Federico Tombari. Batch normalization embeddings for deep domain generalization. CoRR, 2020. URL httpt::///aaat http///xww...gg/a////////// /tt.
12672. 18
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-cxperts layer. In Sth International Conference Learning Representations, ICLR, 2017. URL https:
/ppeneeview net/ ftr.r//mmwee//eeumtmmmm (etm:ndddd11ldpppp. 2, 5, 6, 19, 26
Yuge Shi, Jeffrey Seely, Philip H. S. Torr, N. Siddharth, Awni Y. Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. CoRR, 2021. URL https:
//arxiv. wwg./vvvorgaaaooo aatpaaaas/2000 2104. 09937. 2, 4, 8, 30,
Sarath Sivaprasad, Akshay Goindani, Vaibhav Garg, and Vineet Gandhi. Reappraising domain generalization in neural networks. arXiv preprint arXiv:2 10.07981, 2021. 20
Baochen Sun and Kate Saenko. Deep CORAL: correlation aligmment for deep domain adaptation. In Computer Vision - ECCV 2016 Workshops, 2016. doi: 10.1097/97 3-319-494099-99-- 35. URL
https://doi org 1007 978- 319 49409-9-09211 8_199-8__5... 35. 18, 30
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826, 2016. 31
Jun Tani and Stefano Nolfi. Learning to perceive the world as articulated: approach for hierarchical learning in sensory-motor systems. Neural Networks, 1999. 19
Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou. Fixing the train-test resolution discrepancy. Advances in neural information processing systems, 32, 2019. 31
14

---

Published as a conference paper at ICLR 2023
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transforerrs & distillation through attention. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp- 10347-10357. PMLR, 2021. URL http: htttt::/proceeeiingge mlr. htt.rronnnaa.m......... html. 31
Leslie G Valiant. A theory of the learnable. Communications of the ACM, 27(11):1133-1144 1984. 19
Vladimir Vapnik. Principles of risk minimization for learning theory. Advances in neural information processing systems, 4, 1991. 8, 30
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017. URL https: /ttp://e///cccecigssss// neurips cc wttpscc.ppppppppppeeeeeeeeeeeeee paper/2011
Bashh3f5ee243330 7dee fbd0 (a33:cc444445a00 3aa.aaa222.44.eeeeeoooooocccaaa. Abstract htmm. 19
Hemanth Venkateswara, Jose Fusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pp. 5385-5394. IEEE Computer Society, 2017. doi: 10.1109/CVP.20017..22. URL https: http:::://do...rg..../ 20.09//CP.2.0 2017.511550 572. 28
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-uesd birds-200-2011 dataset. California Institute of Technology, 2011. 8,29
Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorek, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, et al. Image as foreign language: Beit pretraining for all vision and vision-language tasks. arXiv preprint arXiv:2208. 10442, 2022.
Yufei Wang, Haoliang Li, Lap-Pui Chau, and Alex C. Kot. Variational disentanglement for domain generalization. CoRR, 2021. URL https: /t////aree arrivv.vvvv
www.aawwwwwww......rraaaa.t.aaee wwt./oogg///........ 22199 2109. 05826.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Alvise-Rebuffi, Ira Ktena, Taylan Cemgil, et al. A tne-grained analysis on distribution shift. arXiv preprint arXiv:21 2021. 1, 2, 20
Tengyang Xie and Nan Jiang. Batch value-function approximation with only realizability, In International Conference on Machine Learning, pp. 11404-11413, PMLR, 2021. 25
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon Du, Kenichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? In Proceedings of International Conference on Learning Representation, Apr. 2020a. 2, 3, 4, 20
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint 1rXiv:2009.11845 2020b. 4, 19, 25
Haotian Ye, Chuanlong Xie, Tianle Cai, Ruichen Li, Zhenguo Li, and Liwei Wang. Towards theoretical framework
disttribuitiiino generalization. Advances in Neural Information Processing Systems, 34:23519-23533 2021. 19
Xiangyu Yue, Yang Zhang, Sicheng Zhao, Alberto L. Sangiovanni- Vincentelli, Kurt Keutzer, and Boqing Gong. Domain randomization and pyramid consistency: Simulation-to- generalization without accessing target domain data. In 2019 JEEE/CVF International Conference on Computer Vision, 2019. doi: 10. 1001100/C/C222210002222 URL httph:ttt//// http:/di...ogg//////// 1109 ICCV. 201 9. 00219. 9. 19
Sergey Zakharov, Wadim Kehl, and Slobodan Ilie. Deceptionnet: Network-driven domain ran- domization. In 2019 IEEE/CVF International Conference on Computer Vision, 2019. doi: 10.1 10.1109/CCV.221990000 URL.
/ http:://///oooo/////.//////
2006.0002..
00062. 19
15

---

Published as a conference paper at ICLR 2023
Chongzhi Zhang, Mingyuan Zhang, Shanghang Zhang, Daisheng Jin, Qiang Zhou, Zhongang Cai, Haiyu Zhao, Xianglong Liu, and Ziwei Liu. Delving deep into the generalization of vision transformers under distribution shifts. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 7277-7286, 2022. 19
Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, and Aaron C. Courville. Can subnetwork generalization? Proceedings of International Con-
structure be the key ference on Machine Learning. 2021a. URL http::/ mlr. press/v139/ zhang21a a. html.
Hanlin Zhang, Yi-Fan Zhang, Weiyang Liu, Adrian Weller, Bernhard Schölkopf, and Eric Xing. Towards principled disentanglement for domain generalization. CoRR, 2021b. URL https:
//arxivv...////////// org/ aba/ 11811
Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: Learning to adapt domain shift. Advances in Neural Information Processing Systems, 34:23664-2366 2021c. 1, 19
Han Zhao, Remi Tachet des Combes, Kun Zhang, and Geoffrey J. Gordon. On learning invariant representations for domain adaptation. Proceedings of International Conference on Machine Learning, 2019a. URL http: http:::://pcceedeeegeee
www.mprs..es......mmmmmmmm..m..e / / zhaol 9a httml. 18
Sicheng Zhao, Bo Li, Xiangyu Yue, Yang Gu, Pengfei Xu, Runbo Hu, Hua Chai, and Kurt Keutzer. Multi-source domain adaptation for semantic segmentation. In Advances in Neural Information Processing Systems, 2019b. URL htt.h:ttp//// /ptpceees///coceinnggssss neurips co httpppppr/211// hash/ bashhddbaade dbddd5 6c 71 61 4eteeedd2772311111555554 -5000456b------ html. 19
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmenta- tion, arxiv. arXiv preprint arXiv: 04886, 2017. 31
Bolei Zhou, Aditya Khosla, Agata Lapedriza. Aude Oliva, and Antonio Torralba. Object detectors emerge in deep scene cnns. arXiv preprint arXiv: 1412.6856.. 2014. 7,
Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang. Sparse invariant risk minimization. In International Conference on Machine Learning, pp. 27222-27244. PMLR, 2022. 18, 19
Liu Ziyin, Tilman Hartwig, and Masahito Ueda. Neural networks fail to learn periodic functions and how to fix it. Advances Neural Information Processing Systems, 33:1583- 1594, 2020. 19
16

---

Published as a conference paper at ICLR 2023
CONTENTS
1
Introduction
1.1
Motivations
1.2
Contributions
2
Preliminaries
2.1
Notations
2.2
Attribute Factorization
2.3
Algorithmic Alignment
3
On the Importance of Neural Architecture for Domain Generalization
3.1
A Motivating Example: CNNs versus Vision Transformers
3.2
Understanding from Theoretical Perspective
4
Generalizable cexture-of-Experre for Domain Generalization
4.1
-ixtere-of-Experts Layer .ee.r...........................
4.2
Visual Attributes, Conditional Statements, and Sparse MoEs
4.3
Adapting MoE to Domain Generalization
Experimental Results
5.1
DomainBed Results
5.2
GMoE with DG Algorithms
5.3
Single-Source Domain Generalization Results
8
5.4
Model Analysis
Conclusions
Ethics Statements
10
8 Acknowledgments
10
Related Works
18
A.I Domain Generalization
18
A.2 Vision Transformers
19
A.3 Sparse Mixture-of-Experts
19
A.4
Out-of-Distributioo Generalization Theory
19
A.5 Algorithmic Alignment
20
A.6 Empirical Investigations of Distribution Shift
20
Theorems and Proofs
20
B.1 Proof of Theorem
20
17

---

Published as a conference paper at ICLR 2023
B.2
Synthetic Experiments for Theorem
21
B.2.1
Noiseless Version
B.2.2
Noisy Version
22
B.3 Proof of Theorem
24
B.4 Technical Lemmas
24
Implementations
25
C.1 Detailed Architecture
25
C.2 Model Initialization
26
C33 Perturbation
26
C.4 Auxiliary loss
26
Experimental Results
28
D.1 DomainBed Details
D.2 Training Details Hyperparameters
................................
D.3 Leave-one-domain--uu Results
29
D.4 Ablation Study with DG Algorithms
D.5 Single-source Domain Generalization Results
30
D.6 Ablation Study on Model Design
................................
D.7 Ablation Study on ResNet with Strong Data Augmentation
31
D.8 Computational cost comparison
Model Analysis
34
E.1 CUB-DG Results & Visualizations
34
A RELATED WORKS
A.1 DOMAIN GENERALIZATION
Domain Generalization (DG) aims to maintain the good performance of machine learning models in the domains that different from the training (source) domain. The following the categories of mainstream domain generalization research.
(1) Invariant Learning: Aligning domain distributions and finding invariance across domains has been often studied with empirical results and theoretical proof's (Ganin et al., 2016; Zhao et al., 2019a). Specifically, 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., 2019), etc. Besides aligning feature distributions, Arjovsky al. (Arjovsky et al., 2019) proposed IRM to learn an ideal invariant classifier on top of the representation space, which has inspired many follow-up works (Krueger al., 2021; Ahuja et al., 2021: Zhou al., 2022). These studies implement invariant learning via loss function designs. In this paper, show that if we have suitable backbone architecture, ERM is sufficient to find the invariant correlations, which indicates that the backbone architecture is important loss function design.
(2) Ensemble Learning and Meta-learning: Apart from learning invariant features and correlations across the domain, model-specific or domain-specins specific information helps improve the performance DG datasets. The ensemble learning methods combine different models to exploit model-specific information (Mancini al., 2018; Segù et al., 2020; Arpit et al., 2021; Li et al., 2022). In addition,
18

---

Published as a conference paper at ICLR 2023
SWAD (Cha et al., 2021b) inhibits models from being overfit to local sharp minima by averaging model weights below a validation loss threshold. Meta-learning- based approaches (Li et al., 201 8a; 2019; Dou et al., 2019; Zhang et al., 2021c; Bui et al., 2021) leverage domain specifisepccceeee feature to improve the accuracy on each domain.
(3) Data Manipulation: Diverse training data helpful for improving generalization, and re- searchers have proposed different manipilatiov/anmmemttioo techniques (Nazari & Kovashka, 2020; Riemer et al., 2019), domain randomization (Yue et al., 2019; Zakharov al., 2019). Furthermore, line of works (Qiao et al., 2020; Liu et al., 2018: 8: 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 DG settings: 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 find 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.2 VISION TRANSFORMERS
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 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 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 -f-eeeert--esseeeeeeeeeeeee to enhance the performance of ViT models.
A.3 SPARSE CITTUURE-O-EXXPPREE
Eexture-of-Experri models, 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 plethora of methods for a wide ranges of applications (Hu et al., 1997: Tani & Nolfi, 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 on-par with the standalone counterpart. They were considered promising ways to scale up vision models (Riquelme et al., 2021). Some existing studies have applied MoE-type architecture generalization tasks. (Guo al., 2018) considers the muti-source domain adaptation in NLP, where one classifier is trained on each domain and a MoE is adopted to ensemble the classifiers. (Rahaman al., 2021) investigates the systematic generalization problem and proposes to dynamically compose experts instead of selecting experts in MoEs. In this paper, we consider the domain generalization (DG) problem and prove that a backbone more aligned to invariant correlations is more robust to distribution shift. Based the theory, we repurpose MoE for DG tasks and make several modifications for performance enhancement.
A.4 OUT-OF-DIITTRBUUTON GENERALIZATION THEORY
The classic PAC learning (Valiant, 1984) is only applicable to III settings and recently there have been 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
19

---

Published as a conference paper at ICLR 2023
existing works, this paper studies the situation where the invariant correlation is arbitrary nonlinear function, and one of our main focuses is to justify why these assumptions hold in DG of vision tasks.
A.5 ALGORITHMIC ALIGNMENT
The algorithmic alignment framework (Xu et al., 2020a) was originally proposed to understand the effectiveness specialized network structures in reasoning tasks. Specifically, 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 that if they align well, the neural network only needs to learn simple functions, which leads to better sample efficiency. This framework was further extended
Li et al. (2021) to investigate the impact of backbone architectures on the robustness of noisy labels. This paper adopts algorithmic alignment to provide novel perspective of DG and develops a new architecture based on it.
A.6 EMPIRICAL INVESTIGATIONS OF DISTRIBUTION SHIFT
Prior to this work, some empirical studies (Wiles et al., 2021; Sivaprasad et al., 2021) systematically investigate the impact of the data argumentation, the optimizer, the backbone architecture, and DG algorithms on domain generalization datasets, drawing various conclusions. Nevertheless, they did not explain when a specific model will succeed DG and how to design the backbone architecture to improve DG performance. As far we know, our paper is an initial attempt to theoretically investigate these questions, which provides a novel perspective of DG. In this paper, we prove that backbone is more robust to distribution shift if it aligns with the invariant correlation in the dataset and better alignment leads to better performance in DG. A novel architecture is designed based on theory. We believe these results will encourage further research in this direction.
B THEOREMS AND PROOFS
B.1 PROOF OF THEOREM 1
We first state the algorithmic alignment theorem in the IID setting and then extend it the DG setting.
Theorem 3. (Alignment improves IID generalization; (Xu et al., 2020a)) Fix e and 8. Given target function g and a neural network N, suppose { are i...d. samples drawn from some distribution D, and let yı g(æ:).. Assumptions:
(a) Algorithm stability. Let A be a learning algorithm for N Ws Suppose J A((as,yi)) and A((2, v.1). For any x, [lf( - f(x)) S Lo -max; ||x; - Ezll. where D; is the i-th coordinate of x.
f1(( (")) with (b) Sequential learning. We train N; sequentially: N1 has input samples {a (
obtained from the training dataset. For 1, the inputs a for are the outputs from the previous modules, but labels are generated by the correct functions fi-1''' ".. f1 on (1)
(c) Lipschitzness. The learned functions f, satisfy (1f)(z)) - fz()) < Lı||æ xb| for some L1. Under assumptions (a)(b)cc, Aligmmen(N..g.e. S M implies there a learning algorithm A such that
P Pa- Pa-DIIIN(a) (1())) < O(e)) 2 O(8),
where N is the network generated by on the training data {I;; y/ M
Remark 4. (Explanations of Assumptions in Theorem 3) The first and third assumptions are common assumptions in machine learning and are practical. The second assumption is impractical as 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.
20

---

Published as a conference paper at ICLR 2023
Theorem 4. (Impact of Backbone Architecture in DG) Let Let (N2. ... .N.). Assuming train the neural network with ERM, and Assumption 2, 3, we have the following statements:
1. If Alignment( 9 e, 6) & Etrl we have PD. (x)) yl = O(e)) -
2. If Alignment(N 6, 6) S [Err we have Pr.. [[]) wll w()) > 1 1
Proof. The proof is mainly based on Theorem 3. We tackle the distribution shift in with Lemma 1 and analyze the distribution shift in with Assumption 2,
First condition: From Theorem 3,
P Paw (=(V( 0(0) s (08).
By Lemma I, we have
Paw [[)) (c(N(())) & O())] & P..-D...IlVi -x(x)) < 0())) & 0(6). (6)
Combing with the second condition of Assumption 2, we have
D.. [I)) - .(=))l - (e))
PP
- s O(e)] I.D..IIN(()) (c(Ni( s O())
211 - 8) - = - 0(8),
where the first inequality follows Oee) + O(e) = O()), and the second inequality follows the second condition of Assumption 2 and equation 6.
Second condition: Following the proof of equation 6, we obtain
Pp.. (IlV)()))) g. (Mi s O()) 1 O()).
(7)
Combing with the second condition of Assumption 3, we have
PD.. (=)) - s(x))1 > w()]
PP.. (N1 (x)) - -(x) > w(e)] .P... 9. < O())
21 - 8) (1)) -
where the first inequality follows w()) - 0 06)) and the second inequality follows the second condition of Assumption 3 and equation 7. This finishes the proof for Theorem 1.
B.2 SYNTHETIC EXPERIMENTS FOR THEOREM
In this subsection, we set up synthetic experiments to illustrate Theorem
B.2.1 NOISELESS VERSION
We first set up experiments following the assumptions in Theorem 1.
Training dataset generation: Consider data pairs (x, v) (RK)P x (1,-.-, K} in the training and validation datasets generated as follows.
Generate label E (1...., K} uniformly.
Generate K mutually orthogonal feature vectors {cx} such that ck R and ||cx|la 1.
Generate collection of P patches: = (z) .x( E (IR))P
Pixel-level feature. For the first patch, the y-th pixel as 1 and other pixels drawn from (0, 1).
Pettch-eeee feature. Uniformly select one and only one patch given by cy, Random noise. The rest patches are Gaussian noise drawn independently from (O, Ix)-
21

---

Published as a conference paper at ICLR 2023
100
80
Accuracy
Accuracy
Validation Test1
60
Validation
Test1
Test2
Test2
10
30 Training Epoches
20
50
10
20
40
30
Training Epoches
(a) MLP's performance.
(b) FCN's performance.
Figure 4: Performance of the MLP and FCN on the synthetic dataset. ML.P aligns with the invariant correlations test dataset while FCN aligns with the invariant correlations in test dataset
The generation of the validation dataset follows that the training dataset.
Test dataset generation: Two OOD test datasets given. The first dataset differs from the training dataset in the pixel-level feature. Specifically, we set the (y + mod K-th pixel of the first patch as 1 and other pixels in the first patch follow (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)mmod K i.e., making patch-level features spurious.
Backbone architectures: We adopt MLPs and fully convolutional networks (FCNs) as ML.ss align with pixel-level features while FCNs align with patch-level features. It is clear that both MLPs and FCNs can fit 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 three-layer MLP with hidden size (40, 100, 100, 4} and a two-layer FCN with 20 filters.
The experimental results are shown in Fig. 4. The MLP completely fails on the second dataset while the FCN completely fails on the first dataset, which consistent with Theorem
B.2.2 NOISY VERSION
In the previous experiment, both pixel-level features and patch-level features can perfectly fit the label. We give noisy version follows:
1. Pixel-level feature. Given probability value Pr. With probability Pt, the y-th pixel of the first patch is set as 1 and other pixels are drawn from N 1). With probability 1 P1. the y/tth (y/ f y) pixel of the first patch set as and other pixels are drawn from N 1).
2. Patch-level feature. Given probability value p2- With probability P2, uniformly select one and only one patch given by Cy. With probability - P2, uniformly select one and only patch given by Caf' (y/ + v).
The experiments with different P1 and Pa are shown in Fig. 5. 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. 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 and P2 is large, the network will still learn its aligned correlation at the early stage (with above 90% accuracy Testl), but may change when the number of epochs is large. This suggests that early stopping is beneficial to DG as suggested in Cha et al. (2021b) under the condition that the backbone aligns some invariant correlaiinns.
22

---

Published as a conference paper at ICLR 2023
Accursoy
Accurery
250
200
22
150
126
Tanning Spoches
Thaning Epoches
Beiiinn Epoccess
(a) P1 0.9, P2 0.9.
(b) P1 0.8, P2 0.8.
(c) P1 0.7, P2 0.7.
60-
Tex2
Tainnn Epoches
Training Epoches
Traiiin Epoches
(d) P1 0.9, P2 0.95.
(e) P1 0.8, P2 0.85.
(1) p1 0.7, 0.75.
Accursio
Accumbty
Tou
Tainro Epoches
Traniing Epoches
Traling Epaches
111 0.9, P2
(h) P1 0.8, 0.9.
p1 0.7, 0.8.
Figure 5: MLP' performance with different Pl,Pa P2-
23

---

Published as a conference paper at ICLR 2023
B.3 PROOF OF THEOREM 2
Theorem 5. An MoE module in equation 3 with N experts and k = 1 aligns the conditional sentences in Algorithm with
(N (1++)) max (M's M M(G,,,,,,,
if M,
Alignment
(N +1 max
max (M]
M h4t1, MGG, if M,
where M(-:'s'-)) is defined in Definition 1, and 5y is the optimal objective value of the following optimization problem:
S minimize Ey
max (E(1, .N1
M(frFN. (11,,)ez, o hn1) [hylsez,
subject to
U T = (2,3, ... M + +1},
where 11, is the fadicator on interval 1,-
Proof. For N 2 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 2 M: We define the mapping ((x) [ ... 1 E RRM where 1 is the indicator function interval 4, To show the aligmment between MoE and conditional sentences, define the module functions fi, .NN and neural network modules N1, ,NN as
Ho o h1,
i
G, =
St h++1
if f M;= fFFN. fisiss+L, fIIN. O.W.,
0,
O.W.,
Replacing N₁ with ft. the MoE network simulates Algorithm 1. Thus, for > M. MoE algorithmi- cally aligns conditional sentences with
Alignment (N + 1) max (max M,
(8)
where M(.....)) is defined Definition 1.
Case N < M: We define the mapping Ho(x) = [Lu 4, "......,,,,,uy...... 1,, E RVV, where Iz is the solution to the optimization problem equation 5. To show the alignment between MoE and conditional sentences, we define the module functions fi. fN and neural network modules .N, .NN as
fo Ho o h1,
if
N = [G,
if
(11,)ee,, o hT) [hyleez, O.W.,
[FFN. O.W.
Replacing N₁ with ft. the MoE network simulates Algorithm Thus, for N < M, MoE algorithmi- cally aligns with conditional sentences with
Alignment (N 1) max (M (,,,,,,,,)))
where Mi is the optimal objective value to equation 5 and M(-rr.ss.) is defined in Definition
B.4 TECHNICAL LEMMAS
Lemma 1. Under Assumption for given functions f(-) and interval A, PD. (((G A) < 5 implies P P.. [f E A] & C6.
Proof. We denote the inverse mapping as f f-) {slf(s) A). Note that
Pr. IDNG A) = /
Piea.Ni (s)ds
o 1(A)
SC
Purin,NG (s)ds - CPD.. E A]
(A)
This finishes the proof for Lemma
24

---

Published as a conference paper at ICLR 2023
Remark 5. Assumption 1 is similar to j-exploration in batch reinforcement learning (RL), which aims at controlling the distribution shift in RL (Xie & Jiang, 2021). Lemma 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 al. (2020b).
C IMPLEMENTATIONS
C.I DETAILED ARCHITECTURE
For image input X, we first process the images sequence of equal-sized patches Φ, E X using a l-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 T. blocks.
To make fair comparison with other reported methods on DomainBed, we choose the VIT-S/16 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 x 16, 6 heads in multi-head attention layers, and 12 transformer blocks.
We consider two-layer configurations, every-two and last-two, that are widely adopted in Sparse MoF design (Riquelme et al., 2021). The detailed architectures are illustrated in Figure 6. Besides, we also consider evaluating GMoE's s generalization performance with larger scale models (e.g., ViT-Base).
V-MoE (Riquelme et al., 2021) trains 32 experts model on ImageNet-211 dataset. However, consider- ing the number of training data, 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-21 K. 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 experts.
We put the ablation study results on layer configuration and larger size backbone model on Sec. D.6.
Classification Layer
Classification Layer
Transformer Block
Transformer Block
GMoE Block
GMoE Block
Transformer Block
12 Blocks
GMoE Block
Transformer Block
8 BBockk
Transformer Block
GMoE Block
Transformer Block
Patch Embedding
Patch Embedding
Images
Images
(a) Every 2 GMoE.
(b) Last 2 GMoE.
Figure 6: Diagrams of detailed GMoE architecture with two types of layer configuration.
25

---

Published as a conference paper at ICLR 2023
C.2 MODEL INITIALIZATION
To make fair comparison with other algorithms from DomainBed (mostly with ResNet-50 pretrained on ImageNet- IK), also utilize the pretrained ViT models ImageNet- I 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 ImageNeet22IK: Mcanwhile, there is no released pretrained MoE model on ImageNet- I... In Table 4, we see ViT-S/16 and GMoE-S/16 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 by 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 first.
Table 4: Comparison of the different backbone models used in the experiment. ImageNet- I denotes GMoE are initialized with ViT pretrained on ImageNet-JK and copy weights for different experts.
Backbone
Re
ResNet-101
VIT-S/16
GMoE-S/16
Pretrained Dataset
ImageNet-IK
ImageNet-IK
ImageNet-IK
ImageNet- K
Parameters
25.6M
42.9M
21.7M
33.8M
Flops
4.1G
7.9G
4.6G
4.8G
C.3 PERTURBATION
Another specific design for the router adding Gaussian noise to increase stability. Hard routing brings sparsity benefits 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 model to capture the invariant semantic meaning of an object. Existing work (Chen et al., 2022) theoretically claims that adding noise the gating function provides smooth transition of the gating function during training. Empirically, we find adding noise with a standard deviation of N would empirically improve performance and stabilize training dynamics.
C.4 AUXILIARY LOSS
It has been observed that the gating network tends to converge to self-reinforcing imbalanced state, where a few experts are more frequently activated than others. Following Riquelme et al. (2021) I); Shazeer et al. (2017), we define importance loss Limp to encourage balanced usage of experts. The importance of expert e is defined as the normalization of the gating function's 's output in a batch of images:
imp
XX)
Σ g(x)e
(9)
xEX
And the loss Limp is defined by the squared coefficient of variation of imp(X):
(TT((im(X)) MEAN(imp(X))
Cimp
While importance loss seeks that all experts have average similar output routing weights, it is still likely some experts would not routed (see example Table 5).
To address this issue, besides balancing across patches, also need to encourage balanced assign- ment across experts. Ideally, we expect balanced assignment via a load loss C However, the assignment is discrete value; then expect a proxy to make it differentiable in pook-propagation. Following the idea in Riquelme et al. (2021); Shazeer et al. (2017); Mustafa et al. (2022), for each token and each expert e; could compute the probability of e, once being selected in TOP;
26

---

Published as a conference paper at ICLR 2023
Table 5: Simple example TOP; (before SoftMax operation) where average weights across patches (T1,2,3) are balanced. However, Expert has always been ignored.
Patch
Expert
Expert 2
Expert 3
Expert 4
Selected Expert
0.9
0.4
0.1
0.2
Expert
22
0.9
Expert
0.1
0.4
0.2
0.9
Expert 4
after SoftMax operation, and then kept if re-sampling again only the noise expert 6,. other words, we are computing the probability that having expert still being among the TOP; while re-sampling again only the noise of the expert. This probability is formally defined as in Eq. 10, where The the K-th largest entry after SoftMax operation and & is the cumulative distribution function of a Gaussian distribution.
Pe(x) - ①b T] (E" Wx)e
(10)
Then the expert load could be defined as loade(X) Exex Pe (x: and the whole experts load is defined as load(X) (X) {load, J ,The load loss Glood is defined by
(To(d xz)
Lioad (X)
MEAN(load
With the above adaptations for DG tasks, GMoE still works with simple overall loss:
((X)
celussficationn
() X(iimp (X) Loout (asua(() (X)
with hyperparameter balance the main task loss and the auxiliary loss, in which usually we view equal importance of the load and importance loss.
27

---

Published as a conference paper at ICLR 2023
D EXPERIMENTAL RESULTS
D.1 DOMAINBED DETAILS
Following previous DG studies, mainly evaluate our proposed method and report empirical results on DomainBed (Gulrajani & Lopez-Paz, 2021) with 8 benchmark datasets, PACS, VLCS, OfficeHome, Terralncognita, DomainNet, SVIRO, Wilds-Cameliy 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.
Benchmark 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 6: Statistics of DomainBed datasets.
Dataset
PACS
VLCS
OfficeHome
TerraInc
DomainNet
Domains
4
4
6
# Classes
7
5
65
10
345
# Examples
9,991
10,729
15,588
24,788
586,575
Table 7: Statisties of Wilds and CUB-DG. Terralnc stands for Terralncognita, Camelyon stands for Wilds-Cann Cadde-Cmmlyonn and FMOW stands for Wilds-FMOW.
Dataset
SVIRO
W-Camelyon
W-FMOW
CUB-DG
Domains
10
6
4
# Classes
7
2
62
200
# Examples
56,000
455,954
523.846
47,152
In detail, the multi-domain image classification datasets are comprised of:
1. PACS (Li et al., 2017) comprises four domains d E (art, cartoons, photos, sketches). This dataset contains 9, 999 examples dimension (3, 224, 224) and classes.
2. VLCS (Fang al., 2013) comprises photographic domains d (Caltech101, LabelMe,
SUN09, VOC2007). This dataset contains 10, 729 examples of dimension (3, 224, 224) and
5 classes.
3. Office-Home (Venkateswara et al., 2017) includes domains d E [art, clipart, product, real). This dataset contains 15, 588 examples dimension (3. 224, 224) and 65 classes.
4. TerraIncognita (Beery et al., 2018) contains photographs of wild animals taken by camera traps at locations (0,100 L.38, L43,, L46). Our version of this dataset contains 24, 788 examples of dimension (3, 224, 224) and 10 classes.
5. DomainNet (Peng et al., 2019) has domains d E (clipart, infograph, painting, quickdraw, real, sketch). This dataset contains 586, 575 examples of size (3, 224, 224) and 345 classes.
6. SVIRO (Cruz et al., 2020) a Synthetic dataset for Vehicle Interior Rear seat Occupancy across 10 different vehicles, the 10 domains.
The domains are d E {aclass, escape, hilux, i3, lexus, tesla, tiguan, tucson, x5, z0e). This dataset for image classification contains 56,000 image examples of size (3, 224, 224) and classes.
7. 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,954 examples of dimension (3, 224, 224) and classes from hospitals.
8. Wilds-FMOw (Koh et al., 2021) dataset is variant of the functional map of the world dataset (Christic 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).
28

---

Published as a conference paper at ICLR 2023
Table 8: Hyperparameters to reproduce best performance of GMoE each dataset.
Hypen pyperaarametess
PACS
VLCS
OfficeHome
TerraIne
DomainNet
Learning Rate
3 x 10-5
3 x 10-5
1 1 10-5
x 10-5
x 10-5
Weight Decay
1 10
1 10
1 10
9. 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 e [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 has_aial_lll wwwll_ihapeedalleeeeeeeaaaaaaa sa__sshappdaggeree bounding box information. We treat the 15 positional attributes as visual attributes of bird image, and thus we could evaluate the correlation between visual attributes and expert selections.
and
Evaluation protocols We evaluate our proposed method with two protocols on DomainBed (Gul- rajani & 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 overall validation set. Finally, we choose the model maximizing the accuracy on the overall validation set, and report the final accuracy on leave-out domain.
For leave-one-domain-one validation selection, we train the model on all training domains and one domain out the validation set. We choose the model maximizing the accuracy on the Icave-out validation domain, and report accuracy on another leave-out test domain. We should emphasize that the cave-one-domain-
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 reported.
D.2 TRAINING DETAILS & HYPERPARAMETERS
In this section, report the training details of 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 32 for each domain.
Since Gulrajani & Lopez-Paz (2021) only evaluates baselines for 5K iterations on DomainNet, while 5K iterations less than 2 epochs DomainNet. This results in insufficient training on this dataset and it can not fully reveal models' performance with different algorithms. Subsequent state-of-the-a works (Cha et al., 202 (Cha al., 2022) re-evaluated ERM and proposed their methods on 15,000 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 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.3 LEAVE-ONE-DOMAINN -UT RESULTS
We have demonstrated the experimental results in Table with train-validation selection. In this section, we will further provide the results of GMoE on more challenging setting. In this setting, should train GMoE on 5 datasets with in total of 61 individual 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
selection, which is
29

---

Published as a conference paper at ICLR 2023
Table 9: Train-validation selection performance comparison with 5K and 15K training iterations, where train-val stands for train-validation selection and lodo stands for leave-one-domain-ouu selec- tion. The place marked is not applicable stated in Sec. D.3.
Iterations
Algorithms & Models
DomainNet (train-val)
DomainNet (lodo)
ERM (Vapnik, 1991)
41.2 + 0.2
40.6 + 0.2
(Arjovsky et al., 2019)
33.9
+
DANN (Ganin et al., 2016)
41.8 = 0.2
38.2 + 0.2
CORAL (Sun Sacnko, 2016)
41.8 0.2
41.1 + 0.1
5K
MMD (Li et al., 2018b)
39.4 = 0.8
23.4 + 9.4
MLDG (Li et al., 2018a)
41.2 + 0.1
41.0 + 0.2
Fish (Shi al., 2021)
42.7
Fishr (Rame al., 2021)
41.7 = 0.2
ViT-S/16 (Dosovitskiy al., 2021)
42.4 0.0
42.0 + 0.2
GMoE-S/16
44.6 = 0.3
44.7 1 0.2
Swad (Cha et al., 2021a)
46.5 0.1
1SK
MIRO (Cha et al., 2022)
47.4 1 0.0
ViT-S/16 (Dosovitskiy al., 2021)
47.1 0.3
46.1 + 0.4
GMoE-S/16
48.7 + 0.2
48.4 + 0.3
by DomainBed. In the Table 10, we can find that GMoE outperforms previous methods in all datasets.
Table 10: Overall out-of-domain accuracies (%) with leave-one-domain out selection criterion.
Algorithm
PACS
VLCS
OfficeHome
TerraInc
DomainNet
ERM (RexNet50)
83.0 207
77.2 04
65.7 =05
41.4
40.6
IRM
81.5 6.88
76,3
64.3
41.2 436
33.5
DANN
81.0 L.
76.9 1.00
64.9 1 1.2
44.4 1:1.1
38.2 1 020
CORAL
82.6
78.7
68.5
46.3 =
41.1 =
MMD
83.2 0.
77.3 0.05
60,2 5.2
46.5 1:15
23,4 à
MIDG
77.2 0.09
82.9 1.7
66.1 2.5
46.2 2:09
41.0 0:02
VIT-S/16
86.6 1.0.
78.3 1.02
71.9
41.9 1:03
46.1
GMoE-S/16
87.1 0.3
80.0 n.
73.9 0.2
46.7 0.03
48.4
D.4
ABLATION STUDY WITH DG ALGORRITHMS
This section presents ablation study compare GMoE with DG algorithms and VIT with DG algorithms, as shown in Table 11. From the table, we see that GMoE significantly outperforms VIT and ResNet under variety of DG algorithms, showing the generalizability of the proposed backbone architecture.
D.5 SINGLE-SOURGN DOMAIN GENERALIZATION RESULTS
In this section, we will demonstrate the single source domain generalization results on DomainNet's domains. We select different training domains in turn and take the remaining domains as test domains. In Table 12, show the OOD results when models trained only one domain and test on test domains, as well as the IID results on the training domain's validation set. We specify IID and OOD results in different colors.
30

---

Published as a conference paper at ICLR 2023
Table 11: Overall accuracies (%) with train-validation selection criterion.
Algorithms
Backbones
PACS
VL.CS
OfficeHome
TerraIne
DomainNet
Avg
VIT-S/16
88.3
80.4
72.7
46.9
48.8
67.4
Swad
GMoE-S/16
88.6
80.8
74.5
49.1
49.6
68.5
ResNet50
85.5
77.8
68.6
45.1
43.1
64
Fish
VIT-S/16
87.3
78
72.2
42.7
47.3
65.5
GMoE-S/16
87.9
79.7
73.0
46.9
48.8
67.1
ResNet50
85.5
77.8
68.6
47.4
41.7
64.2
Fishr
ViT-S/16
85.6
79.1
71.8
44.1
47.6
65.6
GMoE-S/16
88.3
78.8
72.7
45.6
48.3
66.7
D.6 ABLATION STUDY ON MODEL DESIGN
In order to validate discussions in Sec. 4.3 and Sec C.I, we conduct experiments with different layer configurations, well as adopting larger backbone (e.g.. ViT-Base) model. From the results in Table 13, we see that the last two configuration largely exceeds every two configuration, 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.
D.7 ABLATION STUDY ON RESNET WITH STRONG DATA AUGMENTATION
In this part, we will go into greater depth about the pre-trained model's 's effect on DG performance. As stated in Sec. C.2, GMoE is initialized from the same architecture ViT models, where pre- traitre weights are from Touvron et al. (2021). 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 downstream tasks.
To verify this difference in detail, we consider comparing GMoB-S/16' pre-trained model, VIT-S/16, with stronger pre- cre-trained ResNet-50 V2 model' with significant training tricks data augmentations. The pre-training details on ViT-S/16 and ResNet V2 are listed in Table 14.
Due to the space limit, we use abbreviations for some terms. In detail, they are:
LR Opt. stands for learning 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 al., 2019).
WDT stands for weight decay tuning.
IRT stands for inference resize tuning (Touvron et al., 2019).
The results in Table 15 demonstrate that GMoE and VIT could still remain higher performance than the strong pre-trained model ResNet V2 due to the benefits the backbone architecture.
D.8 COMPUTATIONAL COST COMPARISON
Having shown the effects of GMoE's performance on different domain generalization benchmarks, we now conduct efficiency analysis with respect to iteration time and run-fime memory during training and inference. Since algorithms developed from DomainBed mainly adopt the same architectures
From TORCHVISION's Pretrained Models.
31

---

Published as a conference paper at ICLR 2023
Table 12: Single-source DG results DomainNet. We alternatively choose training domain and the other as domains. We report IID generalization results on training domain's 's validation set
with grey color, and OOD generalization
results on test domains' validation set with light cyan color.
Clipart
Clipart
Info
Paint
Quick
Real
Sketch
ResNet50
68.5
12.3
25.8
9.5
39.6
37.1
ResNet101
70.4
12.2
28.6
12.4
43.0
40.7
ViT-S/16
74.7
16.2
35.3
12.7
50.9
40.4
GMoE-S/16
76.3
16.8
36.2
13.9
Info
Clipart
Info
Paint
Quick
Real
Sketch
ResNe150
27.9
34.0
23.3
2.0
34.3
24.5
ResNet101
29.2
35.6
24.0
3.1
34.1
24.6
VIT-S/16
36.3
40.2
33.6
5.1
47.4
30.7
GMoB-S/16
36.7
41.0
34.2
4.5
48.2
30.8
Paint
Clipart
Info
Paint
Quick
Real
Sketch
ResNet50
37.1
12.9
62.7
2.2
ResNet101
40.5
13.1
63.4
3.1
51.2
35.4
ViT-S/16
42.7
15.9
69.0
5.0
56.4
37.0
GMoE-S/16
43.5
16.1
69.3
5.3
56.4
38.0
Quick
Clipart
Info
Paint
Quick
Real
Sketch
ResNe150
17.3
1.1
2.8
62.4
ResNet101
16.0
1.3
2.7
64.1
6.3
8.5
ViT-S/16
22.1
2.0
6.3
66.2
11.8
11.8
GMoE-S/16
22,8
2.0
6.4
66.8
12.2
12.9
Real
Clipart
Info
Paint
Quick
Real
Sketch
ResNet50
39.5
13.0
40.0
4.4
73.6
30.5
ResNet101
43.3
15.6
43.3
6.8
75.3
34.0
VIT-S/16
46.9
18.0
46.0
5.0
79.5
GMoE-S/16
47.4
18.5
46.9
7.5
79.6
36.4
Sketch
Clipart
Info
Paint
Quick
Real
Sketch
ResNet50
48.9
11.7
31.6
10.8
40.7
63.4
ResNet101
48.4
11.8
31.9
11.3
38.1
64.3
ViT-S/16
52.4
14.4
38.6
13.4
48.2
67.4
GMoE-S/16
52.9
14.1
36.6
13.7
(e.g., ResNet50, Linear Classifier), traditional model complexity measures such flops and model parameters can not truly reflect the difference in efficiency. 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 16, we observe that GMoE achieves relatively low un-time memory and training step time among other competitors.
32

---

Published as a conference paper at ICLR 2023
Table 13: Train-validation selection performance comparison for different GMoE models. All GMoEs (S/16 and B/16) have E = experts and L = 12 blocks. We specify the number of attention heads (H), the patch embedding size (D), layer configuration, and router's type the table header.
Algorithm
Config.
Router
H
D
PACS
VLCS
OfficeHome
TerraIne
DomainNet
GMoE-S/16
Every 2
Linear
6
384
81.8 4
75.0 201
64.0 o4
32.5
46.3 u
GMoF-S/16
Every? 2
Cosine
6 384
81.4 101
74.8 102
62.2 14
40.9
46.4 122
GMoE-S/16
Linear
384
87.8 4:02
0.00
72.7 4:02
46.7 =2
48.3 a.1
GMoB-S/16
Last 2
Cosine
6 384
88.1 +.0.
80.2 #2
74.2 0.04
48.5 04
48.7 #.02
GMo6-B/16
Last 2
Cosine
12
768
89.4 1
81.2
77.2
49.3
51.3
Table 14: Comparison of different training recipes ResNet-50 V2 and ViT-S/16.
Recipe
LR Opt.
TrivialAug.
Ep.
Rand Er.
Label Sm.
FixRes Mt.
WDT
IRT
INIK
ResNet-50 V2
600
80.8
VIT-S/16
4
300
V
K
X
X
79.9
Table 15: Train- Train-validation validation selection performance comparison for VIT, GMoE, and other DG al- gorithms with ResNet-50 V2 as backbone model. On DomainNet, results are reported with 15K iterations.
Algorithm
PACS
VLCS
OfficeHomee
TerraIne
DomainNet
ERM (w/ ResNet-50 V2)
87.2
78.2
68.7
Fishr (w/ ResNet-50 V2)
87.5
77.9
70.4
51.7
47.0
ViT-S/16
86.2
79.7
72.2
42.0
47.1
GMoE-S/16
88.1
80.1
74.2
48.5
48.7
Table 16: Comparison of training/inference iteration time and run-time memory for a mini-batch. mini-batch is formed with 160 images in 224 x 224 resolutions from DomainBed. For both metrics, lower better.
Training
ERM
DANN
IRM
Fish
Fishr
SWAD
VIT-S/16
GMoE-S/16
Step Time (s)
1.01
1.02
L.10
2.79
1.10
1.21
0.90
0.98
Run-time Memory (GB)
13.40
13.42
13.40
3.41
15.25
14.32
11.15
12.28
Inference
FRM
DANN
IRM
Fish
Fishr
SWAD
VIT-S/16
GMoF-S/16
Step Time (8)
0.32
0.33
0.32
0.33
0.34
0.33
0.28
0.30
Run-time Memory (GB)
1.82
1.83
1.82
1.82
1.83
1.84
0.76
1.05
33

---

Published as a conference paper at ICLR 2023
E MODEL ANALYSIS
E. CUB-DG RESULTS & VISUALIZATIONS
Generalization across Image Stylization To evaluate the model performance in generalizing 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 al., 2016), Fish (Shi et al., 2021) and Fishr (Rame et al., 2021) I)). Among them, DANN and Fish are widely discussed and cited papers, and Fishr considered the best algorithm of its kind far.
In this experiment, all three invariant learning methods, including ERM, adopt ResNet-50 as the backbone model. GMoE does not have any specific loss design and constraint. The only difference is the model architecture. In Table 17, we see that GMoF, the new backbone model for DG, is significantly more effective than those methods specifically designed to learn invariance with certain losses and gradient constraints.
Table 17: Out-of-domain accuracy (%) ) in each domain on CUB-DG dataset. The best is in bold.
Algorithm
Candy
Mosaic
Natural
Udnie
ERM (ResNet50)
64.3
20.8
75.3
77.9
DANN
47.0
14.9
76.5
70.6
Fish
62.6
22.4
84.5
77.7
Fishr
62.4
24.3
78.8
73.8
GMoE-S/16
83.1
42.8
89.3
82.7
Experts Selection section.
Following the analysis in Sec. 5.4, we provide more visualization results in this
In Figure 8-9, we provide the expert selection results in natural domain images across different bird classes. We use different colored lines that correspond to different visual attributes, and we connect these lines across images in row. indicates that experts focus the same visual attributes while seeing different bird images from different angles. This suggests that mixture-of-expert have the ability to capture and process visual attributes across class.
In Figure 10-13, we provide the expert selection results for the same image on four domains (with four types of stylization). In order to more vividly display the expert selection across domains, we assign different colors different selected experts. We see that, in most cases, the same areas and visual attributes are handled by the same experts across domains. For example, in Figure 10, expert 1,3 and 4 focus on the bird while expert 0,2 focus the background. And it is consistent across four domains with different stylization. The results reveal that experts are potentially invariant in dealing with visual attributes across domains.
Multi-head Attention Visualization Ideally, the multi-head attention mechanism can be viewed as jointly attending multiple places by ensembling multiple attention heads. Each attention head would focus on its specific attention relationship between all patches and inherently collaborate with each other.
However, recent studies (Cordonnicr et al., 2020) demonstrate that multiple-head attention (MHA) layers could learn redundant key/query projections, i.c. some heads might attend similar fea- tures in input space. This issue demonstrates that two heads head; and head, are computing the same key/query representations up to unitary matrix I E Rdxd such that Wo. = Wo, I and WK. WK, In this case, even though the two heads are computing identical attention scores, i.e. II W'' K the concatenation [Wq Wo,1 E Rdx2d can also be full rank, which indicates that some attention heads would focus the same content but are agnostic each other. We opine that the mixture-of-expeee experess (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
34

---

Published as a conference paper at ICLR 2023
head,
headz
head3
Experty
Patch Embedding
Expertz
Experts
Expert,
Dimensions
Figure 7: Diagram of a GMoE block, where different attention heads attend to different aspects of relations between patches. Different experts handle the learned attention in different patches. The router has been omitted for brevity.
experts. This mechanism extends the network sparsity and hence is able to improve the model generalizability. sum, the MoE layer designed to leverage different experts to capture distinet visual attributes, well their correlations (via attention mechanism) to each other.
To see more about the learned representation multi-head attention layer, we provide the visualization of last layer's attention activations of VIT and GMoE on Figure 14 15 16 17.
35

---

Published as a conference paper at ICLR 2023
Book
Right
Right
Figure 8: Expert selection visualization on block- 10 of GMoE-S/16. Images are from different classes of natural domains on CUB B-D. Different colored lines connect the same type visual attributes across images. Same visual attribute is processed by the same expert, e.g,, beak and tail by expert 3, left/right leg by expert 4.
Wing
Right Wing
Laft Log
Right
Figure 9: Expert selection visualization on block-10 10 of GMoE-S/16. Images are from different classes of natural domains on CUB-DG Different colored lines connect the same type of visual attributes across images. Same visual attribute is processed by the same expert, e.g., beak and tail by expert 3, left/right leg by expert 4.
36

---

Published as a conference paper at ICLR 2023
3 3
2 7
2
2
2
2
2
2
2
0
2
7
7
2
0
2
4
2
2
2
2
2
4
2
2
2
0
2
O
Z
2
0
0
0
0
0
(a) Natural
(b) Candy
3
3
2
3
2
2
2
2
3
2
2
3
2
2
2
2
2
2
2
3
2
2
2
2
2
3
3
3
3
3
2
2
3
2
3
2
2
4
2
2
2
2
2
12
1
3
0
(c) Udnic
(d) Mosaic
Figure 10: Expert selection visualization on block-10 10 of GMoB-S/16. Images are from different domains on CUB-DG.
37

---

Published as a conference paper at ICLR 2023
3
3
0
3
3
3
0
0
4
0
0
0
0
0
0
0
0
3
0
0
0
0
0
(a) Natural
(b) Candy
0
3
3
3
2
2
2
3
0
2
3
2
2
0
3
1
0
2
3
2
3
4
2
2
0
2
0
2
2
2
2
3
2
0
(c) Udnie
(d) Mosaic
Figure 11: Expert selection visualization on block-10 10 of GMoB-S/16. Images are from different
domains on CUB-DG.
38

---

Published as a conference paper at ICLR 2023
2
2
0
0
2
2
0
0
2
7
2
3
0
3
3
3
3
5
2
0
0
0
1
2
(a) Natural
(b) Candy
3
0
3
3
0
0
3
0
3
3
3
4
4
3
3
3
2
4
0
2
2
2
0
0
2
2
(c) Udnic
(d) Mosaic
Figure 12: Expert selection visualization on block-10 of GMoB-S/16. Images are from different
domains on CUB-DG.
39

---

Published as a conference paper at ICLR 2023
0
0
3
3
2
2
3
3
3
2
2
3
3
3
2
2
0
0
2
2
0
0
0
0
2
2
2
(a) Natural
(b) Candy
2
2
2
2
2
0
2
3
3
2
2
3
3
2
3
2
2
0
2
0
3
3
3
2
(c) Udnie
(d) Mosaic
Figure 13: Expert selection visualization on block- 10 of GMoE-S/16. Images are from different
domains on CUB-DG.
40

---

Published as a conference paper at ICLR 2023
Multi-head attention (from head 0-5). Right ones are original images.
VIT
GMoE
ViT
GMoE
VIT
GMoE
ViT
GMoE
ViT
GMoE
Figure 14: Multi-head attention visualization the last block of ViT-S/16 and GMoE-S/16. Images are from Real domain in DomainNet.
41

---

Published as a conference paper at ICLR 2023
Multi-head attention (from head 0-5). Right ones are original images.
VIT
GMoE
ViT
GMoE
VIT
GMoE
ViT
GMoE
ViT
GMoE
Figure 15: Multi-head attention visualization the last block of ViT-S/16 and GMoE-S/16. Images are from Real domain in DomainNet.
42

---

Published as a conference paper at ICLR 2023
Multi-head attention (from head 0-5). Right ones are original images.
VIT
GMoE
ViT
GMOE
VIT
GMoE
ViT
GMoE
ViT
GMoE
Figure 16: Multi-head attention visualization the last block of ViT-S/16 and GMoE-S/16. Images are from Real domain in DomainNet.
43

---

Published as a conference paper at ICLR 2023
Multi-head attention (from head 0-5). Right ones are original images.
VIT
GMoE
ViT
GMoE
VIT
GMoE
ViT
GMoE
ViT
GMoE
Figure 17: Multi-head attention visualization the last block of ViT-S/16 and GMoE-S/16. Images are from Real domain in DomainNet.
44


Reviews:
[{"review_id": "pYbp9hvzcf", "reviewer": "Reviewer_FUc5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper start from an interesting empirical finding that a transformer trained with ERM outperforms CNN trained with domain generalization (DG) algorithms on DG task. They show some experiments and math to prove that. They propose a new GMoE and show its superiority on DG. ", "summary_of_the_review": "See above.", "strength_and_weaknesses": "The paper focus on an interesting aspect and the paper are overall easy to understand. But some important experiments are missing.\n\n1. Missing comparison between transformer with ERM and transformer with DG. The core motivation \"empirical\nrisk minimization (ERM) outperform CNN-based models employing state-of-the-art (SOTA) DG algorithms\" can just because the transformer is a stronger backbone.\n\n2. Why is GMoE superior to a normal MoE. \nThe major difference is that it forces a normalized operation on the input of the gating function and uses a learned embedding dictionary to decide the gating. The ablation study of why this is superior is missing from the paper.\n\n3. The paper states that GMoE relies on visual attributes to have better performance. This is not convinceable from figure 3(b). It seems that e0 and e2 are specialized for background and others (e3,e4, e5) are all activated by almost all the visual attributes. Figure 3(c) is also not convincing to me. It seems no strong relation between experts and attributes. Any network would have similar visualization that some part of it prefers some region in the image.\n\nOverall, the experiment parts seem to be a weakness: not enough experiments, ablations, and convincing visualization.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "WktjBwO8Ff5", "reviewer": "Reviewer_nMjY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a GMoE model for solving classification under domain shifts, specifically in the DG setting. The paper presents results on 8 benchmarks and shows promising performance. ", "summary_of_the_review": "My primary concerns are on missing results and papers and would request the authors to respond to it in the rebuttal period.", "strength_and_weaknesses": "Strengths\n\n- I liked the idea of applying mixture of experts for solving the problem of DG. I resonate with the authors, that it is a bright direction for explorations in DG\n- The experiments are thorough and the results appear promising and t\n\nWeaknesses\n\nMy major concern with the paper is that some important recent papers or results/discussions from cited papers are missing. \n\n(a) For instance, why did the authors ignore the MIRO+SWAD numbers in Cha et al. 2022. It clearly and comprehensively outperforms the proposed method on all the datasets. \n\n(b) Author missed in important recent work by Sivaprasad et al. [A], which makes a similar observation that backbone plays a more crucial role in DG compared to tailored algorithms. They show that ERM with Inception Resnet backbone can achieve 89.11% accuracy on PACS (outperforming the proposed approach on the particular dataset). That clearly dilutes the first contribution of the paper. I suspect more similarities with that work as well in terms of distribution shifts etc. A discussion is warranted.\n\n(c) Finally, were the experiments in Table 1 averaged over multiple runs?\n\n\n[A] Sivaprasad et al. Reappraising Domain Generalization in Neural Networks, arXiv 2021\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "8Y7dGS78zl", "reviewer": "Reviewer_jTSg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tackles the problem of domain generalization by proposing to focus on model architectures that are well aligned to the task of interest. The paper uses the definition of algorithmic alignment of Xu et al. 2020 which defines how easy it is to learn a function that can be decomposed into N subfunctions if your neural network can also decompose into similar subcomponents. First, the paper formally shows that models that align well with the causal correlations in the data will generalize well, while models that align well with the spurious correlations will not generalize. The authors propose a model named Generalizable Mixture-of-Experts (GMoE), similar to other mixture-of-experts Transformer models used primarily in NLP. The authors argue that domain generalization benefits from capturing and effectively combining diverse attributes and for this purpose, some kind of conditional statements are required. It is formally shown that such statements align well with the proposed model, thus the GMoE model should generalize well out-of-distribution. Experimentally, the GMoE is shown to obtain good results compared to recent methods.\n", "summary_of_the_review": "Overall the paper is well-written and has a good formalism that gives theoretical explanations of the good performance of the proposed method. The paper will benefit the community by showing that good architectures that align with the task are to be desired and show a good example of such architecture.\n\n", "strength_and_weaknesses": "**Strengths**:\n\nThe paper tackles an important problem of domain generalization and proposes a good direction to approach it, by designing models that are well aligned with the desired task.\n\nThe formalization of alignment between models and the task is sound and useful for guiding the design of future architectures.\n\nThe proposed model seems to be beneficial for DG tasks. It is good to see that GMoE obtains good results with standard ERM training, but also benefits from DG algorithms. \n\nIt is interesting to see that the experts seem to specialize, as seen by histograms and examples in Section 5.\n\n**Weaknesses**\n\nThe paper did not explain in much detail why conditional statements are good for Domain Generalisation. As this point is crucial in linking the proposed method to DG, it should be better explained. At the high-level Theorem 2 says that the proposed model is good for cases where the true function is composed of similar mechanisms (selection + application of submodule) as the GMoE model (gating + separate FFN). But still, it should be explained in more detail why such functions are preferable for DG. Some intuitions are given in Section 4.2, but expanding them would be good.\n\nThe paper could discuss connections to other works that mixture-of-experts for domain adaptation [A], or use experts in similar ViT models [B].\n\n\n[A] Guo et al. “Multi-Source Domain Adaptation with Mixture of Experts” EMNLP 2018 \n[B] Rahaman et al. \"Dynamic inference with neural interpreters.\" NeurIPS 2021.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7RoX1pmCPt", "reviewer": "Reviewer_aHin", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission considers the problem of improving domain generalization via NN architecture design. The authors propose to adopt an MoE ViT design. The former (MoE) is motivated by its ability to serve conditional computation for easier processing of high-level visual attributed, while the latter (ViT) is motivated by multi-head self-attention properties as low-pass filter with shape bias compared to texture-biased CNNs. Routing scheme (cosine routing) and number of MoE layers are empirically adapted for best generalization performance. Experiments demonstrate superior performance over explicit domain generalization algorithms using more conventional CNN-based architectures.", "summary_of_the_review": "The paper suggests well-grounded architecture choices to improve the SOTA in domain generalization using conditional computation via MoE.", "strength_and_weaknesses": "As a clear strength, the paper provides detailed argumentation to motivate each of its design choices, going beyond a purely empirical study. Most prominently, it motivates the choice of ViT architecture by means of Theorem 1 and the choice of (G)MoE by means of Theorem 2. Both theorems provide a somewhat theoretical grounding of the architectural choices made. The empirical evaluation is extensive and convincing. A minor weakness is that further tweaks in Section 4.3 are made on purely empirical grounds and ablations are relegated to the appendix.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%