Source-linked AI summary
Ensemble of Averages: Improving Model Selection and Boosting Performance in Domain Generalization
Devansh Arpit, Huan Wang, Yingbo Zhou, Caiming Xiong
TL;DR
Domain generalization models can have unstable out-domain performance, complicating reliable selection from in-domain validation data. The paper introduces efficient moving-average modeling and ensembles of moving-average models, explains their gains through a domain-generalization bias-variance decomposition, and reports stronger benchmark performance than vanilla ERM.
Problem
Out-domain performance is unstable across training trajectories and independently trained models, making model selection from in-domain validation data unreliable.
Method
The paper maintains a simple moving average of model parameters during ERM training, uses it for stopping and inference, and ensembles moving-average models from independent runs.
Results
EoA outperforms ensembles of unaveraged models, improves validation-to-test rank correlation, and provides larger gains over ERM as pretraining dataset and model size increase.
Takeaways & Limitations
Moving-average models make out-domain performance more stable and model selection using in-domain validation data more reliable.
Takeaways & Limitations
The theoretical analysis assumes IID training samples from an in-domain distribution and evaluates loss on a different out-domain distribution.
Abstract
from arXiv · showhide
In Domain Generalization (DG) settings, models trained independently on a given set of training domains have notoriously chaotic performance on distribution shifted test domains, and stochasticity in optimization (e.g. seed) plays a big role. This makes deep learning models unreliable in real world settings. We first show that this chaotic behavior exists even along the training optimization trajectory of a single model, and propose a simple model averaging protocol that both significantly boosts domain generalization and diminishes the impact of stochasticity by improving the rank correlation between the in-domain validation accuracy and out-domain test accuracy, which is crucial for reliable early stopping. Taking advantage of our observation, we show that instead of ensembling unaveraged models (that is typical in practice), ensembling moving average models (EoA) from independent runs further boosts performance. We theoretically explain the boost in performance of ensembling and model averaging by adapting the well known Bias-Variance trade-off to the domain generalization setting. On the DomainBed benchmark, when using a pre-trained ResNet-50, this ensemble of averages achieves an average of $68.0\%$, beating vanilla ERM (w/o averaging/ensembling) by $\sim 4\%$, and when using a pre-trained RegNetY-16GF, achieves an average of $76.6\%$, beating vanilla ERM by $6\%$. Our code is available at https://github.com/salesforce/ensemble-of-averages.
1 Introduction
The paper shows that domain-generalization performance is unstable both across independently trained models and along a single model’s optimization trajectory. It proposes simple moving-average models and ensembles of those averages to improve out-domain performance and make validation-based model selection more reliable.
- Domain generalization remains difficult because deep models perform poorly when test distributions differ from training distributions.
- Out-domain performance can vary greatly during training even when in-domain performance changes little, undermining validation-based model selection.
- Simple moving averaging boosts domain-generalization performance and reduces out-domain instability, improving the reliability of model selection.
- The protocol is hyperparameter-free in practice by averaging every iteration from near initialization, which works across multiple datasets and architectures with pretrained initialization.
- Using the moving-average model for early stopping and inference avoids the frequent validation checks required by SWAD, reducing computational requirements.
- An ensemble of moving-average models outperforms an ensemble of unaveraged models and has better validation-to-test rank correlation.
- A bias-variance decomposition explains the gains by attributing individual-model out-domain loss to bias and variance, while averaged or ensemble models mainly retain bias.
- EoA gains over ERM range from 4%−6% and increase with pretraining dataset and model size in the reported benchmark experiments.
2 Model Averaging
The paper maintains a simple moving average of model parameters during training and uses the averaged model for validation-based model selection and inference. Across multiple datasets, model averaging improves the alignment between validation and out-domain performance, making selection more reliable.
- Terminology: The online model is the gradient-optimized network fθ, whose parameters change throughout training.
- Model Averaging Protocol: The moving average model maintains an average of online-model parameters during training, using tail averaging beginning at iteration t0.The paper denotes its parameters by ˆθ and calls this a simple moving average (SMA) model.
- Model Averaging Protocol: For t > t0, the SMA update combines the previous average with the current online-model state using weights determined by the number of averaged iterations.The averaging starts from the online state at t0 and updates through the end of training.
- Model Selection: Validation performance and final test predictions are computed with ˆθt rather than θt, selecting the averaged model with the best validation performance.This protocol is intended to make model selection more reliable for domain generalization.
- Ablation Analysis: Poor rank correlation across independently trained models makes it difficult to identify the best out-domain model using only in-domain validation performance.The paper reports this instability as an implication of the independent-model ablation.
- Ablation Analysis: In most reported cases, model averaging significantly improves Spearman rank correlation between in-domain validation accuracy and out-domain test accuracy.The comparison covers PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet, for both individual models and ensembles of averages.
3 Ensemble of Averages (EoA)
EoA combines independently trained moving-average models to improve out-domain performance stability and accuracy. The paper connects these gains to model averaging’s approximation of ensembling and a domain-generalization bias-variance decomposition.
- EoA Protocol: EoA ensembles independently trained moving-average models selected using in-domain validation performance.The ensemble combines predictions from models trained with different hyperparameters and seeds, using each run’s selected moving-average state.
- 3.1 Analysis: EoA improves out-domain performance stability relative to ensembles of online models during training.The qualitative comparison uses six independently trained models on TerraIncognita with domain L38 held out and ResNet-50 initialization.
- Why does Ensembling and Model Averaging Improve Performance?: The domain-generalization bias-variance decomposition attributes lower expected out-domain error to the strictly positive variance reduction from ensembling.The decomposition separates expected out-domain cross-entropy into an ensemble bias term and a variance term measuring disagreement with the ensemble prediction.
- Why does Ensembling and Model Averaging Improve Performance?: EoA consistently outperforms traditional unaveraged ensembles across ensemble sizes, while larger ensembles typically improve out-domain performance.Figure 3 compares mean and standard error for models with and without moving-average parameters.
- Why does Ensembling and Model Averaging Improve Performance?: Model averaging approximates ensembling because the second-order Taylor term concentrates near zero relative to the moving-average model’s logits.The analysis uses model-output derivatives, and the first-order term vanishes under the moving-average definition.
4 Empirical Results
Across DomainBed datasets, the study benchmarks simple moving averaging and ensembles against ERM, using several pretrained architectures. Model averaging and especially EoA provide larger gains as pretraining data and model size increase.
- DomainBed Benchmarking: The experiments compare ERM, SMA, ensembles, and EoA across PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet.All models use the ERM objective, with moving-average and online-model protocols evaluated under DomainBed procedures.
- DomainBed Benchmarking: Across five DomainBed datasets, EoA outperforms ERM by 4% and SWAD by 1.1% with pretrained ResNet-50.Traditional ensembles and SMA achieve similar performance at 66.8% and 66.5%, respectively.
- Pretraining and Model Scaling: EoA gains over ERM increase from 4% to 5% when comparing ResNet-50 with ResNeXt-50 32x4d, while SMA gains increase from 2.5% to 3.9%.The authors relate these larger gains to the larger pretraining dataset used for ResNeXt-50 32x4d.
- Pretraining and Model Scaling: EoA gains over ERM increase from 5% for ResNeXt-50 32x4d to 6% for RegNetY-16GF, while SMA gains increase from 3.9% to 5%.This comparison changes both model size and pretraining dataset size.
- IID Evaluation: SMA also outperforms ERM without model averaging in the IID setting.
5 Related Work
The paper situates model averaging within theoretical work on bias-variance reduction and regularization, and distinguishes its protocol from SWAD and broader DG approaches. Its focus is complementary: improving out-domain stability and model selection reliability.
- Model Averaging: Tail averaging motivates the protocol because averaging can reduce variance faster than unaveraged optimization, while initialization bias decays with the averaging start iteration.The cited least-squares analysis bounds excess risk using bias and variance terms.
- Model Averaging: Model averaging has a regularization effect related to Tikhonov regularization, connecting it to under-specification in over-parameterized models.
- SWAD: SWAD averages contiguous optimization states selected using validation loss, unlike regular-interval or late-training sampling strategies.
- Domain Generalization Methods: Existing DG methods include domain alignment, regularization, data augmentation, and meta-learning, whereas this work studies model averaging across such methods.
- Model Selection: The paper emphasizes out-domain performance instability during training because it affects the reliability of validation-based model selection.
6 Conclusion
The conclusion reports that hyperparameter-free, efficient averaging improves out-domain performance and model-selection reliability, with EoA outperforming ensembles of unaveraged models. Gains increase with pretrained model and dataset scale, and the authors interpret the improvement through bias-variance trade-offs.
- Conclusion: Model averaging improves out-domain performance over unaveraged ERM, while EoA outperforms ensembles of unaveraged models.
- Conclusion: Model averaging significantly improves rank correlation between in-domain validation accuracy and out-domain test accuracy, supporting more reliable model selection.
- Conclusion: EoA gains over corresponding ERM baselines range from 4% to 6% across pretrained models with increasing pretraining dataset and model size.
- Conclusion: The paper explains the performance boost of EoA by adapting the bias-variance trade-off to domain generalization.
Checklist
The checklist records affirmative responses on contribution and limitation disclosure, theoretical-assumption reporting, reproducibility materials, training details, error bars, compute reporting, and ethical-impact discussion.
- Disclosure: The authors report that they described the work’s limitations and scope, including limitations of their methods and existing DG methods.They refer readers to Appendix E for these discussions.
- Theoretical Reporting: The paper states that theoretical assumptions were provided in Section 3.2, while complete proofs were marked not applicable.
- Reproducibility: The authors report providing code, data, and reproduction instructions, along with training details and hyperparameter-selection information.
- Experimental Reporting: The experiments report error bars across repeated runs and disclose the total compute and resource types used.
- Assets: The paper states that existing assets were cited, but their licenses were not mentioned.
- Societal Impact: The authors state that the work aims to mitigate training-dataset bias effects on learned hypotheses and may have positive societal impact.
B Training and Evaluation Protocols for DomainBed Benchmarking
The benchmarking protocol uses ERM-trained models with specified pretrained architectures and evaluates sensitivity to moving-average sampling frequency on PACS and TerraIncognita.
- The training protocol uses a reduced hyperparameter search space and trains DomainNet models for 15,000 iterations for computational reasons.
- Models use ERM with Adam and pretrained ImageNet ResNet-50 initialization; the final benchmark also includes a pretrained ResNeXt-50 32x4d.
- All experiments were performed on Google Cloud Platform using 24 NVIDIA A100 GPUs.
- Sampling model states every 1000 iterations significantly reduces TerraIncognita performance, while other sampling frequencies have little broad impact.
C Additional Ablation Analysis of Our Model Averaging Protocol (Eq. 1)
Additional ablations examine dataset coverage, the start iteration for averaging, sampling frequency, and rank correlation as factors affecting model averaging and model selection.
- Dataset Details: The experiments use PACS, TerraIncognita, VLCS, OfficeHome, and DomainNet from the DomainBed benchmark, with four domains except DomainNet’s six.
- C.2 Start Iteration: Starting model averaging later reduces average test performance across PACS and TerraIncognita, with a start near initialization performing best.The experiments vary the start iteration across 0, 100, 500, 2000, 4000, 6000, and 8000.
- C.3 Averaging Frequency: Averaging can sample iterates at larger intervals rather than every iteration, and its effect on out-domain performance is evaluated on PACS and TerraIncognita.The tested sampling intervals are 100, 101, 102, and 103 iterations.
- C.4 Rank Correlation: Spearman correlation is used because model selection depends on ranking validation performance during early stopping rather than linear association.
C.5 Instability Reduction: Qualitative Analysis
The qualitative analysis evaluates whether in-domain validation performance supports reliable model selection across training durations and independent runs.
- C.5 Instability Reduction: Qualitative Analysis: If validation and out-domain performance correlate well, training longer should remain manageable because early stopping can address later overfitting.
- C.5 Instability Reduction: Qualitative Analysis: Moving-average models show more stable out-domain performance across training durations than online models on TerraIncognita.The experiment spans 1,000 to 10,000 iterations in 1,000-iteration increments.
- C.5 Instability Reduction: Qualitative Analysis: The cross-run analysis gathers validation and test accuracies from independently seeded or hyperparameter-varied runs at each training iteration.
D Why Does Ensembling Improve Performance?
This section analyzes model averaging and ensembling through training-trajectory behavior, stability of model selection, rank correlation, and a Taylor-expansion-based account of ensemble performance.
- The analysis rewrites Taylor expansion for the model-averaging protocol and omits the first-order term because it is zero.
- Model averaging reduces instability in out-domain performance across training durations, making model selection more reliable than without averaging.
- In most DomainBed cases, model averaging significantly improves Spearman rank correlation between within-run validation accuracy and out-domain test accuracy for individual models and ensembles.
- The appendix histogram experiment evaluates logits and second-order Taylor terms for 1000 test samples across all 10 TerraIncognita classes.Only iterates at t = 300*i for i in {1, 2, ..., 16} are used because of computational constraints.
- Across PACS and TerraIncognita, cross-run rank correlations are poor for both online and moving-average models, limiting reliable selection of a single model from independent runs.
E Discussions and Limitations
The paper identifies boundaries of ensemble-of-averages benefits, including residual bias, limited functional diversity, domain-scale constraints, and dependence on baseline quality. It also discusses computational cost, experimental choices, and information-use limitations.
- Domain Generalization Limitations: Ensembling removes the variance component of expected loss but cannot reduce bias, so further strategies are needed to improve OOD performance.The analysis states that ensembles are dominated by bias after variance is removed.
- Domain Generalization Limitations: The proposal does not use environment IDs, leaving open what other strategies could exploit domain information to improve generalization.The discussion notes that domain alignment is neither necessary nor sufficient according to cited empirical evidence.
- Functional Diversity: Model averaging improves within-run stability, but performance gaps between independent runs remain and its functional diversity is weaker than that of independently trained models.The paper links this limited functional diversity to the smaller-than-expected advantage over traditional ensembles.
- Scalability: Using samples from every training domain in each mini-batch can become prohibitive when the number of domains is very large.Preliminary experiments found that randomly selecting one domain per iteration achieved similar performance.
- Computational Complexity: SMA adds practically negligible overhead, whereas EoA scales linearly with ensemble size when models are trained sequentially but can be parallelized.With sufficient resources, parallel training keeps EoA complexity comparable to vanilla supervised training.
- Experimental Scope: The reported gains depend partly on baseline quality and were obtained with a smaller hyperparameter search space than a prior study.The paper states that stronger pretrained models improved the ERM baseline and produced further gains for model averaging and EoA.