Year: 2021

Abstract:
Generalising robustly to distribution shift is a major challenge that is pervasive across most real-world applications of machine learning. A recent study highlighted that many advanced algorithms proposed to tackle such domain generalisation (DG) fail to outperform a properly tuned empirical risk minimisation (ERM) baseline. We take a different approach, and explore the impact of the ERM loss function on out-of-domain generalisation. In particular, we introduce a novel meta-learning approach to loss function search based on implicit gradient. This enables us to discover a general purpose parametric loss function that provides a drop-in replacement for cross-entropy. Our loss can be used in standard training pipelines to efficiently train robust models using any neural architecture on new datasets. The results show that it clearly surpasses cross-entropy, enables simple ERM to outperform some more complicated prior DG methods, and provides state-of-the-art performance across a variety of DG benchmarks. Furthermore, unlike most existing DG approaches, our setup applies to the most practical setting of single-source domain generalisation, on which we show significant improvement.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD
text: 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.

heading: 3.1 META-LEARNING LOSSES FOR DG
text: 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 cannot 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 θ∗(ω).

heading: 3.2 IMPLICIT GRADIENT
text: The conceptually simplest way to optimise ω is to store all the intermediate iterates generated by the optimiser when training the network in the inner loop, and to then backpropagate through all of these weight updates (Maclaurin et al., 2015). This becomes prohibitively expensive in both memory and computation. Instead, after finding θ∗(ω) we compute the gradient using the Implicit Function Theorem (IFT). The implicit gradient computation takes advantage of the fact that ∂Lω∂θ = 0, because we have found locally optimal model parameters for the inner problem. The gradient we want to compute is given by
∂M ∂ω = ∂M ∂θ ∂θ ∂ω ∣∣∣∣∣ ω,θ∗(ω) , (4)
and the IFT can be used to obtain ∂θ
∂ω = − [ ∂2Lω ∂θ ∂θT︸ ︷︷ ︸ |θ|×|θ| ]−1 × ∂ 2Lω ∂θ ∂ωT︸ ︷︷ ︸ |θ|×|ω| . (5)
The inverse of the Hessian can be rephrased in terms of a Neumann series,[ ∂2Lω ∂θ ∂θT ]−1 = lim i→∞ i∑ j=0 [I − ∂ 2Lω ∂θ ∂θT ]j , (6)
and approximated by truncating the summation to a finite number of terms. In practice, one can make use of vector-Jacobian products to avoid explicitly constructing the Hessian in the summation. Further details can be found in Lorraine et al. (2020), but we provide pseudo-code for computing the implicit gradient in Algorithm 2.
Algorithm 1 IFT-based loss learning for DG. 1: Input: P, ω 2: Output: ω∗ 3: Init ω 4: while not converged or reached max steps do 5: sample p1, ..., pn from P 6: sample D1, ..., Dn from p1, ..., pn 7: Init H = 0 ∈ Rn×|ω| 8: for all Di do 9: Init θi {Get random network weights} 10: Ds = {D1, ..., Dn}/Di, Dt = Di {Construct source/target splits} 11: θ∗i = argminθ Lω(θi, Ds) {Train the network} 12: hi = Hypergradient(Lω,M, (ω, θ∗i ), α) 13: H[i, :] = hi 14: end for 15: h = grad-surgery(H) 16: ω = ω − ηh {Update the loss function} 17: end while Algorithm 2 Computing the hypergradient of the meta-objectiveM, with respect to the loss ω. The grad(·, ·, ·) function from PyTorch computes a Jacobian-vector product when called with a non-scalar first argument. Inspired by Lorraine et al. (2020), we use this to efficiently compute the Hessian required for approximating the Neumann series. Input: Lω,M, (ω, θ∗), α Output:−p∂ 2Lω ∂θ∂ω v = p = ∂M∂θ |(ω,θ∗) for all j = 1, ..., J do v− = α · grad(∂Lω∂θ , θ, v) p += v end for

heading: 3.3 ROBUST GRADIENT ESTIMATION
text: Algorithm 1 summarizes the gradient estimation procedure. To obtain high quality gradient estimates in each outer loop iteration, we employ a leave-one-domain-out strategy. The DG task, P , used for meta-training the loss parameters has m domains associated with it. In each iteration of the outer loop, we train m networks with the prospective loss (i.e., we instantiate m different copies of the
inner loop), where each network has a different target domain and the remainder of the domains are used to train the network. We can then compute a gradient for each of the m networks and aggregate them together in order to perform an update to the loss parameters. Rather than using the mean gradient, we found aggregation using gradient surgery (Yu et al., 2020), which reduces the gradient noise caused by different source/target domain splits in the inner loop, to work better in practice.

heading: 3.4 TAYLOR POLYNOMIAL REPRESENTATION
text: The choice of loss function parameterisation is a crucial factor in our framework. One must balance the ability to represent a sufficiently broad range of loss functions, with the susceptibility to overfitting the data used to learn the loss. The search space we consider is based on the truncated Taylor polynomials used by the evolutionary optimisation loss learning approach of (Gonzalez & Miikkulainen, 2020). This family of loss functions treats the point around which the Taylor polynomial is centred, and also the value of the derivatives at this point, as learnable parameters. In this sense, it is a variational learning method—though it should be stressed it is not a variational Bayesian method. The family of β order multivariate Taylor polynomials has the form
`(z) = β∑ n=0 1 n! ∇n`(c)T (z− c)n, (7)
where c is a fixed point around which function is being expanded. Because c is fixed, the values of the derivative at this point are also fixed. As such, we can replace c and∇n`(c) with meta-learnable parameters. This allows us to parameterize the learned loss function in terms of the gradients it should have at a meta-learnable point. We define our learnable loss as
Lω(ŷ,y) = 1
C C∑ i=1 `ω(ŷi,yi), `ω(ŷi,yi) = β∑ n=0 1 n! ∇n`([ω0, ω1])T ([ŷi,yi]− [ω0, ω1])n, (8)
where each∇n`([ω0, ω1]) can actually be replaced by introducing more meta-parameters to ω. Please see Appendix A.4 for an expanded definition of this loss function.

heading: 3.5 ALGORITHM SUMMARY
text: Meta-train: Given a set of training domains, the loss function search space in Section 3.4, and efficient update strategy in Section 3.2 and Algorithm 1, we are able to train a robust loss function Lω. We conduct such loss function learning only once using a small dataset, and then evaluate the resulting loss on a variety of larger datasets that are unseen during meta-training. Meta-test: Given the learned loss function Lω∗ , we fix it and use it together with the ERM algorithm for novel DG tasks. Each target problem is trained from scratch and has not been seen during loss learning. An overview of the algorithm, and the learning curve of the meta-train phase, are given in Figure 1.

heading: 4 EXPERIMENTS

heading: 4.1 DATASET AND IMPLEMENTATION DETAILS.
text: Meta-train stage: We aim to learn a general purpose loss function that can be used in diverse DG problems, but we first need to select a dataset for initial loss learning. We chose RotatedMNIST (Ghi-
fary et al., 2015), which contains six different domains that are all derived from MNIST (LeCun & Cortes, 2010) but with different rotations: 0%, 15%, 30%, 45%, 60%, and 75%. The leaveone-domain-out strategy for robust implicit gradient estimation, therefore, results in six inner loop instantiations for each outer loop iteration. For efficiency, we use 2-layer MLPs as the base model, which contains 1024-256-10 units from the input layer to the output one with ReLU as the activation function. The learning rates in the inner loop and outer loop are both 0.01, a batch size of 32 is used for the inner loop, and the Neumann series used for approximating the inverse Hessian is truncated at 15 iterations. The result is a set of 12 parameters that define the learned fourth-order polynomial loss function and an example of learned loss is shown in Appendix A.3. The meta-train compute cost is descripted in Appendix A.6.
Meta-test (deployment) stage: We now evaluate our learned ITL using ERM, but with our learned loss instead of cross entropy. Models trained with our method are denoted by ITL-Net. We evaluate the learned loss function on the four common DG benchmarks: VLCS (Fang et al., 2013), PACS (Li et al., 2018a), OfficeHome (Venkateswara et al., 2017), and Terra Incognita (Beery et al., 2018). Two sets of experiments are conducted: (i) We evaluate the conventional PACS benchmark, as it is the most widely used in the DG literature, and enables comparison against the most recent state-of-the-art competitors. (ii) We evaluate all four benchmarks using the recent DomainBed platform, which is designed to enforce fair and consistent hyperparameter tuning across different methods.

heading: 4.2 RESULTS
text: PACS: Setup A pre-trained ResNet18 backbone is used throughout, together with the source and target domain split described in (Li et al., 2018a). We train ResNet-18 with ITL on the training split and perform model selection using the validation set. We use same set of hyperparameters (learning rate of 0.001, weight decay of 0.00001, and batch size of 32 for each domain) for the baseline ERM with Cross-Entropy (ERM +CE) and Binary Cross-Entropy (BCE) to train our model. We compare ITL-Net with the existing state-of-the-art methods on this benchmark, including RSC (Huang et al., 2020), data augmentation-based L2A-OT (Zhou et al., 2021b), Mixstyle (Zhou et al., 2020) including random shuffle (rs) and domain label (dl), regulariser-based Entropy (Zhao et al., 2020), adversarial gradient-based CrossGrad (Shankar et al., 2018), meta learning-based MASF (Dou et al., 2019) and Epi-FCR (Li et al., 2019b) and self-supervision-based JiGen (Carlucci et al., 2019).
PACS: Results We first conduct experiments using the classic PACS protocol to facilitate comparison against many recent competitors that were not evaluated on DomainBed. Table 1 compares our ITL-Net performance vs state-of-the-art methods. From the results we can see that: (i) Simply swapping out the loss in ERM from CE to ITL, leads to a significant 5.5% improvement. (ii) Overall our ITL-Net leads to state-of-the-art performance on this benchmark, surpassing the most recent and sophisticated competitors such as Mixstyle.
DomainBed: Setup We next evaluate ITL using the DomainBed platform, which enforces careful and fair evaluation by ensuring that all competitors use the same hyper-parameter tuning strategy
(random search, driven by source domain validation performance), and the same number of hyperparameter search iterations. We follow the standard DomainBed protocol and use a ResNet-50, with experiments conducted on VLCS, PACS, OfficeHome, and Terra Incognita.
DomainBed: Results The results in Table 2 compare ITL-Net with ERM and some of the most competitive published alternatives: SagNet (Nam et al., 2019), CORAL (Sun & Saenko, 2016), CDANN (Li et al., 2018c) and RSC (Huang et al., 2020) in the original DomainBed paper (Gulrajani & Lopez-Paz, 2021). A detailed comparison is given in Appendix 6. The conclusion of the DomainBed study was that existing methods did not reliably beat ERM under this hyperparameter tuning protocol. In contrast, we can see that ITL-Net provides a clear improvement on ERM and matches or improves on the strongest existing competitor in each case, especially on TerraIncognita and OfficeHome. To formally compare ITL-Net with competitors, we perform significance testing using the Wilcoxon signed-rank test, where the p-value is set as 0.025 and the sample size is number of the Domain datasets applied. For example, when comparing the performance of ITL-Net and ERM, the null hypothesis (H0) is that ITL-Net has equal performance to ERM on DomainBed and the alternative hypothesis (H1) that ITL-Net is statistically significantly better than ERM on DomainBed. The results of the hypothesis tests comparing ITL to each competitor are summarized at the bottom of Table 2, and confirm that ITL-Net outperforms them all.
Single Source DG: Setup Most existing DG methods rely on the availability of multiple source domains in some form: For example to synthesise new domains for data augmentation (Zhou et al., 2021b), or perform feature alignment among training domains (Sun & Saenko, 2016). A unique feature of our ITL-Net is that, since it is only a small modification to ERM, it can be used to learn on a single source domain. Although this setting is not well explored in the literature, it is obviously highly practical as multiple source domains are often not available in practice. To explore this setting, we modify the DomainBed benchmark to train on a single source at a time and average over each source→target combination, rather than training on the conjunction of all sources. Single Source DG: Results From the results in Table 4, we can see that performance drops across the board compared to multi-source training (Table 2), as expected. However, the state of the art alternatives CORAL and SagNet are no longer as competitive compared to ERM as they were in the multi-source case (Table 2)—this is expected as they are designed to exploit cues from multiple source domains. In contrast, our ITL-Net maintains a clear lead over the conventional ERM with cross entropy baseline in this setting. This is a significant achievement as existing work has not produced algorithms that improve robustness under the single-source setting.

heading: 4.3 FURTHER ANALYSIS
text: Qualitative comparison and impact on learning To understand our learned loss, we visually compare it Figure 5 with several other loss functions: CE, SCE (Wang et al., 2019), Label Smoothing (LS) (Pereyra et al., 2017), and FOCAL (Mukhoti et al., 2020). Compared to the standard cross entropy loss, we can see that ITL has softer penalties for severe misclassification, and stronger penalties for moderate misclassification. In addition, it can be noticed that FOCAL has the most similar shape to our ITL. The properties of the FOCAL loss for probability calibration were recently studied in (Mukhoti et al., 2020). Motivated by this, we report the evolution of the target-domain entropy of correctly and incorrectly classified samples during training in Figure 2(left) and Figure 2(right) respectively. Clearly, a byproduct of ITL compared to CE and SCE is a (desirable) increase in uncertainty for misclassified instances. FOCAL has similar behaviour to ITL in general, especially the entropy of prediction distributions for misclassified instances.
Quantitative comparison to other robust losses We next investigate whether the good performance of ITL in cross-domain robustness can be easily replicated by applying existing robust loss functions, or whether our meta-learning pipeline has learned something new in terms of robust model training. SCE (Wang et al., 2019) and GCE (Zhang & Sabuncu, 2018) were designed with label-noise robustness in mind, while Focal (Mukhoti et al., 2020; Lin et al., 2017) was designed for class imbalance and calibration. Label-smooth (LS) (Pereyra et al., 2017) is for improving generalisation and reducing overconfidence. From the results in Table 3, we can see that while some losses improve on CE, ITL leads to the clearest improvement. The somewhat similar behaviour of ITL and FOCAL in terms of induced entropy from Figure 2 does not carry over similar cross-domain recognition accuracy. Note that all experiments in Table 3 were run by us, while competitor performance in Table 2 is taken from (Gulrajani & Lopez-Paz, 2021).
Loss landscape analysis and Perturbation analysis One of the factors that affect model generalisation is the loss landscape at convergence. (Keskar et al., 2017; Chaudhari et al., 2019) observed that flatter loss landscapes lead to good generalisation of the learned model. To this end, we compared a 1D slice through the loss landscape of ITL-Net with that of ERM on both source domain and target domains. Namely, we perturb the converged paraeters by moving it around though gradient direction which generated by the evigenvector of the Hessian matrix . From Figure 4, we can see that for each held out target domain, ITL-Nets have flatter loss landscapes compared with models trained by CE with respect to the source domains. To further analyse the quality of the minimas provided by CE
and ITL, we follow the perturbation analysis routine in (Keskar et al., 2017; Zhang et al., 2018) by adding multiplicative noise to the weights of the converged models. From the results in Figure 3, it can be see that ITL-Nets outperform ERM at every noise ratio.
Repeatability analysis Our message thus far is that a single loss produced by our pipeline can be re-used as a plug-and-play modification to improve vanilla ERM+CE on a wide variety of held-out downstream DG tasks. That said, one might reasonably wonder about the reliability of the loss function learning procedure itself. To investigate this, we repeat our entire pipeline including the meta-train stage five times. We then evaluate the consistency of the resulting five loss functions on the downstream ColoredMNIST task. Furthermore, to evaluate the dependence of our result on the choice of meta-training dataset, we repeat the above experiment on RotatedKMNIST (Clanuwat et al., 2018) to replace the RotatedMNIST used previously. From the results in Table 5, we can see that performance is quite consistent over trials (small standard deviation). It also differs little with choice of pre-training dataset - with both options performing well compared to competitors in Table 2.
Limitations ITL improves ERM+CE for DG tasks in general, but in some cases, the margin over SOTA is small, since other state-of-the-art competitors may beat ERM+CE. It would be more interesting if ITL is highly complementary to SOTA methods based on other architectural, augmentation, or domain-alignment improvements – but this remains for future work to determine.

heading: 5 CONCLUSION
text: We provided the first study of the effect of ERM loss functions on Domain Generalisation. We observe that models trained by ERM with existing robust loss functions can improve performance on Domain Generalisation compared with those trained by Cross-Entropy. To discover the best loss for DG, we perform meta-learning to find a re-usable white-box loss function. This is tractably solved using IFT to obtain gradients of the target domain performance with respect to the source domain loss parameters. This also provides the first demonstration of IFT-based loss learning in the literature. The results show that a simple modification to the standard ERM pipeline improves both multi-source and single source DG, and even surpasses the purpose-designed state of the art models.

heading: A APPENDIX

heading: A.1 THE DETAILED RESULTS FOR DOMAINBED
text: In Table 2 of the main paper, we summarized performance across held out target domains for the standard multi-source DG problem. In Table 6 we now give the detailed results of ITL-Net for each target domain.

heading: A.2 DETAILED RESULTS FOR SINGLE SOURCE DOMAIN EXPERIMENT
text: In Table 4 of the main paper, we reported single-source DG results, summarising performance across choice of source domain and over all target domains. In Table 7 we now give the detailed results of ITL-Net for each choice of source and target domain.

heading: A.3 THE LEARNED LOSS FUNCTION
text: Recall that our experimental design trained a single loss function on RotatedMNIST, and then evaluated it from different perspectives across all our main experiments. The specific parameters of our learned ITL (as visualised in Fig. 5(right)) are given in Table A.3. Then reader can plug-and-play for their own Domain Generalisation problems.
ITL -2.0193, -1.2234, 0.1363, 0.1269, -0.4566, -0.1016, -0.2545, 1.0971, -0.9203, 0.2368, 0.4795, 0.9975

heading: A.4 THE PARAMETERISATION OF THE LEARNABLE LOSS FUNCTION
text: We apply fourth order bi-variate Taylor polynomial to parameterise the learnable loss function. The terms only contain yi are removed from the polynomial since these do not generate gradients with respect to the prediction of the network. The final form, only containing 12 learnable parameters, is given as:

heading: A.5 ROBUST LOSS FUNCITON ILLUSTRATION

heading: A.6 META-TRAIN COMPUTE COST
text: Due to the efficiency of implicit gradient, training our ITL required using PyTorch (Paszke et al., 2017) only required 8 hours on a single V100 GPU to complete 200 gradient descent steps on ω. While the goals and base learning problems are not directly comparable, this is dramatically faster than alternatives that require an entire cluster (Li et al., 2019a), and where even very recent fast methods require about 12 GPU-days (Liu et al., 2021).

heading: A.7 META-TRAINING CONVERGENCE
text: Figure 6 shows the convergence of ITL during meta-training on R-MNIST. The x-axis shows outer loop iterations/loss function updates. The lines graph (i) the inner loop accuracy (mean over all source domains) after model training with the current loss function, and (ii) and the outer loop accuracy (held out domain accuracy). The convergence process is quite smooth.

heading: A.8 FULL LOSS LANDSCAPE PLOTS
text: Figure 7 shows the landscape for all four domains, of which two were shown in Figure 4 of the main manuscript.

Review summaries:
This paper proposed a novel viewpoint in domain generalization, i.e, the loss function search. Specifically, the search procedure is decomposed as a bi-level optimization and solved through implicit function. The author later adopted the Neumann series for memory saving. The extensive empirical results validate the correctness of the proposed approach.  ====== Final Update after rolling discussions  I would appreciate the authors for the detailed responses and I have read the rebuttal. Unfortunately, the concern of rigor is still not addressed after discussion and I will maintain the same rating. Below are my *final* remarks (I fully understand your high-level idea), I hope the author will carefully rethink the paper and improve it in the future version.  - In the rebuttal and paper, many terms are not mathematically/rigorously defined, which made me rather confused. These terms even made me doubtful about the generalization properties since they should be rigorously defined and discussed. (see Note 1 for the detailed examples)  - The motivation for choosing the specific loss parameter families is not properly addressed in the paper. From the rebuttal, I finally got the logic/motivation but it was too late. I do think a major revision for fully justifying the motivation can greatly improve the paper. (see Note 2 for the detailed discussions)  ---------------------  Note 1.  Examples of unclear terms, most of them are not formally formulated/defined. The descriptive words rather than the math formula made the reviewer confused and a bit doubtful about the generalization property since generalization should be rigorously defined.    - "Scale of the domain shift is similar enough between meta-train and deployment datasets." What is the formal definition of the scale of domain shift?   -  “the cardinality of the classification problem is not too different between meta-train and meta-test.” What is the formal definition of the cardinality of classification?    - “But please note that something like MNIST->SVHN transfer is not exactly the problem setting we address.” It is not domain generalization (trained on digit A and deployed on digits B)? What is the formal (or mathematic) definition of a single source setting? It is best to use equations rather than descriptive words to depict the problem settings. In fact, the single source is a confusing word (since it is not mathematically defined.)   -  “large search space is expected to lead to meta-overfitting.” What is the formal definition of meta-overfitting? Overfitting w.r.t. which distribution? The author should be cautious to use unclear words since meta-overfitting has not been mathematically defined.    Note 2. The motivation of choosing the specific loss parameter.  Thanks for the rebuttal. It makes me clear about the motivation. But it requires a major revision of your paper to fully justify it. For example, I would suggest testing the linear combination and showing its results. Also exploring other loss families to fully justify the value of the loss search. In the current version, it is more or less we proposed loss family A (without strong motivation or benefits) and it works in some classification benchmark.  A *comprehensive and systematic* empirical understanding can significantly improve it.    ------------------------------------------------------- 

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.

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.