Diverse Weight Averaging for Out-of-Distribution Generalization

Introduction

Learning robust models that generalize well is critical for many real-world applications . Yet, the classical Empirical Risk Minimization (ERM) lacks robustness to distribution shifts . To improve out-of-distribution (OOD) generalization in classification, several recent works proposed to train models simultaneously on multiple related but different domains . Though theoretically appealing, domain-invariant approaches either underperform or only slightly improve ERM on the reference DomainBed benchmark . The state-of-the-art strategy on DomainBed is currently to average the weights obtained along a training trajectory . argues that this weight averaging (WA) succeeds in OOD because it finds solutions with flatter loss landscapes.

In this paper, we show the limitations of this flatness-based analysis and provide a new explanation for the success of WA in OOD. It is based on WA’s similarity with ensembling , a well-known strategy to improve robustness , that averages the predictions from various models. Based on , we present a bias-variance-covariance-locality decomposition of WA’s expected error. It contains four terms: first the bias that we show increases under shift in label posterior distributions (i.e., correlation shift ); second, the variance that we show increases under shift in input marginal distributions (i.e., diversity shift ); third, the covariance that decreases when models are diverse; finally, a locality condition on the weights of averaged models.

Based on this analysis, we aim at obtaining diverse models whose weights are averageable with our Diverse Weight Averaging (DiWA) approach. In practice, DiWA averages in weights the models obtained from independent training runs that share the same initialization. The motivation is that those models are more diverse than those obtained along a single run . Yet, averaging the weights of independently trained networks with batch normalization and ReLU layers may be counter-intuitive. Such averaging is efficient especially when models can be connected linearly in the weight space via a low loss path. Interestingly, this linear mode connectivity property was empirically validated when the runs start from a shared pretrained initialization . This insight is at the heart of DiWA but also of other recent works , as discussed in Section 6.

In summary, our main contributions are the following:

We propose a new theoretical analysis of WA for OOD based on a bias-variance-covariance-locality decomposition of its expected error (Section 2). By relating correlation shift to its bias and diversity shift to its variance, we show that WA succeeds under diversity shift.

We empirically tackle the covariance term by increasing the diversity across models averaged in weights. In our DiWA approach, we decorrelate their training procedures: in practice, these models are obtained from independent runs (Section 3). We then empirically validate that diversity improves OOD performance (Section 4) and show that DiWA is state of the art on all real-world datasets from the DomainBed benchmark (Section 5).

Theoretical insights

Under the setting described in Section 2.1, we introduce WA in Section 2.2 and decompose its expected OOD error in Section 2.3. Then, we separately consider the four terms of this bias-variance-covariance-locality decomposition in Section 2.4. This theoretical analysis will allow us to better understand when WA succeeds, and most importantly, how to improve it empirically in Section 3.

We consider a neural network (NN) $f(\cdot,\theta):\mathcal{X}\rightarrow\mathcal{Y}$ made of a fixed architecture $f$ with weights $\theta$. We seek $\theta$ minimizing the target generalization error:

2 Weight averaging for OOD and limitations of current analysis

We study the benefits of combining $M$ individual member weights $\{\theta_{m}\}_{m=1}^{M}\triangleq\{\theta(l_{S}^{(m)})\}_{m=1}^{M}$ obtained from $M$ (potentially correlated) identically distributed (i.d.) learning procedures $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$. Under conditions discussed in Section 3.2, these $M$ weights can be averaged despite nonlinearities in the architecture $f$. Weight averaging (WA) , defined as:

is the state of the art on DomainBed when the weights $\{\theta_{m}\}_{m=1}^{M}$ are sampled along a single training trajectory (a description we refine in Remark 1 from Remark 1).

To explain this success, Cha et al. argue that flat minima generalize better; indeed, WA flattens the loss landscape. Yet, as shown in Appendix B, this analysis does not fully explain WA’s spectacular results on DomainBed. First, flatness does not act on distribution shifts thus the OOD error is uncontrolled with their upper bound (see Section B.1). Second, this analysis does not clarify why WA outperforms Sharpness-Aware Minimizer (SAM) for OOD generalization, even though SAM directly optimizes flatness (see Section B.2). Finally, it does not justify why combining WA and SAM succeeds in IID yet fails in OOD (see Figure 7). These observations motivate a new analysis of WA; we propose one below that better explains these results.

3 Bias-variance-covariance-locality decomposition

Given $\{\theta_{m}\}_{m=1}^{M}$ with learning procedures $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$. Denoting \Delta_{L_{S}^{M}}=\max_{m=1}^{M}\mathopen{}\mathclose{{}\left\|\theta_{m}-\theta_{\text{WA}}}\right\|_{2}, $\forall(x,y)\in\mathcal{X}\times\mathcal{Y}$:

This similarity is useful since Equation BV was extended into a bias-variance-covariance decomposition for ENS in . We can then derive the following decomposition of WA’s expected test error. To take into account the $M$ averaged weights, the expectation is over the joint distribution describing the $M$ identically distributed (i.d.) learning procedures $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$.

Equation BVCL decomposes the OOD error of WA into four terms. The bias is the same as that of each of its i.d. members. WA’s variance is split into the variance of each of its i.d. members divided by $M$ and a covariance term. The last locality term constrains the weights to ensure the validity of our approximation. In conclusion, combining $M$ models divides the variance by $M$ but introduces the covariance and locality terms which should be controlled along bias to guarantee low OOD error.

4 Analysis of the bias-variance-covariance-locality decomposition

We now analyze the four terms in Equation BVCL. We show that bias dominates under correlation shift (Section 2.4.1) and variance dominates under diversity shift (Section 2.4.2). Then, we discuss a trade-off between covariance, reduced with diverse models (Section 2.4.3), and the locality term, reduced when weights are similar (Section 2.4.4). This analysis shows that WA is effective against diversity shift when $M$ is large and when its members are diverse but close in the weight space.

\exists\epsilon>0\text{~{}small~{}s.t.~{}}\forall x\in\mathcal{X}_{S},|f_{S}\mathopen{}\mathclose{{}\left(x}\right)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)|\leq\epsilon.

With a bounded difference between the labeling functions $f_{T}-f_{S}$ on $\mathcal{X}_{T}\cap\mathcal{X}_{S}$, under Assumption 1, the bias on domain $T$ is:

This analysis explains why WA fails under correlation shift, as shown on ColoredMNIST in Appendix H. Indeed, combining different models does not reduce the bias. Section 2.4.2 explains that WA is however efficient against diversity shift.

4.2 Variance and diversity shift

Variance is known to be large in OOD and to cause a phenomenon named underspecification, when models behave differently in OOD despite similar test IID accuracy. We now relate OOD variance to diversity shift in a simplified setting. We fix the source dataset $d_{S}$ (with input support $X_{d_{S}}$), the target dataset $d_{T}$ (with input support $X_{d_{T}}$) and the network’s initialization. We get a closed-form expression for the variance of $f$ over all other sources of randomness under Assumptions 2 and 3.

This states that $f$ behaves as a Gaussian process (GP); it is reasonable if $f$ is a wide network . The corresponding kernel $K$ is the neural tangent kernel (NTK) depending only on the initialization. GPs are useful because their variances have a closed-form expression (Section C.4.1). To simplify the expression of variance, we now make Assumption 3.

$\exists(\lambda_{S},\epsilon)\text{~{}with~{}}\linebreak 0\leq\epsilon\ll\lambda_{S}\text{~{}such~{}that~{}}\forall x_{S}\in X_{d_{S}},K(x_{S},x_{S})=\lambda_{S}\text{~{}and~{}}\forall x_{S}^{\prime}\neq x_{S}\in X_{d_{S}},|K(x_{S},x_{S}^{\prime})|\leq\epsilon$.

This states that training samples have the same norm (following standard practice ) and weakly interact . This assumption is further discussed and relaxed in Section C.4.2. We are now in a position to relate variance and diversity shift when $\epsilon\to 0$.

Given $f$ trained on source dataset $d_{S}$ (of size $n_{S}$) with NTK $K$, under Assumptions 2 and 3, the variance on dataset $d_{T}$ is:

The MMD empirically estimates shifts in input marginals, i.e., between $p_{S}(X)$ and $p_{T}(X)$. Our expression of variance is thus similar to the diversity shift formula in : MMD replaces the $L_{1}$ divergence used in . The other terms, $\lambda_{T}$ and $\beta_{T}$, both involve internal dependencies on the target dataset $d_{T}$: they are constants w.r.t. $X_{d_{T}}$ and do not depend on distribution shifts. At fixed $d_{T}$ and under our assumptions, Equation 4 shows that variance on $d_{T}$ decreases when $X_{d_{S}}$ and $X_{d_{T}}$ are closer (for the MMD distance defined by the kernel $K^{2}$) and increases when they deviate. Intuitively, the further $X_{d_{T}}$ is from $X_{d_{S}}$, the less the model’s predictions on $X_{d_{T}}$ are constrained after fitting $d_{S}$.

This analysis shows that WA reduces the impact of diversity shift as combining $M$ models divides the variance per $M$. This is a strong property achieved without requiring data from the target domain.

4.3 Covariance and diversity

The covariance term increases when the predictions of $\{f(\cdot,\theta_{m})\}_{m=1}^{M}$ are correlated. In the worst case where all predictions are identical, covariance equals variance and WA is no longer beneficial. On the other hand, the lower the covariance, the greater the gain of WA over its members; this is derived by comparing Equations BV and BVCL, as detailed in Section C.5. It motivates tackling covariance by encouraging members to make different predictions, thus to be functionally diverse. Diversity is a widely analyzed concept in the ensemble literature , for which numerous measures have been introduced . In Section 3, we aim at decorrelating the learning procedures to increase members’ diversity and reduce the covariance term.

4.4 Locality and linear mode connectivity

To ensure that WA approximates ENS, the last locality term $O(\bar{\Delta}^{2})$ constrains the weights to be close. Yet, the covariance term analyzed in Section 2.4.3 is antagonistic, as it motivates functionally diverse models. Overall, to reduce WA’s error in OOD, we thus seek a good trade-off between diversity and locality. In practice, we consider that the main goal of this locality term is to ensure that the weights are averageable despite the nonlinearities in the NN such that WA’s error does not explode. This is why in Section 3, we empirically relax this locality constraint and simply require that the weights are linearly connectable in the loss landscape, as in the linear mode connectivity . We empirically verify later in Figure 1 that the approximation $f_{\text{WA}}\approx f_{\text{ENS}}$ remains valid even in this case.

DiWA: Diverse Weight Averaging

Our analysis in Sections 2.4.1 and 2.4.2 showed that the bias and the variance terms are mostly fixed by the distribution shifts at hand. In contrast, the covariance term can be reduced by enforcing diversity across models (Section 2.4.3) obtained from learning procedures $\{l_{S}^{(m)}\}_{m=1}^{M}$. Yet, previous methods only average weights obtained along a single run. This corresponds to highly correlated procedures sharing the same initialization, hyperparameters, batch orders, data augmentations and noise, that only differ by the number of training steps. The models are thus mostly similar: this does not leverage the full potential of WA.

Our Diverse Weight Averaging approach seeks to reduce the OOD expected error in Equation BVCL by decreasing covariance across predictions: DiWA decorrelates the learning procedures $\{l_{S}^{(m)}\}_{m=1}^{M}$. Our weights are obtained from $M\gg 1$ different runs, with diverse learning procedures: these have different hyperparameters (learning rate, weight decay and dropout probability), batch orders, data augmentations (e.g., random crops, horizontal flipping, color jitter, grayscaling), stochastic noise and number of training steps. Thus, the corresponding models are more diverse on domain $T$ per and reduce the impact of variance when $M$ is large. However, this may break the locality requirement analyzed in Section 2.4.4 if the weights are too distant. Empirically, we show that DiWA works under two conditions: shared initialization and mild hyperparameter ranges.

2 Approach: shared initialization, mild hyperparameter search and weight selection

The shared initialization condition follows : when models are fine-tuned from a shared pretrained model, their weights can be connected along a linear path where error remains low . Following standard practice on DomainBed , our encoder is pretrained on ImageNet ; this pretraining is key as it controls the bias (by defining the feature support mismatch, see Section 2.4.1) and variance (by defining the kernel $K$, see Section C.4.4). Regarding the classifier initialization, we test two methods. The first is the random initialization, which may distort the features . The second is Linear Probing (LP) : it first learns the classifier (while freezing the encoder) to serve as a shared initialization. Then, LP fine-tunes the encoder and the classifier together in the $M$ subsequent runs; the locality term is smaller as weights remain closer (see ).

As shown in Figure 5, extreme hyperparameter ranges lead to weights whose average may perform poorly. Indeed, weights obtained from extremely different hyperparameters may not be linearly connectable; they may belong to different regions of the loss landscape. In our experiments, we thus use the mild search space defined in Table 7, first introduced in SWAD . These hyperparameter ranges induce diverse models that are averageable in weights.

The last step of our approach (summarized in Algorithm 1) is to choose which weights to average among those available. We explore two simple weight selection protocols, as in . The first uniform equally averages all weights; it is practical but may underperform when some runs are detrimental. The second restricted (greedy in ) solves this drawback by restricting the number of selected weights: weights are ranked in decreasing order of validation accuracy and sequentially added only if they improve DiWA’s validation accuracy.

In the following sections, we experimentally validate our theory. First, Section 4 confirms our findings on the OfficeHome dataset where diversity shift dominates (see Section E.2 for a similar analysis on PACS ). Then, Section 5 shows that DiWA is state of the art on DomainBed .

Empirical validation of our theoretical insights

We consider several collections of weights $\{\theta_{m}\}_{m=1}^{M}$ ($2\leq M<10$) trained on the “Clipart”, “Product” and “Photo” domains from OfficeHome with a shared random initialization and mild hyperparameter ranges. These weights are first indifferently sampled from a single run (every $50$ batches) or from different runs. They are evaluated on “Art”, the fourth domain from OfficeHome.

Figure 1 validates Lemma 1 and that $f_{\text{WA}}\approx f_{\text{ENS}}$. More precisely, $f_{\text{WA}}$ slightly but consistently improves $f_{\text{ENS}}$: we discuss this in Appendix D. Moreover, a larger $M$ improves the results; in accordance with Equation BVCL, this motivates averaging as many weights as possible. In contrast, large $M$ is computationally impractical for ENS at test time, requiring $M$ forwards.

Now we investigate the difference between sampling the weights from a single run or from different runs. Figure 3 first shows that diversity increases when weights come from different runs. Second, in Figure 4, this is reflected on the accuracies in OOD. Here, we rank by validation accuracy the $60$ weights obtained (1) from $60$ different runs and (2) along $1$ well-performing run. We then consider the WA of the top $M$ weights as $M$ increases from $1$ to $60$. Both have initially the same performance and improve with $M$; yet, WA of weights from different runs gradually outperforms the single-run WA. Finally, Figure 5 shows that this holds only for mild hyperparameter ranges and with a shared initialization. Otherwise, when hyperparameter distributions are extreme (as defined in Table 7) or when classifiers are not similarly initialized, DiWA may perform worse than its members due to a violation of the locality condition. These experiments confirm that diversity is key as long as the weights remain averageable.

Experimental results on the DomainBed benchmark

Datasets. We now present our evaluation on DomainBed . By imposing the code, the training procedures and the ResNet50 architecture, DomainBed is arguably the fairest benchmark for OOD generalization. It includes $5$ multi-domain real-world datasets: PACS , VLCS , OfficeHome , TerraIncognita and DomainNet . showed that diversity shift dominates in these datasets. Each domain is successively considered as the target $T$ while other domains are merged into the source $S$. The validation dataset is sampled from $S$, i.e., we follow DomainBed’s training-domain model selection. The experimental setup is further described in Section G.1. Our code is available at https://github.com/alexrame/diwa.

Baselines. ERM is the standard Empirical Risk Minimization. Coral is the best approach based on domain invariance. SWAD (Stochastic Weight Averaging Densely) and MA (Moving Average) average weights along one training trajectory but differ in their weight selection strategy. SWAD is the current state of the art (SoTA) thanks to it “overfit-aware” strategy, yet at the cost of three additional hyperparameters (a patient parameter, an overfitting patient parameter and a tolerance rate) tuned per dataset. In contrast, MA is easy to implement as it simply combines all checkpoints uniformly starting from batch $100$ until the end of training. Finally, we report the scores obtained in for the costly Deep Ensembles (DENS) (with different initializations): we discuss other ensembling strategies in Appendix D.

Our runs. ERM and DiWA share the same training protocol in DomainBed: yet, instead of keeping only one run from the grid-search, DiWA leverages $M$ runs. In practice, we sample $20$ configurations from the hyperparameter distributions detailed in Table 7 and report the mean and standard deviation across $3$ data splits. For each run, we select the weights of the epoch with the highest validation accuracy. ERM and MA select the model with highest validation accuracy across the $20$ runs, following standard practice on DomainBed. Ensembling (ENS) averages the predictions of all $M=20$ models (with shared initialization). DiWA-restricted selects $1\leq M\leq 20$ weights with Algorithm 1 while DiWA-uniform averages all $M=20$ weights. DiWA† averages uniformly the $M=3\times 20=60$ weights from all $3$ data splits. DiWA† benefits from larger $M$ (without additional inference cost) and from data diversity (see Section E.1.3). However, we cannot report standard deviations for DiWA† for computational reasons. Moreover, DiWA† cannot leverage the restricted weight selection, as the validation is not shared across all $60$ weights that have different data splits.

We report our main results in Table 1, detailed per domain in Section G.2. With a randomly initialized classifier, DiWA†-uniform is the best on PACS, VLCS and OfficeHome: DiWA-uniform is the second best on PACS and OfficeHome. On TerraIncognita and DomainNet, DiWA is penalized by some bad runs, filtered in DiWA-restricted which improves results on these datasets. Classifier initialization with linear probing (LP) improves all methods on OfficeHome, TerraIncognita and DomainNet. On these datasets, DiWA† increases MA by $1.3$, $0.5$ and $1.1$ points respectively. After averaging, DiWA† with LP establishes a new SoTA of $68.0\%$, improving SWAD by $1.1$ points.

DiWA with different objectives. So far we used ERM that does not leverage the domain information. Table 2 shows that DiWA-uniform benefits from averaging weights trained with Interdomain Mixup and Coral : accuracy gradually improves as we add more objectives. Indeed, as highlighted in Section E.1.3, DiWA benefits from the increased diversity brought by the various objectives. This suggests a new kind of linear connectivity across models trained with different objectives; the full analysis of this is left for future work.

2 Limitations of DiWA

Despite this success, DiWA has some limitations. First, DiWA cannot benefit from additional diversity that would break the linear connectivity between weights — as discussed in Appendix D. Second, DiWA (like all WA approaches) can tackle diversity shift but not correlation shift: this property is explained for the first time in Section 2.4 and illustrated in Appendix H on ColoredMNIST.

Related work

Generalization and ensemble. To generalize under distribution shifts, invariant approaches try to detect the causal mechanism rather than memorize correlations: yet, they do not outperform ERM on various benchmarks . In contrast, ensembling of deep networks consistently increases robustness and was successfully applied to domain generalization . As highlighted in (whose analysis underlies our Equation BVCL), ensembling works due to the diversity among its members. This diversity comes primarily from the randomness of the learning procedure and can be increased with different hyperparameters , data , augmentations or with regularizations .

Weight averaging. Recent works combine in weights (rather than in predictions) models collected along a single run. This was shown suboptimal in IID but successful in OOD . Following the linear mode connectivity and the property that many independent models are connectable , a second group of works average weights with fewer constraints . To induce greater diversity, used a high constant learning rate; explicitly encouraged the weights to encompass more volume in the weight space; minimized cosine similarity between weights; used a tempered posterior. From a loss landscape perspective , these methods aimed at “explor[ing] the set of possible solutions instead of simply converging to a single point”, as stated in . The recent “Model soups” introduced by Wortsman et al. is a WA algorithm similar to Algorithm 1; yet, the theoretical analysis and the goals of these two works are different. Theoretically, we explain why WA succeeds under diversity shift: the bias/correlation shift, variance/diversity shift and diversity-based findings are novel and are confirmed empirically. Regarding the motivation, our work aims at combining more diverse weights: it may be analyzed as a general framework to average weights obtained in various ways. In contrast, challenges the standard model selection after a grid search. Regarding the task, and our work complement each other: while demonstrate robustness on several ImageNet variants with distribution shift, we improve the SoTA on the multi-domain DomainBed benchmark against other established OOD methods after a thorough and fair comparison. Thus, DiWA and are theoretically complementary with different motivations and applied successfully for different tasks.

Conclusion

In this paper, we propose a new explanation for the success of WA in OOD by leveraging its ensembling nature. Our analysis is based on a new bias-variance-covariance-locality decomposition for WA, where we theoretically relate bias to correlation shift and variance to diversity shift. It also shows that diversity is key to improve generalization. This motivates our DiWA approach that averages in weights models trained independently. DiWA improves the state of the art on DomainBed, the reference benchmark for OOD generalization. Critically, DiWA has no additional inference cost — removing a key limitation of standard ensembling. Our work may encourage the community to further create diverse learning procedures and objectives — whose models may be averaged in weights.

We would like to thank Jean-Yves Franceschi for his helpful comments and discussions on our paper. This work was granted access to the HPC resources of IDRIS under the allocation AD011011953 made by GENCI. We acknowledge the financial support by the French National Research Agency (ANR) in the chair VISA-DEEP (project number ANR-20-CHIA-0022-01) and the ANR projects DL4CLIM ANR-19-CHIA-0018-01, RAIMO ANR-20-CHIA-0021-01, OATMIL ANR-17-CE23-0012 and LEAUDS ANR-18-CE23-0020.

References

Checklist

Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]

Did you describe the limitations of your work? [Yes] In Section 5.2.

Did you discuss any potential negative societal impacts of your work? [Yes] In Appendix A

Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

If you are including theoretical results…

Did you state the full set of assumptions of all theoretical results? [Yes] Assumption 1 discussed in Section C.3.2 and Assumptions 2 and 3 discussed in Section C.4.2.

Did you include complete proofs of all theoretical results? [Yes] In Appendix C

Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] Our code is available at https://github.com/alexrame/diwa.

Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section 5 and Section G.1

Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] Defined by different data splits when possible.

Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Approximately $20000$ hours of GPUs (Nvidia V100) on an internal cluster, mostly for the $2640$ runs needed in Table 1.

If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

If your work uses existing assets, did you cite the creators? [Yes] DomainBed benchmark and its datasets.

Did you mention the license of the assets? [Yes] DomainBed is under “The MIT License”.

Did you include any new assets either in the supplemental material or as a URL? [No]

Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

If you used crowdsourcing or conducted research with human subjects…

Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]

This supplementary material complements the main paper. It is organized as follows:

Appendix A describes the broader impact of our work.

Appendix B points out the limitations of existing flatness-based analysis of WA and shows how our analysis solves these limitations.

Appendix C details all the proofs of the propositions and lemmas found in our work.

Sections C.1 and 1 derive the bias-variance-covariance-locality decomposition for WA (Proposition 1).

Section C.3 establishes the link between bias and correlation shift (Proposition 2).

Section C.4 establishes the link between variance and diversity shift (Proposition 3).

Section C.5 compares WA with one of its member (Lemma 3).

Appendix D empirically compares WA to functional ensembling ENS.

Appendix E presents some additional diversity results on OfficeHome and PACS.

Appendix F ablates the importance of the number of training runs.

Appendix G describes our experiments on DomainBed and our per-domain results.

Appendix H empirically confirms a limitation of WA approaches expected from our theoretical analysis: they do not tackle correlation shift on ColoredMNIST.

Appendix I suggests DiWA’s potential when some target data is available for training .

Appendix A Broader impact statement

We believe our paper can have several positive impacts. First, our theoretical analysis enables practitioners to know when averaging strategies succeed (under diversity shift, where variance dominates) or break down (under correlation shift, where bias dominates). This is key to understand when several models can be combined into a production system, or if the focus should be put on the training objective and/or the data. Second, it sets a new state of the art for OOD generalization under diversity shift without relying on a specific objective, architecture or task prior. It could be useful in medicine or to tackle fairness issues related to under-representation . Finally, DIWA has no additional inference cost; in contrast, functional ensembling needs one forward per member. Thus, DiWA removes the carbon footprint overhead of ensembling strategies at test-time.

Yet, our paper may also have some negative impacts. First, it requires independent training of several models. It may motivate practitioners to learn even more networks and average them afterwards. Note that in Section 5, we restricted ourselves to combining only the runs obtained from the standard ERM grid search from DomainBed . Second, our model is fully deep learning based with the corresponding risks, e.g., adversarial attacks and lack of interpretability. Finally, we do not control its possible use to surveillance or weapon systems.

Appendix B Limitations of the flatness-based analysis in OOD

Consider a set of $N$ covers \mathopen{}\mathclose{{}\left\{\Theta_{k}}\right\}_{k=1}^{N} s.t. the parameter space $\Theta\subset\cup_{k}^{N}\Theta_{k}$ where \operatorname{diam}(\Theta)\triangleq\sup_{\theta,\theta^{\prime}\in\Theta}\mathopen{}\mathclose{{}\left\|\theta-\theta^{\prime}}\right\|_{2},N\triangleq\mathopen{}\mathclose{{}\left\lceil(\operatorname{diam}(\Theta)/\gamma)^{d}}\right\rceil and $d$ is the dimension of $\Theta$. Then, $\forall\theta\in\Theta$ with probability at least $1-\delta$:

\operatorname{Div}(p_{S},p_{T})\triangleq 2\sup_{A}\mathopen{}\mathclose{{}\left|p_{S}(A)-p_{T}(A)}\right| is a divergence between the source and target marginal distributions $p_{S}$ and $p_{T}$: it measures diversity shift.

$v_{k}$ is a VC dimension of each $\Theta_{k}$.

Previous understanding of WA’s success in OOD relied on this upper-bound, where $\mathcal{E}_{d_{S}}^{\gamma}(\theta)$ involves the solution’s flatness. This is usually empirically analyzed by the trace of the Hessian : indeed, with a second-order Taylor approximation around the local minima $\theta$ and $h$ the Hessian’s maximum eigenvalue, $\mathcal{E}_{d_{S}}^{\gamma}(\theta)\approx\mathcal{E}_{d_{S}}(\theta)+h\times\gamma^{2}$.

In the following subsections, we show that this inequality does not fully explain the exceptional performance of WA on DomainBed . Moreover, we illustrate that our bias-variance-covariance-locality addresses these limitations.

B.2 SAM leads to flatter minimas but worse OOD performance

The flatness-based analysis does not explain why WA outperforms other flatness-based methods in OOD. We consider Sharpness-Aware Minimizer (SAM) , another popular method to find flat minima based on minimax optimization: it minimizes the maximum loss around a neighborhood of the current weights $\theta$. In Figure 6, we compare the flatness (i.e., the Hessian trace computed with the package in ) and accuracy of ERM, MA (a WA strategy) and SAM when trained on the “Clipart”, “Product” and “Photo” domains from OfficeHome : they are tested OOD on the fourth domain “Art”. Analyzing the second and the third rows of Figures 6(a) and 6(b), we observe that SAM indeed finds flat minimas (at least comparable to MA), both in training (IID) and test (OOD). However, this is not reflected in the OOD accuracies in Figure 6(c), where MA outperforms SAM. As reported in Table 3, similar experiments across more datasets lead to the same conclusions in . In conclusion, flatness is not sufficient to explain why WA works so well in OOD, because SAM has similar flatness but worse OOD results. In contrast, we highlight in this paper that WA succeeds in OOD by reducing the impact of variance thanks to its similarity with prediction ensembling (see Lemma 1), a privileged link that SAM does not benefit from.

B.3 WA and SAM are not complementary in OOD when variance dominates

We investigate a similar inconsistency when combining these two flatness-based methods. As argued in , we confirm in Figures 6(a) and 6(b) that MA + SAM leads to flatter minimas than MA alone (i.e., with ERM) or SAM alone. Yet, MA does not benefit from SAM in Figure 6(c). showed an even stronger result in Table 3: SWAD + ERM performs better than SWAD + SAM. We recover similar findings in Table 4: DiWA performs worse when SAM is applied in each training run.

This behavior is not explained by Theorem 1, which states that more flatness should improve OOD generalization. Yet it is explained by our diversity-based analysis. Indeed, we observe in Figure 7 that the diversity across two checkpoints along a SAM trajectory is much lower than along a standard ERM trajectory (with SGD). We speculate that this is related to the recent empirical observation made in : “the rank of the CLIP representation space is drastically reduced when training CLIP with SAM”. Under diversity shift, variance dominates (see Equation 4): in this setup, the gain in accuracy of models trained with SAM cannot compensate the decrease in diversity. This explains why WA and SAM are not complementary under diversity shift: in this case, variance is large.

Appendix C Proofs

Given $\{\theta_{m}\}_{m=1}^{M}$ with learning procedures $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$. Denoting \Delta_{L_{S}^{M}}=\max_{m=1}^{M}\mathopen{}\mathclose{{}\left\|\theta_{m}-\theta_{\text{WA}}}\right\|_{2}, $\forall(x,y)\in\mathcal{X}\times\mathcal{Y}$:

to establish the functional approximation, as , it performs Taylor expansion of the models’ predictions at the first order.

to establish the loss approximation, as , it performs Taylor expansion of the loss at the first order.

With a Taylor expansion at the first order of the models’ predictions w.r.t. parameters $\theta$:

Therefore, because $\sum_{m=1}^{M}\Delta_{m}=0$,

With a Taylor expansion at the zeroth order of the loss w.r.t. its first input and injecting Equation 6:

C.2 Bias-variance-covariance-locality decomposition

Our result in Proposition 1 is simplified by leveraging the fact that the learning procedures $L_{S}^{M}=\{l_{S}^{(m)}\}_{m=1}^{M}$ are identically distributed (i.d.). This assumption naturally holds for DiWA which selects weights from different runs with i.i.d. hyperparameters. It may be less obvious why it applies to MA and SWAD . It is even false if the weights $\{\theta(l_{S}^{(m)})\}_{m=1}^{M}$ are defined as being taken sequentially along a training trajectory, i.e., when $0\leq i<j\leq M$ implies that $l_{S}^{(i)}$ has fewer training steps than $l_{S}^{(j)}$. We propose an alternative indexing strategy to respect the i.d. assumption. Given $M$ weights selected by the weight selection procedure, we draw without replacement the $M$ weights, i.e., $\theta(l_{S}^{(i)})$ refers to the $i^{th}$ sampled weights. With this procedure, all weights are i.d. as they are uniformly sampled. Critically, their WA are unchanged for the two definitions.

it follows the bias-variance-covariance decomposition from for functional ensembling. It is tailored to WA by assuming that learning procedures are identically distributed.

it injects the obtained equation into Lemma 1 to obtain the Proposition 1 for WA.

Using $f_{\text{ENS}}\triangleq f_{\text{ENS}}(\cdot,\{\theta(l_{S}^{(m)})\}_{m=1}^{M})\triangleq\frac{1}{M}\sum_{m=1}^{M}f(\cdot,\theta(l_{S}^{(m)}))$ in this decomposition yields,

As $f_{\text{ENS}}$ depends on $L_{S}^{M}$, we extend the bias into:

Under identically distributed $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$,

Thus the bias of ENS is the same as for a single member of the WA.

Under identically distributed $L_{S}^{M}\triangleq\{l_{S}^{(m)}\}_{m=1}^{M}$,

The variance is split into the variance of a single member (divided by $M$) and a covariance term.

C.3 Bias, correlation shift and support mismatch

We first present in Section C.3.1 a decomposition of the OOD bias without any assumptions. We then justify in Section C.3.2 the simplifying Assumption 1 from Section 2.4.1.

This proof is original and based on splitting the OOD bias in and out of $\mathcal{X}_{S}$:

To decompose the first term, we write $\forall x\in\mathcal{X}_{S}$, -\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)=-f_{S}\mathopen{}\mathclose{{}\left(x}\right)+\mathopen{}\mathclose{{}\left(f_{S}(x)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right).

The four terms can be qualitatively analyzed:

The second term is exactly the IID bias, but weighted by the marginal distribution $p_{T}(X)$.

The third term \int_{\mathcal{X}_{T}\cap\mathcal{X}_{S}}2\mathopen{}\mathclose{{}\left(f_{T}\mathopen{}\mathclose{{}\left(x}\right)-f_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right)\mathopen{}\mathclose{{}\left(f_{S}\mathopen{}\mathclose{{}\left(x}\right)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right)p_{T}(x)dx measures to what extent the IID bias compensates the correlation shift. It can be negative if (by chance) the IID bias goes in opposite direction to the correlation shift.

The last term measures support mismatch between test and train marginal distributions. It lead to the “No free lunch for learning representations for DG” in . The error is irreducible because “outside of the source domain, the label distribution is unconstrained”: “for any domain which gives some probability mass on an example that has not been seen during training, then all […] labels for that example” are possible.

C.3.2 Discussion of the small IID bias Assumption 1

the network architecture $f$ which should be able to fit a given dataset $d_{S}$. This is realistic when the network is sufficiently parameterized, i.e., when the number of weights $|\theta|$ is large.

the expected datasets $d_{S}$ which should be representative enough of the underlying domain $S$; in particular the dataset size $n_{S}$ should be large.

the sampled configurations $c$ which should be well chosen: the network should be trained for enough steps, with an adequate learning rate …

For DiWA, this is realistic as it selects the weights with the highest training validation accuracy from each run. For SWAD , this is also realistic thanks to their overfit-aware weight selection strategy. In contrast, this assumption may not perfectlty hold for MA , which averages weights starting from batch $100$ until the end of training: indeed, $100$ batches are not enough to fit the training dataset.

C.3.3 OOD bias when small IID bias

We now develop our equality under Assumption 1.

With a bounded difference between the labeling functions $f_{T}-f_{S}$ on $\mathcal{X}_{T}\cap\mathcal{X}_{S}$, under Assumption 1, the bias on domain $T$ is:

We simplify the second and third terms from Proposition 4 under Assumption 1.

The second term is \int_{\mathcal{X}_{T}\cap\mathcal{X}_{S}}\mathopen{}\mathclose{{}\left(f_{S}\mathopen{}\mathclose{{}\left(x}\right)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right)^{2}p_{T}(x)dx. Under Assumption 1, |f_{S}\mathopen{}\mathclose{{}\left(x}\right)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)|\leq\epsilon. Thus the second term is $O(\epsilon^{2})$.

The third term is \int_{\mathcal{X}_{T}\cap\mathcal{X}_{S}}2\mathopen{}\mathclose{{}\left(f_{T}\mathopen{}\mathclose{{}\left(x}\right)-f_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right)\mathopen{}\mathclose{{}\left(f_{S}\mathopen{}\mathclose{{}\left(x}\right)-\bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right)}\right)p_{T}(x)dx. As $f_{T}-f_{S}$ is bounded on $\mathcal{X}_{S}\cap\mathcal{X}_{T}$, $\exists K\geq 0$ such that $\forall x\in\mathcal{X}_{S}$,

Finally, note that we cannot say anything about \bar{f}_{S}\mathopen{}\mathclose{{}\left(x}\right) when $x\in\mathcal{X}_{T}\setminus\mathcal{X}_{S}$. ∎

To prove the previous equality, we needed a bounded difference between labeling functions $f_{T}-f_{S}$ on $\mathcal{X}_{T}\cap\mathcal{X}_{S}$. We relax this bounded assumption to obtain an inequality in the following Proposition 5.

We follow the same proof as in Proposition 4, except that we now use: $(a+b)^{2}\leq 2(a^{2}+b^{2})$. Then,

C.4 Variance and diversity shift

We prove the link between variance and diversity shift. Our proof builds upon the similarity between NNs and GPs in the kernel regime, detailed in Section C.4.1. We discuss our simplifying Assumption 3 in Section C.4.2. We present our final proof in Section C.4.3. We discuss the relation between variance and initialization in Section C.4.4.

We fix $d_{S},d_{T}$ and denote $X_{d_{S}}=\{x_{S}\}_{(x_{S},y_{S})\in d_{S}}$, $X_{d_{T}}=\{x_{T}\}_{(x_{T},y_{T})\in d_{T}}$ their respective input supports. We fix the initialization of the network. $l_{S}$ encapsulates all other sources of randomness.

C.4.2 Discussion of the same norm and low similarity Assumption 3 on source dataset

Lemma 2 shows that the variance only depends on the input distributions $p(X)$ without involving the label distributions $p(Y|X)$. This formula highlights that the variance is related to shifts in input similarities (measured by $K$) between $X_{d_{S}}$ and $X_{d_{T}}$. Yet, a more refined analysis of the variance requires additional assumptions, in particular to obtain a closed-form expression of $K(X_{d_{S}},X_{d_{S}})^{-1}$. Assumption 3 is useful because then $K(X_{d_{S}},X_{d_{S}})$ is diagonally dominant and can be approximately inverted (see Section C.4.3).

The first part of Assumption 3 assumes that $\exists\lambda_{S}$ such that all training inputs $x_{S}\in X_{d_{S}}$ verify $K(x_{S},x_{S})=\lambda_{S}$. Note that this equality is standard in some kernel machine algorithms and is usually achieved by replacing $K(x,x^{\prime})$ by $\lambda_{S}\frac{K(x,x^{\prime})}{\sqrt{K(x,x)}\sqrt{K(x^{\prime},x^{\prime})}},\forall(x,x^{\prime})\in{(X_{d_{S}}\cup X_{d_{T}})}^{2}$. In the NTK literature, this equality is achieved without changing the kernel by normalizing the samples of $X_{d_{S}}$ such that they lie on the hypersphere; this input preprocessing was used in . This is theoretically based: for example, the NTK $K(x,x^{\prime})$ for an architecture with an initial fully connected layer only depends on $\|x\|,\|x^{\prime}\|,\langle x,x^{\prime}\rangle$ . Thus in the case where all samples from $X_{d_{S}}$ are preprocessed to have the same norm, the value of $K(x_{S},x_{S})$ does not depend on $x_{S}\in X_{d_{S}}$; we denote $\lambda_{S}$ the corresponding value.

The second part of Assumption 3 states that $\exists 0\leq\epsilon\ll\lambda_{S},\text{s.t.~{}}\forall x_{S},x_{S}^{\prime}\in X_{d_{S}}^{2},x_{S}\neq x_{S}^{\prime}\Rightarrow|K(x_{S},x_{S}^{\prime})|\leq\epsilon$, i.e., that training samples are dissimilar and do not interact. This diagonal structure of the NTK , with diagonal values larger than non-diagonal ones, is consistent with empirical observations from at initialization. Theoretically, this is reasonable if $K$ is close to the RBF kernel K_{h}\mathopen{}\mathclose{{}\left(x,x^{\prime}}\right)=\exp(-\mathopen{}\mathclose{{}\left\|x-x^{\prime}}\right\|_{2}^{2}/h) where $h$ would be the bandwidth: in this case, Assumption 3 is satisfied when training inputs are distant in pixel space.

We now provide an analysis of the variance where the diagonal assumption is relaxed. Specifically, we provide the sketch for proving an upper-bound of the variance when the NTK has a block-diagonal structure. This is indeed closer to the empirical observations in at the end of training, consistently with the local elasticity property of NNs . We then consider the dataset $d_{S^{\prime}}\subset d_{S}$ made of one sample per block, to which Assumption 3 applies. As decreasing the size of a training dataset empirically reduces variance , the variance of $f$ trained on $d_{S}$ is upper-bounded by the variance of $f$ trained on $d_{S^{\prime}}$; the latter is given by applying Proposition 3 to $d_{S^{\prime}}$. We believe that the proper formulation of this idea is beyond the scope of this article and best left for future theoretical work.

C.4.3 Expression of OOD variance

Given $f$ trained on source dataset $d_{S}$ (of size $n_{S}$) with NTK $K$, under Assumptions 2 and 3, the variance on dataset $d_{T}$ is:

Our proof is original and is based on the posterior form of GPs in Lemma 2. Given $d_{S}$, we recall Equation 9 that states $\forall x\in\mathcal{X}$:

Denoting $B=K(X_{d_{S}},X_{d_{S}})^{-1}$ with symmetric coefficients $b_{i,j}=b_{j,i}$, then

We fix $x_{T}\in X_{d_{T}}$ and determine the form of $B^{-1}$ in two cases: $\epsilon=0$ and $\epsilon\neq 0$.

We first derive a simplified result, when $\epsilon=0$.

Then, $b_{i,i}=\frac{1}{\lambda_{S}}$ and $b_{i,j}=0$ s.t.

We now relate the second term on the r.h.s. to a MMD distance. As $K$ is a kernel, $K^{2}$ is a kernel and its MMD between $X_{d_{S}}$ and $X_{d_{T}}$ is per :

We recover the same expression with a $O(\epsilon)$ in the general setting where $\epsilon\neq 0$.

where $\|H\|\leq n_{S}\epsilon=O(\epsilon)$. Thus,

Therefore, when $\epsilon$ is small, Equation 10 can be developed into:

C.4.4 Variance and initialization

The MMD depends on the kernel $K$, i.e., only on the initialization of $f$ in the kernel regime per . Thus, to reduce variance, we could act on the initialization to match $p_{S}(X)$ and $p_{T}(X)$ in the RKHS of $K^{2}$. This is consistent with Section 2.4.1 that motivated matching the train and test in features. In our paper, we used the standard pretraining from ImageNet , as commonly done on DomainBed . The Linear Probing initialization of the classifier was shown in to prevent the distortion of the features along the training. This could be improved by pretraining the encoder on a task with fewer domain-specific information, e.g., CLIP image-to-text translation as in .

C.5 WA vs. its members

We validate that WA’s expected error is smaller than its members’ error under the locality constraint.

and the expression of the standard bias-variance decomposition in Equation BV from ,

Appendix D Weight averaging versus functional ensembling

We further compare the following two methods to combine $M$ weights $\{\theta(l_{S}^{(m)})\}_{m=1}^{M}$: $f_{\text{WA}}$ that averages the weights and $f_{\text{ENS}}$ that averages the predictions. We showed in Lemma 1 that $f_{\text{WA}}\approx f_{\text{ENS}}$ when $\max_{m=1}^{M}\|\theta(l_{S}^{(m)})-\theta_{\text{WA}}\|_{2}$ is small.

In particular, when $\{l_{S}^{(m)}\}_{m=1}^{M}$ share the same initialization and the hyperparameters are sampled from mild ranges, we empirically validate our approximation on OfficeHome in Figure 1. This is confirmed on PACS dataset in Figure 9. For both datasets, we even observe that $f_{\text{WA}}$ performs slightly but consistently better than $f_{\text{ENS}}$. The observed improvement is non-trivial; we refer to Equation 1 in for some initial explanations based on the value of OOD Hessian and the confidence of $f_{\text{WA}}$. The complete analysis of this second-order difference is left for future work.

Yet, we do not claim that $f_{\text{WA}}$ is systematically better than $f_{\text{ENS}}$. In Table 5, we show that this is no longer the case when we relax our two constraints, consistently with Figure 5. First, when the classifiers’ initializations vary, ENS improves thanks to this additional diversity; in contrast, DiWA degrades because weights are no longer averageable. Second, when the hyperparameters are sampled from extreme ranges (defined in Table 7), performance drops significantly for DiWA, but much less for ENS. As a side note, the downward trend in this second setup (even for ENS) is due to inadequate hyperparameters that degrade the expected individual performances.

This highlights a limitation of DiWA, which requires weights that satisfy the locality requirement or are at least linearly connectable. In contrast, Deep Ensembles are computationally expensive (and even impractical for large $M$), but can leverage additional sources of diversity. An interesting extension of DiWA for future work would be to consider the functional ensembling of several DiWAs trained from different initializations or even with different network architectures . Thus the Ensemble of Averages (EoA) strategy introduced in is complementary to DiWA and could be extended into an Ensemble of Diverse Averages.

Appendix E Additional diversity analysis

In Section 4, our diversity-based theoretical findings were empirically validated using the ratio-error , a common diversity measure notably used in . In Figure 10, we recover similar conclusions with another diversity measure: the Centered Kernel Alignment Complement (CKAC) , also used in . CKAC operates in the feature space and measures to what extent the pairwise similarity matrices (computed on domain $T$) are aligned — where similarity is the dot product between penultimate representations extracted from two different networks.

E.1.2 Accuracy gain per unit of diversity

E.1.3 Diversity comparison across a wide range of methods

Inspired by , we further analyze in Figure 12 the diversity between two weights obtained from different (more or less correlated) learning procedures.

In the upper part, weights are obtained from a single run. They share the same initialization/hyperparameters/data/noise in the optimization procedure and only differ by the number of training steps (which we choose to be a multiple of $50$). They are less diverse than the weights in the middle part of Figure 12, that are sampled from two ERM runs.

When sampled from different runs, the weights become even more diverse when they have more extreme hyperparameter ranges, they do not share the same classifier initialization or they are trained on different data. The first two are impractical for WA, as it breaks the locality requirement (see Figures 5 and 10(c)). Luckily, the third setting “data diversity” is more convenient and is another reason for the success of DiWA†; its $60$ weights were trained on $3$ different data splits. Data diversity has provable benefits , e.g., in bagging .

Finally, we observe that diversity is increased (notably in features) when two runs have different objectives, for example, Interdomain Mixup and Coral . Thus incorporating weights trained with different invariance-based objectives have two benefits that explain the strong results in Table 2: (1) they learn invariant features by leveraging the domain information and (2) they enrich the diversity of solutions by extracting different features. These solutions can bring their own particularity to WA.

In conclusion, our analysis confirms that “model pairs that diverge more in training methodology display categorically different generalization behavior, producing increasingly uncorrelated errors”, as stated in .

E.1.4 Trade-off between diversity and averageability

We argue in Section 2.4.4 that our weights should ideally be diverse functionally while being averageable (despite the nonlinearities in the network). We know from that models fine-tuned from a shared initialization with shared hyperparameters can be connected along a linear path where error remains low; thus, they are averageable as their WA also has a low loss. In Figure 5, we confirmed that averaging models from different initializations performs poorly. Regarding the hyperparameters, Figure 5 shows that hyperparameters can be selected slightly different but not too distant. That is why we chose mild hyperparameter ranges (defined in Table 7) in our main experiments.

A complete analysis of when the averageability holds when varying the different hyperparameters is a promising lead for future work. Still, Figure 13 is a preliminary investigation of the impact of different learning rates (between learning procedures of each weight). First, we validate that more distant learning rates lead to more functional diversity in Figure 13(a). Yet, we observe in Figure 13(b) that if learning rates are too different, weight averaging no longer approximates functional ensembling because the $O(\Delta_{L_{S}^{M}}^{2})$ term in Lemma 1 can be large.

E.2 On PACS

We perform in Figure 14 on domain “Art” from PACS the same core diversity-based experiments than on OfficeHome in Section 4. We recover the same conclusions.

Appendix F Number of training runs

In our experiments, we train $20$ independent training runs per data split. We selected this value as $20$ is the standard number of hyperparameter trials in DomainBed . In Figure 16 we ablate this choice on the OOD domain “Art” of OfficeHome. We observe that a larger number of runs leads to improved performance and reduced standard deviation. These results are consistent with our theoretical analysis, as the variance is divided per $M$ in Proposition 1. If reducing the training time is critical, one could benefit from significant gains over ERM even with a smaller number of runs: for example, $10$ runs seem sufficient in this case. This analysis complements Figure 4 — where $60$ runs were launched then sorted in increasing validation accuracy.

Moreover, in Table 6 we report DiWA’s results when considering only $5$ runs, with uniform weight selection. Interestingly, it shows that $M=5$ is enough to be competitive against SWAD , the previous state of the art.

Appendix G DomainBed

We now further detail our experiments on the DomainBed benchmark .

DomainBed includes several computer vision classification datasets divided into multiple domains. Each domain is successively considered as the test domain while other domains are used in training. In practice, the data from each domain is split into 80% (used as training and testing) and 20% (used as validation for hyperparameter selection) splits. This random process is repeated with 3 different seeds: the reported numbers are the means and the standard errors over these 3 seeds.

We follow the training protocol from https://github.com/facebookresearch/DomainBed. For each dataset, domain and seed, we perform a random search of $20$ trials on the hyperparameter distributions described in Table 7. Our mild distribution is taken directly from , yet could be adapted by dataset for better results. Even though these distributions are more restricted than the extreme distributions introduced , our ERM runs perform better. It leads to a total amount of $2640$ runs only for Table 1. In Appendix B, the $\rho$ hyperparameter for SAM is sampled from $[0.001,0.002,0.005,0.01,0.02,0.05]$. In Table 2, hyperparameters specific to Interdomain Mixup (“mixup_alpha”) and Coral (“mmd_gamma”) are sampled from the distributions defined in . We use a ResNet50 pretrained on ImageNet, with a dropout layer before the newly added dense layer and fine-tuned with frozen batch normalization layers. The optimizer is Adam . Our classifier is either initialized randomly or with Linear Probing ; in the latter case, we first learn only the classifier (with the encoder frozen) with the default hyperparameters defined in Table 7; the classifier’s weights are then used to initialize all subsequent runs. All runs are trained for 5k steps, except on DomainNet with 15k steps as done in concurrent works . As in , validation accuracy is calculated every $50$ steps for VLCS, $500$ steps for DomainNet and $100$ steps for others.

We consider the training-domain validation set protocol. From each run, we thus take the weights of the epoch with maximum accuracy on the validation dataset — which follows the training distribution. Our restricted weight selection is also based on this training-domain validation set. This strategy is not possible for DiWA† as it averages $M=20\times 3$ weights trained with different data splits: they do not share a common validation dataset. The scores for ERM and Coral are taken from DomainBed . Scores for SWAD and MA are taken from their respective papers. Note that MA and SWAD perform similarly even though SWAD introduced three additional hyperparameters tuned per dataset: “an optimum patient parameter, an overfitting patient parameter, and the tolerance rate for searching the start iteration and the end iteration”. Thus we reproduced MA which was much easier to implement, and closer to our uniform weight selection.

G.2 DomainBed results detailed per domain for each real-world dataset

Tables below detail results per domain for the $5$ multi-domain real-world datasets from DomainBed: PACS , VLCS , OfficeHome , TerraIncognita and DomainNet . Critically, showed that diversity shift dominates in these datasets.

Appendix H Failure of WA under correlation shift on ColoredMNIST

Based on Equation BVCL, we explained that WA is efficient when variance dominates; we showed in Section 2.4.2 that this occurs under diversity shift. This is confirmed by our state-of-the-art results in Table 1 and Section G.2 on PACS, OfficeHome, VLCS, TerraIncognita and DomainNet. In contrast, we argue that WA is inefficient when bias dominates, i.e., in the presence of correlation shift (see Section 2.4.1). We verify this failure on the ColoredMNIST dataset, which is dominated by correlation shift .

Colored MNIST is a colored variant of the MNIST handwritten digit classification dataset where the correlation strengths between color and label vary across domains. We follow the protocol described in Section G.1 except that (1) we used the convolutional neural network architecture introduced in DomainBed for MNIST experiments and (2) we used the test-domain model selection in addition to the train-domain model selection. Indeed, as stated in , “it may be improper to apply training-domain validation to datasets dominated by correlation shift since under the influence of spurious correlations, achieving excessively high accuracy in the training environments often leads to low accuracy in novel test environments”.

In Tables 13 and 14, we observe that DiWA-uniform and MA both perform poorly compared to ERM. Note that DiWA-restricted does not degrade ERM as it selects only a few models for averaging (low $M$). This confirms that our approach is useful to tackle diversity shift but not correlation shift, for which invariance-based approaches as IRM or Fishr remain state-of-the-art.

Appendix I Last layer retraining when some target data is available

The traditional OOD generalization setup does not provide access to target samples (labelled or unlabelled). The goal is to learn a model able to generalize to any kind of distributions. This is arguably the most challenging generalization setup: under these strict conditions, we showed that DiWA outperforms other approaches on DomainBed. Yet, in real-world applications, some target data is often available for training; moreover, last layer retraining on these target samples was shown highly efficient in . The complete analysis of DiWA for this new scenario should be properly addressed in future work; yet, we now hint that a DiWA strategy could be helpful.

Specifically, in Table 15, we consider that after a first training phase on the “Clipart”, “Product” and “Photo” domains, we eventually have access to some samples from the target “Art” domain (20% or 80% of the whole domain). Following , we re-train only the last layer of the network on these samples before testing. We observe improved performance when the (frozen) feature extractor was obtained via DiWA (from the first stage) rather than from ERM. It suggests that features extracted by DiWA are more adapted to last layer retraining/generalization than those of ERM. In conclusion, we believe our DiWA strategy has great potential for many real-world applications, whether some target data is available for training or not.