Fishr: Invariant Gradient Variances for Out-of-distribution Generalization 

1 INTRODUCTION . The success of deep neural networks in supervised learning ( Krizhevsky et al. , 2012 ) relies on the crucial assumption that the train and test data distributions are identical . In particular , the tendency of networks to rely on simple features ( Kalimeris et al. , 2019 ; Valle-Perez et al. , 2019 ; Geirhos et al. , 2020 ) is generally a desirable behavior reflecting Occam ’ s razor . However , in case of distribution shift , this simplicity bias deteriorates performance since more complex features are needed ( Tenenbaum , 2018 ; Shah et al. , 2020 ) . For example , in the fight against Covid-19 , most of the deep learning methods developed to detect coronavirus from chest scans were shown useless for clinical use ( DeGrave et al. , 2021 ; Roberts et al. , 2021 ) : indeed , networks exploited simple bias in the training datasets such as patients ’ age or body position rather than ‘ truly ’ analyzing medical pathologies . To better generalize under distribution shifts , most works such as Blanchard et al . ( 2011 ) or Muandet et al . ( 2013 ) assume that the training data is divided into different training domains in which there is a constant underlying causal mechanism ( Peters et al. , 2016 ) . To remove the domain-dependent explanations , different invariance criteria across those training domains have been proposed . ( Ganin et al. , 2016 ; Sun et al. , 2016 ; Sun & Saenko , 2016 ) enforce similar feature distributions , others ( Arjovsky et al. , 2019 ; Krueger et al. , 2021 ) force the classifier to be simultaneously optimal across all domains . Yet , despite the popularity of this research topic , none of these methods perform significantly better than the classical Empirical Risk Minimization ( ERM ) when applied with controlled model selection and restricted hyperparameter search ( Gulrajani & Lopez-Paz , 2021 ; Ye et al. , 2021 ) . These failures motivate the need for new ideas . To foster the emergence of a shared mechanism with consistent generalization properties , our intuition is that learning should progress consistently and similarly across domains . Besides , the learning procedure of deep neural networks is dictated by the distribution of the gradients with respect to the network weights ( Yin et al. , 2018 ; Sankararaman et al. , 2020 ) — usually backpropagated in the network during gradient descent . Thus , we seek distributional invariance across domains in the gradient space : domain-level gradients should be similar , not only in average direction , but most importantly in statistics such as covariance and dispersion . In this paper , we propose the Fishr regularization for out-of-distribution generalization in classification — summarized in Fig . 1 . We match the domain-level gradient variances , i.e. , the second moment of the gradient distributions . In contrast , previous gradient-based works such as Fish ( Shi et al. , 2021 ) only match the domain-level gradients means , i.e. , the first moment . Moreover , our strategy is also motivated by the close relations between the gradient variance , the Fisher Information ( Fisher , 1922 ) and the Hessian . This explains the name of our work , Fishr , using gradients as in Fish and related to the Fisher Information Matrix . Notably , we will study how Fishr forces the model to have similar domain-level Hessians and promotes consistent explanations — by generalizing the inconsistency formalism introduced in Parascandolo et al . ( 2021 ) . To reduce the computational cost , we justify an approximation that only considers the gradients in the classifier . This is simple to implement with the BackPACK ( Dangel et al. , 2020 ) package . We summarize our contributions as follows : • We introduce the Fishr regularization that brings closer the domain-level gradient variances . • Based on the relation between the gradient covariance , the Fisher Information and the Hessian , we show that Fishr matches domain-level Hessians and improves generalization by reducing inconsistencies across domains . • We justify a simple and scalable implementation . Empirically , we first validate that Fishr tackles distribution shifts on the synthetic Colored MNIST ( Arjovsky et al. , 2019 ) . Then , we show that Fishr performs best on the DomainBed benchmark ( Gulrajani & Lopez-Paz , 2021 ) with the ‘ Oracle ’ model selection method and third with the ‘ Training ’ model selection when compared with state-of-the-art counterparts . Critically , Fishr is the only method to perform better ( on VLCS , OfficeHome , TerraIncognita and DomainNet ) or similarly ( on PACS ) than ERM with both selection methods on all ‘ real ’ datasets . 2 CONTEXT AND RELATED WORK . We first describe our task and provide the notations used along our paper . Then we remind some important related works to understand how our Fishr stands in a rich literature . Problem definition and notations We study out-of-distribution ( OOD ) generalization for classification . Our model is a deep neural network ( DNN ) fθ ( parametrized by θ ) made of a deep features extractor Φφ on which we plug a dense linear classifier wω : fθ = wω ◦Φφ and θ = ( φ , ω ) . In training , we have access to different domains E : for each domain e ∈ E , the datasetDe = { ( xie , y i e ) } ne i=1 contains ne i.i.d . ( input , labels ) samples drawn from a domain-dependent probability distribution . Combined together , the datasets { De } e∈E are of size n = ∑ e∈E ne . Our goal is to learn weights θ so that fθ predicts well on a new test domain , unseen in training . As described in Koh et al . ( 2020 ) and Ye et al . ( 2021 ) , most common distribution shifts are diversity shifts — where the training and test distributions comprise data from related but distinct domains — or correlation shifts — where the distribution of the covariates at test time differs from the one during training . To generalize well despite these distribution shifts , fθ should ideally capture an invariant mechanism across training domains . Following standard notations , ‖M‖2F denotes the Frobenius norm of matrix M ; ‖v‖ 2 2 denotes the euclidean norm of vector v ; 1 is a column vector with all elements equal to 1 . The standard Expected Risk Minimization ( ERM ) ( Vapnik , 1999 ) framework simply minimizes the average empirical risk over all training domains , i.e. , 1|E| ∑ e∈E Re ( θ ) where Re ( θ ) = 1ne ∑ne i=1 ` ( fθ ( xie ) , yie ) and ` is the loss , usually the negative log-likelihood . Many approaches try to exploit some external source of knowledge ( Xie et al. , 2021 ) , in particular the domain information . As a side note , these partitions may be inferred if not provided ( Creager et al. , 2021 ) . Some works explore data augmentations to mix samples from different domains ( Wang et al. , 2020 ; Wu et al. , 2020 ) , some re-weight the training samples to favor underrepresented groups ( Sagawa et al. , 2020a ; b ; Zhang et al. , 2021 ) and others include domain-dependent weights ( Ding & Fu , 2017 ; Mancini et al. , 2018 ) . Yet , most recent works promote invariance via a regularization criterion and only differ by the choice of the statistics to be matched across training domains . They can be categorized into three groups : these methods enforce agreement either ( 1 ) in features ( 2 ) in predictors or ( 3 ) in gradients . First , some approaches aim at extracting domain-invariant features and were extensively studied for unsupervised domain adaptation . The features are usually aligned with adversarial methods ( Ganin et al. , 2016 ; Gong et al. , 2016 ; Li et al. , 2018b ; c ) or with kernel methods ( Muandet et al. , 2013 ; Long et al. , 2014 ) . Yet , the simple covariance matching in CORAL ( Sun et al. , 2016 ; Sun & Saenko , 2016 ) performs best on various tasks for OOD generalization ( Gulrajani & Lopez-Paz , 2021 ) . With Zije the j-th dimension of the features extracted by Φφ for the i-th example xie of domain e ∈ E = { A , B } , CORAL minimizes ‖Cov ( ZA ) − Cov ( ZB ) ‖ 2 F where Cov ( Ze ) = 1 ne−1 ( Z > e Ze − 1ne ( 1 > Ze ) > ( 1 > Ze ) ) is the feature covariance matrix . CORAL is more powerful than mere feature matching ∥∥∥ 1nA1 > ZA − 1nB 1 > ZB∥∥∥22 as in Deep Domain Confusion ( DDC ) ( Tzeng et al. , 2014 ) . Yet , Johansson et al . ( 2019 ) and Zhao et al . ( 2019 ) show that these approaches are insufficient to guarantee good generalization . Motivated by arguments from causality ( Pearl , 2009 ) and the idea that statistical dependencies are epiphenomena of an underlying causal structure , Invariant Risk Minimization ( IRM ) ( Arjovsky et al. , 2019 ) explains that the predictor should be invariant ( Peters et al. , 2016 ; Rojas-Carulla et al. , 2018 ) , i.e. , simultaneously optimal across all domains . Among many suggested improvements ( Chang et al. , 2020 ; Idnani & Kao , 2020 ; Teney et al. , 2020 ; Ahmed et al. , 2021 ) , Risk Extrapolation ( V-REx ) ( Krueger et al. , 2021 ) argues that training risks from different domains should be similar and thus penalizes |RA − RB |2 when E = { A , B } . These ideas have been applied in semi-supervised learning ( Li et al. , 2021 ) . Yet , recent works point out pitfalls of IRM ( Javed et al. , 2020 ; Guo et al. , 2021 ; Kamath et al. , 2021 ) , that does not provably work with non-linear data ( Rosenfeld et al. , 2021 ) and could not improve over ERM when hyperparameter selection is restricted ( Koh et al. , 2020 ; Gulrajani & Lopez-Paz , 2021 ; Ye et al. , 2021 ) . A third and most recent line of work promotes agreements between gradients with respect to θ. Gradient agreements help batches from different tasks to cooperate , and have been previously employed for multitasks ( Du et al. , 2018 ; Yu et al. , 2020 ) , continual ( Lopez-Paz & Ranzato , 2017 ) , meta ( Finn et al. , 2017 ; Zhang et al. , 2020 ) and reinforcement ( Zhang et al. , 2019 ) learning . In OOD generalization , Koyama & Yamaguchi ( 2020 ) ; Parascandolo et al . ( 2021 ) ; Shi et al . ( 2021 ) try to find minimas in the loss landscape that are shared across domains . Specifically , these works tackle the domain-level expected gradients : ge = E ( xe , ye ) ∼De∇θ ` ( fθ ( xe ) , ye ) . ( 1 ) When E = { A , B } , IGA ( Koyama & Yamaguchi , 2020 ) minimizes ||gA − gB ||22 ; Fish ( Shi et al. , 2021 ) increases gA · gB ; AND-mask ( Parascandolo et al. , 2021 ) and others ( Mansilla et al. , 2021 ; Shahtalebi et al. , 2021 ) update weights only when gA and gB point to the same direction . Along with the increased computation cost , the main limitation of these gradient-based methods is the per-domain batch averaging of gradients , that removes more granular statistics ; in particular , this averaging removes the information from pairwise interactions between gradients from samples in a same domain . In opposition , our new regularization for OOD generalization keeps extra information from individual gradients and matches across domains the domain-level gradient variances . In a nutshell , Fishr is similar to the covariance-based CORAL ( Sun et al. , 2016 ; Sun & Saenko , 2016 ) but in the gradient space rather than in the feature space .

Summary: This paper introduces a simple technique for out-of-distribution generalization, called Fishr. Intuitively, given the success of CORAL, Fishr enforces consistency through mean regularization between the element-wise variances of the gradients across domains. It is efficiently solvable using the existing BackPACK package. Several explanations have been provided for Fishr. Then, the paper concludes with applying Fishr to standard OOD benchmarks showing superior performance of Fishr in several different problems.