SWAD: Domain Generalization by Seeking Flat Minima
Authors: Junbum Cha
Sanghyuk Chun
Kyungjae Lee
Han-Cheol Cho
Seunghyun Park
Yunsung Lee
Sungrae Park
Venue: NeurIPS 2021 Poster (2021)
Keywords: domain generalization
robustness
data distribution shift
domain shift
flatness minimization
sharpness minimization
generalization

Abstract: 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.

SWAD: Domain Generalization by Seeking Flat Minima
Junbum Chalt
Sanghyuk Chun Seunghyun Park
Kyungjae Lee
Han-Cheol Cho
Yunsung Lee
Sungrae Park
Kakao Brain
NAVER AI Lab
Chung-Ang University 6 Upstage AI Research
NAVER Clova
Korea University
Abstract
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 heen 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. Un- fortunately, simply solving ERM complex, non-convex loss function easily lead to sub-optimal generalizability by seeking sharp minima. In this paper, we theoretically show that finding flat minima results in smaller domain generaliza- tion gap. We also propose 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 dense and overfit-aware stochastic weight sampling strategy. SWAD shows state-of-the-art performances on five DG benchmarks, namely PACS, VLCS, OfficeHome, Tnroanncogni and DomainNet, with consistent and large margins of +1.6% averagely out- of-domain accuracy. We also compare SWAD with conventional generalization methods, such as data augmentation and consistency regularization methods, ver- ify that the remarkable performance improvements are originated from by seeking flat minima, from better in-domain generalizability. Last but 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:///tt t/tbbbb.tttttttt/// thub. com/ khanrc/swad
Introduction
Independent and identically distributed (i.id.) condition is the underlying assumption of machine learning experiments. However, this assumption may not hold in real-world scenarios, í.e., the training and the test data distribution may differ significantly by distribution shifts. For example, a self-driving car should adapt to adverse weather day-to day-to-night to--ttt shifts 2 Even in simple image recognition scenario, systems rely on wrong cues for their prediction, e-g., geographic distribution (3], demographic statistics 4 texture (5), backgrounds (6). Consequently, practical system should require generalizability to distribution shift, which is yet often failed by traditional approaches.
Domain generalization (DG) aims to address domain shift simulated by training and evaluating on different domains. DG tasks assume that both task labels and domain labels are accessible. For example, PACS dataset (7) has seven task labels (e.g., "dog", "horse") and four domain labels (e.g., "photo", "sketch"). Previous approaches explicitly reduced domain gaps in the latent space g
"Equal contribution Part of work done while at NAVER Clova Correspondence Junbum Cha <junbum.che cha @ kakkabbrain.com comm, Sungrae Park <sungrae, park @ @uustage.ai>
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Table 1: Comparisons with SOTA. The proposed SWAD outperforms other stale-of-the-ari DG methods on five different DG benchmarks with significant gaps (+ (+1.6pp 66pp in the average).
PACS
VLCS
OfficeHome
TerraInc
Domainliet
Avg.
ERM 29
85.5
77.5
66.5
46.1
40.9
63.3
Best SOTA competitor
86.6
30
78.8 [3
68.7
B1
48.6
32
43.6
[ 33
65.3
SWAD (proposed)
88.1
79.1
70.6
50.0
46.5
66.9
Previous SOTA 31 SWAD
88.3
78.9
71.3
51.0
46.8
67.3
12.. obtained well transffranfeeeeee model parameters by the meta-learning framework [3 data
augmentation 17-19] or capturing causal relation 20 21 Despite numerous previous attempts for a decade, Gulrajani and Lopez-Paz [22 showed that a simple empirical risk minimization (ERM) approach works comparably or even outperforms the previous attempts on diverse DG benchmarks
under a fair evaluation protocol, called "DomainBed".
Unfortunately, although ERM showed surprising empirical success on DomainBed, simply mini- mizing the empirical loss a complex and non-convex loss landscape is typically not sufficient to arrive at good generalization [23]26] In particular, the connection between the generalization gap and the flatness of loss landscapes has been actively discussed under the i.i.d. condition [23 Izmailov et al. (25] argued that seeking flat minima will lead to robustness against the loss landscape shift between training and test datasets, while simple ERM converges to the boundary of a wide flat minimum and achieves insufficient generalization. In the DG scenario, because training and test loss landscapes differ more drastically due to the domain shift, we conjecture that the generalization gap between flat and sharp minima is larger than expected in the i.i.d. scenario.
To show that flatter minima generalize better unseen domains, we formulate robust risk minimiza- tion (RRM) problem defined by the worst-case empirical risks within neighborhoods in parameter We theoretically show that the generalization gap of DG, i.e., the error on the target domain, is upper bounded by RRM, i.e., a flat optimal solution. Based on our theoretical observation, we modify stochastic weight averaging (SWA) (25], one of the popular existing flatness- aatarss-aware aware solvers, by introducing dense and overfit-aware stochastic weight sampling strategy. First, we suggest to sample weights densely, i.e,, for every iteration. Also, we search the start and end iterations for averaging by considering the validation loss to avoid overfitting. We empirically show that the proposed Stochastic Weight Averaging Densely (SWAD) finds flatter minima than the vanilla SWA does, resulting in better generalization to unseen domains.
space 26 34
Contribution. Our main contribution is introducing flatness into DG, and showing remarkably outperforming performances against existing DG methods. As shown in Table ! our SWAD improves the average DG performances by 3.6pp against the ERM bascline and 1.6pp against the existing best methods. Furthermore, by combining SWAD and previous SOTA [3I]. we even achieve 0.4pp improvements against the vanilla SWAD results. We also empirically show that while popular in- domain generalization methods without considering flatness, e.g., Mixup [35 or CutMix (36), are not effective to oft-oindomain generalization (Table flatness-aware methods, e.g., SWA 25 or SAM (26), are only effective methods to both in-domain and out-of-domain generalization.
2 A Theoretical Relationship between Flatness and Domain Generalization
Let D :0:020) ( be a set of training domains, where D; is distribution over input space X, and / is the total number of domains. From each domain, we observe n training data points which consist of input x and target label y. (af, vf) ~ D,. We also define a set of target domain T := (7.), similarly, where the number of target domains T is usually set to one. For the sake of simplicity, unlike Ben-David et al. (37), we assume that there exists a global laheling function h((z) that generates target label for multiple domains, Le., y h(x;)) for all and j. Domain generalization (DG) aims to find a model parameter E θ which generalizes well over both multiple training domains 'D and unseen target domain T. More specifically, let us consider a bounded instance loss function é : y x y + [0. c). such that E(y1, V2) = 0 holds if and only if y1 labels. For simplicity, we set to one in our proofs, but we note that E(:-)) can be generalized for any bounded loss function. Then, we can define population loss over multiple domains by
where y is a set of
Ep t E'= ER (((f(at; 0), y'))], where f(; 0) is model parameterized by 0. Formally,

---

the goal of DG is to find model which minimizes both ED( and (T by only minimizing an empirical risk (0) in Ej= ((((x) v²)) over training domains D.
In practice, ERM, Le., arg ming ED can have multiple solutions that provide similar values of the training losses but significantly different generalizability on (D( and (0). Unfortunately, the typical optimization methods, such as SGD and Adam [38). often lead sub-optimal generalizability as finding sharp and narrow minima even under the i.i.d. assumption (23-28] In the DG scenario, the generalization gap between empirical loss and target domain loss becomes even worse due to domain shift. Here, we provide a theoretical interpretation of the relationship between finding a flat minimum and minimizing the domain generalization gap. inspired by previous studies (23
We consider robust empirical loss function defined by the worst-case loss within neighborhoods in the parameter space as ED (0)= == max| ( + A), where || denotes the L2 norm and 7 is a radius which defines neighborhoods of 9. Intuitively, if is sufficiently larger than the "radius" of a sharp optimum 02 of (0). 0, is no longer an optimum of 8 (0) as well as its neighborhoods within the y-ball. On the other hand, if an optimum 8 has larger "radius" than y. there exists a local optimum within -- - See Figure |I Hence, solving the robust risk minimization (RRM), i.e., arg ming (#), will find near solution of flat optimum showing better generalizability (26] 34). However, as domain shift worsen the generalization gap by breaking the i.i.d. assumption, it is not trivial that RRM will find an optimum with better DG performance. To answer the question, we first show the generalization bound between E and ET as follows:
Empirical risk
Rob ust rak
Mirimum of the empirical risk
Mirimum rebust
Flat minimum
minimum
Figure 1: Robust risk minimization
(RRM) and flat minima. With proper RRM will find flat minima.
Theorem 1. Consider a set of N covers {0x) such that the parameter space O C U O2 where diam( ()) supo,oree |10 1221 N = (diam(()/7)) and d is dimension of O. Let Ug be a VC dimension of each Øx Then, for any E θ the following bound holds with probability at least 1 - 8,
1
Uk In (m//wg) + In(N/8)
(T) < E3) 2T
Σ Div(D,, i=1
max RE[I,N]
(I)
m
where m = nl is the number of the training samples and Div T) = 2sup PP (T()) is a divergence between two distributions.
Proof can be done similarly as 37 and 33 In Theorem the test loss (O) is bounded by three
terms: (1) the robust empirical loss (0). (2) the discrepancy between training distribution and test distribution, i.e. the quantity of domain shift, and (3) a confidence bound related to the radius y and the number of the training samples m. Our theorem is similar to Ben-David et al. [37], while our theorem does not have the term related to the difference in labeling functions across the domains. It is because simply assume there is no difference between labeling functions for each domain for simplicity. If one assumes a different labeling function, the dissimilarity term can be derived casily because it is independent and compatible with our main proof. More details of Theorem 1 including proof and discussions on the confidence bound, are in Appendix C. and C.2.
From Theorem I one can conjure that minimizing the robust empirical loss is directly related to the generalization performances on the target distribution. We show that the domain generalization gap on the target domain T by the optimal solution of RRM, 87, is upper bounded as follows:
Theorem 2. Let in denote the optimal solution of the RRM, i.e., 07 := arg ming E (f), and let v be VC dimension of the parameter space Q Then, the gap between the optimal test loss, ming: ET (0"), and the test loss of 0", E(()), has the following bound with probability at least 8.
E7 min Er (8') < E7(07) - min Eo + I Σ Div(D.. T)
2k In (m/ /v//)) + In (2N/5)
v eln (m/v) + In (2/8)
+ max kE[L.N]
m
m
3

---

minimum
trrsshod
stochastic weights
Vol dation lass
Average sparsely
Average densely
dationtion
The number iteratioons
The number of iterations
(a) SWA
(h) SWAD (proposed)
Figure 2: Comparison between SWA and SWAD. (a) SWA collects stochastic weights for every K epochs from pre-defined Ku epochs the final epoch. (b) Our SWAD collects stochastic weights densely, i.e., for every iteration, obtain sufficiently many weights. SWAD collects the weights from start iteration t, to the end iteration f,, where lg and le are obtained by monitoring the validation loss (overfit-aware scheduling).
Proof is in Appendix C.3. It implies that if we find the optimal solution of the RRM (i.e., (7), then the generalization gap in the test domain (i.e., {6)) ming Eạ (0")) is upper bounded by the gap between the RRM and ERM (i.e.. (07) ming & (0")). Other terms in Theorem 2 are the discrepancy between the train domains 'D and the target domain T, and the confidence bounds caused by sample means. We remark that if we choose a proper y the optimal solution of the RRM will find point near a flat optimum of ERM as shown in Figure I] Hence, Theorem and the intuition from Figure Sourr[impe iimply that seeking a flat minimum
ERM lead to better domain generalization gap.
3 SWAD: Domain Generalization by Seeking Flat Minima
We have shown that flat minima will bring better domain generalization. In this section, we propose Stochastic Weight Averaging Densely (SWAD) algorithm, and provide empirical quantitative and qualitative analyses on SWAD and flatness understand why SWAD works better than ERM.
3.1 A baseline method: stochastic weight averaging
Since the importance of flatness in loss landscapes has emerged [23/28] several methods have been proposed to find flat minima [25] 26 (39] We select stochastic weight averaging (SWA) [25] as baseline, which finds flat minima by a weight ensemble approach. More specifically, SWA updates pretrained model (namely, a model trained with sufficiently enough training epochs, Kn) with cyclical [40] or high constant learning rate scheduling. SWA gathers model parameters for every K epochs during the update and averages them for the model ensemble. SWA finds an ensembled solution of different local optima found by a sufficiently large learning rate to escape a local minimum. Izmailov al. 25 empirically showed that SWA finds flatter minima than ERM. We also considered sharpness- aharpness-aware aware minimization (SAM) [26). which is another popular flatness-aware solver, but SWA finds flatter minima than SAM (See Figure 3). We illustrate overview of SWA in Figure 2a
3.2
Dense and overfit-aware stochastic weight sampling strategy
Despite its advantages, directly applying SWA to DG task has two problems. First, SWA averages few weights (usually less than ten) by sampling weights for every K epochs, results in an inaccurate approximation of flat minima on a high-dimensional parameter space (e.g. 23M for ResNet-50 (41). Furthermore, common DG benchmark protocol uses relatively small training epochs (e.g., Gulrajani and Lopez-Paz (22) trained with less than two epochs for DomainNet benchmark), resulting in insufficient stochastic weights for SWA. From this motivation, propose "dense" sampling strategy for gathering sufficiently enough stochastic weights.
In addition, widely used DG datasets, such as PACS (> 10K images, 7 classes) and VLCS (~ 11K images, 5 classes), are relatively smaller than large-scale datasets, such as ImageNet 42 (x= 1.2M images, 1K classes). In this case, we observe that simple ERM approach is rapidly reached to a local optimum only within few epochs, and easily suffers from the overfitting issue, i.e., the validation loss is increased after a few training epochs. It implies that directly applying the vanilla SWA will suffer from the overfitting issue by averaging sub-optimal solutions (i.e., overfitted parameters). Hence, we need an "overfit-aware" sampling scheduling to omit the sub-optimal solutions for SWA.

---

[cyclcc
(constard)
SWAD
0.20
0.00
30
40
20
50
(a) Average train flatness
(b) Average test flatness
1.00
Train
22
6.10
0.0
0.00
250
0.50
0.00
Test
0.4
346
2.2
0.20
332
Art painting (test)
Cartoon (test)
Photo (test)
Sketch (test)
(c) Flatness for each target domain
Figure 3: Local flatness comparisons. We plot the local flatness via loss gap. ie.. (0) =101==111~1EE(1) E())], of ERM, SAM, SWA, and SWAD by varying radius on different domains of PACS dataset. For each figure, Y-axis indicates the flatness (0) and X-axis indicates the radius y. We measure the train flatness F3 (0) on seen domains and the test flatness F (0) on unseen domain. Each point computed by Monte-Can Carlo approximation with 100 random samples. This comparisons show SWAD finds flatter minima than only ERM but also SAM and SWA.
The main idea of Stochastic Weight Averaging Densely (SWAD) is dense and overfit-aware stochastic weight gathering strategy. First, instead of collecting weights for every K epochs, SWAD collects weights for every iteration. This dense sampling strategy easily collects sufficiently many weights than the sparse one. We also employ overfi--aware sampling scheduling by considering traces of the validation loss. Instead of sampling weights from Kg pretraining epochs the final epoch, we search the start iteration (when the validation loss achieves a local optimum for the first time) and the end iteration (when the validation loss no longer decreased, but keep increasing). More specifically, we introduce three parameters: an optimum patient parameter N., an overfitting patient parameter Ne, and the tolerance rate for searching the start iteration Lg and the end iteration 4, First, we
search t, which satisfies (i.., val - c(,++) = g(t,) where E(dl denotes the validation loss at iteration Simply. t, the first iteration where the loss value is no longer decreased during N, iterations. Then, we find f. satisfying jo................. c(t++i) > rttt.) In other words, te is the first iteration where the validation loss values exceed the tolerance during Ne iterations.
We illustrate the overview of SWAD and the comparison of SWAD to SWA in Figure 2 Detailed pseudo code is provided in Appendix B.4. We compare SWAD with other possible SWA strategies in $4.3 and show that our design choice works better for DG tasks.
3.3 Empirical analysis of SWAD and flatness
Here, we analyze solutions found by SWAD in terms of flatness. We first verify that the SWAD solution is flatter than those of ERM, SWA, and SAM. Our loss surface visualization shows that the SWAD solution located on the center of the flat region, while ERM finds boundary solution. Finally, we show that the sharp boundary solutions by ERM are not generalized well, resulting in sensitivity to the model selection. All following empirical analyses are conducted on PACS dataset, validating by all four domains (art painting, cartoon, photo, and sketch).
Local flatness anaylsis.
begin with, we quantify the local flatness model parameter by assuming that flat minima will have smaller changes of loss value within its neighborhoods than sharp minima. For the given model parameter 0, we compute the expected loss value changes between and parameters on the sphere surrounding with radius 7, i.e., (0) = ||
(((())))))) [E('') E(O)]. In
5

---

2.00
LOG0
Train
CRn
LOM
Test
1369
1288
cim
600
Art painting (test)
Cartoon (test)
Photo (test)
Sketch (test)
Figure 4: Loss surfaces model parameters in PACS dataset for each target domain. The three triangles indicate model weights chosen the end training phase with equal intervals. Each plane defined the three weights and losses upon the plane are visualized with contours. The center cross mark averaged point of three weights. The first and second rows show the averaged training loss and the test loss surfaces, respectively.
Art painting
Cartoon
Photo
Sketch
100
100
1:0
130
M
26
34
24
2003
2000
3050
4000
5000
20
1003
2020
3050
4000
1000 5000
2000
3050
4000
5000
1000
2000
3000
4000
5002
Art painting (lest)
(b) Cartoon (test)
(c) Photo (lest)
(d) Sketch (test)
Figure 5: Validation accuracies for in-domains. The X- and Y-axis indicate the training iterations and accuracy, respectively, about the validation domains (legend) and the test domain (caption). The vertical dot lines represent start and iterations, ts and le, identified by overfit-awe aware sampling strategy SWAD.
practice, (0) is approximated by Monte-Carlo sampling with 100 samples. Note that the proposed local flatness F, (0) is computationally efficient than measuring curvature using the Hessian- bassed quantities. Also, F. (0) has an unbiased finite sample estimator, while the worst-case loss value, i.e., || (E(6') E()) has no unbiased finite sample estimator.
In Figure 3 we compare (0) of ERM, SAM, SWA with cyclic learning rate, SWA with constant learning rate, and SWAD by varying radius y,. SAM and SWA find solutions with lower local flatness than ERM average. SWAD finds the most flat minimum in every experiment.
Loss surface visualization. We visualize the loss landscapes by choosing three model weights on the optimization trajectory (01, 02. 0g.7 and computing the loss values by linear combinations of 9,, 02 0 as (25] More details are in Appendix B.5. In Figure 4 we observe that for all cases, FRM solutions are located at the boundary of flat minimum of training loss, resulting in poor generalizability in test domains, that is aligned with our theoretical analysis and empirical flatness analysis. Since ERM solutions are located on the boundary of a flat loss surface, we observe that ERM solutions are very sensitive to model selection. In Figure 5 we illustrate the validation accuracies for each train-test domain combination of PACS by ERM, over training iterations (one epoch is equivalent to 83 iterations). We first observe that ERM rapidly reaches the best accuracy within only a few training epochs, namely less than 6 epochs. Furthermore, the ERM validation accuracies fluctuate lot, and the final performance is very sensitive to the model selection criterion.
On the other hand, we observe that SWA solutions are located on the center of the training loss surfaces as well as of the test loss surfaces (Figure 4.. Also, our overfit-aware stochastic weight gathering strategy (denoted as the vertical dot lines in Figure prevents the ensembled weight from overfitting and makes SWAD model selection-free.
We choose weights iteration 2500, 3500, 4500 during the training.
Each point is defined by and computed by 02 O, and (0--81) B1) (82 88
10g (02 01)

---

Table Comparison with domain generalization methods and SWAD. Out-of-domain accuracies on five domain generalization benchmarks are shown. We highlight the best results and the second best results. Note that ERM (reproduced), Mixstyle are reproduced numbers, and other numbers are from the original literature and Guirajani and Lopez-Pat Pazez-Paz [2
(denoted with t). Our experiments repeated three times.
Algorithm
PACS
VLCS
OfficeHome
TerraInc
DomainNet
Avg.
MASF
14
82.7
DMG 3
83.4
43.6
MetaReg
15
83.6
43.6
ER
85.3
PAdaIN
47
85.4
EISNet
48
85.8
DSON
30
86.6
ERM
85.5
77.5
66.5
46.1
40.9
63.3
ERM (reproduced)
84.2
77.3
67.6
47.8
44,0
64.2
IRM
83.5
78.6
47.6
33.9
GROu
49
84.4
76.7
66.0
43.2
33.3
60.7
M-Mixup 50
84.6
77.4
47.9
39.2
MIDG
84.9
77.2
66.8
47.8
41.2
63.6
CORAL 31
86.2
78.8
68.7
47.7
41.5
64.5
MMD 53
84.7
77.5
66.4
42.2
23.4
58.8
DANN
83.7
78.6
65.9
46.7
38.3
62.6
CDANN
10
82.6
77.5
65.7
45.8
38.3
62.0
MTL
64
84,6
77.2
66.4
45.6
40.6
62.9
SagNet
86.3
77.8
68.1
48.6
40.3
64.2
ARM 16
85.1
77.6
64.8
45.5
35.5
61.7
VREx
21
84.9
46.4
33.6
61.9
RSC
55
85.2
77.1
65.5
46.6
38.9
62.7
Mixstyle
85.2
77.9
60.4
44.0
34.0
60.3
SWAD (ours)
88.1
79.1
70.6
50.0
46.5
66.9
(20.1)
(=0.1)
(2:0.2
(20.3)
(=0011)
4 Experiments
4.1 Evaluation protocols
Dataset and optimization protocol. Following Gulrajani and Lopez-Paz 12 we exhaustively evaluate our method and comparison methods on various benchmarks: PACS (9.91 images, classes, and 4 domains), VLCS 43] (10,729 images, 5 classes, and domains), OfficeHome [44] (15,588 images, 65 classes, and domains), TerraIncognita (45 (24,78 788 images, 10 classes, and 4 domains), and DomainNet 46) (586,575 images, 345 classes, and domains).
For fair comparison, we follow training and evaluation protocol by Guirajani and Lopez-Pazz 22 including the dataset splits, hyperparameter (HP) search and model selection (while SWAD does need it) on the validation set, and optimizer HP,, except the HP search space and the number of iterations for DomainNet. We use reduced HP search space to reduce the computational costs. also tripled the number of iterations for DomainNet from 5, 5,000 15,000 because we observe that 5,000 is not sufficient to convergence. We re-evaluate ERM with 15,000 iterations, and observe 3.1pp average performance improvement (40.9% + 44.0%) in DomainNet. For training, we choose domain as the target domain and use the remaining domains the training domain where 20% samples are used for validation and model selection. ImageNet [2] trained ResNet-50 41 is employed the initial weight, and optimized by Adam (38] optimizer with learning rate of 5e-5. We construct mini batch containing all domains where each domain has 32 images. We set SWAD HPs N, 3, N₂ 6, and 1.2 for VLCS and 1.3 for the others by HP search on the validation sets. Additional implementation details, such as other HPs, are given in Appendix B.
Evaluation metrics. We report out-of-domain accuracies for each domain and their average, i.e., model trained and validated on training domains and evaluated on the unseen target domain. Each out-of-domain performance is an average of three different runs with different train-validate validation splits.

---

4.2
Main results
Comparison with domain generalization methods. We report the full out-of-domain performances on five DG benchmarks in Table The full tables including out- of-domain accuracies for each domain are in Appendix E. In all experiments, our SWAD achieves significant performance gain against ERM as well as the previous best results: +2.6pp PACS, +0.3pp in VLCS, +1.4pp in TerraIncognita, +1.9pp in officeHome and +2.9pp in DomainNet comparing to the previous best results. We observe that SWAD provides two practical advan- tages comparing to previous methods. First, SWAD training objectives
Table 3:
Comparison between generaliza- tion methods on PACS. The scores are aver- aged over all settings using different target do- mains. (f) (4) indicate statistically signifi- cant improvement and degradation from ERM.
Out-of-domain In-domain
ERM
85.3+0.4
96,6+0.0
EMA
8-
97.0=0.1(t)
SAM
85.5=0.(-)
97.4=0.1(1)
Mixup
84.8=0.3()
97.3=0.1(1)
CutMix
83.8=0.4()
97.6=0.(1))
does not need any modification or model architecture, ie., it universally applicable to any other methods. As an example, we show that SWAD actually improves the performances of other DG meth- ods, such as CORAL (31 in Table I Moreover, we discussed before, SWAD is free to the model selection, resulting in stable performances (t.e., small standard errors) on various benchmarks. Note that only compare results with ResNet-50 backbone for a fair comparison. We describe the implementation details of each comparison method and the hyperparameter search protocol in Appendix B.
VAT
85.42006-)
96.9-0.2(1)
IT-model
83.5+0.5())
96.8=0.2(1))
SWA
85.910.1(1)
97.1
SWAD
87.110.2(1)
97.7
Comparison with conventional generalization methods. We also compare SWAD with other conventional generalization methods to show that the remarkable domain generalization gaps by SWAD is not achieved by better generalization, but by seeking flat minima. The comparison methods include flatness-aware optimization methods, such as SAM (26). ensemble methods, such as EMA (56), data augmentation methods, such as Mixup 335 and CutMix (36), and consistency regularization methods, such as VAT [57] and II-model [58]. We also split in io-ammiin datasets into training (60%), validation (20%). and test (20%) splits, while no in-domain test set used for Table 2 Every experiment is repeated three times.
The results are shown in Table 3
We observe that all conventional methods helps in-domain generalization, i.e., performing better than ERM on in-domain test set. However, their out-of-domain performances are similar to or even worse than ERM. For example, CutMix and IT-model improve in-domain performances by 1.0pp and 0.2pp but degrade out-of-domain performances by 1.5pp and 1.8pp. SAM, another method for seeking flat minima, slightly increases both in-domain and out-of-domain performances but the out-of-domain performance is not statistically significant. We will discuss performances of SAM in other benchmarks later. contrast, the vanilla SWA and our SWAD significantly improve both in-domain and oft-of-domain performances. SWAD improves the performances by SWA with statistically significantly gaps: 1.2pp on the out-of domann and 0. 0.6pp on the in-domain. Further comparison between SWA and SWAD is provided in $4.
Table 4: Combination of SWAD and other methods. The scores averaged every target domain case. The performances of ERM, CORAL, and SAM are optimized by HP searches of DomainBed. contrast, for the SWAD combination cases, CORAL and SAM use default HPs without additional HP search. We additionally compare SWAD to SWAw Note that ERM SWAD is same as "SWAD" in Table
PACS
VLCS
OfficeHome
TerraInc
DomainNet
Avg. (A)
ERM
85.5 40.2
77.5 40.4
66.5 40.3
46.1 +1.8
40.9 #0.1
63.3
ERM SWAwi coust
86.9 20.2
76.6 +0.1
69.3 +0.3
49.2 #1.2
45.9 #0.0
65.6 (+2.3)
ERM + SWAD
88.1 +0.1
79.1 +0.1
70.6 +0.2
50.0 #0.3
46.5 +0.1
66.9 (+3.6)
CORAL
86.2 +0.3
78.8 +0.6
68.7 +0.3
47.6 +100
41.5 +0.1
64.5
CORAL + SWAD
88.3 +0.1
78.9 40.1
71.3 40.1
51.0 40.1
46.8 10.0
67.3 (+2.8)
SAM
85.8 40.2
79.4 40.1
69.6 +0.1
43.3 40.7
44.3 10.0
64.5
SAM SWAD
87.1 +0.2
78.5 +0.2
69.9 #0.1
45.3 40.9
46.5 10.1
65.5 (+1.0)
8

---

Combinations with other methods. Since SWAD does not require any modification on training procedures and model architectures. SWAD is universally applicable to any other methods. Here, we combine SWAD with ERM, CORAL 31 and SAM 26 Results are shown in Table 4 Both CORAL and SAM solely show better performances than HRM with +1.2pp average out-of-domain accuracy gap. Note that SAM is not DG method but a sharpness-
optimization method to find flat minima. It supports our theoretical motivation: DG can be achieved by seeking flat minima.
By applying SWAD on the baselines, the performances are consistently improved by 3.6pp on ERM, 2.8pp on CORAL, and 1.0pp on SAM. Interestingly. CORAL + SWAD show the best performances with both incorporating different advantages of utilizing domain labels and seeking flat minima. We also observe that SAM+ SWAD shows worse performance than ERM SWAD, while SAM performs better than ERM. We conjecture that it is because the objective control by SAM restricts the model parameter diversity durinig training, reducing the diversity for SWA ensemble. However, applying SWAD on SAM still leads to better performances than the sole SAM. The results demonstrate that the application of SWAD on other baselines is simple yet effective method for DG.
4.3 Ablation study
Table 5: Ablation studies of the stochastic weights selection strategies on PACS and VLCS. In the configura- tion, "t,", "t,". "Ir", and "interval" indicate start and end iterations of sampling. learning rate schedule, and a stochastic weight sampling interval, respectively. "Opt" and "Overfit" indicate the start and iterations identified by our overfit-aware sampling strategy, and "Val" means the start and end itcrations whose averaging shows the best accuracy on the validation set. "Cyclic" and "Const" represent cyclic and constant Icarning rate schedules. All experiments repeated three times.
Configuration
Out-of-domain VLCS
In-domain
L.
Le
Ir interval
PACS
Avg.
PACS
VLCS
Avg.
S SWA cyclic
4000
5000
Cyclic
100
85.9 20.1 76.6 10.1
81.2
97.1 #0.1
85.0 20.2
91.0
SWAw const
4000
5000
Const
100
86.5 =0.3
76.7 #0.2
81.6
97.3 #0.1
85.0 =0.2
91.1
SWAD wio Dense
Opt
Overfit
Const
100
86.5 =0.4 78.0 10.7
82.2
97.6 #0.1
85.8 =0.4
91.7
SWAD wii Ope Overiit
4000
5000
Const
86.6 =0.6 76.9 +0.3
81.7
97.5 #0.1
85.2 =0.1
91.3
SWAD WADwii Overiix
Opt
5000
Const
87.1 10.3
77.6 10.1
82.4
97.7 10.1
85.8 10.3
91.8
SWAD D Di-cn-vul
Val
Val
Const
86.2 =0.2
78.6
#0.1
82.4
97.5
#0.2
85.8
=0.3
91.7
SWAD (proposed)
Opt
Overfit
Const
1
87.1 =0.2
78.9
+0.2
83.0
97.7 #0.1
86.1 =0.5 91.9
Table provides ablative studies on the starting and ending iterations for averaging, the learning rate schedule, and the sampling interval. syete (SWA in Table and SWAw constant are vanilla SWAs with fixed sampling positions. We also report SWAD by eliminating three factors: the dense sampling strategy, and searching the start iteration, searching the end iteration. The dense sampling strategy lets SWAD estimate more accurate approximation degeneration in the average out-of-domain accuracy (SWAD w/o Dense). When take average from to to the final iteration, the out-of-domail performance degrades by 0.6pp (S (SWAD Overti). Similarly, fixed scheduling without the overfit- aware scheduling only shows very marginal improvements from the vanilla SWA (SWAD Overit) We also evaluate SWADmi-n--a that uses the range achieving the best performances on the validation set, but it becomes overfitted to the validation, results in lower performances than SWAD. The results demonstrate the benefits of combining "dense" and "overfit-awan sampling strategies of SWAD.
flat minima: showing 0.8pp
4.4 Exploring the other applications: ImageNet robustness
Table 6: ImageNet robustness benchmarks. We show the ImageNet generalization performances on Iet-c, background challenge (BGC), and ImageNet-R.
Method
ImageNet (%) 1
ImageNet-C (mCE) 4
BGC (%) f
ImageNet-R (%)
ERM
76.5
57.6
8.7
36.7
SWA
76.9
56.8
10.9
37.5
SWAD (ours)
77.0
55.7
11.8
38.8

---

Since SWAD does not rely on domain labels, it can be applied to other robustness tasks not containing domain labels. Table 6 show the generalizability of SWAD on ImageNet 442 and its shifted benchmarks, namely, ImageNet-Ci [59], ImageNer-R [60], and background challenge (BGC) [6]] SWAD consistently improves robustness performances against the ERM baseline and the SWA baseline. These results support that our method is robustly and widely applicable to improve both in-domain and out-of-domain generalizability. The detailed setup is provided in Appendix B.6.
5 Discussion and Limitations
Despite many benefits from SWAD, such as the significant performance improvements, model selection-free property, working plug-and-play manner for various methods, there are some potential limitations. Here, we discuss the limitations of SWAD for further improvements.
Confidence error in Theorem I. While the confidence error in Theorem I tells the effect of y on generalization error bound, there exists limitation in that the confidence error term shows improper behavior with respect to y if y is close to zero, The behavior we expect is that the confidence error of RRM converges to the confidence error of ERM decreases to zero, however, the current theorem does not show such tendency since the confidence bound diverges to infinity when y goes to zero. However, would like note that this limitation is not drawback of RRM, but it is caused by the looseness of the union bound which is a mathematical technique used to derive the confidence error of RRM. Our RRM formulation has similarity to previous works 26 34 and we note that the counter-intuitive behavior of the confidence bound and y also appears in Foret et al. [26]-
SWAD is not perfect flatness-aware optimization method. Note that SWAD is not perfect and theoretically guaranteed solver for flat minima, but heuristic approximation with empirical benefits. However, even if a better flatness-aware optimization method is proposed, our theoretical contribution still holds: showing the relationship between flat minima and DG.
SWAD does not strongly utilize domain-specifie information. In Theorem the domain gener- alization gap is bounded by three factors: flat minima, domain discrepancy, and confidence bound. Most of the existing approaches focus on domain discrepancy, reducing the difference between the source domains and the target domain by domain invariant learning 88 [ SWAD focuses on the first factor, the flat minima. While the domain labels are used to construct a mini-batch, SWAD does not strongly utilize dopain-specifie information. implies that if one can consider both flatness and domain discrepancy, better domain generalization can be achievable. Table 4 gives us clue: the combination of CORAL (utilizing domain-specific information) and SWAD (seeking flat minima) shows the best performance among all comparison methods. As future research direction, we encourage studying a method that can achieve both flat optima and small domain discrepancy.
6 Concluding Remarks
In this paper, theoretically and empirically demonstrate that domain generalization (DG) is achievable by seeking flat minima. We propose SWAD that captures flatter minima than the vanilla SWA does. The extensive experiments on five DG benchmarks show superior performances of SWAD compared with existing DG methods. In addition, combinations of SWAD and existing DG methods even show better performances than the vanilla SWAD. We theoretically and empirically observe that seeking flat minima can achieve better generalizability to both in-domain and out-of-domain, while strong in-domain generalization methods without consideration of flatness, e.g., Mixup CutMix, cannot guarantee to achieve out-of-domain generalizability in both theory and practice. This study first brings the concept of flatness into DG tasks, and shows strong empirical performances not only in DG but also in ImageNet benchmarks. We hope that this study promotes a new research direction of seeking flat minima for domain generalization and other robustness tasks.
Acknowledgments and Disclosure of Funding
NAVER Smart Machine Learning (NSML) 662 and Kakao Brain Cloud platform have been used in experiments. This work was supported by IITP grant funded by the Korea government (MSIT) (No. 2021-0-01341, AI Graduate School Program, CAU).
10

---

References
[1] Dengxin Dai and Luc Van Gool. Dark model adaptation: Semantic image segmentation from daytime to nighttime. In 2018 21st International Conference Intelligent Transportation Systems (ITSC), pages 2819938244 IEEE, 2018.
[2] Claudio Michaelis, Benjamin Mitzkus, Robert Geirhos, Evgenia Rusak, Oliver Bringmann, Alexander S Feker, Matthias Bethge, and Wieland Brendel. Benchmarking robustness in object detection: Autonomous driving when winter is coming, arXiv preprint arXiv:1907.077
2019,
[3] Terrance de Vries, Ishan Misra, Changhan Wang, and Laurens van der Maaten. Does object recognition work for everyone? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 52-59, 2019.
[4] Kaiyu Yang. Klint Qinami, Li Fei-Fei, Jia Deng, and Olga Russakovsky. Towards fairer datasets: Filtering and balancing the distribution of the people subtree the imagenet hierarchy. In Proceedings of the 2020 Conference Fairness, Accountability, and Transparency, pages 547-558, 2020.
[5] Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. mmagenet-traineds cmns are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations, 2019.
[6] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision (ECCV), pages 456-473, 2018.
[7] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In JEEE International Conference on Computer Vision, pages 5542- 5550, 2017.
[8] Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pages 10-18. PMLR, 2013.
[9] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarin training of neural networks. Journal of machine learning research, 17(1):2096 2030, 2016.
[10] Ya Li, Mingming Gong, Xinmei Tian, Tongliang Liu, and Dacheng Tao. Domain generaliza- tion via conditional invariant representations. In AAAI Conference on Artificial Intelligence, volume 32, 2018.
[11] Hyojin Bahng, Sanghyuk Chun, Sangdoo Yun, Jaegul Choo, and Scong Joon Oh. Learning de-biased representations with biased representations. In International Conference on Machine Learning (ICMI). 2020.
[12] Shanshan Zhao, Mingming Gong, Tongliang Liu, Huan Fu, and Dacheng Tao. Domain general- ization via entropy regularization. Neural Information Processing Systems, 33, 2020.
[13] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Meta- learning for domain generalization. In AAAI Conference on Artificial Intelligence, volume 32, 2018.
[14] Qi Dou, Daniel C Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. Neural Information Processing System, 2019,
[15] Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain generalization using meta-regulaarzzttonn Neural Information Processing Systems, 31:998-1008. 2018.
[16] Marvin Zhang, Henrik Marklund, Abhishek Gupta, Sergey Levine, and Chelsea Fino. Adap- tive risk minimization: meta-learning approach for tackling group shift. arXiv preprint 20X.022077.02931, 2020.

---

[17] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In International Conference on Learning Representations, 2021.
[18] Shiv Shankar, Vihari Piratla, Soumen Chakrabarti, Siddhartha Chaudhuri, Preethi Jyothi, and Sunita Sarawagi. Generalizing across domains via cross-gradient training. In International Conference on Learning Representations, 2018.
[19] Fabio M Carlucci, Antonio D D'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain generalization by solving jigsaw puzzles. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2229-2233 2019.
[20] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprini arXiv: arXiv: 1907.02893 2019.
[21] David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Remi Le Priol, and Aaron Courville. outtrf--distibution generalization via risk extrapolation (rex). arXiv preprint arXiv:2003.00 00680 2020.
[22] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021.
[23] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[24] Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. In Neural Information Processing Systems, 2018.
[25] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. Conference on Uncertainty in Artificial Intelligence, 2018.
[26] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awar min- imization for efficiently improving generalization. In International Conference Learning Representations, 2021.
[27] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Uncertainty in Artificial Intelligence, 2017.
[28] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02/777 2019.
[29] V Vapnik. Statistical learning theory. NY: Wiley, 1998.
[30] Seonguk Sco, Yumin Suh, Dongwan Kim, Jongwoo Han, and Bohyung Han. Learning to optimize domain specific normalization for domain generalization. European Conference Computer Vision, 2020.
[31] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pages 443-450 Springer, 2016
[32] Hyeonseob Nam, HyunJae Lee, Jongchan Park, Wonjun Yoon, and Donggeun Yoo. Reducing domain gap by reducing style bias. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8690-8699, 2021.
[33] Prithvijit Chattopadhyay, Yogesh Balaji, and Judy Hoffman. Learning to balance specificity and invariance for in and out domain generalization. In European Conference on Computer Vision, pages 301-318. Springer, 2020.
[34] Matthew Norton and Johannes Royset. Diametrical risk minimization: Theory and computa- tions. arXiv preprint arXiv:1910.10840 2019.
12

---

[35] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. International Conference on Learning Representations, 2018.
[36] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In International Conference on Computer Vision (ICCV), 2019.
[37] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jen- nifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1): 151-175, 2010.
[38] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- lional Conference on Learning Representations, 2015.
[39] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12): 124018, 2019.
[40] Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pages 464 472. IEEE, 2017.
[41] Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pattern Recognition, 2016.
[42] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211-255) 2015.
[43] Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In IEEE International Conference on Computer Vision, pages 1657-1664, 2013.
[44] Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In IEEE conference on computer vision and pattern recognition, pages 5018-5001 00--5027,2220 2017.
[45] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In European Conference on Computer Vision, pages 456-473, 2018.
[46] Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In IEEE/CVF International Conference on Computer Vision. pages 1406-1415, 2019.
[47] Oren Nuriel, Sagie Benaim, and Lior Wolf. Permuted adain: Reducing the bias towards global statistics in image classification. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
[48] Shujun Wang, Lequan Yu, Caizi Li, Chi-Wing Fu, and Pheng-Ann Heng. Learning from extrinsic and intrinsic supervisions for domain generalization. In European Conference on Computer Vision, pages 159-176. Springer, 2020.
[49] Shiori Sagawa* Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
[50] Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang Adversarial domain adaptation with domain mixup. In AAAI Conference on Artificial Intelli- gence, volume 34, pages 6502-6509, 2020.
[51] Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training, arXiv preprint vrXiv:2001.000 2020.
[52] Yufei Wang, Haoliang Li, and Alex C Kot. Heterogeneous domain generalization via domain mixup. In IEEE International Conference on Acoustics, Speech and Signal Processing, pages 3622-3626. IEEE, 2020.
13

---

[53] Haoliang Li. Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adversarial feature learning, IEEE Conference on Computer Vision and Pattern Recognition. pages 5400-5409, 2018.
[54] Gilles Blanchard, Aniket Anand Deshmukh, Urun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. Journal of Machine Learning Research, 22(2):1-555 2021.
[55] Zeyi Huang, Haohan Wang, Eric P Xing, and Dong Huang. Self-challenging improves cross- domain generalization. European Conference on Computer Vision, 2, 2020.
[56] Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, (0(4):838-855 1992.
[57] Takeru Miyato, Shin-ichi Maeda, Masanori Koyama, and Shin Ishii Virtual adversarial training: a regularization method for supervised and semi-supervised learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):19791199, 2018.
[58] Samuli Laine and Timo Aila. Temporal ensembling for semi-supervised learning. In Interna- tional Conference on Learning Representations, 2017.
[59] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Learning Representations, 2018. 8.
[60] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: critical analysis of out-of cut-of-distribution generalization. arXiv preprint arXiv:2006. 16241, 2020.
[61] Kai Yuanqing Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander Madry. Noise or signal: The role of image backgrounds in object recognition. International Conference on Learning Representations, 2020.
[62] Hanjoo Kim, Minkyu Kim, Dongjoo Seo, Jinwoong Kim, Heungseok Park, Soeun Park, Hyun- woo Jo, KyungHyun Kim, Youngil Yang, Youngkwan Kim, et al. Nsml: Meet the mlaas platform with real-world case study. arXiv preprint arXiv: /10.099/0. 2018.
14


Reviews:
[{"review_id": "OEUIkBI_712", "reviewer": "Reviewer_1Nc7", "summary": "This paper extends the idea of Stochastic Weighting Averaging (SWA) to the context of Domain Generalization (DG). Theoretically, the authors argues that finding a flat minima improves domain generalization. Algorithmically, the authors modify SWA by increasing the number of samples being averaged (e.g. per iteration) and by selecting the interval from where the samples are taken. The proposed algorithm  is evaluated in  the DomainBed benchmark and shows compelling results.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n- The paper tackles an important and relevant problem to the ML community.\n- The proposed algorithm is relatively simple to implement in practice and perform well experimentally.\n\nCons:\n\n- My major concern is wrt to the novelty of the work and the connection/motivation between the proposed algorithm SWAD, domain generalization and the derived theoretical insights.\n\n- __Theory__.  While I appreciate the  provided theoretical insights, they directly follow from standard results in domain adaptation [1,2,3]  and Norton and Royset [29] . Most importantly, I believe some details  have been overseen in the technical proof.  Specifically, I believe the proof misses a term  (line 86 eq. 20 in suppl material).  Lemma 1 is between any two h, but the risk is between h and the labeling function. If I understand correct, this is overseen in line 86.  The author should rewrite lemma 1 before directly using it in the proof. Please refer to [1] and [2] for details. The second part of the proof and Lemma 2 follow from Norton and Royset [5].  This should be clearly mentioned and if they are different the authors should highlight main differences.  I would appreciate if the authors can clarify on this.\n\n    Moreover, in practice,  the provided upper bound can be easily dominated by the divergence between  source and target.  Theorem 1, does not capture either the dissimilarity of the labeling functions (as a result of the missing term), which can also simply dominate the performance on target and being small is necessary for adaptation [2,3,4].  Even if these were all negligible and the result correct, the provided bound only shows that a flat minimal can  improve generalization but this is not necessarily novel.   It is hard to see a clear connection between these results and the proposed algorithm.  In other words, the theoretical results do not necessarily motivate the use of SWAD vs SWA. I would appreciate if the authors can provide further explanation on this aspect. \n\n- __Algorithm__. The proposed algorithm is based on a heuristic on top of SWA  that does not necessarily exploit the current scenario (DG). I believe this heuristic improves SWA in general and it might be algorithmically incremental vs vanilla SWA. Moreover, it introduces 3 new hyperparameters (well tuned in the experiments).  A better motivation than the one provided in 134-147 can improve the presentation/motivation of the algorithm.\n\n- __Experiments and Analysis__.  This is one of the strongest sections of the paper.  Just a few questions:\n    - Why SWA is not shown in Table 2?\n    - Table 4 mentioned HP search was not applied but ERM+SWAD includes HP search (This is the same result from Table 2).\n\n__Suppl__\n\n- Lemma 1 is a well know result also from [1,2,3] and should be cited appropriately.\n- Lemma 2 follows [5] and show be cited appropriately.\n\n\n\n\n\n\n**Originality**: The use of SWA in the context of DG is new.  However, algorithmically the proposed method is an extension (based on a heuristic) of a well known technique.  Theoretically, the presented results also follow from existing work and it is not clear how they motivate the need of SWA vs SWAD (see theory and algorithm above for more details).  \n\n**Quality**:  I believe the proposed algorithm should be better motivated as it is not clear the connection/motivation between  SWAD, domain generalization and the derived theoretical insights.  Moreover,  I believe several  small details where overseen in the technical proofs and those should be corrected to make them technically sound. (theory above for more details)\n\n**Clarity**:  The submission is clearly written and well organized. The main issue in this aspect is with respect to the motivation and connection among the two different parts (theoretical results and algorithms).  (see algorithm above for more details)\n\n**Significance**:  The empirical section is one of the strongest part of the paper and the results advance SoTA.  The method is relatively simple to implement in practice. However, it introduces a few extra hyperparams. \n\nAs expressed above, my major concern is wrt to the novelty of the work and the connection between the proposed algorithm SWAD, domain generalization and the derived theoretical insights. I believe the paper could significantly benefit  from a major revision.\n\n\n__References__\n\n[1] Ben-David et al.  A theory of learning from different domains Springer 2010\n\n[2] Impossibility Theorems for Domain Adaptation PMLR,2010.\n\n[3] Acuna et al. f-Domain-Adversarial Learning: Theory and Algorithms. ICML 2021\n\n[4] Zhao et al On Learning Invariant Representations for Domain Adaptation, ICML 2019\n\n[5] Norton et al Diametrical Risk Minimization: Theory and Computations, 2019 (cited)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "11"}, "extra_text": {}}, {"review_id": "LdgdXgOYePp", "reviewer": "Reviewer_75KP", "summary": "This paper studies how to improve the model performance on domain generalization problems from the perspective of loss landscape. Specifically, the authors propose to apply weight averaging (WA) to find flat minima and demonstrate this could lead to improve model generalization on unseen domains. Theoretically, the authors provide generalization bounds for the proposed method. Empirically, the proposed method achieves better performance on benchmark datasets compared with previous methods.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper tries to tackle the domain generalization problem by finding flat minima. It is interesting to investigate the loss landscape of the DNNs on domain generalization problems. \n\n**Originality**: The method proposed in this paper (i.e., SWAD) is mainly built on previous approaches such as SWA, and the authors mention this in the submission.\n\n**Quality**: This submission is a complete piece of work and the claims are well supported by the empirical results. The evaluation is reasonable. \n\n**Clarity**: This submission is well-written and easy to follow.\n\n**Significance**: This paper provides interesting empirical observations and improves the state-of-the-art by applying their proposed new approach.\n\nPros:\n1. The proposed method consistently improves the model performance on domain generalization problems, and the proposed method is flexible and can be easily integrated with other training approaches for better performance.\n\nCons:\n1. The generalization bound does not provide intuition for robust risk minimization (RRM), e.g. the theoretical results do not suggest performing RRM is better than standard ERM.\n2. The proposed method is heavily built on the previous approach SWA, which only differs in terms of how frequently the averaging is performed and which iteration is being selected for averaging. \n\nSuggestions:\n1. It would be better to put the algorithm (Algorithm 1 in appendix B) in the main body for describing the proposed method.\n\nMissing reference:\n1. Reference on generalization bound, [1].\n\nOverall, I think this paper provides a simple approach that can improve model (out-of-domain) performance on domain generalization problems. I would suggest `Marginally above the acceptance threshold'.\n\n[1]. A theory of learning from different domains. Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, Jennifer Wortman Vaughan. Machine learning, 2010.\n\n\n======================================[updates]======================================\n\nThanks for the authors' response. I have read the authors' responses as well as the other reviews. I will stick with my score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N.A.", "review_text": "", "extra_scores": {"ethical_concerns": "N.A.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "5tGXR4V_Dhq", "reviewer": "Reviewer_aAQf", "summary": "This paper studies the domain generalization (DG) problem setting and proposes a modification to the stochastic weight averaging (SWA) method that leads to better performance on DG testbeds. The proposed method, named SWAD, uses the same basic idea as SWA but samples parameters more densely and uses additional techniques for preventing the incorporation of overfit parameters. Some theoretical analysis supports the general thrust of seeking flat minima for generalization. The focus of the paper, however, is empirical, as SWAD demonstrates improved performance compared to many prior methods on several DG testbeds. Furthermore, SWAD can be combined with these prior methods, such as CORAL, for even better performance.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality\n---\nThe paper is, perhaps by design, not particularly original. Rather, it makes the observation that an existing set of techniques with a particular goal (seeking flat minima for improved generalization) also has applications to a different goal (DG). This is still, to the best of my knowledge, a novel observation and contribution, and furthermore the modifications made to SWA to arrive at SWAD also appear novel. Nevertheless, these modifications are relatively small, so the method can still be viewed as a version of the general idea of SWA, and thus originality is not the paper's strong suit.\n\n\nQuality\n---\nThe paper is of relatively high quality, primarily due to the empirical results. I did not carefully check the theory for correctness, but I am unsure as to the significance of the theorems. They appear rather similar to prior theoretical results on how flat minima can lead to better generalization, but with some additional divergence terms thrown in as we are dealing with the DG setting. So the high level conclusion seems to be: standard (in distribution) generalization also helps with DG, so long as the domain divergence is not too large. This is not an interesting result. It would be nice to have stronger theory that more directly connects flat minima to DG, but the current theory does not do this.\n\nTwo more minor comments about the theory. First, the reliance on VC dimension is not desirable, since this could be in the billions for the models used and, in particular for DG testbeds, many orders of magnitude larger than the number of training points. Second, I am unsure about the argument in L114-120. \"It implies that...\" seems to ignore several terms on the RHS of Eq (2) that could be rather large. The final sentence (\"Hence, Theorem 2 and...\") also seems to require quite a bit more support for such a general statement.\n\nNevertheless, as mentioned, I am appreciative of the empirical results that the paper provides, including a detailed study on several testbeds and an ablation study that verifies all of the proposed components of the method. I am inclined to view this paper as an empirical contribution, and a fairly nice one that further supports the notion that strong results can be achieved on existing DG testbeds without paying much attention to the domain labels themselves (instead focusing on standard generalization).\n\n\nClarity\n---\nThe paper is overall well written and easy to follow. See my comments above for some suggestions regarding the clarity of the theory. My other comments are more minor. First, though the technical description of the overfit-aware sampling strategy in L151-160 is nice, it would be made more clear with additional intuition and descriptive text, e.g., what is each parameter doing specifically and how did this particular instantiation come about? Second, there is a minor typo on L157 (remove the \"\\mathcal{E}_{val}^{(t_s)} =\"). Third, it is not immediately obvious what the significance is of Figure 5.\n\n\nSignificance\n---\nAs noted, I view this paper as a nice (significant) empirical contribution. Perhaps the significance could be further extended by exploring problems beyond DG. Indeed, there seems to be nothing central to the proposed method that ties it to DG, so could it perhaps be applied to other problems? A discussion of this point or, even better, some experiments would further improve the paper.\n\n---\nEdit after reading the author response\n---\nThanks for your additional comments. My scores remains unchanged and I continue to recommend acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have adequately addressed limitations and societal impacts in Section 5 and the Appendix.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "3q3xgRZkqZ", "reviewer": "Reviewer_74mu", "summary": "This paper provides some theoretical justification that so-called flat solutions exhibit better domain generalization along with a strongly performing training methodology, SWAD, for improving domain generalization. They first prove theorems bounding the domain generalization performance in terms of robust risk minimization. Then they propose an approach to finding flat solutions related to previous work on in distribution generalization, and they thoroughly explore this method with regard to the flatness of solutions found, its domain generalization performance, and also perform ablations of the method.\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The connection between the theoretical work in the first part of the paper and the method developed and tested is fairly tenuous. I appreciate the motivation for finding flatter minima but the paper would be tighter if there were a more direct connection between the flatness term in eq (1) and the algorithm. In particular, does SWAD encourage flatness w.r.t. the first term in eq. (1) more than other measures of flatness? My apologies if I missed this.\n\nHow novel/different is Theorem 1 compared to classic results from e.g. Ben-David (2010)? Of course, flatness doesn’t enter into these results but the rest seems somewhat similar.\n\nThe performance of the SWAD algorithm seems quite impressive, and I think this makes the paper a solid submission.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "A thoughtful limitations section was included but societal impacts was not. However, I don’t see particular societal impact concerns specific to this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes Stochastic Weight Averaging Densely (SWAD) to improve domain generalization. As the name suggests, it is a modification of Stochastic Weight Averaging (SWA) by averaging more densely (every iteration).

Reviewers mostly agree on strength and weaknesses of this work. They all believe that this work is a solid empirical contribution. However, there are two major shortcomings: 1- The novelty is very limited given the minimal changes applied to SWA to get SWAD. 2- The theoretical results seem very unrelated to reasons for the success of SWAD (in particular, as opposed to SWA). I highly recommend rewriting the paper with less emphasize on the theoretical results (maybe moving it to appendix?).

Overall, given the empirical contributions, I'm slightly inclined to accept the paper. I hope authors would take reviewers' suggestion into account to improve the paper.
Decision: Accept (Poster)