Diverse Weight Averaging for Out-of-Distribution Generalization

Alexandre Ramé; Matthieu Kirchmeyer; Thibaud Rahier; Alain Rakotomamonjy; Patrick Gallinari; Matthieu Cord

Source-linked AI summary

Diverse Weight Averaging for Out-of-Distribution Generalization

Alexandre Ramé, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, Matthieu Cord

arXiv:2205.09739v2May 19, 2022cs.CVcs.AIcs.LG

TL;DR

Out-of-distribution generalization remains difficult under distribution shifts, and existing flatness-based accounts do not fully explain weight averaging’s success. The paper introduces DiWA, which averages weights from independently trained, diverse models using a bias-variance-covariance-locality analysis. DiWA improves DomainBed state-of-the-art performance without additional inference cost.

Problem
Existing flatness-based analysis does not fully explain why weight averaging succeeds for out-of-distribution generalization.
Method
DiWA averages weights from independently trained models, motivated by a bias-variance-covariance-locality decomposition and the goal of reducing covariance through diversity.
Results
DiWA improves the state of the art on DomainBed, with LP-initialized DiWA† reaching 68.0% and improving SWAD by 1.1 points.
Takeaways & Limitations
Weight averaging can exploit diverse models to improve OOD generalization without the additional inference cost of standard ensembling.
Takeaways & Limitations
DiWA cannot use diversity that breaks linear connectivity and, like other weight-averaging methods, does not address correlation shift.

Abstract

from arXiv · show

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.

1 Introduction

The paper challenges flatness-based explanations of weight averaging and proposes DiWA, which averages weights from independent runs to increase model diversity while preserving averageability.

OOD generalization is difficult because ERM lacks robustness to distribution shifts, while domain-invariant methods only slightly improve or underperform ERM on DomainBed.
The paper reinterprets weight averaging through its similarity to ensembling and introduces a bias-variance-covariance-locality decomposition of expected error.
The proposed analysis motivates increasing diversity to reduce covariance while maintaining the locality needed for successful weight averaging.
DiWA averages weights from independent training runs because their differing training procedures produce more diverse models than a single trajectory.
The paper’s contributions include a new theoretical analysis, diversity-focused training procedures, and empirical validation on DomainBed.

2 Theoretical insights

The paper formulates OOD generalization as minimizing target error under source-target distribution shifts, then explains weight averaging through a bias-variance-covariance-locality decomposition. The analysis identifies diversity shift as a setting where averaging can reduce error, provided models remain diverse yet close in weight space.

Problem formulation: The target objective is minimizing expected loss on target distribution pT despite training only on related but different source data pS.
Weight averaging: Weight averaging combines M member weights and is effective on DomainBed when members are sampled along a single training trajectory.
Limits of flatness: The flatness-based explanation is insufficient because it does not control distribution-shifted OOD error or explain WA’s differing behavior with SAM.
BVCL decomposition: WA’s expected OOD error decomposes into bias, variance, covariance, and locality terms, with variance divided by M but covariance and locality requiring control.
Shift analysis: The analysis predicts that bias dominates under correlation shift, whereas variance dominates under diversity shift; WA is therefore effective when models are diverse but close in weight space.
Diversity shift: Combining M models reduces diversity-shift variance by M without requiring target-domain data.
Covariance and diversity: Lower prediction covariance increases WA’s gain over individual members, motivating functionally diverse models and decorrelated learning procedures.

3 DiWA: Diverse Weight Averaging

DiWA reduces covariance by averaging models from diverse independent runs while using shared initialization and mild hyperparameter variation to preserve weight locality. Weight selection can filter harmful runs, and experiments report state-of-the-art performance across DomainBed datasets.

Independent runs: Unlike prior approaches that average highly correlated models from one trajectory, DiWA uses many runs with diverse learning procedures.
DiWA motivation: DiWA targets the covariance term in WA’s OOD error by decorrelating the learning procedures used to obtain averaged models.
Diverse procedures: DiWA’s implementation trains models from shared initialization and varies hyperparameters, batch orders, augmentations, stochastic noise, and training duration.
Shared initialization: Shared pretrained initialization supports low-loss linear connectivity, while linear probing keeps weights closer and reduces the locality term.
Hyperparameter search: Mild hyperparameter ranges are used because extreme differences can place weights in regions that are not linearly connectable.
Weight selection: Restricted weight selection ranks models by validation accuracy and can remove detrimental runs that hurt uniform averaging.

4 Empirical validation of our theoretical insights

The experiments test whether model diversity improves weight averaging and whether independently trained weights outperform weights sampled along one run. Results show that diversity improves WA, especially with more models, while gains depend on weights remaining averageable.

WA versus prediction averaging: WA accuracy slightly but consistently exceeds prediction averaging, and larger model sets improve results.Weight averaging also avoids the inference cost of M forward passes required by prediction averaging.
Diversity and accuracy: WA’s accuracy gain over individual members increases with prediction diversity, with the relationship becoming stronger as M grows.The reported gain is measured using the ratio-error diversity between different and simultaneous errors.
Different runs: Weights sampled from different runs become more diverse than weights sampled along one run and eventually produce higher OOD accuracy as M increases.Both approaches initially improve similarly, but different-run WA gradually outperforms single-run WA.
Averageability conditions: The different-run advantage holds for mild hyperparameter ranges and shared initialization, but extreme ranges or differing initializations can make WA underperform its members.These settings can violate the locality condition needed for weights to remain averageable.

5 Experimental results on the DomainBed benchmark

On DomainBed, DiWA improves weight-averaging results by combining independently trained models, with restricted selection helping when some runs are detrimental. With linear-probing classifier initialization, DiWA† reaches a new state of the art at 68.0%.

Evaluation setup: DiWA uses multiple runs instead of retaining only one grid-search run, while ERM and MA select the highest-validation-accuracy model among 20 runs.Each DiWA run selects its highest-validation-accuracy epoch, and results average across three data splits.
Results: DiWA†-uniform is best on PACS, VLCS, and OfficeHome with random classifier initialization, while DiWA-restricted helps on TerraIncognita and DomainNet by filtering bad runs.Uniform averaging can be hurt by detrimental runs; restricted selection adds weights only when validation accuracy improves.
Results: 1.3, 0.5, and 1.1 points are DiWA†’s improvements over MA on OfficeHome, TerraIncognita, and DomainNet, respectively, with linear-probing initialization.The reported gains are dataset-specific and use the DiWA† variant.
Results: 68.0% establishes a new state of the art for DiWA† with linear probing, improving SWAD by 1.1 points.This is the paper’s aggregate DomainBed result after weight averaging.
Different objectives: DiWA-uniform improves as weights trained with Interdomain Mixup and Coral are added, consistent with benefits from increased solution diversity.The paper suggests, but leaves for future work, a linear-connectivity explanation across models trained with different objectives.

6 Related work

Related work contrasts prediction ensembling and weight averaging for out-of-distribution generalization. DiWA differs by providing a theoretical diversity-based explanation and averaging independently trained models.

Generalization and ensemble: Ensembling deep networks consistently improves robustness and has been applied to domain generalization, whereas invariant approaches do not outperform ERM on various benchmarks.The paper positions its analysis near prior work connecting ensembling with robustness.
Weight averaging: Prior weight-averaging methods commonly combine models collected along a single run, while other approaches relax constraints or explicitly increase diversity.Examples include high learning rates, weight-space volume, cosine-similarity minimization, and tempered posteriors.
Weight averaging: DiWA differs from Model Soups by explaining why weight averaging succeeds under diversity shift and by linking correlation shift to bias and diversity shift to variance.The paper describes these theoretical links and diversity findings as novel and empirically confirmed.

7 Conclusion

The paper explains weight averaging’s OOD success through its ensembling nature and a bias-variance-covariance-locality decomposition. DiWA then averages independently trained models, improves DomainBed’s state of the art, and preserves single-model inference cost.

Contributions: The proposed decomposition relates bias to correlation shift and variance to diversity shift, identifying diversity as important for generalization.This analysis motivates the design of DiWA.
Contributions: DiWA averages weights from independently trained models to exploit the diversity highlighted by the theoretical analysis.The method targets diversity without changing inference cost.
Implications: DiWA improves the state of the art on DomainBed without additional inference cost, avoiding a key limitation of standard ensembling.The conclusion frames this as a practical advantage of weight averaging.

Checklist

The checklist reports that the paper documents its assumptions, proofs, implementation details, evaluation procedures, limitations, and potential societal impacts. It also highlights a theoretical account of when weight averaging succeeds and broader benefits and costs of DiWA.

Reproducibility and theory: The paper states that it provides complete assumptions and proofs for its theoretical results.Assumption 1 is discussed in Appendix C.3.2, Assumptions 2 and 3 in Appendix C.4.2, and complete proofs are included in Appendix C.
Reproducibility and theory: The authors report releasing code and specifying training details, data splits, hyperparameters, error bars, and compute resources.The reported compute is approximately 20,000 Nvidia V100 GPU hours, mostly for the 2,640 runs in Table 1.
Impact and limitations: DiWA is presented as improving OOD generalization without additional inference cost, while requiring independent training of several models and retaining deep-learning risks.The broader-impact discussion mentions adversarial attacks, lack of interpretability, and possible misuse in surveillance or weapon systems.
Theoretical analysis: The paper’s analysis connects correlation shift to bias and diversity shift to variance, explaining why weight averaging is effective when diversity shift dominates.The decomposition also identifies covariance and locality as relevant terms, with effectiveness depending on sufficiently many diverse but nearby weight-averaged models.

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

The analysis links variance to input-distribution shifts and examines kernel assumptions that simplify its characterization. It also contrasts weight averaging with functional ensembling, identifying conditions under which their predictions are similar and limitations when locality fails.

The variance depends on input distributions and shifts in input similarities measured by the kernel, not on label distributions.
Assumption 3 requires equal self-similarity for source inputs and sufficiently small pairwise similarities, yielding a diagonally dominant kernel matrix.Equal self-similarity can arise from normalized inputs, while low pairwise similarity corresponds to weak interaction among training samples.
A relaxed block-diagonal analysis uses one sample per block to approximate the variance for datasets whose neural tangent kernel is not strictly diagonal.
The paper derives the variance under these assumptions and relates its input-distribution term to an MMD in the RKHS of K^2.
Weight averaging approximates prediction ensembling when models share initialization and use mild hyperparameter ranges, with slight empirical gains on OfficeHome and PACS.
DiWA degrades when classifier initializations vary or hyperparameters are sampled from extreme ranges because weights become less averageable, whereas ensembling retains additional diversity.
For M = 9, accuracy gain increases by 0.297 per unit of prediction diversity and by 0.179 per unit of feature diversity; the slope increases with M.

E.1.3 Diversity comparison across a wide range of methods

Diversity increases when models come from independent or methodologically different training procedures, but diversity must remain compatible with weight averageability. The resulting trade-off explains why some sources of diversity help DiWA while others undermine it.

Weights from different runs are more diverse than weights sampled along a single run, even when the latter use extreme hyperparameters.
Different data, classifier initializations, or extreme hyperparameter ranges increase diversity across runs, but the first two practical constraints can break locality and averageability.
Different training objectives increase feature diversity, and combining invariant objectives can both learn invariant features and extract different features.
More distant learning rates increase functional diversity, but sufficiently different rates make weight averaging diverge from functional ensembling.
The PACS analysis recovers the same core diversity-based conclusions reported for OfficeHome.

F Number of training runs

Increasing the number of independent training runs improves DiWA’s performance and reduces variability, while five runs are already competitive with the prior state of the art.

Larger numbers of training runs improve OOD performance and reduce standard deviation when uniformly averaging DiWA weights.
Five runs with uniform weight selection are enough for DiWA to be competitive with SWAD.

G.1 Description of the DomainBed benchmark

DomainBed evaluates cross-domain generalization across several multi-domain classification datasets using repeated train, validation, and test splits. The paper reports strong WA results under diversity shift but a failure under correlation shift on ColoredMNIST.

DomainBed contains multiple computer-vision classification datasets whose domains are alternately held out as test domains while the others provide training data.
Each domain is split into 80% training/testing and 20% validation data, with results averaged over three random seeds.
The experiments use 20 random-search trials per dataset, domain, and seed under the stated hyperparameter distributions.
Model selection chooses the epoch with maximum training-domain validation accuracy, while DiWA† cannot use a shared validation dataset across its different data splits.
The benchmark datasets include PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet, where diversity shift dominates according to prior analysis.
DiWA achieves state-of-the-art results on these diversity-shift benchmarks but performs poorly relative to ERM on ColoredMNIST, where correlation shift dominates.

I Last layer retraining when some target data is available

When target samples become available, DiWA can be used as a feature extractor before last-layer retraining. Preliminary results indicate improved performance over ERM, while full analysis remains future work.

The traditional OOD setting assumes no access to labelled or unlabelled target samples, whereas this scenario permits target data for last-layer retraining.
After source-domain training, experiments retrain only the last network layer using 20% or 80% of the target Art domain.
DiWA-derived frozen feature extractors improve performance after last-layer retraining compared with ERM-derived features.The comparison uses OfficeHome’s Art domain after training on Clipart, Product, and Photo.
The authors present this result as evidence that DiWA features may be better adapted to last-layer retraining and generalization, while calling for fuller analysis in this setting.