SWAD: Domain Generalization by Seeking Flat Minima

1 Introduction . Independent and identically distributed ( i.i.d . ) condition is the underlying assumption of machine learning experiments . However , this assumption may not hold in real-world scenarios , i.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 or day-to-night shifts [ 1 , 2 ] . Even in a simple image recognition scenario , systems rely on wrong cues for their prediction , e.g. , geographic distribution [ 3 ] , demographic statistics [ 4 ] , texture [ 5 ] , or backgrounds [ 6 ] . Consequently , a 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 [ 8– ∗Equal contribution †Part of work done while at NAVER Clova Correspondence to : Junbum Cha < junbum.cha @ kakaobrain.com > , Sungrae Park < sungrae.park @ upstage.ai > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 12 ] , obtained well-transferable model parameters by the meta-learning framework [ 13–16 ] , 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 minimizing the empirical loss on a complex and non-convex loss landscape is typically not sufficient to arrive at a 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–28 ] . Izmailov et al . [ 25 ] argued that seeking flat minima will lead to robustness against the loss landscape shift between training and test datasets , while a 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 to unseen domains , we formulate a robust risk minimization ( RRM ) problem defined by the worst-case empirical risks within neighborhoods in parameter space [ 26 , 34 ] . 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-aware solvers , by introducing a 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 . Contribution . Our main contribution is introducing flatness into DG , and showing remarkably outperforming performances against existing DG methods . As shown in Table 1 , our SWAD improves the average DG performances by 3.6pp against the ERM baseline and 1.6pp against the existing best methods . Furthermore , by combining SWAD and previous SOTA [ 31 ] , we even achieve 0.4pp improvements against the vanilla SWAD results . We also empirically show that while popular indomain generalization methods without considering flatness , e.g. , Mixup [ 35 ] or CutMix [ 36 ] , are not effective to out-of-domain generalization ( Table 3 ) , 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 : = { Di } Ii be a set of training domains , where Di is a distribution over input space X , and I is the total number of domains . From each domain , we observe n training data points which consist of input x and target label y , ( xij , y i j ) n j=1 ∼ Di . We also define a set of target domain T : = { Ti } T i 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 labeling function h ( x ) that generates target label for multiple domains , i.e. , yij = h ( x i j ) for all i and j. Domain generalization ( DG ) aims to find a model parameter θ ∈ Θ 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 × Y → [ 0 , c ] , such that ` ( y1 , y2 ) = 0 holds if and only if y1 = y2 where Y is a set of labels . For simplicity , we set c to one in our proofs , but we note that ` ( · , · ) can be generalized for any bounded loss function . Then , we can define a population loss over multiple domains by ED ( θ ) = 1I ∑I i=1 Exi∼Di [ ` ( f ( xi ; θ ) , yi ) ) ] , where f ( · ; θ ) is a model parameterized by θ . Formally , the goal of DG is to find a model which minimizes both ED ( θ ) and ET ( θ ) by only minimizing an empirical risk ÊD ( θ ) : = 1In ∑I i=1 ∑n j=1 ` ( f ( x i ; θ ) , yi ) ) over training domains D. In practice , ERM , i.e. , arg minθ ÊD ( θ ) , can have multiple solutions that provide similar values of the training losses but significantly different generalizability on ED ( θ ) and ET ( θ ) . 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–28 ] . We consider a robust empirical loss function defined by the worst-case loss within neighborhoods in the parameter space as ÊγD ( θ ) : = max‖∆‖≤γ ÊD ( θ + ∆ ) , where ‖ · ‖ denotes the L2 norm and γ is a radius which defines neighborhoods of θ . Intuitively , if γ is sufficiently larger than the “ radius ” of a sharp optimum θs of ÊD ( θ ) , θs is no longer an optimum of ÊγD ( θ ) as well as its neighborhoods within the γ-ball . On the other hand , if an optimum θf has larger “ radius ” than γ , there exists a local optimum within γ-ball – See Figure 1 . Hence , solving the robust risk minimization ( RRM ) , i.e. , arg minθ ÊγD ( θ ) , will find a near solution of a 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̂γD and ET as follows : Theorem 1 . Consider a set of N covers { Θk } Nk=1 such that the parameter space Θ ⊂ ∪Nk Θk where diam ( Θ ) : = supθ , θ′∈Θ ‖θ− θ′‖2 , N : = ⌈ ( diam ( Θ ) /γ ) d ⌉ and d is dimension of Θ . Let vk be a VC dimension of each Θk . Then , for any θ ∈ Θ , the following bound holds with probability at least 1− δ , ET ( θ ) < ÊγD ( θ ) + 1 2I I∑ i=1 Div ( Di , T ) + max k∈ [ 1 , N ] √ vk ln ( m/vk ) + ln ( N/δ ) m , ( 1 ) where m = nI is the number of the training samples and Div ( Di , T ) : = 2 supA |PDi ( A ) −PT ( A ) | is a divergence between two distributions . Proof can be done similarly as [ 37 ] and [ 34 ] . In Theorem 1 , the test loss ET ( θ ) is bounded by three terms : ( 1 ) the robust empirical loss ÊγD ( θ ) , ( 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 γ 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 we 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 easily 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.1 and C.2 . From Theorem 1 , 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 , θ̂γ , is upper bounded as follows : Theorem 2 . Let θ̂γ denote the optimal solution of the RRM , i.e. , θ̂γ : = arg minθ ÊγD ( θ ) , and let v be a VC dimension of the parameter space Θ . Then , the gap between the optimal test loss , minθ′ ET ( θ′ ) , and the test loss of θ̂γ , ET ( θ̂γ ) , has the following bound with probability at least 1− δ . ET ( θ̂γ ) −min θ′ ET ( θ′ ) ≤ ÊγD ( θ̂ γ ) −min θ′′ ÊD ( θ′′ ) + 1 I I∑ i=1 Div ( Di , T ) + max k∈ [ 1 , N ] √ vk ln ( m/vk ) + ln ( 2N/δ ) m + √ v ln ( m/v ) + ln ( 2/δ ) m ( 2 ) Proof is in Appendix C.3 . It implies that if we find the optimal solution of the RRM ( i.e. , θ̂γ ) , then the generalization gap in the test domain ( i.e. , ET ( θ̂γ ) − minθ′ ET ( θ′ ) ) is upper bounded by the gap between the RRM and ERM ( i.e. , ÊγD ( θ̂γ ) −minθ′′ ÊD ( θ′′ ) ) . 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 γ , the optimal solution of the RRM will find a point near a flat optimum of ERM as shown in Figure 1 . Hence , Theorem 2 and the intuition from Figure 1 imply that seeking a flat minimum of ERM will lead to a better domain generalization gap .

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.