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 proposes a new regularization method for OOD generalization. The central idea is to align the covariances of gradients between different domains. Specifically, - Authors first categorize different works in this area into three groups:     1. Methods that match **features** across domains such as those in adversarial training and most notably CORAL which considers the covariance of the **features**.     2. Methods that state that **the classifier on the top** should be invariant. Such as IRM.     3. Those that align **gradients** such as IGA.  - Authors then propose Fishr which falls into the 3rd category. It aligns gradients between the domain but with the important different that it aligns the *covariance* of the gradients not the *average* of them.  - Authors establish equivalences between the the covariance, the fisher information matrix, and the hessian which in fact describes the loss landscape.  - As a practical approach, authors propose to match the **variances** instead of the **covariances** and also only the penultimate linear layer.  - Finally, authors evaluate their approach on standard OOD benchmarks in two different setups:     1. When the hyper-parameters are tuned using the test set.     2. When the hyper-parameters are tuned using the training set.