Diverse Weight Averaging for Out-of-Distribution Generalization
Authors: Alexandre Rame
Matthieu Kirchmeyer
Thibaud Rahier
Alain Rakotomamonjy
patrick gallinari
Matthieu Cord
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Learning
Computer Vision
Out-of-Distribution Generalization
Domain Generalization
Distribution Shifts
Ensembling
Weight Averaging

Abstract: Standard neural networks struggle to generalize under distribution shifts in computer vision. Fortunately, combining multiple networks can consistently improve out-of-distribution generalization. In particular, weight averaging (WA) strategies were shown to perform best on the competitive DomainBed benchmark; they directly average the weights of multiple networks despite their nonlinearities. In this paper, we propose Diverse Weight Averaging (DiWA), a new WA strategy whose main motivation is to increase the functional diversity across averaged models. To this end, DiWA averages weights obtained from several independent training runs: indeed, models obtained from different runs are more diverse than those collected along a single run thanks to differences in hyperparameters and training procedures. We motivate the need for diversity by a new bias-variance-covariance-locality decomposition of the expected error, exploiting similarities between WA and standard functional ensembling. Moreover, this decomposition highlights that WA succeeds when the variance term dominates, which we show occurs when the marginal distribution changes at test time. Experimentally, DiWA consistently improves the state of the art on DomainBed without inference overhead.

Diverse Weight Averaging for Out-of-Distribution Generalization
Alexandre Ramé",, Matthieu Kirchmeyer!.2
Thibaud Rahier?, Alain Rakotomamonjy Patrick Gallinari'2 Matthieu Cord Sorbonne Université, CNRS, ISIR, F-75005 Paris, France CCrito AI Lab, Paris, France aaleoo Paris, France *Unvversité de Rouen, LITIS, France
'Equal contribution
Abstract
Standard neural networks struggle to generalize under distribution shifts in com- puter vision. Fortunately, combining multiple networks can consistently improve out-of-distributio generalization. In particular, weight averaging (WA) strategies were shown to perform best on the competitive DomainBed benchmark; they di- rectly average the weights of multiple networks despite their nonlinearities. In this paper, we propose Diverse Weight Averaging (DiWA), a new WA strategy whose main motivation is to increase the functional diversity across averaged models. To this end, DiWA averages weights obtained from several independent training runs: indeed, models obtained from different runs are more diverse than those collected along a single run thanks to differences in hyperparameters and training procedures. We motivate the need for diversity by a new meas-variance-oovarian locality decomposition of the expected error, exploiting similarities between WA and standard functional ensembling. Moreover, this decomposition highlights that WA succeeds when the variance term dominates, which we show occurs when the marginal distribution changes at test time. Experimentally: DIWA consistently improves the state of the art on DomainBed without inference overhead.
Introduction
Learning robust models that generalize well is critical for many real-world applications 2]. Yet,
the classical Empirical Risk Minimization (ERM) lacks robustness to distribution shifts [3, 4, 5]. To improve out-of- --sttribbbtiitioon (OOD) generalization in classification, several recent works proposed to train models simultancously on multiple related but different domains [6]. Though theoretically appealing, domain-invariant approaches [7] either underperform [8, 9] or only slightly improve [10, 11] ERM on the reference DomainBed benchmark [12]. The state-of-the-ae strategy on DomainBed is currently average the weights obtained along training trajectory [13]. [14] 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 Hatness-based analysis and provide new explanation for the success of WA in OOD. It is based on WA's similarity with ensembling [15], a well-known strategy to improve robustness [16, 17]. that averages the predictions from various models. Based WA's expected error. contains four terms: first the bias that show increases under shift in label posterior distributions (i.e., correlation shift [19]); second, the variance that we show increases under shift in input marginal distributions (i.e., diversity shift [19]); third, the covariance that decreases when models are diverse; finally, a locality condition on the weights of averaged models.
on [18], we present a bias-variance- www.nncc-oovriincceeleeee
decomposition
Based on this analysis, aim at obtaining diverse models whose weights are averageable with our Diverse Weight Averaging (DiWA) approach. In practice, DiWA averages weights the models
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

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 [20, 21]. Yet. averaging the weights of independently trained networks with batch normalization [22] and ReI.U layers [23] may be counter-intuitive. Such averaging efficient especially when models can be connected linearly in the weight space via low loss path. Interestingly, this linear mode connectivity property [24] was empirically validated when the runs start from a shared pretrained initialization [25]. This insight is at the heart of DiWA but also of other recent works [26, 27, 28], as discussed in Section 6.
In summary, our main contributions are the following:
We propose new theoretical analysis of WA for OOD based on a bias variiance-covvrriance 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 [12] (Section 5).
2 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- -arranai- 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.
2.1 Notations and problem definition
Notations. We denote X' the input space of images, )' the label space and f: y - R. a loss function. S is the training (source) domain with distribution ps, and T is the test (target) domain with distribution pr- For simplicity, we will indistinctly use the notations ps and pr to refer to the joint, posterior and marginal distributions of (X, Y).. We note fs, fr X' + the source and target labeling functions. We assume that there is no noise in the data: then fs is defined on Xss A {i E X/ps(x) > 0} by v(I,y) ~ ps, fs(r) = y and similarly fr is defined on XY A {x X/PT(x)) > 0} by v(x, 8) ~ pr, fz) = y.
We consider a neural network (NN) f(, 0) X + y made of fixed architecture f with weights 0. We seek 0 minimizing the target generalization error:
Problem.
(0) ((xx))p [e(f(x,)) (zz,0).2) 8)]-
(I)
[(; should approximate fr on Xy. However, this is complex in the OOD setup because we only
have data from domain S in training, related yet different from T. The differences between S and T
are due to distribution shifts (i.e., the fact that ps((, Y) * prp (X) which are decomposed per [19] into diversity shift (a.k.a. covariate shift), when marginal distributions differ (i.c., pX()) * (X)).) and correlation shift (a.k.a. concept shift), when posterior distributions differ (i.e., (Y * pr (YlX) and fs * fr). The weights are typically learned on training dataset ds from S (composed of rig i.i.d. samples from ps (X, Y)) with a configuration c, which contains all other sources of randomness in learning (e.g., initialization, hyperparameters, training stochasticity, epochs, etc.). We call Is = {dg, c} a learning procedure on domain S, and explicitly write ((ls) to refer to the weights obtained after stochastic minimization of 1/ng ((ay)eds ((f(x, 0). y) w.r.t. 0 under lg.
2.2 Weight averaging for OOD and limitations of current analysis
Weight averaging. We study the benefits of combining M individual member weights (0mm m==1 (0(1g))) ) obtained from M (potentially correlated) identically distributed (i.d.) learning pro- cedures LA & { -m=1- Under conditions discussed in Section 3.2, these M weights can be averaged despite nonlinearities in the architecture f. Weight averaging (WA) [13], defined as:
Swx 2 Owx), where Own e Own (L 2 1/M > m=1
0m:
(2)
2

---

is the state of the art [14, 29] on DomainBed [12] when the weights {0m Jm are sampled along a single training trajectory (a description we refine in Remark from Appendix C.2).
Limitations of the flatness-based analysis. To explain this success, Cha et al. [14] 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 Appendix B.I). Second, this analysis does not clarify why WA outperforms Sharpness- Aware Minimizer (SAM) [30] for OOD generalization, even though SAM directly optimizes flatness (see Appendix B.2). Finally, it does not justify why combining WA and SAM succeeds in IID [31] yet fails in OOD (see Appendix B.3). These observations motivate new analysis of WA; we propose one below that better explains these results.
2.3 Biss--aarancee-ovaaiaanc-llllll decomposition
We now introduce bias- wwwpneee--oovvaine--eeeeeee
decomposition which extends the bias- variance decomposition [32] to WA. In the rest of this theoretical section, E is the Mean Squared Error for simplicity: yet, results may be extended to other losses in [33]. In this case, the expected error of a model with weights ((ss) w.r.t. the learning procedure I5 was decomposed in [32] into:
(r((s) = (=p (baas)))ppiiiss"((( [bias"(x; y) + var(z)),
(BV)
where bia y). var(x) are the bias and variance of the considered model w.r.t. sample (1, y). defined later in Equation (BVCL). To decompose WA's error, we leverage the similarity (al- ready highlighted in [13)) between WA and functional ensembling (ENS) [15, 34), more tra- ditional way to combine collection of weights. More precisely, ENS averages the predictions, fENS A fENS( (-, (0. Mm & 1/M Σ Em=1 (m). Lemma I establishes that fwA is first-order approximation of fENS when {8.m) 1 are close in the weight space.
Lemma 1 (WA and ENS. Proof in Appendix C.I. Adapted from [13, 28].). Given {0m Jim with learning procedures L
(1(m)) 1m=1- Denoting Lg max mm0m Owall ((z Xx))
JMM
fwa(z) fx) + O(47.)) and ((fwa(z). 9) ((fens(z),)) O(4
This similarity is useful since Equation (BV) was extended into bias slus-variance-oovaree sition for ENS in [18, 35J. 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 A (1(m) )
decompo-
Jm=1-
Proposition (Bias-variance) verrance--ovvrrianc--laalii decomposition of the expected generalization error of WA in OOD. Proof in Appendix C.2.). Denoting fs(x) [f 0(1s))}, under identically
distributed learning procedures L & (15T) (1(m)) Jm=1, the expected generalization error on domain T of (T A 1
M Em= Om over the joint distribution of T. is:
Lg (Owa () E(2.o><r biias?zz,,)) + M
var(x)) + M -
-ov((=) + +(A)),
M
where bias(z, y) y fs(())
and var var(z) Eis [1F(x, ((((x,0(ls)) -((s)) Fs(
(BVCL)
and cov (I) Eisi''s (f 0(1s)) - fs())) (1 0(1))) - Js(x)))
and A Ey 47, with A Ex m=1 max
16mm Owa 112-
Cov is the prediction covariance between two member models whose weights are averaged. The locality term 2 is the expected squared maximum distance between weights and their average.
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 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.
3

---

2.4 Analysis of the biss--ataacc--cvarineee--bliii 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 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.
2.4.1 Bias and correlation shift (and support mismatch)
We relate OOD bias correlation shift [19] under Assumption 1, where fs(x) - Eig (f 0(1s))] As discussed in Appendix C.3.2, Assumption is reasonable for a large NN trained on a large dataset representative of the source domain S. It is relaxed in Proposition 4 from Appendix C.3.
Assumption (Small IID bias). 3e small s.t. Vr E Xs, (fs()) fs(x))
Proposition 2 (OOD bias and correlation shift. Proof in Appendix C.3). With bounded difference between the labeling functions fT - fs on X7 n Xg, under Assumption 1, the bias on domain is:
B (( Ibias (r, u)] = Correlation shift + Support mismatch + O(t),
where Correlation shift- J Jxnnxs
(sr(x) - f=(=s)(())))))))))))})))))) pr(x)dx, (s(=))"pr())ddr
(3)
and Support mismatch
/ Xr
(fr(x) f=(())))
pr(x)dr. pr(
We analyze the first term by noting that fr(x)) Epp x] and fs(x) E.g [YX =), Vr E Ar 0 Xx. This expression confirms that our correlation shift term measures shifts in posterior distributions between source and target, as in [19]. It increases in presence of spurious correlations: e.g., on ColoredMNIST [8] where the color/label correlation is reversed at test time. The second term is caused by support mismatch between source and target. It was analyzed in [36] and shown irreducible in their "No free lunch for learning representations for DG", Yet, this term can be tackled if we transpose the analysis in the feature space rather than the input space. This motivates encoding the source and target domains into a shared latent space, e.g., by pretraining the encoder on task with minimal nomain-specifn information as in [36].
This analysis explains why WA fails under correlation shift, as shown on ColoredMNIST Ap- pendix H. Indeed, combining different models does not reduce the bias. Section 2.4.2 explains that WA is however efficient against diversity shift.
2.4.2 Variance and diversity shift
Variance is known to be large in OOD [5] and to cause phenomenon named underspecification, when models behave differently in OOD despite similar test IID accuracy. We now relate OOD variance to diversity shift [19] in a simplified setting. We fix the source dataset ds (with input support Xaş). ), the target dataset dr (with input support X ) and the network's 's initialization. We get a closed-form expression for the variance of over all other sources of randomness under Assumptions and 3.
Assumption 2 (Kernel regime). f is in the kernel regime [37, 38].
This states that f behaves as a Gaussian process (GP); it is reasonable if f is a wide network [37, 39J. The corresponding kernel K is the neural tangent kemel (NTK) [37] depending only on the initialization. GPs are useful because their variances have a closed-form expression (Appendix C.4.1). To simplify the expression of variance, we now make Assumption 3.
ds). 3(As, e) with 0 < c < Äs such that V E Xas. K(Ts, = Xs and Va's + ug E Xa., (xss, e.
Assumption 3 (Constant norm and low intra-sample similarity
This states that training samples have the same norm (following standard practice [39, 40, 41, 42)) and weakly interact [43, 44]. This assumption is further discussed and relaxed in Appendix C.4.2. We now in position to relate variance and diversity shift when 0.

---

Proposition 3 (OOD variance and diversity shift. Proof in Appendix C.4). Given f trained on source dataset dy (of size ns)) with NTK K, under Assumptions and 3, the variance on dataset dr is:
Ex [var(xT)] =
n18 MMD² (Xdg: Xdy) + XT 2Xg
ng
BT + +()),
(4)
2Xs
where MMD is the empirical Maximum Mean Discrepancy in the RKHS of (² = (K(x,)) ;xr & ExyeXay ((Fr, and fr - () STtE, (xTT, x4) are the em- pirical mean similarities respectively measured between identical (wars. K) and different (w.rt. K²) samples averaged over Xdy-
The MMD empirically estimates shifts in input marginals, i.e., between ps (X) and (X). Our expression of variance is thus similar the diversity shift formula in [19]: MMD replaces the L divergence used in [19]. The other terms, Xr and By, both involve internal dependencies on the target dataset dr: they are constants w.r.t. Xd, and do not depend on distribution shifts. At fixed dr and under our assumptions, Equation (4) shows that variance on d7 decreases when Xás and Xd, are closer (for the MMD distance defined by the kernel K2) and increases when they deviate. Intuitively, the further Xdr is from Xdy. the less the model's predictions on Xdy are constrained after fitting ds-
This analysis shows that WA reduces the impact of diversity shift as combining M models divides the variance per M. This is strong property achieved without requiring data from the target domain.
2.4.3 Covariance and diversity
The covariance term increases when the predictions of {f pm 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), detailed in Appendix C.5. It motivates tackling covariance by encouraging members to make different predictions, thus to be functionally diverse. Diversity is widely analyzed concept in the ensemble literature [15], for which numerous measures have been introduced [45, 46, 47]. In Section 3, we aim at decorrelating the learning procedures to increase members' diversity and reduce the covariance term.
2.4.4 Locality and linear mode connectivity
To ensure that WA approximates ENS, the last locality term 3( 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, locality. In practice, we consider that the main goal of this locality term is to ensure that the weights 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 [24]. We empirically verify later in Figure that the approximation fwa 2 fENS remains valid even in this case.
thus seek good trade-off between diversity and
3 DiWA: Diverse Weight Averaging
3.1 Motivation: weight averaging from different runs for more diversity
Limitations of previous WA approaches. 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 { ( jM Yet, previous methods [14, 29] only average weights obtained along a single run. This corresponds 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.
DIWA. Our Diverse Weight Averaging approach seeks to reduce the OOD expected error in Equa- tion (BVCL) by decreasing covariance across predictions: DiWA decorrelates the learning procedures (m) M m= Our weights are obtained from M > 1 different runs, with diverse learning procedures:
5

---

Algorithm DiWA Pseudo-code
Require: Oo pretrained encoder and initialized classifier; (lm ) hyperparameter configurations
Training: Vm = 1 to H, 0m & FineTTune(6a hm)
Weight selection:
Uniform: M {1. H). Restricted: Rank {0.m m== by decreasing ValAcc(Om)) M t- 0.
for m to H do
If ValAco(f MU)(M) 2 ValAcc(OMM)
M+M U {m}
Inference: with f(, BM), where 8 Emem lm/lM)
these have different hyperparameters (learning rate, weight decay and dropout probability). batch or- ders, 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 [21] 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.
3.2 Approach: shared initialization, mild hyperparameter search and weight selection
Shared initialization. The shared initialization condition follows [25]: when models are fine-tuned from shared pretrained model, their weights can be connected along a linear path where error remains low [24]. Following standard practice on DomainBed [12], our encoder is pretrained on ImageNet [48]; this pretraining is key it controls the bias (by defining the feature support mismatch, see Section 2.4.1) and variance (by defining the kernel K see Appendix C.4.4). Regarding the classifier initialization,
test two methods. The first is the random initialization, which may distort the features [49]. The second is Linear Probing (LP) [49]: it first learns the classifier (while freezing the encoder) to serve 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 [49]).
Mild hyperparameter search. 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, thus the mild search space defined in Table 7, first introduced in SWAD [14]. These hyperparameter ranges induce diverse models that are averageable in weights.
Weight selection. The last step of our approach (summarized in Algorithm 1) is to choose which weights average among those available. We explore two simple weight selection protocols, as in [28]. The first uniform equally averages all weights; it is practical but may underperform when some runs are detrimental. The second restricted (greedy in [28]) 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 [50] where diversity shift dominates [19] (see Appendix E.2 for similar analysis on PACS [51]). Then, Section 5 shows that DIWA is state of the art on DomainBed [12].
4 Empirical validation of our theoretical insights
JM We consider several collections of weights {0.m 1m=1 (2 S M < 10) trained on the "Clipart", "Product" and "Photo" domains from OfficeHome [50] with a shared random initialization and mild hyperparameter ranges. These weights are first indifferently sampled from single run (every 50 batches) or from different runs. They are evaluated on "Art", the fourth domain from OfficeHome.
WA vs. ENS. Figure validates Lemma and that Swa as SENS- More precisely, fww slightly but consistently improves fENs: discuss this in Appendix D. Moreover, larger A4 improves the

---

M-2 (slape: 0.116)
0.88
M=2
M-3
M=3
AcciOua) -8 Accil0m)
M-4 Islape: 0.196)
M-4
M-5 Islope: 0.208)
M-6 Islope: 0.235)
Acc(0yal
0.64
6
M=B tslope: 0.282)
M=7 Islope: 0.259)
M-7
M=# tslope: 0.297)
M-8
0.2
M
0.04
0.00
0.50
0.60
0.92
0.64
6.66
0.00
077
Acc(Eers)
Prediction diversity
Figure 1: Each dot displays the accuracy (f) of weight averaging (WA) vs. accuracy (f) of prediction averaging (ENS) for M models.
Figure Each dot displays the accuracy (t) gain of WA over its members vs. the prediction diversity [46] (t) for M models.
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.
Diversity and accuracy. We validate in Figure that fwa benefits from diversity. Here, we measure diversity with the ratio-error [46], i.e., the ratio Nair/ Nsimul between the number of different errors Nair and of simultaneous errors Nsimal in test for a pair in (ʃ(: (m,0m)))s mm)) A higher average over the (2) pairs means that members are less likely to eIr on the same imputs. Specifically, the gain of Acc(Ows) over the mean individual accuracy 37 Em=1 Acc(Om) increases with diversity. Moreover, this phenomenon intensifies for larger M: the linear regression's 's slope (i.e., the accuracy gain per unit of diversity) increases with M. This consistent with the (M 1)/M factor of cov(r) in Equation (BVCL), as further highlighted in Appendix 1.2. Finally, in Appendix E.1.1, we show that the conclusion also holds with CKAC [47], another established diversity measure.
Increasing diversity thus accuracy via different runs. Now we investigate the difference between sampling the weights from single run or from different runs. Figure first shows that diversity increases when weights come from different runs. Second, in Figure 4, this is reflected the accuracies OOD. Here, we rank by validation accuracy the 60 weights obtained (1) from 60 different runs and (2) along well-perioo well-performing
We then consider the WA of the top M weights as M increases from 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 similarly
hyperparameter distributions are extreme defined in Table 7) when classifiers 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.
Weights from differant Weights from one single run
0.68
Acci0mal (ccet....))))
12
0.2
Frequency (%6
Acc10ge))
-0.-
0.4
Ours
Weights from different runs Weights from single
Extrome hyperparams
2-
0.2
-03
Difterent classitier ints
Prediction diversity
M
do
Prediction diversity
Figure 3: Frequencies of predic- Figure WA accuracy (t) as M tion diversities (1) [46] across 2 weights obtained along a single are obtained along a single run or from different runs.
Figure 5: Each dot displays the
increases, when the M weights
accuracy (f) gain of WA over its members vs. prediction diversity for M 10 models.
or from different runs.
7

---

Experimental results on the DomainBed benchmark
Datasets. We now present our evaluation on DomainBed [12]. By imposing the code, the training
procedures and the ResNet50 [52] architecture, DomainBed is arguably the fairest benchmark for OOD generalization. It includes multi-domain real-world datasets: PACS [51]. VLCS [53], OfficeHome [50], TerraIncognita [54] and DomainNet [55]. [19] 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-doman model selection. The experimental setup is further described Appendix G.I. Our code is available at https: : /tt///ggttttt com/ alexr alexrame//ee ame/ /diwa.
Baselines. ERM the standard Empirical Risk Minimization. Coral [10] is the best approach based on domain invariance. SWAD (Stochastic Weight Averaging Densely) [14] and MA (Moving Average) [29] average weights along one training trajectory but differ in their weight selection strategy. SWAD [14] 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 [29] 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 [29] for the costly Deep Ensembles (DENS) [15] (with different initializations): we discuss other ensembling strategies in Appendix
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. practice, we sample 20 configurations from the hyperparameter distributions detailed in Table 7 and report the mean and standard deviation across 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 < M < 20 weights with Algorithm while DiWA-uniform averages M = 20 weights. DiWAf averages uniformly the M 3 x 20 = 60 weights from all 3 data splits. DiWA benefits from larger M (without additional inference cost) and from data diversity (see Appendix 1.3). However, we cannot report standard deviations for DiWAI 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.
5.1 Results on DomainBed
We report our main results in Table 1, detailed per domain in Appendix G.2. With randomly initialized classifier, DiWAI -uNN--r is the best on PACS, VLCS and OfficeHome: DiWA-uniform the second best on PACS and OfficeHome. On Terralncognita and NemainNet, DiWA is penalized by some bad runs, filtered DiWA -es-rrstticttd which improves results on these datasets. Classifier initialization with linear probing (LP) [49] 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 new SoTA of 68.0%, improving SWAD by 1.1 points.
Table 1: Accuracy (%, 1) DomainBed with ResNet50 (best in bold and second best underlined).
Algorithm
Weight selection
Init
PACS
VLCS
OfficeHome
TerraInc
DomainNet
Avg
ERM
85.5
66.5 0.3
46.1
1
63.3
Coral [10]
N/A
86.2 + 0.3
78.8 006
68.7 077+0.
47.6 1.0
41.5 + 0.1
64.6
SWAD [14]
Overfit-aware
Random
88.1
79.1
70.6
50.0
46.5 0.1
66.9
MA [29]
Uniform
87.5 0.2
0.2
70.6 0.1
50,3 0.5
46.0 + 0.1
66.5
DENS [15, 29]
Uniform: M
87.6
78.5
70.8
49,2
47.7
66.8
ERM
85.5 0.5
77.6
67.4
4..
1
64.6
MA [29]
Uniform
87.9 7.9-4000 0.1
78.4 4 0.1
70.3 0.1
49.9 0.2
46.4 + 0.1
66.6
ENS
Unifor::
78.7
7..
1.5
47.4
0.2
67.1
Random
DIWA
Restricted: M x 20
87.9 + 0.2
79.2
+ 0.1
70.5 4.1
50.5
46.7
4 0.1
67.0
DIWA
Uniform:
88.8
79.1
0.2
7.10
48,9
46.1 + 0.1
66.8
DIWA!
Uniform: 60
89.0
79.4
71.6
49,0
67.1
ERM
N/A
85.9 0.6
18.1
6.2
50.4 1.8
44.3 + 0.2
65.6
MA [29]
Uniform
87.8 0.3
78.5
0.4
0.3
51.4
46.6 0.0
67.1
ENS
Uniform: M = 20
88.1 0.3
78.5
0.1
7..7
50.8
47.0 0.2
67.2
LP
DIWA
Restrieted: 20
88.0 0.3
78.5
0.1
71.5
$1.6 0.9
47.7 - 0.1
67.5
DIWA
Unifom: M = 20
88.7 1 0.2
78.4 0.2
72.1 0.2
51.4
47.4 1 0.2
67.6
DIWA!
Uniform: M = 60
89.0
78.6
72.8
$1.9
47.7
68.0
8

---

DiWA with different objectives. So far we used ERM that does not leverage the domain information. Table 2 shows that DiWA-uniform benefits from aver- aging weights trained with Interdomain Mixup [56] and Coral [10]: accuracy gradually improves as we add more objectives. Indeed, as highlighted in Ap-
Table 2: Accuracy (%, f) on OfficeHome domain "Art" with various objectives.
Algorithm
No WA
MA
DIWA
DIWA
ERM
62.9 + 1.3
65.0 +022
67.3 0.2
67.7
63.1 0.7.
66.2 4.33
67.8
68.4
Coral
64.4
64.4
+.04
67.7+02
68.2
pendix 1.3, DiWA benefits from the increased diver- sity 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.
ERMIComa
ERM/MIHU
N/A
NIA
NIA
68.1 18.11.03
68.7
68.9
ERMIMiupNCoom
N/A
NIA
68.4 4 0.4
69.1
5.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.
6 Related work
Generalization and ensemble. To generalize under distribution shifts, invariant approaches [8, 9, 11, 10, 57, 58] try to detect the causal mechanism rather than memorize correlations: yet, they do not outperform ERM on various benchmarks [12, 19, 59]. In contrast, ensembling of deep networks [15, 60, 61] consistently increases robustness [16] and was successfully applied to domain generalization [29, 62, 63, 64, 65, 66]. As highlighted in [18] (whose analysis underlies our Equation (BVCL)), ensembling works due the diversity among its members. This diversity comes primarily from the randomness of the learning procedure [15] and can be increased with different hyperparameters [67], data [68, 69, 70], augmentations [71, 72] or with regularizations [73, 65, 66, 74, 75].
Weight averaging. Recent works [13, 76, 77, 78] combine in weights (rather than in predictions) models collected along single run. This was shown suboptimal in IID [17] but successful in OOD
[14, 29]. Following the linear mode connectivity [24, 79] and the property that many independent models are connectable [80], second group of works average weights with fewer constraints [26, 27, 28, 81, 82, 83]. To induce greater diversity, [84] used a high constant learning rate; [80] explicitly encouraged the weights to encompass more volume in the weight space; [83] minimized cosine similarity between weights; [85] used a tempered posterior. From a loss landscape perspective [20], these methods aimed at "explorfing] the set of possible solutions instead of simply converging to single point", as stated in [84]. The recent "Model soups" introduced by Wortsman et al. [28] is
WA algorithm similar 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, [28] challenges the standard model selection after a grid search. Regarding the task, [28] and our work complement each other: while [28] demonstrate robustness on several ImageNet variants with distribution shift, we improve the So'T on the multi-domain DomainBed benchmark against other established OOD methods after a thorough and fair comparison. Thus, DiWA and [28] are theoretically complementary with different motivations and applied successfully for different tasks.
7 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 new bias-variance-covarie wore--ccciiiie-oooaliit 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 BedainBed, 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.

---

Acknowledgements
We would like thank Jvan-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 AD01 101 1953 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-00 DIACLIM ANR- ANR-19-CON 19-CHIA-001 8-01, RAIMO ANR-20-CHIA-00011-00 OATMIL ANR-17-CE23- 0012 and LEAUDS ANR- ANR--8-CE223 18-CE23-0020.
and the ANR projects
References
[1] John R. Zech, Marcus A. Badgeley, Manway Liu, Anthony B. Costa, Joseph Titano, and Eric Karl Oermann. Variable generalization performance a deep learning model to detect pneumonia in chest radiographs: A cross- seos-secctionl study. PLOS Medicine, 2018. (pp. and 17)
[2] Alex DeGrave, Joseph D Janizek, and Su-In Lee. Ai for radiographic covid-19 detection selects shorteuts signal. Nature Machine Intelligence, 2021. (pp. 1 and 17)
[3] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In ICLR, 2019. (p.
[4] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In NeurlPS, 2020. (p. 1)
[5] Alexander D' Amour, Katherine Heller, Dan Moldovan, Ben Adlam, Babak Alipanahi, Alex Beutel, Christina Chen, Jonathan Deaton, Jacob Eisenstein, Matthew D Hoffman, et al. Un- derspecification presents challenges for credibility modern machine learning. JMLR, 2020. (pp. and 4)
[6] Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML, 2013. (p. 1)
[7] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. JSTOR, 2016. (p. 1)
[8] Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint, 2019, (pp. 1,4, 9, and 36)
[9] David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In ICML, 2021. (pp. and
[10] Baochen Sun, Jiashi Feng, and Kate Saenko. Return of frustratingly easy domain adaptation. In AAAI, 2016. (pp. 8, 9, 29, 33, 34, 35, and 36)
[11] Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of- -ut-of-distribution generalization. In ICML, 2022. (pp. 1, 9, and 36)
[12] Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. In ICLR, 2021. (pp. 2, 6, 8, 9, 16, 17, 18, 27, 32, 33, and 36)
[13] Pavel Izmailov, Dmitrii Podoprikhin. Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In UAI, 2018.
(pp. 1, 3, 9, and 20)
[14] Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. In NewrlPS, 2021. (pp. 3. 5, 6, 8, 9, 17, 18, 19, 20, 23, 32, 33, 34, and 35)
[15] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. NeurlPS, 2017. (pp. 3, 9, 18, 28, 34, and 35)
10

---

[16] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In NeurIPS, 2019. (pp. 1 and 9)
[17] Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry Vetrov. Pitfalls of in-domain uncertainty estimation and ensembling in deep learning. In ICLR. 2020. (pp. I and 9)
[18] Naonori Ueda and Ryohei Nakano. Generalization error of ensemble estimators. In ICNN, 1996. (pp. 1, 3, 9, and 21)
[19] Nanyang Ye, Kaican Li, Lanqing Hong, Haoyue Bai, Yiting Chen, Fengwei Zhou, and Zhenguo Li. Ood-bench: Benchmarking and understanding -ut-of-distribution generalization datasets and algorithms. CVPR, 2022. (pp. 1, 4, 6, 8, 9, 23, 33, and 36)
[20] Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape perspective. arXiv preprint, 2019. (pp. 2 and
[21] Raphael Gontijo-Lppes, Yann Dauphin, and Ekin Dogus Cubuk. No one representation to rule them all: Overlapping features of training methods. In ICLR, 2022. (pp. 2, 6, 29, and 30)
[22] Sergey loffe and Christian Szegedy, Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015. (p. 2)
[23] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint, 2018.
(p. 2)
[24] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In ICML, 2020. (pp. 5, 6, and 9)
[25] Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learning? In NeurlPS, 2020. (pp. 2, 6, 28, and 30)
[26] Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Hanna Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In CVPR, 2022. (pp. 9, and 28)
[27] Michael Matena and Colin Raffel. Merging models with fisher-weighted averaging. In NeurlPS, 2022. (pp. 2 and
[28] Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari S. Morcos, Hongseok Namkoong. Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In ICML, 2022. (pp. 2, 6, 20, and 28)
[29] Devansh Arpit, Huan Wang, Yingbo Zhou, and Caiming Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In NeurIPS, 2022. (pp. 3, 8, 18, 19, 20, 23, 28, 33, 34, 35, and 36)
[30] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- minimization for efficiently improving generalization. In ICLR, 2021. (pp. 3, 18, and 19)
[31] Jean Kaddour, Linging Liu, Ricardo Silva, and Matt J. Kusner. When do flat minima optimizers work? NeurIPS, 2022. (pp. 19)
[32] Ron Kohavi, David H Wolpert, et al. Bias plus variance decomposition for zero-one loss functions. In ICML, 1996. (pp. 21, and 27)
[33] Pedro Domingos. A unified bias-variance decomposition. In ICML, 2000. (p. 3)
[34] Thomas G Dietterich. Ensemble methods in machine learning. In MCS, 2000. (p.
[35] Gavin Brown, Jeremy Wyatt, and Ping Sun. Between two extremes: Examining decompositions of the ensemble objective function. In MCS, 2005. (pp. and 21)

---

[36] Yangjun Ruan, Yann Dubois, and Chris J. Maddison. Optimal representations for covariate shift. In ICLR, 2022. (pp. 23, and 27)
[37] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemel: Convergence and generalization in neural networks. In NeurIPS, 2018. (pp. 4, 25, and 27)
[38] Amit Daniely. Sgd learns the conjugate kernel class of the network. NeurIPS, 2017. (p. 4)
[39] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks gaussian processes. In ICLR, 2017. (pp. and 25)
[40] Julien Ah-Pine. Normalized kernels as similarity indices. In PAKDD, 2010. (pp. 4 and 25)
[41] Benyamin Ghojogh, Ali Ghodsi, Fakhri Karray, and Mark Crowley. Reproducing kernel hilbert space, mercer's theorem, cigenfunctions, nystrom method, and use of kernels in machine learning: Tutorial and survey, arXiv preprint, 2021. (pp. and 25)
[42] Jason Rennie. How to normalize a kernel matrix. MIT Computer Science Artificial Intelligence Lab Tech Rep, 2005. (pp.4 and 25)
[43] Hangfeng He and Weijie Su. The local elasticity of neural networks. In ICLR, 2020. (pp. and 25)
[44] Mariia Seleznova and Gitta Kutyniok. Neural tangent kernel beyond the infinite- width limit: Effects of depth and initialization. ICML, 2022. (pp. 4 and 25)
[45] Ludmila I Kuncheva and Christopher J Whitaker. Measures of diversity in classifier ensembles and their relationship with the ensemble accuracy. Machine learning. 2003. (p. 5)
[46] Matti Aksela. Comparison of classifier selection methods for improving committee perfor- mance. In MCS, 2003. (pp. 5, 19, 28, 29, 30, and 31)
[47] Simon Korbllith, Mohammad Norouzi, Honglak Lee, and Geoffrey E. Hinton. Similarity of neural network representations revisited. In ICML,, 2019. (pp. 28, 29, and 30)
[48] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In NeurlPS, 2012. (pp. 6 and 27)
[49] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distributtin. ICLR, 2022. (pp. 6, 8, 27, 32, 33, 34, and 35)
[50] Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In CVPR, 2017. (pp. 6, 8, 18, and 33)
[51] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In ICCV, 2017. (pp. 6, 8, and 33)
[52] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016. (pp. 8 and 33)
[53] Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In ICCV, 2013. (pp. and 33)
[54] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In ECCV, 2018. (pp. and 33)
[55] Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV, 2019. (pp. 33, and 36)
[56] Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint, 2020. (pp. 29, and 33)
12

---

[57] Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In ICLR, 2020. (pp. and 17)
[58] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks, JMLR, 2016. (p. 9)
[59] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, 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: benchmark of in -ild distribution shifts. In JCML, 2021. (p. 9)
[60] Lars Kai Hansen and Peter Salamon. Neural network ensembles. IEEE transactions on pattern analysis and machine intelligence, 1990. (p.
[61] Anders Krogh and Jesper Vedelsby. Neural network ensembles, cross validation, and active learning. In NeurlPS, 1995. (p. 9)
[62] Kowshik Thopalli, Sameeksha Katoch, Jayaraman J. Thiagarajan, Pavan K. Turaga, and Andreas Spanias. Multi-domain ensembles for domain generalization. In NeurIPS Workshop,
2021. (p.9) 9)
[63] Yusuf Mesbah, Youssef Youssry Ibrahim, and Adil Mehood Khan. Domain generalization using ensemble learning. ISWA, 2022. (p. 9)
[64] Ziyue Li, Kan Ren, Xinyang Jiang, Bo Li, Haipeng Zhang, and Dongsheng Li. Domain generalization using pretrained models without fine-tuning. arXiv preprint, 2022. (p. 9)
[65] Yoonho Lee, Huaxiu Yao, and Chelsea Finn. Diversify and disambiguate: Learning from underspecified data, arXiv preprint, 2022. (p. 9)
[66] Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy. Agree to disagree: Diversity through disagreement for better transferability. arXiv preprint, 2022. (p. 9)
[67] Florian Wenzel, Jasper Snoek, Dustin Tran, and Rodolphe Jenatton. Hyperparameter ensembles for robustness and uncertainty quantification. In NeurIPS, 2020. (p. 9)
[68] Leo Breiman. Bagging predictors. Machine learning, 1996. (pp. and 29)
[69] Jeremy Nixon, Balaji Lakshminarayanan, and Dustin Tran. Why are bootstrapped deep ensembles not better? In NeurIPS Workshop, 2020. (p. 9)
[70] Teresa Yeo, Oguzhan Fatih Kar, and Amir Roshan Zamir. Robustness via cross-domain ensembles. In JCCV, 2021. 9)
[71] Yeming Wen, Ghassen Jerfel, Rafael Muller, Michael W Dusenberry, Jasper Snoek, Balaji Lakshminarayanan, and Dustin Tran. Combining ensembles and data augmentation can harm your calibration. In ICLR, 2021. (p. 9)
[72] Alexandre Rame, Remy Sun, and Matthieu Cord. Mixmo: Mixing multiple inputs for multiple outputs via deep subnetworks. In ICCV, 2021. (pp. 9 and 28)
[73] Alexandre Rame and Matthieu Cord. Dice: Diversity in deep ensembles via conditional redundancy adversarial estimation. In ICLR, 2021. (pp. and 28)
[74] Tianyu Pang, Kun Xu, Chao Du, Ning Chen, and Jun Zhu. Improving adversarial robustness via promoting ensemble diversity. In ICML, 2019. (p. 9)
[75] Damien Teney, Ehsan Abbasnejad, Simon Lucey, and Anton van den Hengel. Evading the simplicity bias: Training diverse set of models discovers solutions with superior ood generalization. arXiv preprint, 2021. (p. 9)
[76] Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In ICML, 2018. (p.9) 9)
13

---

[77] Hao Guo, Jiyong Jin, and Bin Liu. Stochastic weight averaging revisited. arXiv preprint, 2022.
(p. 9)
[78] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey Hinton. Lookahead optimizer:
steps forward, step back. NeurIPS, 32, 2019, (p.
[79] Vaishnavh Nagarajan and Zico Kolter. Uniform convergence may be unable explain generalization in deep learning. NeurIPS, 2019. (p.9)
[80] Gregory Benton, Wesley Maddox, Sanae Lotfi, and Andrew Gordon Wilson. Loss surface simplexes for mode connecting volumes and fast ensembling. In ICML, 2021. (p. 9)
[81] Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Large-batch training that generalizes well. ICLR, 2020. (p. 9)
[82] Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining. arXiv preprint, 2022. (p. 9)
[83] Mitchell Wortsman, Maxwell Horton, Carlos Guestrin, Ali Farhadi, and Mohammad Rastegari. Learning neural network subspaces. ICML, 2021. (p.9) 9)
[84] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In NeurlPS, 2019. (p. 9)
[85] Pavel Izmailov, Wesley Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Subspace inference for bayesian deep learning. In UAI, 2019. (p. 9)
[86] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In ICML SCIS Workshop, 2022. (pp. 17
and 37)
[87] Su Lin Blodgett, Lisa Green, and Brendan O'Connor. Demographic dialectal variation in social media: A case study of african- aaririn-american english. In EMNLP, 2016. (p. 17)
[88] Solon Barocas and Andrew D Selbst. Big data's 's disparate impact. Calif. L. Rev. 2016. (p. 17)
[89] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In ICMI,, 2017. (p. 18)
[90] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley, Relative flatness and generalization. In NeurlPS, 2021. (p. 18)
[91] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In Big Data, 2020. (p. 18)
[92] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint, 2022. (p. 19)
[93] Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, 2003. (p. 24)
[94] Fernando Pérez-Cruz, Steven Van Vaerenbergh, Juan José Murillo-Fuentes, Miguel Lázaro- Gredilla, and Ignacio Santamaria. Gaussian processes for nonlinear signal processing: An overview of recent advances. EEE Signal Process. Mag., 2013. (p. 25)
[95] Greg Yang and Hadi Salman. fine-grained spectral perspective on neural networks. arXiv preprint, 2019. (p. 25)
[96] Damien Brain and Geoffrey Webh. On the effect of data set size bias and variance in classification learning. AKAW, 1999. (p. 25)
[97] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012. (p. 26)
14

---

[98] Jan R Magnus and Heinz Neudecker. Matrix differential calculus with applications in statisties and econometrics. John Wiley & Sons, 2019. (p. 26)
[99] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. ICML, 2021. (p. 27)
[100] Saurabh Singh, Derek Hoiem, and David Forsyth. Swapout: Learning an ensemble of deep architectures. In NeurlPS, 2016. (p. 28)
[101] Bradley Efron. Bootstrap methods: another look the jackknife. In Breakthroughs in statistics. 1992. (p. 29)
[102] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015. (p. 33)
[103] Yann LeCun, Corinna Cortes, and Chris Burges. Mnist handwritten digit database, 2010. 36)
[104] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Domain adjusted regression
Erm may already learn features sufficient for out-of-distribution generalization. 2022. (p. 37)
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Section 5.2.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] In Appendix A
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] Assumption discussed in Appendix C.3.2 and Assumptions 2 and 3 discussed in Appendix C.4.2.
(b) Did you include complete proofs of all theoretical results? [Yes] In Appendix C
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as a URL)? [Yes] Our code available at https: /tte:://giii /gggttbt...//g com/ wwweaaaeeeeeeeaaaaaa.aaaaaaaaeee /diwa.
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] See Section 5 and Appendix G.1
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes] Defined by different data splits when possible.
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, cloud provider)? [Yes] Approximately 20000 hours of GPUs (Nvidia V100) on an internal cluster, mostly for the 2640 runs needed in Table 1.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] DomainBed bench- mark [12] and its datasets.
(b) Did you mention the license of the assets? [Ys]] DomainBed is under "The MIT License".
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're 'r using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "elf1erCJeni", "reviewer": "Reviewer_b9sv", "summary": "The paper studies why and when averaging the weights of a model improve out-of-domain generalization by way of a new a bias-variance-covariance-locality decomposition of the expected target domain error. They show how diversity is needed and thus propose Diverse Weight Averaging (DiWA) wherein weights of N independent trained models from shared initialization are averaged together (in an optionally greedy fashion). They show how DiWA outperforms weight averaging across N model checkpoints of a single training run and other baselines on the public DomainBed benchmark.\n", "questions": "I would be curious to know how DiWA and WA would fare when the Sharpness-Aware Minimization (SAM) is used in each training run. Perhaps, the benefits of DiWA or WA would be diminish in the presence of SAM?\n", "limitations": "Yes.\n* \"extreme hyperparameter ranges lead to weights whose average may perform poorly\"\n* \"diversity is key as long as the weights remain averageable\" -- linear connectivity in weight space is needed.\n", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The paper is very well written -- exposition is clear, well-organized, high quality (free of typos), and each section is well-motivated. For example, \"Limitations of the flatness-based analysis\" discusses why the current explanation for weight averaging's performance is insufficient and why a different type of analysis (like their 4 term decomposition) is needed.\n* The proposed bias-variance-covariance-locality decomposition (Proposition 1) is novel and aids our understanding of weight averaging. Furthermore, it directly motivates their proposed DiWA.\n* DiWA is evaluated on DomainBed, which is the standard public benchmark for out-of-domain generalization right now. Experimental setup is sound and the comparisons with baselines fair.\n\nWeaknesses:\n* Despite DiWA outperforming weight averaging checkpoints of a single run (WA), it is unlikely to be adopted for some use cases. For example, for large language models, there is no extra cost for doing WA but training the model multiple times as needed by DiWA is not practical or feasible. The extra compute is probably better spent scaling the model size or training for longer.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TJIvTfNGLdV", "reviewer": "Reviewer_ySFH", "summary": "In this paper, the authors propose an approach to mitigating the generalization problem in computer vision, under the common scenarios of either a difference in the data generating marginals (covariate shift) or a difference in the class-conditional distributions (which they call concept shift or correlation shift).  They connect previous work on the errors incurred by weight averaging and standard ensembling, and offer a new decomposition that better explains when weight averaging fails under correlation shift or covariate shift.  Their proposed method of averaging diverse sets of weights leads to improved results on the DomainBed benchmark suite of datasets.  \n", "questions": "One thing I found inconsistent, was the combination of the results of Lemma 1 and the evidence in Figure 1.  Lemma 1 relates the loss of a weight-averaged model and a prediction-averaged model, showing  that the loss of the weight averaged model is the same as the prediction averaged model, plus a (non-negative) term based on the max of the 2-norm between the weight average parameters and an individual member.  Yet if that is so, I would expect prediction averaged models to outperform weight-averaged ones.  Figure 1 establishes that in fact the opposite is true.  Is it the case that I am misinterpreting Lemma 1?\n\nAn answer might lie at the end of section 4 on lines 256-266, as well as part of Appendix D: \"These experiments confirm that diversity is key as long as the weights remain averageable.\".  So, in effect, weight averaging performs well when the members being averaged lie within  a low-error subspace that is compact.  I would like to see this concept of \"averagable\" explored further.  There seems to be a practical friction between training diverse members of the ensemble (say by varying the initializations) and ensuring that the members do not stray too far apart so as to no longer stay averagable.  A closer investigation of when this happens (say from Figure 5) in terms of violations of the linear mode connectivity would be illuminating.", "limitations": "Yes, they have.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths:\n\nThe paper is well written.  The sections are subdivided cleanly and they preview the arguments established in each subsection, culminating in a larger point.  The authors are clearly well versed in the literature of ensembling and transfer learning.  \n\nIn particular, section 2.4 and its subsections develop the argument for why weight averaged ensembles are robust to covariate shift (or diversity shift) quite convincingly.  \n\nAlso, the analysis in section 2.2 of why flatness-based minimizers do not have any effect on OOD error was illuminating.\n\n### Weaknesses:\n\nThere are a few small issues that I'd like to see addressed:\n\nAt the outset in the abstract, the authors state that the best current approaches for out of distribution generalization are derived from averaging the weights along a training run.  This is not quite so, as there are multiple other approaches (SWAG, Deep Ensembles, LNNS, Loss surface simplexes for mode connecting volumes and fast ensembling). These all introduce greater member diversity though different mechanisms, so it would be more correct to broaden the related work section to include them, and to acknowledge them in the abstract.\n\nOn lines 185-186 the authors state that hey follow the objective of decorrelating the learning procedures of ensemble members, citing the DICE paper.  It’s an exaggeration to say that the DICE objective is followed in section 3.  The authors of DICE perform considerable work to implement a variational approximation to the conditional entropy bottleneck between two learners.  The authors of this work do not go to such lengths to decorrelate the features extracted from the inputs by ensemble members.\n\n### Comments:\n\nI think the paper would benefit from a perspective that comments upon the difficulty of transfer learning based on the presence of neural collapse, or learning neural network subspaces, especially as the authors note that the size of the space spanned by the furthest member from the barycentre is important.\n\nIt's hard to believe that this paper's claims of increased diversity by new initializations are an actual advance, this is present in all the modern ensembling literature (cf. papers descended from the DeepEnsembles literature).  In Table 1,  there does not seem to be a comparison against simple deep ensembles ($f_{ENS}$ in the notation of the paper.)\n\nA question I have with the setup in 2.1 is why would we not want to compare against gains made by fine-tuning in the target distribution.  In most practical scenarios, there is limited target data available to train on.  Certainly in almost all modern NLP tasks, this is a baseline to consider.  Including this baseline, where a model (or ensemble of models) that is trained on S is allowed to be fine-tuned in increments (of data or learning iterations) on data drawn from T before being evaluated, would qualify the degree to which DiWA is a feasible solution to practical transfer learning.\n\nOverall, despite some small quibbles, I find that there is a lot to offer in this paper and the community would benefit from seeing it in a first class venue.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IlOzSzvzMU", "reviewer": "Reviewer_xhCM", "summary": "This paper proposes Diverse Weight Averaging (DiWA), which averages weights obtained from different training runs starting from the same initial parameters and mildly different hyperparameters. The paper decomposes the expected OOD error into (bias, variance, covariance, locality) and claims that weight averaging is most effective when the covariance term dominates, verified empirically. They show improved OOD generalization performance on the DomainBed benchmark.", "questions": "What is the performance of ENS using the models obtained by DiWA?\n\nWhat is the rough training cost of obtaining 60 models for the final results? Were you unable to obtain error bounds in table 1 because of the computation costs?\n\nMinor comments\nLine 89: “…why WA outperforms in OOD Sharpness-Aware Minimizer…” I’m unsure whether the grammar here is technically wrong, but I took a while to parse this sentence.", "limitations": "To my knowledge, this work has no potential negative societal impact other than what was already present in the existing literature.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I think this is a solid paper with a conceptually simple method and clean writing.\n\nOver ENS, WA has the benefit of requiring only one feedforward at test time. Both Lemma 1 and figure 1 show that WA is a close approximation to ENS in the training setting of DiWA. In fact, WA seems to perform slightly better than ENS.\n\nDiWA is simple and effective, and the experiments show that it can be combined with previous advances for more benefits: LP initialization, leveraging diversity from ERM/Mixup/Coral algorithms, etc.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "C3TzFNAgT1", "reviewer": "Reviewer_P5Dq", "summary": "In this paper, the authors proposed a simple method to improve the out-of-domain generalization, i.e., by averaging weights from different training runs rather than one run. Some theoretical analysis is also given for explaining weight averaging (WA) and the proposed method. Experiments on the DomainBed benchmark show the performance of the proposed method. ", "questions": "1. How to decide the total number of runs? Is the number larger usually mean the final performance is better?  \n2. Is there any intuition of why using MMD in Proposition 3? How about other divergences?", "limitations": "The authors have addressed the limitations of their proposed method. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The research question of when WA succeeds and how to improve it, is interesting and relatively less explored.\n2. The paper provides theoretical insights for WA and the proposed method.\n\nWeaknesses:\n1. The writings of the paper may have some overclaim issues. In the abstract, the authors claim that \"for out-of-distribution generalization in computer vision, the best current approach averages the weights along a training run\" (line 2-3). But from the introduction, it seems that such a claim is only based on the observation from the DomainBed benchmark (line 22-24). I'm wondering does the observation on one benchmark could represent the whole situation in computer vision? \n2. The contribution of this paper is quite minimal. Especially regarding the experimental results, it seems there is no large performance improvement of the proposed method when compared with MA [29] in Table 1; under the random initialization, the performance of the proposed method looks similar to the performance of SWAD [14].   \n3. The proposed method needs to ensemble weights from different runs (e.g., >=20 runs in the experiments), which may cause unnecessary additional computational cost.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept