Diverse Weight Averaging for Out-of-Distribution Generalization

TL;DR

This work tackles out-of-distribution generalization in vision by scrutinizing weight averaging (WA) and proposing Diverse Weight Averaging (DiWA). It introduces a bias-variance-covariance-locality decomposition of WA's OOD error, linking correlation shift to bias and diversity shift to variance, and identifies covariance as a bottleneck mitigated by model diversity. DiWA mitigates covariance by training multiple independent runs with shared initialization and mild hyperparameter variation, achieving state-of-the-art results on DomainBed without extra test-time cost. The findings demonstrate that increasing diversity across WA members yields tangible OOD gains, while also acknowledging limitations against correlation shift and the need to preserve linear connectability. Overall, DiWA offers a practical, scalable path to robust OOD generalization through diverse but weight-averageable models.

Abstract

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

Figures (17)

• Figure 1: Each dot displays the accuracy ($\uparrow$) of weight averaging (WA) vs. accuracy ($\uparrow$) of prediction averaging (ENS) for $M$ models.
• Figure 2: Each dot displays the accuracy ($\uparrow$) gain of WA over its members vs. the prediction diversity aksela2003comparison ($\uparrow$) for $M$ models.
• Figure 3: Frequencies of prediction diversities ($\uparrow$) aksela2003comparison across $2$ weights obtained along a single run or from different runs.
• Figure 4: WA accuracy ($\uparrow$) as $M$ increases, when the $M$ weights are obtained along a single run or from different runs.
• Figure 5: Each dot displays the accuracy ($\uparrow$) gain of WA over its members vs. prediction diversity ($\uparrow$) for $2\leq M < 10$ models.

Theorems & Definitions (22)

• Lemma 1: WA and ENS. Proof in \ref{['app:wa_loss']}. Adapted from izmailov2018Wortsman2022ModelSA.
• Proposition 1: Bias-variance-covariance-locality decomposition of the expected generalization error of WA in OOD. Proof in \ref{['app:proof_bvc']}.
• Proposition 2: OOD bias and correlation shift. Proof in \ref{['app:bias_correlation']}
• Proposition 3: OOD variance and diversity shift. Proof in \ref{['app:var_diversity']}
• Theorem 1: Equation 21 from cha2021wad, simplified version of their Theorem 1
• Lemma : \ref{['lemma:wa_ensembling']}
• proof
• Remark 1
• Proposition : \ref{['prop:b_var_cov']}
• proof