Diverse Weight Averaging for Out-of-Distribution Generalization
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-covariancelocality 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 [1, 2].
Yet, the classical Empirical Risk Minimization (ERM) lacks robustness to distribution shifts [3, 4, 5].
To improve out-of-distribution (OOD) generalization in classification, several recent works proposed to train models simultaneously 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-art strategy on DomainBed is currently to average the weights obtained along a 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 flatness-based analysis and provide a 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 on [18], 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 [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.
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 [20, 21].
Yet, averaging the weights of independently trained networks with batch normalization [22] and ReLU layers [23] 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 [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 a new theoretical analysis of WA for OOD based on a bias-variance-covariancelocality 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).
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-variancecovariance-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, Y the label space and ℓ : Y 2 → R + a loss function.
S is the training (source) domain with distribution p S , and T is the test (target) domain with distribution p T .
For simplicity, we will indistinctly use the notations p S and p T to refer to the joint, posterior and marginal distributions of ( X, Y ) .
We note f S , f T : X
→ Y the source and target labeling functions.
We assume that there is no noise in the data: then f S is defined on X S ≜ { x ∈ X /p S ( x ) > 0 } by ∀ ( x, y ) ∼ p S , f S ( x ) =
y and similarly f T is defined on
X T ≜ { x ∈ X /p T ( x ) > 0 } by ∀ ( x, y ) ∼ p T , f T ( x ) = y .
Problem.
We consider a neural network (NN) f ( · , θ ) :
X → Y made of a fixed architecture f with weights θ .
We seek θ minimizing the target generalization error:
[Equation]
f ( · , θ ) should approximate f T on X T .
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 p S ( X, Y ) ̸
= p T ( X, Y ) ) which are decomposed per [19] into diversity shift (a.k.a. covariate shift), when marginal distributions differ (i.e., p S ( X ) ̸ = p T ( X ) ), and correlation shift (a.k.a. concept shift), when posterior distributions differ (i.e., p S ( Y | X ) ̸ = p T ( Y | X ) and f S ̸ = f T ).
The weights are typically learned on a training dataset d S from S (composed of n S i.i.d.
samples from p S ( 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 l S = { d S , c } a learning procedure on domain S , and explicitly write θ ( l S ) to refer to the weights obtained after stochastic minimization of 1 /n S (cid:80) (
x,y ) ∈
d S ℓ ( f ( x, θ ) , y ) w.r.t.
θ under l S .
2.2 Weight averaging for OOD and limitations of current analysis
Weight averaging.
We study the benefits of combining M individual member weights { θ m } M m =1 ≜ { θ ( l ( m ) S ) } Mm =1 obtained from M (potentially correlated) identically distributed (i.d.) learning procedures L MS ≜ { l ( m ) S } Mm =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: f WA ≜ f ( · , θ WA ) , where θ WA ≜ θ WA ( L MS ) ≜ 1 /M (cid:88) M m =1 θ m , (2) is the state of the art [14, 29] on DomainBed [12] when the weights { θ m } Mm =1 are sampled along a single training trajectory (a description we refine in Remark 1 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.1).
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 a new analysis of WA; we propose one below that better explains these results.
2.3 Bias-variance-covariance-locality decomposition
We now introduce our bias-variance-covariance-locality decomposition which extends the biasvariance decomposition [32] to WA.
In the rest of this theoretical section, ℓ is the Mean Squared Error for simplicity: yet, our results may be extended to other losses as in [33].
In this case, the expected error of a model with weights θ ( l S ) w.r.t.
the learning procedure l S was decomposed in [32] into:
[Equation]
where bias( x, y ) , var( x ) are the bias and variance of the considered model w.r.t.
a sample ( x, y ) , defined later in Equation (BVCL).
To decompose WA’s error, we leverage the similarity (already highlighted in [13]) between WA and functional ensembling (ENS) [15, 34], a more traditional way to combine a collection of weights.
More precisely, ENS averages the predictions, f ENS ≜ f ENS ( · , { θ m } Mm =1 ) ≜ 1 /M (cid:80)
Mm =1 f ( · , θ m ) .
Lemma 1 establishes that f WA is a first-order approximation of f ENS when { θ m } Mm =1 are close in the weight space.
Lemma 1 (WA and ENS.
Proof in Appendix C.1.
Adapted from [13, 28].) .
Given { θ m } Mm =1 with learning procedures L MS ≜ { l ( m ) S } Mm =1 .
Denoting ∆ L MS = max Mm =1 ∥ θ m − θ WA
∥ 2 , ∀ ( x, y ) ∈
X ×Y : f WA ( x ) =
f ENS ( x ) + O (∆ 2 L MS ) and ℓ ( f WA ( x ) , y ) =
ℓ ( f ENS ( x ) , y )
+ O (∆ 2 L MS ) .
This similarity is useful since Equation (BV) was extended into a bias-variance-covariance decomposition for ENS in [18, 35].
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 MS ≜ { l ( m ) S } Mm =1 .
Proposition 1 (Bias-variance-covariance-locality decomposition of the expected generalization error of WA in OOD.
Proof in Appendix C.2.) .
Denoting ¯ f S ( x ) =
E l S [ f ( x, θ ( l S ))] , under identically distributed learning procedures L MS ≜ { l ( m ) S } Mm =1 , the expected generalization error on domain T of θ WA ( L MS ) ≜ 1 M (cid:80)[-Mm =
-]
Mm =1 θ m over the joint distribution of L MS is:
[Equation]
where bias( x, y ) =
y − ¯ f S ( x ) , and var( x ) =
E l S (cid:104)(cid:0) f
( x, θ ( l S ))
− ¯ f S ( x ) (cid:1) 2 (cid:105) , and cov( x ) =
E l S ,l ′ S (cid:2)(cid:0) f
( x, θ ( l S ))
− ¯ f S ( x ) (cid:1)(cid:0) f
( x, θ ( l ′ S )))
− ¯ f S ( x ) (cid:1)(cid:3) , and ¯∆ 2 =
E L MS ∆ 2 L MS with ∆ L MS =
M max m =1 ∥ θ m − θ WA ∥ 2 .
(BVCL) 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 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.
2.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 .
2.4.1 Bias and correlation shift (and support mismatch)
We relate OOD bias to correlation shift [19] under Assumption 1, where ¯ f S ( x ) ≜ E l S
[ f ( x, θ ( l S ))] .
As discussed in Appendix C.3.2, Assumption 1 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.
[Equation]
Proposition 2 (OOD bias and correlation shift.
Proof in Appendix C.3) .
With a bounded difference between the labeling functions f T − f S on X T ∩ X S , under Assumption 1, the bias on domain T is:
[Equation]
[Equation]
Correlation shift + Support mismatch
[Equation]
where Correlation shift = (cid:90)
[Equation]
and Support mismatch =
(cid:90)
[Equation]
We analyze the first term by noting that f T ( x ) ≜ E p T [ Y | X = x ] and f S ( x ) ≜ E p S
[ Y | X
= x ] , ∀ x ∈ X T ∩ X S .
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 a task with minimal domain-specific information as in [36].
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.
2.4.2 Variance and diversity shift
Variance is known to be large in OOD [5] 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 [19] 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.
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, 39].
The corresponding kernel K is the neural tangent kernel (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.
Assumption 3 (Constant norm and low intra-sample similarity on d S ) .
∃ ( λ S , ϵ ) with 0 ≤ ϵ ≪ λ S such that ∀ x S ∈ X d S , K ( x S , x S ) =
λ S and ∀ x ′ S ̸
= x S ∈ X d S , | K ( x S , x ′ S ) |≤ ϵ .
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 are now in a 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 d S (of size n S ) with NTK K , under Assumptions 2 and 3, the variance on dataset d T is:
[Equation]
where MMD is the empirical Maximum Mean Discrepancy in the RKHS of K 2 ( x, y )
= ( K ( x, y )) 2 ; λ T ≜ E x T ∈
X
dT K ( x T , x T ) and β T ≜ E ( x T ,x ′ T ) ∈
X 2 dT ,x T ̸ = x ′ T K 2 ( x T , x ′ T ) are the empirical mean similarities respectively measured between identical (w.r.t.
K ) and different (w.r.t.
K 2 ) samples averaged over X d T .
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 [19]: MMD replaces the Ldivergence used in [19].
The other terms, λ T and β 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 .
2.4.3 Covariance and diversity
The covariance term increases when the predictions of { f ( · , θ m ) } Mm =1 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 Appendix 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 [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 O ( ¯∆ 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 [24].
We empirically verify later in Figure 1 that the approximation f WA
≈ f ENS remains valid even in this case.
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 { l ( m ) S } Mm =1 .
Yet, previous methods [14, 29] 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.
DiWA.
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 ( m ) S } Mm =1 .
Our weights are obtained from M ≫ 1 different runs, with diverse learning procedures:
Algorithm
Require: θ 0 pretrained encoder and initialized classifier; { h m } Hm =1 hyperparameter configurations.
Training: ∀ m = 1 to H , θ m ≜
FineTune( θ 0 , h m )
Weight selection:
Uniform: M = { 1 , · · · , H } .
Restricted: Rank { θ m } Hm =1 by decreasing ValAcc( θ m ) .
M ← ∅ .
for m = 1 to H do
If ValAcc( θ M∪{ m } ) ≥ ValAcc( θ M )
M ← M ∪ { m }
Inference : with f ( · , θ M ) , where θ M = (cid:80) m ∈M θ m / |M| .
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 [21] and reduce the impact of variance when M is large.
However, this may break the locality requirement analyzed in Section [-2.4.
-]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 a 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 as 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, we 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 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 [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, we thus use 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 to 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 a similar analysis on PACS [51]).
Then, Section 5 shows that DiWA is state of the art on DomainBed [12].
Empirical validation of our theoretical insights
We consider several collections of weights { θ m } Mm =1 ( 2 ≤ 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 a single run (everybatches) or from different runs.
They are evaluated on “Art”, the fourth domain from OfficeHome.
WA vs. ENS.
Figure 1 validates Lemma 1 and that f WA ≈ f ENS .
More precisely, f WA slightly but consistently improves f 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.
Diversity and accuracy.
We validate in Figure 2 that f WA benefits from diversity.
Here, we measure diversity with the ratio-error [46], i.e., the ratio N diff /N simul between the number of different errors
N diff and of simultaneous errors N simul in test for a pair in { f ( · , θ m ) }
Mm =1 .
A higher average over the (cid:0) M 2 (cid:1) pairs means that members are less likely to err on the same inputs.
Specifically, the gain of
Acc( θ WA ) over the mean individual accuracy 1 M (cid:80)
Mm =
1 Acc( θ m ) increases with diversity.
Moreover, this phenomenon intensifies for larger M : the linear regression’s slope (i.e., the accuracy gain per unit of diversity) increases with M .
This is consistent with the ( M − 1) /M factor of cov( x ) in
Equation (BVCL), as further highlighted in Appendix E.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 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) fromdifferent 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 [12].
By imposing the code, the training procedures and the ResNet50 [52] architecture, DomainBed is arguably the fairest benchmark for OOD generalization.
It includes 5 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-domain model selection.
The experimental setup is further described in Appendix G.1.
Our code is available at https://github.com/alexrame/diwa .
Baselines.
ERM is 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 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 ≤ M ≤ 20 weights with
Algorithm 1 while DiWA-uniform averages
all M = 20 weights.
DiWA † averages uniformly the
M = 3 × 20 = 60 weights from all 3 data splits.
DiWA † benefits from larger M (without additional inference cost) and from data diversity (see Appendix 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.
5.1 Results on DomainBed
We report our main results in Table 1, detailed per domain in Appendix 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) [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 a new SoTA of 68 . 0% , improving SWAD by 1 . 1 points.
[Table] [Table]
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 [56] and Coral [10]: accuracy gradually improves as we add more objectives.
Indeed, as highlighted in Appendix 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.
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.
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 to 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 a 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], a 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 “explor[ing] the set of possible solutions instead of simply converging to a single point”, as stated in [84].
The recent “Model soups” introduced by Wortsman et al.
[28] 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, [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 SoTA 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.
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.
Acknowledgements
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 AD 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-CE230012 and LEAUDS ANR-18-CE23-0020.
Appendices
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.
• Appendices C.1 and C.2 derive the bias-variance-covariance-locality decomposition for WA (Proposition 1).
• Appendix C.3 establishes the link between bias and correlation shift (Proposition 2).
• Appendix C.4 establishes the link between variance and diversity shift (Proposition 3).
• Appendix 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 [86].
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 [1, 2] or to tackle fairness issues related to under-representation [57, 87, 88].
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 [12].
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.
B Limitations of the flatness-based analysis in OOD
Theorem 1 (Equation 21 from [14], simplified version of their Theorem 1) .
Consider a set of N covers { Θ k } Nk =1 s.t.
the parameter space Θ
⊂ ∪ Nk Θ k where diam(Θ) ≜ sup θ,θ
′ ∈
Θ ∥ θ − θ ′ ∥ 2 , N ≜ (cid:6) (diam(Θ) /γ )
d (cid:7) and d is the dimension of Θ .
Then, ∀ θ ∈ Θ with probability at least 1 − δ :
[Equation]
where:
• E T ( θ ) ≜ E ( x,y ) ∼ p T ( X,Y )
[ ℓ ( f θ ( x ); y )] is the expected risk on the target domain,
• Div( p S , p T ) ≜ 2 sup A | p S ( A )
− p T ( A ) | is a divergence between the source and target marginal distributions p S and p T : it measures diversity shift.
• E S ( θ ) ≜ E ( x,y ) ∼ p S ( X,Y )
[ ℓ ( f θ ( x ); y )] is the expected risk on the source domain, • E γd S ( θ ) ≜ max ∥ ∆ ∥≤ γ E d S ( θ + ∆) (where E d S ( θ + ∆)
≜ E ( x,y ) ∈
d S [ ℓ ( f θ +∆ ( x ); y )] ) is the robust empirical loss on source training dataset d S from S of size n S , • v k is a VC dimension of each Θ k .
Previous understanding of WA’s success in OOD relied on this upper-bound, where E γd S ( θ ) involves the solution’s flatness.
This is usually empirically analyzed by the trace of the Hessian [89, 90, 91]: indeed, with a second-order Taylor approximation around the local minima θ and h the Hessian’s maximum eigenvalue, E γd S ( θ ) ≈ E d S ( θ ) + h × γ 2 .
In the following subsections, we show that this inequality does not fully explain the exceptional performance of WA on DomainBed [12].
Moreover, we illustrate that our bias-variance-covariancelocality addresses these limitations.
B.1 Flatness does not act on distribution shifts
The flatness-based analysis is not specific to OOD.
Indeed, the upper-bound in Equation (5) sums up two noninteracting terms: a domain divergence Div( p S , p T ) that grows in OOD and E γd S ( θ ) that measures the IID flatness.
The flatness term can indeed be reduced empirically with WA: yet, it does not tackle the domain gap.
In fact, Equation (5) states that additional flatness reduces the upper bound of the error similarly no matter the strength of the distribution shift, thus as well OOD than IID.
In contrast, our analysis shows that variance (which grows with diversity shift, see Section 2.4.2) is tackled for large M : our error is controlled even under large diversity shift.
This is consistent with our experiments in Table 1.
Our analysis also explains why WA cannot tackle correlation shift (where bias dominates, see Appendix H), a limitation [14] does not illustrate.
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) [30], another popular method to find flat minima based on minimax optimization: it minimizes the maximum loss around a neighborhood of the current weights θ .
In Figure 6, we compare the flatness (i.e., the Hessian trace computed with the package in [91]) and accuracy of ERM, MA [29] (a WA strategy) and SAM [30] when trained on the “Clipart”, “Product” and “Photo” domains from OfficeHome [50]: they are tested OOD on the fourth domain “Art”.
Analyzing the second and the third rows of Figures 6a and 6b, 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 6c, where MA outperforms SAM.
As reported in Table 3, similar experiments across more datasets lead to the same conclusions in [14].
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 [15] (see
Lemma 1), a privileged link that SAM does not benefit from.
[Table] [Figure]
[Table]
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 [31], we confirm in
Figures 6a and 6b 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 6c.
[14] 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 [92]: “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.
C Proof C.1 WA loss derivation
Lemma (1) .
Given { θ m } Mm =1 with learning procedures L MS ≜ { l ( m ) S } Mm =1 .
Denoting ∆ L MS = max Mm =1 ∥ θ m − θ WA
∥ 2 , ∀ ( x, y ) ∈
X × Y : f WA ( x ) =
f ENS ( x ) + O (∆ 2 L MS ) and ℓ ( f WA ( x ) , y ) =
ℓ ( f ENS ( x ) , y )
+ O (∆ 2 L MS ) .
Proof.
This proof has two components:
• to establish the functional approximation, as [13], it performs Taylor expansion of the models’ predictions at the first order.
• to establish the loss approximation, as [28], it performs Taylor expansion of the loss at the first order.
Functional approximation With a Taylor expansion at the first order of the models’ predictions w.r.t.
parameters θ :
[Equation]
Therefore, because (cid:80) Mm =1 ∆ m = 0 , f ENS − f WA = O (cid:0) ∆ 2 (cid:1) where ∆ =
M max m =1 ∥ ∆ m ∥ 2 .
Loss approximation With a Taylor expansion at the zeroth order of the loss w.r.t.
its first input and injecting Equation (6):
[Equation]
C.2 Bias-variance-covariance-locality decomposition
Remark 1.
Our result in Proposition 1 is simplified by leveraging the fact that the learning procedures
L MS = { l ( m ) S } Mm =1 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 [29] and SWAD [14].
It is even false if the weights { θ ( l ( m ) S ) }
Mm =1 are defined as being taken sequentially along a training trajectory, i.e., when 0 ≤ i < j ≤ M implies that l ( i ) S has fewer training steps than l ( j ) S .
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., θ ( l ( i ) S ) 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.
Proposition (1) .
Denoting ¯ f S ( x ) =
E l S [ f ( x, θ ( l S ))] , under identically distributed learning procedures L MS ≜ { l ( m ) S } Mm =1 , the expected generalization error on domain T of θ WA ( L MS ) ≜
1 M (cid:80) Mm =
θ m over the joint distribution of L MS is: E
[Equation]
MS E T ( θ WA
( L MS )) = E ( x,y ) ∼ p T (cid:104) bias 2 ( x, y ) + 1 M var( x ) + M − 1 M cov( x ) (cid:105)
+ O ( ¯∆ 2 ) , where bias( x, y ) =
y − ¯ f S ( x ) , and var( x ) =
E l S (cid:104)(cid:0) f
( x, θ ( l S ))
− ¯ f S ( x ) (cid:1) 2 (cid:105) , and cov( x ) =
E l S ,l ′ S (cid:2)(cid:0) f
( x, θ ( l S ))
− ¯ f S ( x ) (cid:1)(cid:0) f
( x, θ ( l ′ S )))
− ¯ f S ( x ) (cid:1)(cid:3) , and ¯∆ 2 =
E L MS ∆ 2 L MS with ∆ L MS =
M max m =1 ∥ θ m − θ WA ∥ 2 .
(BVCL) 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.
Proof.
This proof has two components:
• it follows the bias-variance-covariance decomposition from [18, 35] 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.
BVC for ensembling with identically distributed learning procedures With ¯ f S ( x ) =
E l S [ f ( x, θ ( l S ))] , we recall the bias-variance decomposition [32] (Equation (BV)):
[Equation]
, where bias( x, y ) =
Bias { f | ( x, y ) } =
y − ¯ f S ( x ) , and var( x ) =
Var { f | x } =
E l S (cid:104)(cid:0) f
( x, θ ( l S ))
[Equation]
S ( x ) (cid:1) 2 (cid:105) .
Using f ENS ≜
f ENS ( · , { θ ( l ( m ) S ) }
Mm =1 ) ≜ 1 M (cid:80)
Mm =1 f ( · , θ ( l ( m ) S ))
in this decomposition yields,
[Equation]
As f ENS depends on L MS , we extend the bias into:
[Equation]
Under identically distributed L MS ≜ { l ( m ) S } Mm =1 ,
M
M (cid:88)
[Equation]
Thus the bias of ENS is the same as for a single member of the WA.
Regarding the variance:
[Equation]
(cid:35)(cid:33) 2   .
Under identically distributed
[Equation]
The variance is split into the variance of a single member (divided by M ) and a covariance term.
Combination with Lemma 1 We recall that per Lemma 1,
[Equation]
Then we have:
[Equation]
We eventually obtain the result:
[Equation]
C.3 Bias, correlation shift and support mismatch
We first present in Appendix C.3.1 a decomposition of the OOD bias without any assumptions.
We then justify in Appendix C.3.2 the simplifying Assumption 1 from Section 2.4.1.
C.3.1 OOD bias
Proposition 4 (OOD bias) .
Denoting ¯ f S ( x ) =
E l S [ f ( x, θ ( l S ))] , the bias is:
E ( x,y ) ∼ p T [bias 2 ( x, y )]
= (cid:90)
X T ∩X S ( f T ( x ) − f S ( x ))
2 p T ( x ) dx ( Correlation shift ) + (cid:90)
X T ∩X S (cid:0) f S ( x ) − ¯ f S ( x ) (cid:1) 2 p T ( x ) dx ( Weighted IID bias ) + (cid:90)
X T ∩X S 2( f T ( x ) − f S ( x )) (cid:0) f S ( x ) − ¯ f S ( x ) (cid:1) p T ( x ) dx ( Interaction IID bias and corr.
shift )
+ (cid:90)
X T \X S (cid:0) f T ( x ) − ¯ f S ( x ) (cid:1) 2 p T ( x ) dx.
( Support mismatch )
Proof.
This proof is original and based on splitting the OOD bias in and out of X S :
[Equation]
To decompose the first term, we write ∀ x ∈ X S , − ¯ f S ( x ) = − f S ( x ) + (cid:0) f S ( x ) − ¯ f S ( x ) (cid:1) .
[Equation]
The four terms can be qualitatively analyzed:
• The first term measures differences between train and test labelling function.
By rewriting ∀ x ∈ X T ∩ X S , f T ( x ) ≜ E p T [ Y | X = x ] and f S ( x ) ≜ E p S
[ Y | X
= x ] , this term measures whether conditional distributions differ.
This recovers a similar expression to the correlation shift formula from [19].
• The second term is exactly the IID bias, but weighted by the marginal distribution p T ( X ) .
• The third term (cid:82) X T ∩X
S 2( f T ( x ) − f S ( x )) (cid:0) f S ( x ) − ¯ f S ( x ) (cid:1) 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 [36].
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
Assumption 1 states that ∃ ϵ > 0 small s.t. ∀ x ∈ X S , | f S ( x ) − ¯ f S ( x ) |≤ ϵ where ¯ f S ( x ) =
E l S [ f ( x, θ ( l S ))] .
¯ f S is the expectation over the possible learning procedures l S =
{ d S , c } .
Thus
Assumption 1 involves:
• 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 | θ | 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 [14], this is also realistic thanks to their overfit-aware weight selection strategy.
In contrast, this assumption may not perfectlty hold for MA [29], 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.
Proposition (2.
OOD bias when small IID bias) .
With a bounded difference between the labeling functions f T − f S on X T ∩ X S , under Assumption 1, the bias on domain T is:
[Equation]
shift + Support mismatch + O ( ϵ ) , where Correlation shift = (cid:90)
[Equation]
, and Support mismatch = (cid:90)
[Equation]
Proof.
We simplify the second and third terms from Proposition 4 under Assumption 1.
The second term is (cid:82) X T ∩X S (cid:0) f S ( x ) − ¯ f S
( x ) (cid:1) 2 p T ( x ) dx .
Under Assumption 1, | f S ( x ) − ¯ f S ( x ) |≤ ϵ .
Thus the second term is O ( ϵ 2 ) .
The third term is (cid:82) X T ∩X
S 2( f T ( x ) − f S ( x )) (cid:0) f S ( x ) − ¯ f S ( x ) (cid:1) p T ( x ) dx .
As f T − f S is bounded on
X S ∩ X T , ∃ K ≥ 0 such that ∀ x ∈ X S ,
[Equation]
Thus the third term is O ( ϵ ) .
Finally, note that we cannot say anything about ¯ f S ( x ) when x ∈ X T \ X S .
To prove the previous equality, we needed a bounded difference between labeling functions f T − f S on X T ∩ X S .
We relax this bounded assumption to obtain an inequality in the following Proposition 5.
Proposition 5 (OOD bias when small IID bias without bounded difference between labeling functions) .
Under Assumption 1,
[Equation]
Proof.
We follow the same proof as in Proposition 4, except that we now use: ( a + b ) 2 ≤ 2( a 2 + b 2 ) .
Then,
[Equation]
[Table]
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 Appendix C.4.1.
We discuss our simplifying Assumption 3 in Appendix C.4.2.
We present our final proof in Appendix C [Table]
We discuss the relation between variance and initialization in Appendix C.4. [Table]
C.4.1 Neural Networks as Gaussian Processes
We fix d S [Table] d T and denote X d S
= { x S } ( x S ,y S ) ∈
d S , X d T = { x T } ( x T ,y T )
∈
d T their respective input supports.
We fix the initialization of the network.
l S encapsulates all other sources of randomness.
Lemma 2 (Inspired from [93]) .
Given a NN f ( · , θ ( l S )) under Assumption 2, we denote K its neural tangent kernel and K ( X d S , X d S ) ≜ ( K ( x S , x ′ S ))
x S ,x ′ S ∈
X 2 dS ∈ R n S × n S .
Given x ∈ X , we denote
K ( x, X d S ) ≜
[ K ( x, x S )] x S ∈ X dS ∈ R n S .
Then:
[Equation]
Proof.
Under Assumption 2, NNs are equivalent to GPs.
var( x ) is the formula of the variance of the
GP posterior given by Eq.
(2.26) in [93], when conditioned on d S .
This formula thus also applies to the variance f ( · , θ ( l S )) when l S varies (at fixed d S and initialization).
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 Appendix C.4.3).
The first part of Assumption 3 assumes that ∃ λ S such that all training inputs x S ∈ X d S verify
K ( x S , x S ) =
λ S .
Note that this equality is standard in some kernel machine algorithms [40, 41, 42] and is usually achieved by replacing K ( x, x ′ ) by λ S
[Equation]
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 [39].
This is theoretically based: for example, the NTK K ( x, x ′ ) for an architecture with an initial fully connected layer only depends on ∥ x ∥ , ∥ x ′ ∥ , ⟨ x, x ′ ⟩ [95].
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 ∈ X d S ; we denote λ S the corresponding value.
The second part of Assumption 3 states that ∃ 0 ≤ ϵ ≪ λ S , s.t. ∀ x S , x ′ S ∈
X 2 d S , x S ̸
= x ′ S ⇒ | K ( x S , x ′ S )
|≤ ϵ , i.e., that training samples are dissimilar and do not interact.
This diagonal structure of the NTK [37], with diagonal values larger than non-diagonal ones, is consistent with empirical observations from [44] at initialization.
Theoretically, this is reasonable if K is close to the RBF kernel
K h ( x, x ′ ) = exp( −∥ x − x ′ ∥
22 /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 [44] at the end of training, consistently with the local elasticity property of NNs [43].
We then consider the dataset d S ′ ⊂
d S made of one sample per block, to which Assumption 3 applies.
As decreasing the size of a training dataset empirically reduces variance [96], the variance of f trained on d S is upper-bounded by the variance of f trained on d S ′ ; the latter is given by applying Proposition 3 to d S ′ .
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
Proposition (3) .
Given f trained on source dataset d S (of size n S ) with NTK K , under Assumptionsand 3, the variance on dataset d T is:
[Equation]
(4) with MMD the empirical Maximum Mean Discrepancy in the RKHS of K 2 ( x, y )
=
( K ( x, y )) 2 ; λ T ≜
E x T ∈ X dT K ( x T , x T ) and β T ≜
E ( x T ,x ′ T ) ∈
X 2 dT ,x T ̸ = x ′ T K 2 ( x T , x ′ T ) the empirical mean similarities resp.
measured between identical
Proof.
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 ∀ x ∈ X :
[Equation]
Denoting B = K ( X d S , X d S ) − 1 with symmetric coefficients
b i,j = b j,i , then var( x
[Equation]
Assumption 3 states that K ( X d S , X d S ) =
A + H where A = λ S
I n S and H =
( h ij ) 1 ≤ i ≤ n S 1 ≤ j ≤ n S with h i,i = [-
-]0 and max i,j | h i,j |≤ ϵ .
We fix x T ∈ X d T and determine the form of B − 1 in two cases: ϵ = 0 and ϵ ̸ = 0 .
Case when ϵ
=
We first derive a simplified result, when ϵ = 0 .
Then, b i,i = 1 λ S and b i,j = 0 s.t.
[Equation]
We can then write:
[Equation]
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 [97]:
[Equation]
Finally, because ϵ = 0 , E {+x S ̸ = x ′ S ∈ X 2 dS K 2 ( x S , x ′ S ) = 
 s.t.+}
[Equation]
We recover the same expression with a O ( ϵ ) in the general setting where ϵ ̸ = 0 .
Case when ϵ ̸ =
We denote I :
[Equation]
GL
[Equation]
A
[Equation]
the inversion function defined on
GL n S ( R ) , the set of invertible matrices of M n S ( R ) .
The function I is differentiable [98] in all A ∈ GL n S ( R ) with its differentiate given by the linear application dI A : (cid:26) M n S ( R ) → M n S ( R )
H (cid:55)→ − A − 1 HA − 1 .
Therefore, we can perform a Taylor expansion of I at the first order at A :
[Equation]
where ∥ H ∥≤ n S ϵ
=
O ( ϵ ) .
Thus,
[Equation]
λ [Equation] S + H [Equation]
=
[Equation]
λ S
I n S
[Equation] − ( [Equation] Therefore, when ϵ is small, Equation (10) can be developed into:
[Equation]
Following the derivation for the case ϵ = 0 , and remarking that under Assumption 3 we have E x S ̸
= x ′ S ∈ X 2 dS K 2 ( x S , x ′ S )
=
O ( ϵ 2 ) , yields:
[Equation]
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 [37].
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 [48], as commonly done on
DomainBed [12].
The Linear Probing [49] initialization of the classifier was shown in [49] 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 [99] image-to-text translation as in [36].
C.5 WA vs. its members
We validate that WA’s expected error is smaller than its members’ error under the locality constraint.
Lemma 3 (WA vs. its members.) .
[Equation]
Proof.
The proof builds upon Equation (BVCL):
[Equation]
and the expression of the standard bias-variance decomposition in Equation (BV) from [32],
[Equation]
The difference between the two provides:
[Equation]
Cauchy Schwartz inequality states | cov( Y, Y ′ )
|≤ (cid:112) var( Y )var( Y ′ ) , thus cov( x ) ≤ var( x ) .
Then:
[Equation]
D Weight averaging versus functional ensembling
We further compare the following two methods to combine M weights { θ ( l ( m ) S ) }
Mm =1 : f WA that averages the weights and f ENS [15] that averages the predictions.
We showed in Lemma 1 that f WA ≈ f ENS when max Mm =1 ∥ θ ( l ( m ) S ) − θ WA ∥ 2 is small.
In particular, when { l ( m ) S } Mm =1 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 WA performs slightly but consistently better than f ENS .
The observed improvement is non-trivial; we refer to Equation 1 in [28] for some initial explanations based on the value of OOD Hessian and the confidence of f WA .
The complete analysis of this second-order difference is left for future work.
Yet, we do not claim that f WA is systematically better than f 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 [15] 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 [100].
Thus the
Ensemble of Averages (EoA) strategy introduced in [29] is complementary to DiWA and could be extended into an Ensemble of Diverse Averages.
[Table]
E Additional diversity analysis E.1 On OfficeHome E.1.1 Feature diversity
In Section [Table] our diversity-based theoretical findings were empirically validated using the ratioerror [46], a common diversity measure notably used in [73, 72].
In Figure 10, we recover similar conclusions with another diversity measure: the Centered Kernel Alignment Complement (CKAC) [47], also used in [25, 26].
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
In Figures 2 and 10a [Table] we indicated the slope of the linear regressions relating diversity to accuracy gain at fixed M (between 2 and 9 ).
For example, when M = 9 weights are averaged, the accuracy
[Figure]
[Figure] [Figure]
gain increases by 0 .
297 per unit of additional diversity in prediction [46] (see Figure 2) and by 0 .per unit of additional diversity in features [47] (see Figure 10a).
Most importantly, we note that the slope increases with M .
To make this more visible, we plot slopes w.r.t.
M in Figure 11.
Our observations are consistent with the ( M − 1) /M
factor in front of cov( x ) in Equation (BVCL).
This shows that diversity becomes more important for large M .
Yet, large M is computationally impractical in standard functional ensembling, as one forward step is required per model.
In contrast,
WA has a fixed inference time which allows it to consider larger M .
Increasing M from 20 to 60 is the main reason why DiWA † improves DiWA.
E.1.3 Diversity comparison across a wide range of methods
Inspired by [21], 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 , 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 10c).
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 [101], e.g., in bagging [68].
Finally, we observe that diversity is increased (notably in features) when two runs have different objectives, for example, Interdomain Mixup [56] and Coral [10].
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 [21].
[Figure]
First , it confirms Figures 3 and 10b: weights obtained from two different runs are more different than those sampled from a single run (even with extreme hyperparameters).
Second , this shows that weights from two runs are more diverse when the two runs have different hyperparameters/data/classifier initializations/training objectives.
Domain “Art” on OfficeHome.
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 [25] 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 13a.
Yet, we observe in Figure 13b that if learning rates are too different, weight averaging no longer approximates functional ensembling because the O (∆ 2 L MS ) term in Lemma 1 can be large.
[Figure]
Considering M = 2 weights obtained from two learning procedures with learning rates lrand lr 2 (sampled from the extreme distribution in Table 7), we plot in Figure 13a the prediction diversity for these M = 2 models vs. | lr 1 − lr 2 | .
Then, in Figure 13b, we plot the accuracy differences
Acc( DiWA ) − Acc( ENS ) vs. | lr 1 − lr 2 | .
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.
[Figure]
[Figure]
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 [12].
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
[Table]
in Table 6 we report DiWA’s results when considering only 5 runs
[Table]
with uniform weight selection. Interestingly, it shows that M = 5 is enough to be competitive against SWAD [14], the previous state of the art.
[Table]
G DomainBed G.1 Description of the DomainBed benchmark
We now further detail our experiments on the DomainBed benchmark [12].
Data.
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 withdifferent seeds: the reported numbers are the means and the standard errors over these 3 seeds.
Training protocol.
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 [14], yet could be adapted by dataset for better results.
Even though these distributions are more restricted than the extreme distributions introduced [12], our ERM runs perform better.
It leads to a total amount of 2640 runs only for Table 1.
In
Appendix B, the ρ hyperparameter for SAM is sampled from [0 . 001 , 0 . 002 , 0 .
In Table 2, hyperparameters specific to Interdomain Mixup [56] (“mixup_alpha”) and Coral [10] (“mmd_gamma”) are sampled from the distributions defined in [12].
We use a ResNet50 [52] 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 [102].
Our classifier is either initialized randomly or with Linear Probing [49]; 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 [14, 29].
As in [14], validation accuracy is calculated everysteps for VLCS, 500 steps for DomainNet and 100 steps for others.
[Table]
Model selection and scores.
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 × 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 [12].
Scores for SWAD [14] and MA [29] 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 [29] 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 [51], VLCS [53], OfficeHome [50], TerraIncognita [54] and DomainNet [55].
Critically, [19] showed that diversity shift dominates in these datasets.
[Table] [Table]
[Table] [Table]
[Table]
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 Appendix 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 [8] dataset, which is dominated by correlation shift [55].
Colored MNIST is a colored variant of the MNIST handwritten digit classification dataset [103] where the correlation strengths between color and label vary across domains.
We follow the protocol described in Appendix G.1 except that (1) we used the convolutional neural network architecture introduced in DomainBed [12] for MNIST experiments and (2) we used the test-domain model selection in addition to the train-domain model selection.
Indeed, as stated in [19], “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 [8] or Fishr [11] remain state-of-the-art.
[Table] [Table]
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 [86, 104].
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 [86], 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.
[Table]