PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Mixture-of-Experts are Domain Generalizable Learners

Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, Ziwei Liu

arXiv:2206.04046v6cs.CVcs.AIcs.LG

TL;DR

Machine-learning models struggle to generalize to out-of-distribution visual data, while DG research has mainly emphasized loss design. This paper analyzes backbone alignment with dataset correlations and proposes GMoE; experiments report strong DomainBed performance for GMoE trained with ERM and further gains with DG algorithms.

  • Problem

    Domain generalization seeks to improve model resilience to distribution shifts, but prior approaches have mainly implemented DG ideas through specialized loss functions.

    Source: §1.1 · pp. 1–2

  • Method

    The paper analyzes how backbone alignment with invariant and spurious correlations relates to robustness, then proposes GMoE, a sparse mixture-of-experts vision transformer designed to capture visual attributes.

    Source: §1.1 · p. 2, §4.2 · p. 6

  • Results

    GMoE trained with ERM outperforms counterparts on almost all reported datasets, and adding DG algorithms improves its performance over ERM training.

    Source: §4.2 · p. 6, §5.1 · p. 8

  • Takeaways & Limitations

    The paper concludes that backbone architecture is a useful direction for DG and demonstrates GMoE’s superior performance on DomainBed.

    Source: §6 · p. 9

    The theoretical analysis assumes that training-feature support covers test-feature support and defines the target function using noise-free labels.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

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

1 INTRODUCTION

The paper shifts domain generalization (DG) research from loss-function design toward backbone architecture, motivated by a ViT trained with ERM outperforming CNNs trained with state-of-the-art DG algorithms. It develops an alignment-based account of robustness and proposes GMoE, which shows strong DomainBed performance with and without DG algorithms.

1.1 MOTIVATIONS

  • DG targets resilience to distribution shifts such as changes in background, lighting, texture, shape, and geographic or demographic attributes.Source: §1.1 · p. 1
  • DG methods commonly implement approaches such as domain alignment and invariant prediction through specialized loss functions.The paper presents backbone architecture design as an orthogonal direction to this loss-focused work.Source: §1.1 · p. 1, §1.2 · p. 2
  • A ViT-S/16 trained with ERM outperforms ResNet-50 trained with state-of-the-art DG algorithms on DomainNet, OfficeHome, and VLCS despite similar parameter counts and close in-distribution performance.The paper uses this observation to motivate studying backbone architecture alongside loss design.Source: §1.1 · p. 2

1.2 CONTRIBUTIONS

  • The paper argues that networks aligned with invariant correlations are more robust to distribution shifts, whereas alignment with spurious correlations is associated with lower robustness.The theorems are verified on synthetic and real datasets.Source: §1.2 · p. 2
  • Generalizable Mixture-of-Experts (GMoE) combines sparse mixture-of-experts with vision transformers and is designed to align better with invariant correlations.The paper reports a theory-guided architecture enhancement for DG.Source: §1.2 · p. 2, §1.1 · p. 2
  • GMoE with ERM achieves state-of-the-art performance on 7 of 8 DomainBed datasets in train-validation selection and all 8 in leave-one-domain-out evaluation.GMoE trained with DG algorithms performs better than GMoE trained with ERM.Source: §1.2 · p. 2

2 PRELIMINARIES

The preliminaries describe distribution shifts through changing attribute distributions and introduce algorithmic alignment as a measure of how a backbone matches a decomposed target task. They also explain how this alignment framework motivates the paper’s DG analysis.

2.1 NOTATIONS

  • The notation section identifies Etr, Dtr, Ete, and Dte as the training dataset, training distribution, test dataset, and test distribution.Source: §2.1 · p. 2

2.2 ATTRIBUTE FACTORIZATION

  • The attribute-factorization account describes inputs and labels through attributes, with labels depending on one or more attributes.In classic computer vision, these attributes are called visual attributes.Source: §2.2 · p. 3
  • Attribute factorization models distribution shifts as changes in attribute distributions while retaining a shared conditional generative process.The example uses color as a training shortcut for shape classification, which fails when the test attribute distribution changes.Source: §2.2 · p. 3

2.3 ALGORITHMIC ALIGNMENT

  • Algorithmic alignment measures how a network’s modules match a target function decomposed into component functions.The framework relates easier component subtasks to improved IID generalization and is extended to DG later in the paper.Source: §2.3 · p. 3
  • The alignment value uses a sample-complexity measure for learning each component function at a specified precision and failure probability.Source: §2.3 · p. 3

3 ON THE IMPORTANCE OF NEURAL ARCHITECTURE FOR DOMAIN GENERALIZATION

The section begins with a comparison suggesting that backbone choice can matter for DG, then develops an alignment-based account of robustness under distribution shifts. Its analysis links robustness to whether the architecture aligns with invariant rather than spurious correlations.

3 ON THE IMPORTANCE OF NEURAL ARCHITECTURE FOR DOMAIN GENERALIZATION

  • The section frames its analysis as a progression from a motivating comparison to a formal framework for studying architecture’s effect on DG.Source: §3 · p. 3

3.1 A MOTIVATING EXAMPLE: CNNS VERSUS VISION TRANSFORMERS

  • A ViT trained with ERM already outperforms CNNs trained with state-of-the-art DG algorithms on several datasets, motivating a theoretical investigation of backbone architecture.Figure 1 also compares ViT-S/16 with 21.8M trainable parameters against ResNet-50 with 25.6M.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 UNDERSTANDING FROM A THEORETICAL PERSPECTIVE

  • The theorem assumes that training features cover the support of test features and that labels are noise-free for the target-function definition.The authors describe feature-support coverage as practical in DG settings, particularly with pretrained models.Source: §3.2 · p. 4
  • Under the theorem’s assumptions, alignment with an invariant correlation yields low test error, whereas alignment with a spurious correlation yields higher test error with high probability.The authors report that synthetic-data experiments match the theorem and that the result is verified on real datasets.Source: §3.2 · p. 5
  • In OfficeHome and DomainNet domains where shape is invariant and texture is spurious, the authors connect ViT’s shape bias and CNNs’ texture bias to the observed performance comparison.The passage attributes the shape and texture biases to multi-head attention and convolution, respectively.Source: §3.2 · p. 5
  • The authors argue that functional visual attributes are more fundamental than shape and motivate designing backbones to capture them.Source: §3.2 · p. 5

4 GENERALIZABLE MIXTURE-OF-EXPERTS FOR DOMAIN GENERALIZATION

GMoE adapts sparse mixture-of-experts to domain generalization by aligning expert computation with visual attributes and conditional rules. Its design combines attribute-oriented routing with high-level MoE placement.

4.1 MIXTURE-OF-EXPERTS LAYER

  • A ViT MoE layer replaces its feed-forward network with expert FFNs and uses a top-k gate to select experts.The gate applies a softmax to its learned projection; top-k retains the largest k values.Source: §4.1 · p. 5

4.2 VISUAL ATTRIBUTES, CONDITIONAL STATEMENTS, AND SPARSE MOES

  • Theorem 2 connects sparse MoE with one selected expert to conditional IF/ELSE computations for recognizing and combining visual attributes.The paper argues that conditional rules can capture attributes and combine them for domain-general recognition.Source: §4.2 · p. 6
  • Grouping similar conditional statements lowers the alignment objective, and experiments find that sparse MoE experts handle groups of similar visual attributes.The analysis says MoE alignment with conditional statements is better than with the original FFN.Source: §4.2 · p. 6

4.3 ADAPTING MOE TO DOMAIN GENERALIZATION

  • The cosine router matches normalized patch representations to learned embeddings, treating them as filters for visual attributes rather than relying only on domain similarity.The authors interpret the learned embeddings as a visual-attribute codebook and the normalized dot product as a matched filter.Source: §4.3 · p. 7, Figure 2 · p. 7
  • GMoE places MoE layers in the last two even layers; this outperforms every-two placement with fewer computations.The paper argues that processing visual attributes involves high-level conditional sentences.Source: §4.3 · p. 7
  • GMoE uses perturbation and load-balance loss to train diverse experts.These techniques follow Riquelme et al. (2021).Source: §4.3 · p. 7

5 EXPERIMENTAL RESULTS

GMoE trained with ERM performs strongly across DomainBed, while additional experiments examine its compatibility with DG algorithms, single-source generalization, and expert behavior.

5.1 DOMAINBED RESULTS

  • GMoE achieves the best train-validation results on PACS, VLCS, OfficeHome, and DomainNet, ranks third on TerraIncognita, and outperforms the other two baselines on three large-scale datasets.The large-scale datasets are SVIRO, Wilds-Camelyon, and Wilds-FMOW.Source: Table 1 · p. 8, §5.1 · p. 8

5.2 GMOE WITH DG ALGORITHMS

  • Replacing ResNet-50 with GMoE significantly promotes accuracy for Fish and SWAD, showing that GMoE can complement DG algorithms.Fish modifies the loss function, while SWAD uses model ensembling.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 SINGLE-SOURCE DOMAIN GENERALIZATION RESULTS

  • In single-source DomainNet experiments, GMoE’s OOD generalization gain over ResNet or ViT is much larger than its IID gain.Models are trained on Paint and tested on Paint validation data and the other five domains; the table reports accuracy and FLOPs.Source: §5.3 · p. 8, Table 3 · p. 9

5.4 MODEL ANALYSIS

  • Without visual-attribute supervision, GMoE experts specialize in distinct attributes, and similar attributes can share an expert, such as both wings.The analysis matches each attribute to its nearest nine patches and correlates those patches with selected experts.Source: §5.4 · p. 9
  • Router visualizations show experts 0 and 2 assigned to background patches, expert 3 to wings, beak, and tail, and expert 4 to both legs.Figure 3(c) visualizes block-10 top-1 router selections for patches in a natural-domain CUB-DG image.Source: §5.4 · p. 9, Figure 3 · p. 9

6 CONCLUSIONS

The paper presents backbone alignment with invariant correlations as a route to robustness under distribution shifts and proposes GMoE based on that analysis. It frames this work as an initial step toward architecture design for domain generalization.

6 CONCLUSIONS

  • The authors prove that architectures aligned with invariant correlations are more robust to distribution shifts, verify this on synthetic and real datasets, and propose GMoE, which performs well on DomainBed.They describe the paper as an initial step in exploring backbone architecture’s impact on domain generalization.Source: §6 · p. 9

A.1 DOMAIN GENERALIZATION

  • Domain-generalization approaches include invariant learning, ensemble and meta-learning, data manipulation, and neural architecture search; the paper positions backbone design as complementary to these approaches.The authors state that GMoE may gain performance when combined with the approaches discussed.Source: §A.1 · pp. 18–19

A.2 VISION TRANSFORMERS

  • Prior work reports empirical robustness of vision transformers over CNNs to adversarial noise and distribution shifts, while this paper offers theoretical validation and uses the analysis to motivate MoE.Source: §A.2 · p. 19

A.3 SPARSE MIXTURE-OF-EXPERTS

  • Sparse MoE routing selects only a few experts to keep inference time on par with a standalone counterpart; this paper repurposes MoE for domain generalization.The authors state that they modify the architecture for performance enhancement.Source: §A.3 · p. 19

A.4 OUT-OF-DISTRIBUTION GENERALIZATION THEORY

  • Unlike cited OOD theories focused on linear invariant correlations or differing train-test supports, this work studies arbitrary nonlinear invariant correlations.The paper also focuses on why the assumptions hold in vision-task domain generalization.Source: §A.4 · p. 19

A.5 ALGORITHMIC ALIGNMENT

  • The paper adopts algorithmic alignment, which relates generalization bounds to alignment between an algorithm and a network computation graph, to study domain generalization.The framework’s stated intuition is that good alignment lets a network learn simpler functions, improving sample efficiency.Source: §A.5 · p. 20

A.6 EMPIRICAL INVESTIGATIONS OF DISTRIBUTION SHIFT

  • Earlier empirical studies examined data augmentation, optimizers, backbones, and DG algorithms but did not explain when models succeed or how backbone design could improve DG.This paper presents its theoretical investigation as an initial attempt to address those questions.Source: §A.6 · p. 20

B.1 PROOF OF THEOREM 1

Theorem 1 links architecture–correlation alignment to robustness under distribution shift: alignment with invariant correlations supports generalization, while alignment with spurious correlations predicts poor test performance. Synthetic experiments test this relationship in noiseless and noisy settings.

B.1 PROOF OF THEOREM 1

  • The proof’s generalization result is conditional on assumptions including auxiliary labels, which the authors call impractical, though they note the same pattern appears in end-to-end experiments.The proof also assumes algorithm stability and sequential learning.Source: §B.1 · p. 20
  • Under the theorem’s assumptions, invariant-correlation alignment yields small test error, whereas spurious-correlation alignment yields large test error under distribution shift.The result assumes ERM and conditions on the corresponding alignment values being at most the training-set size.Source: §B.1 · pp. 20–21

B.2.1 NOISELESS VERSION

  • In the noiseless synthetic tests, the MLP fails on test dataset 2 and the FCN fails on test dataset 1, matching their different alignment patterns.The figure identifies the MLP with invariant correlations in test dataset 1 and the FCN with invariant correlations in test dataset 2.Source: §B.2.1 · p. 22, Figure 4 · p. 22

B.2.2 NOISY VERSION

  • With noisy pixel- and patch-level signals, the MLP favors its aligned pixel-level feature when signal-to-noise ratios are similar, but may change features after longer training when p1 and p2 differ greatly.The MLP initially exceeds 90% accuracy on Test1; the passage says early stopping can benefit DG when the backbone aligns with invariant correlations.Source: §B.2.2 · p. 22

B.3 PROOF OF THEOREM 2

Theorem 2 relates sparse MoE architecture to conditional statements through algorithmic alignment. The technical discussion also bounds the assumed distribution shift: training and test support must match.

B.3 PROOF OF THEOREM 2

  • An MoE network can simulate the conditional-statement algorithm; with at least as many experts as functions, it assigns one expert per function, while fewer experts group similar functions.The theorem statement concerns an MoE module with k = 1.Source: §B.3 · pp. 22, 24

B.4 TECHNICAL LEMMAS

  • The support condition is the largest distribution shift the authors assume, because nonlinear out-of-support generalization is impossible under the cited result.The lemma’s implication compares train and test probabilities for the same support.Source: §B.4 · p. 24

C.1 DETAILED ARCHITECTURE

GMoE adapts a vision-transformer backbone with sparse experts and routing, using design choices intended for domain generalization and smaller training datasets. The section also specifies its initialization, auxiliary losses, benchmarks, evaluation protocols, and training schedules.

C.1 DETAILED ARCHITECTURE

  • GMoE uses a ViT patch-embedding pipeline and places MoE layers in every-two or last-two configurations, with six experts per block and TOP-2 routing for each patch.The authors choose six experts because DG datasets are typically 1–2 orders of magnitude smaller than ImageNet-21K, so more experts could be inadequately trained.Source: §C.1 · p. 25

C.2 MODEL INITIALIZATION

  • GMoE is initialized from ImageNet-1K-pretrained ViT weights; the reported comparison says ViT-S/16 and GMoE-S/16 achieve strong generalization with fewer parameters and a smaller architecture.The stated initialization is intended to avoid attributing gains to a stronger pretrained model such as V-MoE pretrained on ImageNet-21K.Source: §C.2 · pp. 25–26, Table 4 · p. 26

C.3 PERTURBATION

  • Adding Gaussian noise to the router empirically improves performance and stabilizes training dynamics.The routing noise addresses discontinuities that can occur when the top two routing outputs are close.Source: §C.3 · p. 26

C.4 AUXILIARY LOSS

  • GMoE combines importance and load losses to encourage balanced expert usage and assignment across patches and experts.The load-loss proxy estimates an expert’s chance of remaining in TOPk when its noise is resampled; the total loss balances auxiliary and task losses.Source: §C.4 · pp. 26–27

D.1 DOMAINBED DETAILS

  • The evaluation covers eight DomainBed datasets and uses both train-validation selection and leave-one-domain-out validation selection.For the latter protocol, one domain is held out for validation and another for testing.Source: §D.1 · pp. 28–29

D.2 TRAINING DETAILS & HYPERPARAMETERS

  • DomainNet results use 15K iterations for comparison with later SOTA methods, while GMoE is also trained for 5K iterations for a fair comparison across settings.The authors note that 5K iterations amount to fewer than two DomainNet epochs; other DomainBed datasets use 5K iterations.Source: §D.2 · p. 29

D.3 LEAVE-ONE-DOMAIN-OUT RESULTS

GMoE outperforms previous methods across the reported leave-one-domain-out datasets, and its advantage over ViT and ResNet remains under a variety of DG algorithms.

D.3 LEAVE-ONE-DOMAIN-OUT RESULTS

  • GMoE outperforms previous methods on all 5 datasets in the more challenging leave-one-domain-out setting.The section summarizes 61 individual experiments across the five datasets; Table 10 reports overall out-of-domain accuracies.Source: §D.3 · pp. 29–30, Table 10 · p. 30

D.4 ABLATION STUDY WITH DG ALGORITHMS

  • Across a variety of DG algorithms, GMoE significantly outperforms both ViT and ResNet, supporting the generalizability of the proposed backbone architecture.The comparison is reported as an ablation study in Table 11.Source: §D.4 · p. 30

D.5 SINGLE-SOURCE DOMAIN GENERALIZATION RESULTS

The section evaluates single-source domain generalization by training on one DomainNet domain and testing on the remaining five, with IID validation results also reported.

D.5 SINGLE-SOURCE DOMAIN GENERALIZATION RESULTS

  • The evaluation alternates the training domain, tests on the other five DomainNet domains, and reports IID validation performance on the training domain.IID and OOD results are distinguished by color in Table 12.Source: §D.5 · p. 30
  • Table 11 reports overall accuracies (%) using a train-validation selection criterion.Source: Table 11 · p. 31

D.6 ABLATION STUDY ON MODEL DESIGN

The appendices test GMoE design choices, compare it with a strongly pretrained ResNet, and assess computational cost. The reported results favor higher-level configurations and show relatively low runtime memory and training-step time among competitors.

D.6 ABLATION STUDY ON MODEL DESIGN

  • The last two layer configurations outperform the other two, while a cosine router and larger model also improve generalization.The comparison includes GMoE-S/16 and GMoE-B/16 configurations with six experts and 12 blocks.Source: §D.6 · p. 31, Table 13 · p. 33

D.7 ABLATION STUDY ON RESNET WITH STRONG DATA AUGMENTATION

  • GMoE and ViT retain higher performance than ResNet-50 V2 despite its stronger pretraining recipe and training tricks.The comparison reports train-validation selection performance for ViT, GMoE, and DG algorithms using ResNet-50 V2 as backbone.Source: §D.7 · p. 31, Table 14 · p. 33, Table 15 · p. 33

D.8 COMPUTATIONAL COST COMPARISON

  • GMoE achieves relatively low runtime memory and training-step time among competitors.Table 16 compares training and inference iteration time and runtime memory for 160-image mini-batches; lower is better.Source: §D.8 · p. 32, Table 16 · p. 33

E.1 CUB-DG RESULTS & VISUALIZATIONS

On CUB-DG, GMoE is reported as more effective than invariant-learning DG methods, while visualizations show experts attending to recurring attributes across bird classes and stylizations. The paper also presents MoE as a way to distribute patch processing across experts.

E.1 CUB-DG RESULTS & VISUALIZATIONS

  • GMoE is reported as significantly more effective than methods designed to learn invariance through losses and gradient constraints on CUB-DG.Table 17 reports out-of-domain accuracy by domain.Source: §E.1 · p. 34, Table 17 · p. 34
  • Across bird classes and viewing angles, the same experts focus on the same visual attributes, such as beaks and legs.Figures 8–9 illustrate expert selection across natural-domain images.Source: §E.1 · p. 34, Figure 8 · p. 36, Figure 9 · p. 36
  • Across four stylizations of an image, the same regions and attributes are usually handled by the same experts.The examples show experts 1, 3, and 4 focusing on the bird and experts 0 and 2 on the background.Source: §E.1 · p. 34, Figure 10 · p. 37, Figure 11 · p. 38
  • MoE may alleviate redundant attention heads; the visualization contrasts attention signals across heads with MoE’s patch-wise information handling.The paper describes this as potentially increasing disentanglement.Source: §E.1 · p. 34, Figure 14 · p. 41, Figure 15 · p. 42
  • The MoE layer assigns patches to different experts to handle visual attributes and their correlations through attention.The GMoE block diagram depicts experts processing patches, with the router omitted.Source: Figure 7 · p. 35, §E.1 · p. 35

Related papers

Loading 2206.04046v6…