SWAD: Domain Generalization by Seeking Flat Minima

By Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, Sungrae Park. 2021

606 citations recorded by the source corpus.

Domain generalization (DG) methods aim to achieve generalizability to an unseen target domain by using only training data from the source domains. Although a variety of DG methods have been proposed, a recent study shows that under a fair evaluation protocol, called DomainBed, the simple empirical risk minimization (ERM) approach works comparable to or even outperforms previous methods. Unfortunately, simply solving ERM on a complex, non-convex loss function can easily lead to sub-optimal generalizability by seeking sharp minima. In this paper, we theoretically show that finding flat minima results in a smaller domain generalization gap. We also propose a simple yet effective method, named Stochastic Weight Averaging Densely (SWAD), to find flat minima. SWAD finds flatter minima and suffers less from overfitting than does the vanilla SWA by a dense and overfit-aware stochastic weight sampling strategy. SWAD shows state-of-the-art performances on five DG benchmarks, namely PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet, with consistent and large margins of +1.6% averagely on out-of-domain accuracy. We also compare SWAD with conventional generalization methods, such as data augmentation and consistency regularization methods, to verify that the remarkable performance improvements are originated from by seeking flat minima, not from better in-domain generalizability. Last but not least, SWAD is readily adaptable to existing DG methods without modification; the combination of SWAD and an existing DG method further improves DG performances. Source code is available at https://github.com/khanrc/swad.

Subfields: Artificial Intelligence, Computational Mechanics, Computer Vision and Pattern Recognition, Representation & Self-Supervised Learning, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

In Computer Vision and Pattern Recognition

In Representation & Self-Supervised Learning

In Optimization & Training

Citations in the corpus

A Comprehensive Survey of Continual Learning: Theory, Method and Application (2023, methodology): “The alternatives include using an approximation of the Hessian [89], [309] or calculating δ only along the trajectory of forward and backward parameter changes [58], [173], [182], [295], [308].”
CoSCL: Cooperation of Small Continual Learners is Stronger than a Big One (2022, neutral): “Inspired by the PAC-Bayes theory [26] and previous work in domain generalization [5,4], we first present the upper bounds of these two errors as follows (please see a complete proof in Appendix A): Proposition 1.”
Domain Generalization: A Survey (2021, methodology): “Following [43], Khosla et al. [46] targeted the cross-dataset generalization problem in classiﬁcation and detection tasks, and proposed to learn domain-speciﬁc bias vectors and domain-agnostic weight vectors based on support vector machine (SVM) classiﬁers.”
Generalizing to Unseen Domains: A Survey on Domain Generalization (2021, methodology): “Additionally, Cha et al. [166] observed the ﬂat minima is important to DG and they designed a simple stochastic weight averaging densely method to ﬁnd the ﬂat minima.”
Towards Out-Of-Distribution Generalization: A Survey (2021, background): “[29] proposed Stochastic Weight Averaging Densely (SWAD) as a means of identifying flat minima solutions, which have demonstrated strong OOD generalization performance across five domain generalization benchmarks.”
Questions for Flat-Minima Optimization of Modern Neural Networks (2022, methodology): “Cha et al. (2021) introduce a SWA variant for domain generalization tasks.”
TIES-Merging: Resolving Interference When Merging Models (2023, background): “Therefore, merging fine-tuned models can improve performance on a single target task [30, 23, 82, 9], improving out-of-domain generalization [31, 29, 7, 4, 60, 59], creating a multitask models from different tasks [31, 29, 38], for federated learning [46, 41], compression [39], multimodal merging…”
BioMistral: A Collection of Open-Source Pretrained Large Language Models for Medical Domains (2024, background): “Recent studies suggest merging pre-trained models to enhance performance and out-of-domain generalization (Cha et al., 2021; Arpit et al., 2022; Wortsman et al., 2022; Jin et al., 2023; Ilharco et al., 2023).”