L OSS F UNCTION L EARNING FOR D OMAIN G ENERALIZATION BY I MPLICIT G RADIENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 efﬁciently 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 signiﬁcant improvement.
I NTRODUCTION
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), overﬁtting (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 deﬁne 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 modiﬁcation 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-speciﬁc—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 ﬂexible enough to include interesting new losses, but simple enough to generalise across tasks without overﬁtting 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 ﬁrst 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 ﬁrst study on the signiﬁcance of supervised loss function choice in DG (ii)
We demonstrate the ﬁrst efﬁcient 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.
R ELATED W ORK
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 efﬁciency (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 speciﬁc (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 ofﬂine (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, efﬁcient 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 ofﬂine/long inner loop meta-learning.
To make this meta-optimisation tractable, we exploit the Implicit Function Theorem, to efﬁciently generate accurate hypergradients of the validation domain performance with respect to the training domain loss function parameters.
Besides being the ﬁrst demonstration of loss learning for DG, to our knowledge it is also the ﬁrst demonstration of any implicit gradient-based loss learning.
M ETHOD
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
[Equation]
∼ p
[Equation] [Equation]
} Kj =1 , and use the union of all these sets as the full training set, D ( s ) = (cid:83) ni
=1 D ( s ) i .
Empirical Risk Minimisation (ERM) then simply ﬁnds the model parameters, θ , that minimise the loss measured on this training set,
[Equation]
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 ) =
(cid:83) mi =1 D ( t ) i .
Standard evaluation metrics such as accuracy can then be computed using this data.
3. M ETA -L EARNING L OSSES 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 ﬁnd 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 efﬁciently 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
[Equation]
where M 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 θ ∗ ( ω ) .
3. I MPLICIT G RADIENT
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 ﬁnding θ ∗ ( ω ) 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
[Equation]
and the IFT can be used to obtain
[Equation]
The inverse of the Hessian can be rephrased in terms of a Neumann series,
[Equation]
and approximated by truncating the summation to a ﬁnite 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.
but we provide pseudo-code for computing the implicit gradient in Algorithm 2.
1: Input
2: Output
3: Init
4: while not converged or reached max steps do 5: sample p 1 , ..., p n from P 6: sample D 1 , ..., D n from p 1 , ..., p n 7:
Init H = 0 ∈ R n ×| ω | 8: for all D i do 9:
Init θ i {Get random network weights} 10:
D s =
{ D 1 , ..., D n } /D
i ,
D t = D i {Construct source/target splits} 11: θ ∗ i =
arg
min θ L ω ( θ i , D s ) {Train the network} 12: h i =
Hypergradient ( L ω , M , ( ω, θ ∗ i ) , α ) 13:
H [ i, :]
= h i 14: end for 15: h = grad-surgery ( H ) 16: ω = ω − ηh {Update the loss function} 17: end while
PyTorch computes a Jacobian-vector product when called with a non-scalar ﬁrst argument.
Inspired by Lorraine et al.
(2020), we use this to efﬁciently compute the Hessian required for approximating the Neumann series.
Input:
[Equation]
Output:
[Equation]
p
[Equation]
for all j = 1 , ..., J do
[Equation]
end for
3. R OBUST G RADIENT E STIMATION
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
[Figure]
𝜃 ∗ 𝜔 ∗ = arg min
# 𝑀(𝑓 " ∗ (#) )
Source
Target
Deploy Loss ℒ # ∗ 𝜃 ∗ = arg min " ℒ #
Inner Loop Meta-test 𝜃 ∗
Source
Target 𝑓 " ∗
Guitar 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.
3. T AYLOR P OLYNOMIAL R EPRESENTATION
The choice of loss function parameterisation is a crucial factor in our framework.
One must balance the ability to represent a sufﬁciently broad range of loss functions, with the susceptibility to overﬁtting 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
[Equation]
(7) where c is a ﬁxed point around which function is being expanded.
Because c is ﬁxed, the values of the derivative at this point are also ﬁxed.
As such, we can replace c and ∇ n (cid:96) ( 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 deﬁne our learnable loss as
[Equation]
where each ∇ n (cid:96) ([ ω 0 , ω 1 ]) can actually be replaced by introducing more meta-parameters to ω .
Please see Appendix A.4 for an expanded deﬁnition of this loss function.
3. A LGORITHM S UMMARY
Meta-train:
Given a set of training domains, the loss function search space in Section 3.4, and efﬁcient 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 ﬁx 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.
E XPERIMENTS 4. D ATASET AND IMPLEMENTATION DETAILS .
Meta-train stage: We aim to learn a general purpose loss function that can be used in diverse DG problems, but we ﬁrst need to select a dataset for initial loss learning.
We chose RotatedMNIST (Ghi-
[Table]
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 efﬁciency, 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 deﬁne 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), OfﬁceHome (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
. 4. R ESULTS 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 shufﬂe (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 ﬁrst 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 signiﬁcant 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
[Table] [Table]
(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, OfﬁceHome, 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 OfﬁceHome.
To formally compare ITL-Net with competitors, we perform signiﬁcance 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 ( H 0 ) is that ITL-Net has equal performance to ERM on DomainBed and the alternative hypothesis ( H 1 ) that ITL-Net is statistically signiﬁcantly 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 conﬁrm 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 modiﬁcation 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 signiﬁcant achievement as existing work has not produced algorithms that improve robustness under the single-source setting.
[Table]
[Figure]
4. F URTHER A NALYSIS
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)
[Figure]
and FOCAL (Mukhoti et al., 2020).
Compared to the standard cross entropy loss, we can see that ITL has softer penalties for severe misclassiﬁcation, and stronger penalties for moderate misclassiﬁcation.
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 classiﬁed 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 misclassiﬁed instances.
FOCAL has similar behaviour to ITL in general, especially the entropy of prediction distributions for misclassiﬁed 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 overconﬁdence.
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 ﬂatter 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 ﬂatter 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
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
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.
[Figure]
Repeatability analysis
Our message thus far is that a single loss produced by our pipeline can be re-used as a plug-and-play modiﬁcation 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
[Figure]
we repeat our entire pipeline including the meta-train stage ﬁve times.
We then evaluate the consistency of the resulting ﬁve 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.
C ONCLUSION
We provided the ﬁrst 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 ﬁnd 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 ﬁrst demonstration of IFT-based loss learning in the literature.
The results show that a simple modiﬁcation to the standard ERM pipeline improves both multi-source and single source DG, and even surpasses the purpose-designed state of the art models.
R
A A PPENDIX A. T HE DETAILED RESULTS FOR D OMAIN B ED
In Table 2 of the main paper
[Table]
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
[Table] [Table]
A. D ETAILED R ESULTS FOR S INGLE S OURCE D OMAIN E XPERIMENT
In Table 4 of the main paper
[Table]
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.
[Table]
Sourceset
Artistic
Clipart
[Table]
A. T HE LEARNED LOSS FUNCTION
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 speciﬁc 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.
[Table]
A. T HE PARAMETERISATION OF THE LEARNABLE LOSS FUNCTION
We apply fourth order bi-variate Taylor polynomial to parameterise the learnable loss function.
The terms only contain y i are removed from the polynomial since these do not generate gradients with respect to the prediction of the network.
The ﬁnal form, only containing 12 learnable parameters, is given as:
[Equation]
A. R OBUST LOSS FUNCITON ILLUSTRATION
[Figure] [Figure]
A. M ETA - TRAIN COMPUTE COST
Due to the efﬁciency 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)
[Figure]
and where even very recent fast methods require about 12 GPU-days (Liu et al., 2021).
A. M ETA -T RAINING CONVERGENCE
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.
A. F ULL LOSS LANDSCAPE PLOTS
Figure 7 shows the landscape for all four domains, of which two were shown in Figure 4 of the main manuscript.
[Figure] [Figure]