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: Authors introduce a bi-level optimization procedure aiming to learn parametric loss functions that result in predictors able to better generalize to new data sources, unseen during training. Specifically, an outer optimization loop iteratively updates a parametric loss function to minimize the empirical risk (under a standard loss) on a left-out data source, not presented to the model during the inner optimization loop, which trains the model on a set of domains using the current version of the parametric loss. The proposed approach is computationally expensive, but authors empirically showed that losses learned in low-scale tasks directly transfer to other cases, rendering the proposal practical.