Gradient Matching for Domain Generalization

1 INTRODUCTION The goal of domain generalization is to train models that performs well on unseen , out-of-distribution data , which is crucial in practice for model deployment in the wild . This seemingly difficult task is made possible by the presence of multiple distributions/domains at train time . As we have seen in past work ( Arjovsky et al. , 2019 ; Gulrajani and Lopez-Paz , 2020 ; Ganin et al. , 2016 ) , a key aspect of domain generalization is to learn from features that remain invariant across multiple domains , while ignoring those that are spuriously correlated to label information ( as defined in Torralba and Efros ( 2011 ) ; Stock and Cisse ( 2017 ) ) . Consider , for example , a model that is built to distinguish between cows and camels using photos collected in nature under different climates . Since CNNs are known to have a bias towards texture ( Geirhos et al. , 2018 ; Brendel and Bethge , 2019 ) , if we simply try to minimize the average loss across different domains , the classifier is prone to spuriously correlate “ cow ” with grass and “ camels ” with desert , and predict the species using only the background . Such a classifier can be rendered useless when the animals are placed indoors or in a zoo . However , if the model could recognize that while the landscapes change with climate , the biological characteristics of the animals ∗Work done during internship at Facebook AI Research . †Now at Zoom . ( e.g . humps , neck lengths ) remain invariant and use those features to determine the species , we have a much better chance at generalizing to unseen domains . Similar intuitions have already motivated many approaches that consider learning “ invariances ” across domains as the main challenge of domain generalization . Typically , a lot of these work focus on learning invariant representations directly by removing the domain information ( Ganin et al. , 2016 ; Sun and Saenko , 2016 ; Li et al. , 2018 ) . In this work , we propose an inter-domain gradient matching ( IDGM ) objective . Instead of learning invariant features by matching the distributions of representations from different domains , our approach does so by encouraging consistent gradient directions across domains . Specifically , our IDGM objective augments the loss with an auxiliary term that maximizes the gradient inner product between domains , which encourages the alignment between the domain-specific gradients . By simultaneously minimizing the loss and matching the gradients , IDGM encourages the optimization paths to be the same for all domains , favouring invariant predictions . Figure 1 illustrates a motivating example described in Section 3.2 : given 2 domains , each containing one invariant feature ( orange cross ) and one spurious feature ( yellow and red cross ) . While empirical risk minimization ( ERM ) minimizes the average loss between these domains at the cost of learning spurious features only , IDGM aligns the gradient directions and is therefore able to focus on the invariant feature . While the IDGM objective achieves the desirable learning dynamic in theory , naive optimization of the objective by gradient descent is computationally costly due to the second-order derivatives . Leveraging the theoretical analysis of Reptile , a meta-learning algorithm ( Nichol et al. , 2018 ) , we propose to approximate the gradients of IDGM using a simple first-order algorithm , which we name Fish . Fish is simple to implement , computationally effective and as we show in our experiments , functionally similar to direct optimization of IDGM . Our contribution is a simple but effective algorithm for domain generalization , which exhibits state-ofthe-art performance on 13 datasets from recent domain generalization benchmark WILDS ( Koh et al. , 2020 ) and DOMAINBED ( Gulrajani and Lopez-Paz , 2020 ) . The strong performance of our method on a variety of datasets demonstrates that it is broadly applicable in different applications/subgenres of domain generalization tasks . We also perform detailed analysis in Section 4.4 to explain the effectiveness of our proposed algorithm . 2 RELATED WORK . Domain Generalization In domain generalization , the training data is sampled from one or many source domains , while the test data is sampled from a new target domain . We will now discuss the five main families of approaches to domain generalization : 1 . Distributional Robustness ( DRO ) : DRO approaches minimize the worst-case loss over a set of data distributions constructed from the training domains . Rojas-Carulla et al . ( 2015 ) proposed DRO to address covariate shift ( Gretton et al. , 2009a ; b ) , where P ( Y |X ) remains constant across domains but P ( X ) changes . Later work also studied subpopulation shift , where the train and test distributions are mixtures of the same domains , but the mixture weights change between train and test ( Hu et al. , 2018 ; Sagawa et al. , 2019 ) ; 2 . Domain-invariant representation learning : This family of approaches to domain generalization aims at learning high-level features that make domains statistically indistinguishable . Prediction is then based on these features only . The principle is motivated by a generalization error bound for unsupervised domain adaptation ( Ben-David et al. , 2010 ; Ganin et al. , 2016 ) , but the approach readily applies to domain generalization ( Gulrajani and Lopez-Paz , 2020 ; Koh et al. , 2020 ) . Algorithms include penalising the domain-predictive power of the model ( Ganin et al. , 2016 ; Wang et al. , 2019 ; Huang et al. , 2020 ) , aligning domains through contrastive loss ( Motiian et al. , 2017 ) , matching mean and variance of feature distributions across domains ( Sun and Saenko , 2016 ) , learning useful representations by solving Jigsaw puzzles ( Carlucci et al. , 2019 ) , using the maximum mean discrepancy to match the feature distributions ( Li et al. , 2018b ) or introducing training constraints across domains using mixup formulation ( Yan et al. , 2020 ) . 3 . Invariant Risk Minimization ( IRM ) : IRM is proposed by Arjovsky et al . ( 2019 ) , which learns an intermediate representation such that the optimal classifiers ( on top of this representation ) of all domains are the same . The motivation is to exploit invariant causal effects between domains while reducing the effect of domain-specific spurious correlations . From an optimization perspective , when IRM reaches its optimal , all the gradients ( for the linear classifier ) has to be zero . This is why IRM ’ s solution won ’ t deviate from ERM when ERM is optimal for every domain , which is not the case for our proposed IDGM objective due to the gradient inner product term . 4 . Data augmentation : More recently , approaches that simulates unseen domains through specific types of data augmentation/normalization has been gaining traction . This includes work such as Zhou et al . ( 2020 ) ; Volpi and Murino ( 2019 ) ; Ilse et al . ( 2021 ) , as well as Seo et al . ( 2019 ) which utilises ensemble learning . 5 . Gradient alignment : Two concurrent work – Koyama and Yamaguchi ( 2021 ) and Parascandolo et al . ( 2021 ) – utilise similar gradient-alignment principle for domain generalization . Koyama and Yamaguchi ( 2021 ) proposes IGA , which learns invariant features by minimizing the variance of inter-domain gradients . The key difference between IGA and our objective is that IGA is completely identical to ERM when ERM is the optimal solution on every training domain , since the variances of the gradients will be zero . While they achieve the best performance on the training set , both IGA and ERM could in some cases , completely fail when generalizing to unseen domains ( see Section 3.2 for such an example ) . Our method , on the contrary , biases towards non-ERM solutions as long as the gradients are aligned , and is therefore able to avoid this issue . Parascandolo et al . ( 2021 ) on the other hand , proposes to mask out the gradients that have opposite signs for different domains . Unlike their work that prunes gradients that are inconsistent , our approach actively encourage gradients from different domains to be consistent by maximizing the gradient inner product . Additionally , in Lopez-Paz and Ranzato ( 2017 ) we also see the application of gradientalignment , however in this case it is applied under the continual learning setting to determine whether a gradient update will increase the loss of the previous tasks . Apart from these algorithms that are tailored for domain generalization , a well-studied baseline in this area is ERM , which simply minimizes the average loss over training domains . Using vanilla ERM is theoretically unfounded ( Hashimoto et al. , 2018 ; Blodgett et al. , 2016 ; Tatman , 2017 ) since ERM is guaranteed to work only when train and test distributions match . Nonetheless , recent benchmarks suggest that ERM obtains strong performance in practice , in many case surpassing domain generalization algorithms ( Gulrajani and Lopez-Paz , 2020 ; Koh et al. , 2020 ) . Our goal is to fill this gap , using an algorithm significantly simpler than previous approaches . Connections to meta-learning There are close connections between meta-learning ( Thrun and Pratt , 1998 ) and ( multi-source ) domain adaptation . In fact , there are a few works in domain generalization that are inspired by the meta-learning principles , such as Li et al . ( 2018a ) ; Balaji et al . ( 2018 ) ; Li et al . ( 2019 ) ; Dou et al . ( 2019 ) . Specifically , Li et al . ( 2020 ) also proposes to adapt Reptile for domain generalization tasks , however they study their method under the sequential learning setting , whereas our method can be trained on all domains and therefore learns faster , especially when the number of domains is large . In Ren et al . ( 2018 ) , we also see the leveraging of gradient inner product in meta-learning , where it is used to determine the importance weight of training examples . We discuss the connection between our proposed algorithm to meta-learning in more details in Appendix A.1 . Note that our proposed algorithm Fish is similar to the Mean Teacher method ( Tarvainen and Valpola , 2017 ) , where a teacher model ( equivalent to θ in Algorithm 1 ) is computed using a moving average of the student model ( equivalent to θ̃ in Algorithm 1 ) . 3 METHODOLOGY . 3.1 GOALS . Consider a training dataset Dtr consisting of S domains Dtr = { D1 , · · · , DS } , where each domain s is characterized by a dataset Ds : = { ( xsi , ysi ) } ns i=1 containing data drawn i.i.d . from some probability distribution . Also consider a test dataset Dte consisting of T domains Dte = { DS+1 , · · · , DS+T } , where Dtr ∩ Dte = ∅ . The goal of domain generalization is to train a model with weights θ that generalizes well on the test dataset Dte such that : arg min θ ED∼Dte E ( x , y ) ∼D [ l ( ( x , y ) ; θ ) ] , ( 1 ) where l ( ( x , y ) ; θ ) is the loss of model θ evaluated on ( x , y ) . A naive approach is to apply ERM , which simply minimizes the average loss on Dtr , ignoring the discrepancy between train and test domains : Lerm ( Dtr ; θ ) = ED∼Dtr E ( x , y ) ∼D [ l ( ( x , y ) ; θ ) ] . ( 2 ) The ERM objective does not exploit the invariance across different domains inDtr and could perform arbitrarily poorly on test data . We demonstrate this effect with the following simple linear example .

Summary: This paper proposed inter-domain gradient matching for domain generalization. They also approximated the proposed model with a simple first-order algorithm to avoid costly second-order computations.  The performance on the WILDs and DomainBed seems better than the ERM algorithm.