Loss Function Learning for Domain Generalization by Implicit Gradient

1 INTRODUCTION . Deep learning is highly successful when the training and testing samples meet the i.i.d . assumption . However , this assumption is violated in many practical applications of machine learning from medical imaging to earth observation imaging ( Koh et al. , 2021 ) . This has led a large number of studies to investigate approaches to training models with increased robustness to distribution shift at testingtime , a problem setting known as Domain Generalisation ( DG ) . Despite the volume of research in this area ( Zhou et al. , 2021a ) , a recent careful benchmarking exercise , DomainBed ( Gulrajani & Lopez-Paz , 2021 ) showed that simple empirical risk minimisation ( ERM ) on a combination of training domains is a very strong baseline when properly tuned . State-of-the-art alternatives based on sophisticated architectures , regularisers , and data augmentation schemes failed to reliably beat ERM ( Gulrajani & Lopez-Paz , 2021 ) . Rather than propose an alternative to ERM for DG , we investigate a previously unstudied hyperparameter of ERM , namely the choice of loss function—which has been ubiquitously taken to be standard cross-entropy ( CE ) in prior DG work . Loss function choice has been shown to impact calibration ( Mukhoti et al. , 2020 ) , overfitting ( Gonzalez & Miikkulainen , 2019 ) , and label-noise robustness ( Wang et al. , 2019 ) in standard supervised learning , so it is intuitive that it would impact robustness to domain-shift . However , it has not yet been studied in this context . Our preliminary experiments showed that equipping ERM with some recent robust loss functions in place of CE does lead to improvements in DG performance where sophisticated alternatives have failed ( Gulrajani & Lopez-Paz , 2021 ) . This raises the question : can one design a loss function specialised for DG ? To answer this question , we define a meta-learning algorithm to learn a parametric ( white-box ) loss function suitable for DG . Our desiderata are : ( 1 ) Performing ERM with this loss on a source domain should lead to good performance when tested on out-of-domain target data ; and ( 2 ) It should provide a ‘ plug-and-play ’ drop-in replacement for cross-entropy that , once learned , can be used without further modification or computational expense with any new dataset or model architecture . While there has been growing interest in meta-learning for loss function design ( Li et al. , 2019a ) , they mostly fail to meet these criteria . They learn problem-specific—rather than re-usable—losses . If applied to DG , this would imply replacing simple ERM learning with sophisticated meta-learning pipelines to train a loss on a per-problem basis . In contrast , our discovered loss provides a drop in replacement for CE that leads standard training pipelines to produce more robust models . To train a general purpose robust loss function we need a search space that is flexible enough to include interesting new losses , but simple enough to generalise across tasks without overfitting to the problem used for loss learning . We choose a 12-dimensional space of fourth order Taylor polynomials . Furthermore , we need a loss that is suitable for all stages of training . This precludes the majority of approaches based on online meta-learning which update the loss and base model iteratively ( Li et al. , 2019a ; c ) , and also suffer from short-horizon bias ( Wu et al. , 2018 ) . Evolutionary methods ( Gonzalez & Miikkulainen , 2019 ) and reinforcement-learning ( Li et al. , 2019a ) could support loss learning in principle , but are too slow to be feasible . Therefore we develop the first implicit-gradient based approach to loss learning . This allows us to tractably compute meta-gradients of the target recognition performance with respect to the loss used for training in the source domain . We use a simple DG task ( RotatedMNIST ) to train our robust loss , termed Implicit Taylor Loss ( ITL ) , to replace CE in ERM . Subsequent experiments show that ERM with ITL surpasses CE across a range of DG benchmarks , and leads to state of the art performance , despite being much simpler and faster than competitor DG methods . While the majority of existing DG methods require multiple source domains to conduct data augmentation or feature alignment strategies , ITL improves single-source domain generalisation , a crucial problem setting which has been minimally studied thus far . To summarise our contributions : ( i ) We provide the first study on the significance of supervised loss function choice in DG ( ii ) We demonstrate the first efficient solution to loss-learning based on meta-gradients computed by the Implicit Function Theorem . ( iii ) Empirically , we show that our learned ITL loss enhances simple ERM and achieves state of the art DG performance across a range of benchmarks , including the challenging single-source DG scenario . 2 RELATED WORK . Domain Generalisation : Domain Generalisation aims to learn a model using data from one or more source domains , but with the further requirement that it is robust to testing on novel target domain data—without accessing target data during training . DG is now a well studied ( Zhou et al. , 2021a ) area with diverse approaches including data augmentation ( Shankar et al. , 2018 ; Zhou et al. , 2021b ) , robust training algorithms such as domain alignment objectives ( Li et al. , 2018b ) , and other regularisers ( Li et al. , 2019c ; Balaji et al. , 2018 ) . Most DG studies have assumed the multi-source setting , which enables new data-augmentation strategies ( Zhou et al. , 2021b ) , and allows generalisation-promoting design features to be tuned by domain-wise cross-validation . In particular , a few studies ( Li et al. , 2019c ; Balaji et al. , 2018 ) have considered meta-learning based DG , where a regulariser applied in a training domain is tuned by meta-gradients from the resulting validation-domain performance . The resulting model is then deployed to the true target domain within the same family . These methods require regulariser meta-learning for each given multi-source DG problem family . In contrast , we propose to learn a simple loss function once , which then provides a drop-in replacement for CE in any single- , or multi-source DG problem . A recent criticism of the DG literature showed that no method consistently outperformed a well tuned ERM baseline on the carefully designed DomainBed benchmark ( Gulrajani & Lopez-Paz , 2021 ) . Rather than competing with ERM , we simply enhance the ERM loss function and this leads to a clear improvement on DomainBed . Loss Function learning : Loss Function learning aims to discover new losses that improve model optimisation from various perspectives including conventional generalisation performance , ( Gonzalez & Miikkulainen , 2019 ; Liu et al. , 2021 ) , optimisation efficiency ( Li et al. , 2019a ; Gonzalez & Miikkulainen , 2019 ; Wang et al. , 2020 ; Bechtle et al. , 2020 ) , and noise robustness ( Li et al. , 2019a ) . Key dichotomies are in the search space of black box ( neural ) ( Bechtle et al. , 2020 ; Li et al. , 2019c ) vs white-box ( human-readable ) ( Li et al. , 2019a ; Gonzalez & Miikkulainen , 2019 ; Wang et al. , 2020 ) losses ; whether learned losses are problem specific ( Li et al. , 2019a ; Wang et al. , 2020 ) or reusable ( Gonzalez & Miikkulainen , 2019 ) ; the meta-optimisation algorithm used—evolution ( Liu et al. , 2021 ; Gonzalez & Miikkulainen , 2019 ) , RL ( Li et al. , 2019a ; Wang et al. , 2020 ; Bechtle et al. , 2020 ) , or gradient ( Li et al. , 2019c ) ; and whether the loss is updated offline ( Liu et al. , 2021 ; Gonzalez & Miikkulainen , 2019 ) ( long inner loop , typically intractable ) , or online ( Li et al. , 2019a ; Wang et al. , 2020 ) ( short inner loop , efficient but suffers from short-horizon bias ( Wu et al. , 2018 ) ) . No studies have thus far investigated loss learning for domain-shift robustness . In order to learn a reusable robust loss we use a white-box loss search space of Taylor polynomials proposed in Gonzalez & Miikkulainen ( 2020 ) , and offline/long inner loop meta-learning . To make this meta-optimisation tractable , we exploit the Implicit Function Theorem , to efficiently generate accurate hypergradients of the validation domain performance with respect to the training domain loss function parameters . Besides being the first demonstration of loss learning for DG , to our knowledge it is also the first demonstration of any implicit gradient-based loss learning . 3 METHOD . The need for Domain Generalisation arises when one is using machine learning to build a model where the available training data is not representative of the data that will be observed by the model once it has been deployed . In particular , it is assumed that there is an underlying distribution over domains , P , from which we can sample several source domain distributions , { p ( s ) 1 , ... , p ( s ) n ∼ P } , to make use of during training . We can construct a training set for each of these source domain distributions by sampling K data points , D ( s ) i = { ( x ( s , j ) i , y ( s , j ) i ) ∼ p ( s ) i } Kj=1 , and use the union of all these sets as the full training set , D ( s ) = ⋃n i=1D ( s ) i . Empirical Risk Minimisation ( ERM ) then simply finds the model parameters , θ , that minimise the loss measured on this training set , min θ 1 n n∑ i=1 1 K K∑ j=1 L ( fθ ( x ( j ) i ) , y ( j ) i ) , ( 1 ) where L ( · , · ) is a loss function ( typically cross entropy ) measuring how well the predicted labels match the ground truth labels . One can empirically check the resulting model ’ s robustness to domain shift by sampling one or more target domain distributions , { p ( t ) 1 , ... , p ( t ) m ∼ P } , from the same distribution over domains that was used to generate the training data . Data can then be sampled for each of these target domains , yielding a test dataset D ( t ) = ⋃m i=1D ( t ) i . Standard evaluation metrics such as accuracy can then be computed using this data . 3.1 META-LEARNING LOSSES FOR DG . Our goal is to replace the standard CE loss typically used in ERM with a learned loss function . We are motivated by recent work showing that learned losses can enable models to perform better for a variety of other problem settings , such as training with label noise ( Wang et al. , 2019 ) and improving calibration ( Mukhoti et al. , 2020 ) . We formulate the task of learning the parameters , ω , of a loss function , Lω , as a bilevel optimisation problem . The outer objective is to find the ω that maximises the performance of a model evaluated on the target domain data , and the inner problem is to train a model to minimise the value of Lω measured on the source domain data . The loss parameters are optimised using gradient-based methods that take advantage of the implicit function theorem to efficiently compute gradients for the outer optimisation problem . Crucially , once the optimal loss function ω∗ has been found , new DG problems can be solved via ERM on the Lω∗ loss . The bilevel optimisation we use to formalise the meta-learning process is given by ω∗ = argmin ω 1 m m∑ i=1 1 K K∑ j=1 M ( fθ∗ ( ω ) ( x ( t , j ) i ) , y ( t , j ) i ) ( 2 ) s.t . θ∗ ( ω ) = argmin θ 1 n n∑ i=1 1 K K∑ j=1 Lω ( fθ ( x ( s , j ) i ) , y ( s , j ) i ) ( 3 ) whereM is a loss function used to measure the performance of the model on the target domains , typically chosen to be cross entropy . Optimising ω is challenging due to the need to backpropagate through the long inner loop optimisation of θ . Existing approaches for learning loss functions typically resort to slow evolutionary or reinforcement learning updates ( Li et al. , 2019a ; Gonzalez & Miikkulainen , 2019 ; Wang et al. , 2020 ) in the outer loop , or to an online approximation based on alternating steps on ω and θ ( Li et al. , 2019a ; Wang et al. , 2020 ) . The latter approach leads to a loss function ω∗ that can not be transferred to new tasks , as it suffers from a short-horizon bias ( Wu et al. , 2018 ) . To solve this problem , we use the Implicit Function Theorem ( IFT ) to compute ∂M∂ω without truncating the inner optimisation problem to approximate θ∗ ( ω ) .

Summary: This paper proposes a loss for domain generalization. The loss is learned by meta-learning on based on RotatedMNIST. The paper evaluate the loss's generalization ability in several other datasets.