G RADIENT
M
ATCHING FOR D OMAIN G ENERALIZA -
TION
Anonymous authors
Paper under double-blind review
A BSTRACT
Machine learning systems typically assume that the distributions of training and test sets match closely.
However, a critical requirement of such systems in the real world is their ability to generalize to unseen domains.
Here, we propose an inter-domain gradient matching objective that targets domain generalization by maximizing the inner product between gradients from different domains.
Since direct optimization of the gradient inner product can be computationally prohibitive — it requires computation of second-order derivatives —- we derive a simpler ﬁrst-order algorithm named Fish that approximates its optimization.
We perform experiments on the W ILDS benchmark, which captures distribution shift in the real world, as well as the D OMAIN B ED benchmark that focuses more on syntheticto-real transfer.
Our method produces competitive results on both benchmarks, demonstrating its effectiveness across a wide range of domain generalization tasks.
Anonymized code is available at this url 1 .
I NTRODUCTION
[Figure]
The goal of domain generalization is to train models that performs well on unseen, out-of-distribution data, which is crucial in practice for model deployment in the wild.
This seemingly difﬁcult task is made possible by the presence of multiple distributions/domains at train time.
As we have seen in past work (Arjovsky et al., 2019; Gulrajani and Lopez-Paz, 2020; Ganin et al., 2016), a key aspect of domain generalization is to learn from features that remain invariant across multiple domains, while ignoring those that are spuriously correlated to label information (as deﬁned in Torralba and Efros (2011); Stock and Cisse (2017)).
Consider, for example, a model that is built to distinguish between cows and camels using photos collected in nature under different climates.
Since CNNs are known to have a bias towards texture (Geirhos et al., 2018; Brendel and Bethge, 2019), if we simply try to minimize the average loss across different domains, the classiﬁer is prone to spuriously correlate “cow” with grass and “camels” with desert, and predict the species using only the background.
Such a classiﬁer can be rendered useless when the animals are placed indoors or in a zoo.
However, if the model could recognize that while the landscapes change with climate, the biological characteristics of the animals (e.g. humps, neck lengths) remain invariant and use those features to determine the species, we have a much better chance at generalizing to unseen domains.
Similar intuitions have already motivated many approaches that consider learning “invariances” across domains as the main challenge of domain generalization.
Typically, a lot of these work focus on learning invariant representations directly by removing the domain information (Ganin et al., 2016; Sun and Saenko, 2016; Li et al., 2018).
In this work, we propose an inter-domain gradient matching (IDGM) objective.
Instead of learning invariant features by matching the distributions of representations from different domains, our approach does so by encouraging consistent gradient directions across domains.
Speciﬁcally, our IDGM objective augments the loss with an auxiliary term that maximizes the gradient inner product between domains, which encourages the alignment between the domain-speciﬁc gradients.
By simultaneously minimizing the loss and matching the gradients, IDGM encourages the optimization paths to be the same for all domains, favouring invariant predictions.
Figure 1 illustrates a motivating example described in Section 3.2: given 2 domains, each containing one invariant feature (orange cross) and one spurious feature (yellow and red cross).
While empirical risk minimization (ERM) minimizes the average loss between these domains at the cost of learning spurious features only, IDGM aligns the gradient directions and is therefore able to focus on the invariant feature.
While the IDGM objective achieves the desirable learning dynamic in theory, naive optimization of the objective by gradient descent is computationally costly due to the second-order derivatives.
Leveraging the theoretical analysis of Reptile, a meta-learning algorithm (Nichol et al., 2018), we propose to approximate the gradients of IDGM using a simple ﬁrst-order algorithm, which we name
Fish.
Fish is simple to implement, computationally effective and as we show in our experiments, functionally similar to direct optimization of IDGM.
Our contribution is a simple but effective algorithm for domain generalization, which exhibits state-ofthe-art performance on 13 datasets from recent domain generalization benchmark W ILDS (Koh et al., 2020) and D OMAIN B ED (Gulrajani and Lopez-Paz, 2020).
The strong performance of our method on a variety of datasets demonstrates that it is broadly applicable in different applications/subgenres of domain generalization tasks.
We also perform detailed analysis in Section 4.4 to explain the effectiveness of our proposed algorithm.
R ELATED W ORK
Domain Generalization In domain generalization, the training data is sampled from one or many source domains, while the test data is sampled from a new target domain.
We will now discuss the ﬁve main families of approaches to domain generalization:
Distributional Robustness (DRO) : DRO approaches minimize the worst-case loss over a set of data distributions constructed from the training domains.
Rojas-Carulla et al.
(2015) proposed DRO to address covariate shift (Gretton et al., 2009a;b), where P ( Y | X ) remains constant across domains but P ( X ) changes.
Later work also studied subpopulation shift , where the train and test distributions are mixtures of the same domains, but the mixture weights change between train and test (Hu et al., 2018; Sagawa et al., 2019); 2.
Domain-invariant representation learning : This family of approaches to domain generalization aims at learning high-level features that make domains statistically indistinguishable.
Prediction is then based on these features only.
The principle is motivated by a generalization error bound for unsupervised domain adaptation (Ben-David et al., 2010; Ganin et al., 2016), but the approach readily applies to domain generalization (Gulrajani and Lopez-Paz, 2020; Koh et al., 2020).
Algorithms include penalising the domain-predictive power of the model (Ganin et al., 2016; Wang et al., 2019; Huang et al., 2020), aligning domains through contrastive loss (Motiian et al., 2017), matching mean and variance of feature distributions across domains (Sun and Saenko, 2016), learning useful representations by solving Jigsaw puzzles (Carlucci et al., 2019), using the maximum mean discrepancy to match the feature distributions (Li et al., 2018b) or introducing training constraints across domains using mixup formulation (Yan et al., 2020).
Invariant Risk Minimization (IRM) : IRM is proposed by Arjovsky et al.
(2019), which learns an intermediate representation such that the optimal classiﬁers (on top of this representation) of all domains are the same.
The motivation is to exploit invariant causal effects between domains while reducing the effect of domain-speciﬁc spurious correlations.
From an optimization perspective, when IRM reaches its optimal, all the gradients (for the linear classiﬁer) has to be zero .
This is why IRM’s solution won’t deviate from ERM when ERM is optimal for every domain, which is not the case for our proposed IDGM objective due to the gradient inner product term.
Data augmentation : More recently, approaches that simulates unseen domains through speciﬁc types of data augmentation/normalization has been gaining traction.
This includes work such as Zhou et al.
; Volpi and Murino (2019); Ilse et al.
(2021), as well as Seo et al.
(2019) which utilises ensemble learning.
Gradient alignment : Two concurrent work – Koyama and Yamaguchi (2021) and Parascandolo et al.
(2021) – utilise similar gradient-alignment principle for domain generalization.
Koyama and Yamaguchi (2021) proposes IGA, which learns invariant features by minimizing the variance of inter-domain gradients.
The key difference between IGA and our objective is that IGA is completely identical to ERM when ERM is the optimal solution on every training domain, since the variances of the gradients will be zero.
While they achieve the best performance on the training set, both IGA and ERM could in some cases, completely fail when generalizing to unseen domains (see Section 3.2 for such an example).
Our method, on the contrary, biases towards non-ERM solutions as long as the gradients are aligned, and is therefore able to avoid this issue.
Parascandolo et al.
(2021) on the other hand, proposes to mask out the gradients that have opposite signs for different domains.
Unlike their work that prunes gradients that are inconsistent, our approach actively encourage gradients from different domains to be consistent by maximizing the gradient inner product.
Additionally, in Lopez-Paz and Ranzato (2017) we also see the application of gradientalignment, however in this case it is applied under the continual learning setting to determine whether a gradient update will increase the loss of the previous tasks.
Apart from these algorithms that are tailored for domain generalization, a well-studied baseline in this area is ERM, which simply minimizes the average loss over training domains.
Using vanilla
ERM is theoretically unfounded (Hashimoto et al., 2018; Blodgett et al., 2016; Tatman, 2017) since
ERM is guaranteed to work only when train and test distributions match.
Nonetheless, recent benchmarks suggest that ERM obtains strong performance in practice, in many case surpassing domain generalization algorithms (Gulrajani and Lopez-Paz, 2020; Koh et al., 2020).
Our goal is to ﬁll this gap, using an algorithm signiﬁcantly simpler than previous approaches.
Connections to meta-learning There are close connections between meta-learning (Thrun and
Pratt, 1998) and (multi-source) domain adaptation.
In fact, there are a few works in domain generalization that are inspired by the meta-learning principles, such as Li et al.
(2018a)
; Balaji et al.
; Li et al.
Dou et al.
In Ren et al.
(2018), we also see the leveraging of gradient inner product in meta-learning, where it is used to determine the importance weight of training examples.
We discuss the connection between our proposed algorithm to meta-learning in more details in Appendix A.1.
Note that our proposed algorithm Fish is similar to the Mean Teacher method (Tarvainen and Valpola, 2017), where a teacher model (equivalent to θ in Algorithm 1) is computed using a moving average of the student model (equivalent to ˜ θ in Algorithm 1).
However with Fish, the "student model" is replaced by the "teacher model" at the beginning of each inner-loop, whereas with Mean Teacher the two models are kept separated.
M ETHODOLOGY 3. G OALS
Consider a training dataset D tr consisting of S domains D tr = {D 1 , · · · , D S } , where each domain s is characterized by a dataset D s :
= { ( x s i , y s i ) } n s i =
1 containing data drawn i.i.d.
from some probability distribution.
Also consider a test dataset D te consisting of T domains D te =
{D S +1 , · · · , D S + T } , where D tr ∩ D te = ∅ .
The goal of domain generalization is to train a model with weights θ that generalizes well on the test dataset D te such that:
[Equation]
where l (( x, y ); θ ) is the loss of model θ evaluated on ( x, y ) .
A naive approach is to apply ERM, which simply minimizes the average loss on D tr , ignoring the discrepancy between train and test domains:
[Equation]
The ERM objective does not exploit the invariance across different domains in D tr and could perform arbitrarily poorly on test data.
We demonstrate this effect with the following simple linear example
. 3. T HE PITFALL OF ERM: A LINEAR EXAMPLE
Consider a binary classiﬁcation setup where data ( x, y ) ∈
B 4 × B , and a data instance is denoted x = [ f 1 , f 2 , f 3 , f 4 ] , y .
The train domains are {D 1 , D 2 } , and test domain is D 3 .
The goal is to learn a linear model
Wx + b = y, W ∈ R 4 , b ∈ R on the train data, such that the error on the test domain is minimized.
The setup and dataset of this example is illustrated in Figure 2.
As we can see in Figure 2, f 1 is the invariant feature in this dataset, since the correlation between fand y is stable across different domains.
The relationships between y and f 2 , f 3 and f 4 changes for
D 1 , D 2 , D 3 , making them the spurious features .
Importantly, if we consider one domain only, the spurious features f 2 , f 3 and f 4 are a more accurate indicator of the label than the invariant feature f 1 .
For instance, using f 2 to predict y can give 97% accuracy on D 1 , while using f 1 only achieves 93% accuracy.
[Table]
The performance of ERM on this simple example is shown in Table 1 (ﬁrst row).
From the trained parameters W and b , we see that the model places most of its weights on spurious features f 2 and f 3 .
While this achieves the highest train accuracy ( 97% ), the model cannot generalize to unseen domains and performs poorly on test accuracy ( 57% ).
3. I NTER - DOMAIN G RADIENT M ATCHING (IDGM)
To mitigate the problem with ERM, we need an objective that learns from features that are invariant across domains.
Let us consider the case where the train dataset consists of S = 2 domains
D tr = {D 1 , D 2 } .
Given model θ and loss function l , the expected gradients for data in the two domains is expressed as
[Equation]
The direction, and by extension, inner product of these gradients are of particular importance to our goal of learning invariant features.
If G 1 and G 2 point in a similar direction, i.e. G 1 · G 2 > 0 , taking a gradient step along G 1 or G 2 improves the model’s performance on both domains, indicating that the features learned by either gradient step are invariant across {D 1 , D 2 } .
This invariance cannot be guaranteed if G 1 and G 2 are pointing in opposite directions, i.e. G 1 · G 2 ≤ 0 .
To exploit this observation, we propose to maximize the gradient inner product (GIP) to align the gradient direction across domains.
The intended effect is to ﬁnd weights such that the input-output correspondence is as close as possible across domains.
We name our objective inter-domain gradient matching (IDGM), and it is formed by subtracting the inner product of gradients between domains (cid:98) G from the original ERM objective.
For the general case where S ≥ 2 , we can write
[Equation]
where γ is the scaling term for (cid:98) G .
Note that GIP can be computed in linear time as (cid:98) G
= || (cid:80) i G i
|| 2 − (cid:80)
i ||
G
i || 2 (ignoring the constant factor).
We can also compute the stochastic estimates of Equation (4) by replacing out the expectations over the entire dataset by minibatches.
We test this objective on our simple linear dataset, and report results in the second row of Table 1.
Note that to avoid exploding gradient we use the normalized GIP during training.
The model has lower training accuracy compared to ERM ( 93% ), however its accuracy remains the same on the test set, much higher than ERM.
The trained weights W reveal that the model assigns the largest weight to the invariant feature f 1 , which is desirable.
The visualization in Figure 1 also conﬁrms that by maximizing the gradient inner product, IDGM is able to focus on the feature that is common between domains, yielding better generalization performance than ERM.
3. O PTIMIZING IDGM WITH F ISH
The proposed IDGM objective, although effective, requires computing the second-order derivative of the model’s parameters due to the gradient inner product term, which can be computationally prohibitive.
To mitigate this, we propose a ﬁrst-order algorithm named Fish 2 that approximates the optimization of IDGM with inner-loop updates.
In Algorithm 1 we present Fish.
As a comparison, we also present direct optimization of IDGM using SGD in Algorithm 2.
Algorithm 1 Fish.
1: for iterations = 1 , 2 , ··· do 2: (cid:101) θ ← θ 3: for D i ∈ permute ( {D 1 , D 2 , ··· , D S } ) do 4:
Sample batch d i ∼ D i 5: (cid:101) g i =
E d i (cid:34) ∂l (( x,y ); (cid:101) θ ) ∂ (cid:101) θ (cid:35) //Grad
wrt (cid:101) θ 6:
Update (cid:101) θ ← (cid:101) θ − α (cid:101) g
i 7: end for 8:
Update θ ← θ + (cid:15) ( (cid:101) θ − θ ) 9: end for
1: for iterations = 1 , 2 , ··· do 2: (cid:101) θ ← θ 3: for D i ∈ permute ( {D 1 , D 2 , ··· , D S } ) do 4:
Sample batch d i ∼ D i 5: g i =
E d i (cid:34) ∂l (( x,y ); θ )
∂θ (cid:35) //Grad
wrt θ 6: end for 7: ¯ g = 1 S
S (cid:88) s =
g s , (cid:98) g
=
GIP (batch) (cid:122) (cid:125)(cid:124) (cid:123)
S ( S − 1) i (cid:54) =
j (cid:88) i,j ∈ S g i · g j 8:
Update θ ← θ − (cid:15) (¯ g − γ ( ∂ (cid:98) g/∂θ )) 9: end for
Fish performs S inner-loop ( l3-l7 ) update steps with learning rate α on a clone of the original model (cid:101) θ , and each update uses a minibatch d i from the domain selected in step i .
Subsequently, θ is updated by a weighted difference between the cloned model and the original model (cid:15) ( (cid:101) θ − θ ) .
To see why Fish is an approximation to directly optimizing IDGM, we can perform Taylor-series expansion on its update in l8 , Algorithm 1.
Doing so reveals two leading terms: 1) ¯ g : averaged gradients over inner-loop’s minibatches (effectively the ERM gradient); 2) ∂ (cid:98) g/∂θ : gradient of the minibatch version of GIP.
Observing l8 of Algorithm 2, we see that ¯ g and (cid:98) g are actually the two gradient components used in direct optimization of IDGM.
Therefore, Fish implicitly optimizes
IDGM by construction (up to a constant factor), avoiding the computation of second-order derivative ∂ (cid:98) g/∂θ .
We present this more formally for the full gradient G in Theorem 3.1.
Theorem 3.1 Given twice-differentiable model with parameters θ and objective l .
Let us deﬁne the following:
[Equation]
where ¯ G = 1 S (cid:80) Ss =1 G s and is the full gradient of ERM.
Then we have
[Equation]
Note that the expectation in G f is over the sampling of domains and minibatches.
Theorem 3. indicates that when α is sufﬁciently small, if we remove the scaled ERM gradient component ¯ G from
Fish’s update, we are left with a term G f that is in similar direction to the gradient of maximizing the
GIP term in IDGM, which was originally second-order.
Note that this approximation comes at the cost of losing direct control over the GIP scaling γ — we therefore also derived a smoothed version of Fish that recovers this scaling term, however we ﬁnd that changing the value of γ does not make much difference empirically.
See Appendix B for more details.
The proof to Theorem 3.1 can be found in Appendix A.
We follow the analysis from Nichol et al.
(2018), which proposes Reptile for model-agnostic meta-learning (MAML), where the relationship between inner-loop update and maximization of gradient inner product was ﬁrst highlighted.
Nichol et al.
(2018) found the GIP term in their algorithm to be over minibatches from the same domain , which promoted within-task generalization; in Fish we construct inner-loop using minibatches over different domains – it therefore instead encourages across-domain generalization.
We compare the two algorithms in further details in Appendix A.1.
We also train Fish on our simple linear dataset, with results in Table 1, and see it performs similarly to IDGM – the model assigns the most weight to the invariant feature f 1 , and achieves 93% accuracy on both train and test dataset.
E XPERIMENTS 4. C D S PRITES -N
Dataset We propose a simple shape-color dataset C D S PRITES -N based on the D S PRITES dataset (Matthey et al., 2017), which contains a collection of white 2D sprites of different shapes, scales, rotations and positions.
C D S PRITES -N contains N domains.
The goal is to classify the shape of the sprites, and there is a shape-color deterministic matching that is speciﬁc per domain.
This way we have shape as the invariant feature and color as the spurious feature.
See Figure 3 for an illustration.
To construct the train split of C D S PRITES -N, we take a subset of D S PRITES that contains onlyshapes (square and oval).
We make N replicas of this subset and assign 2 colors to each, with every color corresponding to one shape (e.g. yellow block in Figure 3a, pink → squares, purple → oval).
For the test split, we create another replica of the D S PRITES -N subset, and randomly assign one of the 2 N colors in the training set to each shape in the test set.
We design this dataset with CNN’s texture bias in mind (Geirhos et al., 2018;
Brendel and Bethge, 2019).
If the value of N is small enough, the model can simply memorize the N colors that correspond to each shape, and make predictions solely based on colors, resulting in poor performance on the test set where color and shape are no longer correlated.
Our dataset allows for precise control over the features that remains stable across domains and the features that change as domains change; we can also change the number of domains N easily, making it possible to examine the effect N has on the performance for domain generalization.
Results We train the same model using three different objectives including Fish, dicrect optimization of IDGM and ERM on this dataset with number of domains N ranging from 5 to 50 .
Again, for direct optimization of IDGM, we use the normalized gradient inner product to avoid exploding gradient.
We plot the average train, test accuracy for each objective over 5 runs against the number of domains
N in Figure 4.
We can see that the train accuracy is always 100% for all methods regardless of N (Figure 4a), while the test performance varies: Figure 4b shows that direct optimization of IDGM (red) and Fish (blue) obtain the best performances, with the test accruacy rising to over 90% when
N ≥ 10 and near 100% when N ≥ 20 .
The predictions of ERM (yellow), on the other hand, remain nearly random on the test set up until N = 20 , and reach 95% accuracy only for N ≥ 40 .
This experiment conﬁrms the following: 1) the proposed IDGM objective have much stronger domain generalization capabilities compared to ERM; 2) Fish is an effective approximation of IDGM, with similar performance to its direct optimization.
We also plot the gradient inner product progression of
Fish vs. ERM during training in Figure 9a, showing clearly that Fish does improve the gradient inner product across domain while ERM does not; 3) we also observe during training that Fish is abouttimes faster than directly optimizing IDGM, demonstrating its computational efﬁciency.
[Figure] [Figure]
4. W ILDS
Datasets We evaluate our model on the W ILDS benchmark (Koh et al., 2020), which contains multiple datasets that capture real-world distribution shifts across a diverse range of modalities.
We report experimental results on 6 challenging datasets in W ILDS
[Table]
and ﬁnd Fish to outperform all baselines on most tasks.
A summary of the W ILDS datasets can be found in Appendix C. For hyperparameters including learning rate, batch size, choice of optimizer and model architecture, we follow the exact conﬁguration as reported in the W ILDS benchmark.
Importantly, we also use the same model selection strategy used in W ILDS to ensure a fair comparison.
See details in Appendix D.
[Table]
Results See a summary of results in Table 2, where we use the metrics recommended in W ILDS for each dataset.
Again, following practices in W ILDS , all results are reported over 3 random seed runs, apart from C AMELYON 17 which is reported over 10 random seed runs.
We included additional results as well as a in-depth discussion on each dataset in Appendix C, and an ablation studies on
Fish’s hyperparameters in Appendix F and Appendix
E. We make the following observations:
Strong performance across datasets: Considering results on all 6 datasets, Fish is the best performing algorithm on W ILDS .
It outperforms all baseline on 4 datasets and achieves similar level of performance to the best method on the other 2 (A MAZON and I W ILD C AM ).
Fish’s strong performance on different types of data and architectures such as R ES N ET (He et al., 2016), D ENSE N ET (Huang et al., 2017) and BERT (Devlin et al., 2018) demonstrated it’s capability to generalize to a diverse variety of tasks; 2.
Strong performance on different domain generalization tasks: We make special note the C IVIL C OMMENTS dataset captures subpopulation shift problems, where the domains in test are a subpopulation of the domains in train, while all other W ILDS datasets depicts pure domain generalization problems, where the domains in train and test are disjointed.
As a result, the baseline models for C IVIL C OMMENTS selected by the W ILDS benchmark are different from the methods used in all other datasets, and are tailored to avoiding systematic failure on data from minority subpopulations.
We see that Fish works well in this setting too without any changes or special sampling strategies (used for baselines on C IVIL C OMMENTS , see more in Table 10), demonstrating it’s capability to perform in different domain generalization scenarios; 3.
Failure mode of domain generalization algorithms: We noticed that on I W ILD C AM and A MAZON , ERM is the best algorithm, outperforming all domain generalization algorithms except for Fish on A MAZON .
We believe that these domain generalization algorithms failed due to the large number of domains in these two datasets — 324 for I W ILD C AM and 7,676 for A MAZON .
This is a common drawback of current domain generalization literature and is a direction worth exploring.
[Table]
4. D OMAIN B ED
Datasets While W ILDS is a challenging benchmark capturing realistic distribution shift, to test our model under the synthetic-to-real transfer setting and provide more comparisons to SOTA methods, we also performed experiments on the D OMAIN B ED benchmark (Gulrajani and Lopez-Paz, 2020).
See a summary of D OMAIN B ED in Appendix H.
Results Following recommendations in D OMAIN B ED , we report results using training domain as validation set for model selection.
See results in Table 3, reported over 5 random trials.
Averaging the performance over all 7 datasets, Fish ranks second out of 10 domain generalization methods.
It performs only marginally worse than Coral (0.1 % ), and is one of the three methods that performs better than ERM.
This showcases Fish’s effectiveness on domain generalization datasets with stronger focus to synthetic-to-real transfer, which again demonstrates its versatility and robustness on different domain generalization tasks.
4. A NALYSIS
We show extensively through empirical evaluation that Fish is very effective for a variety domain generalization tasks.
In this section, we perform analysis to validate that Fish’s strong performance is due to inter-domain gradient inner product maximization.
Does Fish maximize gradient inner product (GIP) empirically?
In Figure 5, we plot the progression of GIP during training using different objectives.
We train both Fish (blue) and ERM (yellow) on
C D S PRITES -N until convergence while tracking the normalized GIP between minibatches from different domains used in each inner-loop.
To ensure a fair comparison, we use the exact same sequence of data for Fish and ERM (see Appendix I for more details).
From Figure 5, it is clear that during training, the normalized GIP of Fish increases, while that for ERM stays at the same value.
The observations here shows that Fish is indeed effective in increasing/maintaining the level of inter-domain GIP.
We conduct the same GIP tracking experiments for the W ILDS datasets we studied as well to shed some lights on its efﬁciency — see Appendix G for results.
Does maximizing GIP between random batches help domain generalization?
In this part of the analysis, we examine the effectiveness of another element of our algorithm – the construction of minibatches.
We conduct experiments where data are grouped randomly instead of by domain for the inner-loop update.
By doing so, we are still maximizing the inner product between minibatches, however not strictly between domain.
We therefore expect the results to be slightly worse than Fish, and the bigger the domain gap is, the more advantage Fish has against the random grouping strategy.
We show the results for random grouping ( Fish, RG ) in Table 4.
As expected, the random grouping strategy performs worse than Fish on all datasets.
The experiment demonstrated that our algorithm does beneﬁt from the domain grouping strategy, and that maximising GIP between random batches of data, while still achieving strong results, does not achieve the same domain generalization performance as Fish.
[Table]
Differences between Fish and ERM We have shown through empirical evaluation that Fish and
ERM are sufﬁciently different.
Most notably, they converge to completely distinct minimum in our linear toy example and C D SPRITES experiment, resulting in large decrepencies in test accuracy.
However, some may notice that Fish as a ﬁrst-order method shares algorithmic similarities to ERM.
In fact
[Table]
if we ﬁx the meta learning rate of Fish (cid:15) = 1 and use domain sampling for ERM, the two algorithms are equivalent.
Although the optimal (cid:15) s found through hyperparameter search are much smaller than 1 , one might wonder if we can offset the increased meta learning rate (cid:15) by reducing the learning rate α proportionally, and achieve similar performance to Fish using simply ERM with domain grouping.
In Table 4 ( ERM, DG ), we verify that this is infeasible empirically.
We train an ERM model with domain grouping strategy (equivalent to Fish with (cid:15) = 1 ) using 3 different learning rates: 1) prop.
lr , learning rate α is lowered proportionally to keep α(cid:15) constant; 2) 0.1 × prop.
lr and 3) 0.001 × prop.
lr which further reduces the learning rate.
We see that on all datasets, all three conﬁgurations result in worse performance than ERM and Fish.
This shows that one cannot improve the results of ERM to match that of Fish by simply adjusting the learning rate/adopting domain grouping strategy, and that
Fish’s gain on domain generalization tasks is not due to its proximity to ERM.
In Appendix A.2, we also show that setting (cid:15) = 1 is equivalent to setting S as the total number of iterations during training, which causes the effect of maximizing GIP to diminish.
This provide further support to our empirical ﬁnding in this analysis.
Does maximising GIP help domain generalization?
In Parascandolo et al.
, authors show that by masking out the gradients that have opposite signs in different domains, they are able to trade-off some “learning speed” for prioritizing learning the invariances.
Our principle for proposing the IDGM objective is identical, except in this case we actively encourage these gradients to be of consistent signs.
This unfortunately, does result in computing the costly second-order derivative, which makes it difﬁcult for us to empirically evaluate the IDGM objective on non-toy dataset 3 .
We believe that the large performance gain of IDGM over ERM on the linear and C D SPRITES datasets, along with the detailed discussion on the relationship between gradient sign consistency and invariant representation, is enough to justify for IDGM’s ability to improve domain generalization.
C ONCLUSION
In this paper we presented inter-domain gradient matching (IDGM) for domain generalization.
To avoid costly second-order computations, we approximated IDGM with a simple ﬁrst-order algorithm,
Fish.
We demonstrated our algorithm’s capability to learn from invariant features (as well as ERM’s failure to do so) using simple datasets such as C D S PRITES -N and the linear example.
We then evaluated the model’s performance on W ILDS and D OMAIN B ED , demonstrating that Fish performs well on different subgenres of domain generalization, and surpasses baseline performance on a diverse range of vision and language tasks using different architectures such as DenseNet, ResNet-50 and
BERT.
Our experiments can be replicated with 1500 GPU hours on NVIDIA V100.
Despite its strong performance, similar to previous work on domain generalization, when the number of domains is large Fish struggles to outperform ERM.
We are currently investigating approaches by which Fish can be made to scale to datasets with orders of magnitude more domains and expect to report on this improvement in our future work.
E THICS S TATEMENT
We believe there are no ethical concerns within this work.
None of the datasets used involve human identities, and our motivation in proposing this work in no way concerns any surveillance/military use.
As with many other domain generalization methods, our algorithm aims at improving the performance of machine learning systems deployed in the wilds, which can be used in many ways that will beneﬁt the society.
A T AYLOR E XPANSION OF R EPTILE AND F ISH I NNER - LOOP U PDATE
In this section we provide proof to Theorem 3.1.
We reproduce and adapt the proof from Nichol et al.
(2018) in the context of Fish, for completeness.
We demonstrate that when the inner-loop learning rate α is small, the direction of G f aligns with that of G g , where
[Equation]
Expanding G g
G g is the gradient of maximizing the gradient inner product (GIP).
[Equation]
Expanding G f
To write out G f , we need to derive the gradient update of Fish, θ − ˜ θ .
Let us ﬁrst deﬁne some notations.
For each inner-loop with S steps of gradient updates, we assume a loss functions l as well as a sequence of inputs { d i } Si =1 , where d i :=
{ x b , y b } Bb =1 denotes a minibatch at step i randomly drawn from one of the available domains in {D 1 , · · · , D S } .
For reasons that will become clearer later, take extra note that the subscript i here denotes the index of step, rather than the index of domain.
We also deﬁne the following:
[Equation]
sequence of parameters
[Equation]
point
[Equation]
In the following analysis we omit the expectation E d i and input ( x, y ) to l and instead denote the loss at step i as l i .
Performing second-order Taylor approximation to (cid:101) g
i yields:
[Equation]
Applying ﬁrst-order Taylor approximation to (cid:101) g
j gives us
[Equation] [Equation]
plugging this back to Equation (15) yields:
[Equation] [Equation]
For simplicity reason
[Equation]
let us consider performing two steps in inner-loop updates,
i
[Equation]
e. S
= 2 .
We can then write the gradient of Fish θ − ˜ θ as
[Equation] [Equation]
+ O ( α 3 ) .
Furthermore, taking the expectation of θ − ˜ θ under minibatch sampling gives us (assuming independence between g 1 and g 2 )
[Equation]
interchanging indices
[Equation]
averaging
[Equation] [Equation]
Note that the only reason we can interchange the indices in 2 is because the subscripts represent steps in the inner loop rather than index of domains.
Plugging 1 , 2 in Equation (19) yields:
[Equation]
We can also expand this to the general case where S ≥ 2 :
[Equation]
The second term in Equation (5) is ¯ G , which is the full gradient of ERM deﬁned as follow:
[Equation]
Plugging Equation (21) and Equation (22) to Equation (5) yields
[Equation]
[Equation]
Comparing Equation (7) to Equation (24), we have:
[Equation]
A. F ISH AND R EPTILE : D IFFERENCES AND C ONNECTIONS
As we introduced, our algorithm Fish is inspired by Reptile, a model agnostic meta-learning (MAML) algorithm.
Meta-learning aims at reducing the sample complexity of new, unseen tasks.
A popular school of thinking in meta-learning is MAML, ﬁrst proposed in Finn et al.
Andrychowicz et al.
The key idea is to backpropagate through gradient descent itself to learn representations that can be easily adapted to unseen tasks.
There are close connections between meta-learning (Thrun and Pratt, 1998) and (multi-source) domain adaptation. In fact, there are a few works in domain generalization that are inspired by the meta-learning principles, such as Li et al. (2018a); Balaji et al.
; Li et al.
Dou et al.
In Ren et al. (2018), we also see the leveraging of gradient inner product in meta-learning, where it is used to determine the importance weight of training examples. Even though meta learning and domain generalization both study N -way, K -shot problems, there are some distinct differences that set them apart. The most prominent one is that in meta learning, some examples in the test dataset will be made available at test time ( K > 0 ), while in domain
1: for i = 1 , 2 , · · · do 2: ˜ θ ← θ 3:
Sample task D t ∼ {D 1 , · · · , D T } 4: for s ∈ { 1 , · · · , S } or D t ∈ {D 1 , · · · , D T } do 5:
Sample batch d t ∼ D t 6: g t = ∂
L ( d t ; ˜ θ ) /∂
˜
θ 7:
Update ˜ θ ← ˜ θ − αg t 8: end for 9:
Update θ ← θ + (cid:15) (˜ θ − θ ) 10: end for
[Figure]
generalization no example in the test dataset is seen by the model ( K = 0 ); another important difference is that while domain generalization aims to train models that perform well on an unseen distribution of the same task , meta-learning assumes multiple tasks and requires the model to quickly learn an unseen task using only K examples.
Due to these differences, it does not make sense in general to use MAML framework in domain generalization.
As it turns out however, the idea of aligning gradients to improve generalization is relevant to both methods — The fundamental difference here that MAML algorithms such as Reptile aligns the gradients between batches from the same task Nichol et al.
(2018), while Fish aligns those between batches from different tasks.
To see how this is ahiceved, let us have a look at the algorithmic comparison between Fish (blue) and Reptile (green) in Algorithm 3.
As we can see, the key difference between the algorithm of Fish and Reptile is that Reptile performs its inner-loop using minibatches from the same task, while Fish uses minibatches from different tasks ( l4-8 ).
Based on the analysis in Nichol et al.
(2018) (which we reproduce in Appendix A), this is why Reptile maximizes the within-task gradient inner products and
Fish maximizes the across-task gradient inner products.
A natural question to ask here is – how does this affect their empirical performance?
In Figure 6, we show the train and test performance of Fish (blue) and Reptile (green) on C D S PRITES -N.
We can see that despite the algorithmic similarity between Fish and Reptile, the two methods behave very differently on this domain generalization task: while Fish’s test accuracy goes to 100% at
N = 10 , Reptile’s test performance is always 50% regardless of N .
Moreover, we observe a dip in
Reptile’s training performance early on, with the accuracy plateaus at 56% when N > 20 .
Reptile’s poor performance on this dataset is to be expected since its inner-loop is designed to encourage within-domain generalization, which is not helpful for learning what’s invariant across domains.
A. W HEN (cid:15) =
When we set the meta learning rate of Fish (cid:15) = 1 , Fish reduces to ERM with a data-loader that samples from only one domain in each minibatch.
Note that with this change, the inner-loop step S now represents the total training iterations, and as a result it increases from a constant value ( ≤in our experiments) to orders of magnitude larger (up to 10 5 for some datasets that we use).
We demonstrate here that for this reason, the effect of maximizing inter-domain gradient inner product is signiﬁcantly less prominent when (cid:15)
= 1 .
To see this, let us revisit Equation (21), where we demonstrate that the expectation to Fish’s update
E [ θ − ˜ θ ] can be written as the sum of the following three terms via Taylor series expansion,
[Equation]
where (1) is the sum of gradients over S steps of SGD updates (equivalent to ERM), (2) the gradient inner product over S steps of gradients and (3) is the higher order terms of Taylor series expansion.
For both Fish and Reptile, the higher order terms are ignored under two conditions: 1) constant/noninﬁnite inner-loop step S and 2) small learning rate α .
However, when (cid:15) = 1 , inner-loop step S → ∞ .
As a result, the higher order terms cannot be ignored, and we can no longer conclude that gradient inner product plays an important rule in the model’s updates.
B S MOOTH F ISH : A MORE GENERAL ALGORITHM B. D ERIVATION
We conclude in Appendix A that a component of Fish’s update G f =
E [ θ − ˜ θ ] − αS · ¯ G is in the same direction as the gradient of GIP, G g .
It is therefore possible to have explicit control over the scaling of the GIP component in Fish, similar to the original IDGM objective, by writing the following:
[Equation]
By introducing the scaling term γ , we have better control on how much the objective focus on inner product vs average gradient.
Note that γ = 1 recovers the original Fish gradient, and when γ =
the gradient G sm is equivalent to ERM’s gradient with learning rate αS .
We name the resulting algorithm
SmoothFish.
See Algorithm 4.
B. R ESULTS
We run experiments on 4 datasets in W ILDS using SmoothFish, with γ ranging in [0 , 0 . 1 , 0 . 2 , 0 . 5 , 0 . 8 , 1 , 2 , 10] .
Note that when γ = 0 , SmoothFish is equivalent to ERM and when γ = 1 it is equivalent to Fish.
See results in Figure 7.
The other hyperparameters including α , meta steps, (cid:15) used here are the same as the ones used in our main experiments section.
1: for iterations = 1 , 2 , · · · do 2: (cid:101) θ ← θ 3: for D i ∈ permute ( {D 1 , D 2 , · · · , D S } ) do 4: Sample batch d i ∼ D i 5: g i = E d i (cid:20) ∂l (( x, y ); θ ) ∂θ (cid:21) //Grad wrt θ 6: (cid:101) g
i =
E d i (cid:34) ∂l (( x, y ); (cid:101) θ ) ∂ (cid:101) θ (cid:35) //Grad
wrt (cid:101) θ 7:
Update (cid:101) θ ← (cid:101) θ − α (cid:101) g i 8: end for 9: ¯ g = 1 S
S (cid:88) s =
g i , g sm =
αS ¯ g + γ (cid:16) ( (cid:101) θ − θ ) − αS ¯ g (cid:17) 10:
Update θ ← θ + (cid:15)g sm 11: end for
[Figure]
C D ISCUSSIONS AND R ESULTS ON W ILDS
We provide a more detailed summary of each dataset in Table 5.
Some entries in # Domains are omitted because the domains in each split overlap.
Note that in this paper we report the results on W ILDS v1 — the benchmark has been updated since with slightly different dataset splits.
We are currently working on updating our results to v2 of W ILDS .
[Table]
[Table]
[Table]
C. P OVERTY M AP -W ILDS
[Table] : Asset index prediction (real-valued)
[Table] countries
The task is to predict the real-valued asset wealth index of an area, given its satellite imagery.
Since the number of domains considered here is large (23 countries), instead of looping over all S domains in each inner-loop, we sample N << S domains in each iteration and perform inner-loop updates using minibatches from these domains only to speed up computation.
For this dataset we choose
N = 5 by hyper-parameter search.
Evalutaion: Pearson Correlation (r).
Following the practice in W ILDS benchmark, we compare the results by computing Pearson correlation (r) between the predicted and ground-truth asset index over 3 random seed runs.
Results: We train the model using a ResNet-18 (He et al., 2016) backbone.
See Table 6.
We see that Fish obtains the highest test performance, with the same validation performance as the best baseline.
The performance is more stable between validation and test, and the standard deviation is smaller than for the baselines.
We also report the results of ERM models trained in our environment as “ERM (ours)”, which shows similar performance to the canonical results reported in the W ILDS benchmark itself (“ERM”).
[Equation]
The C AMELYON 17-W ILDS dataset contains 450,000 lymph-node scans from 5 hospitals.
Due to the size of the dataset, instead of training with Fish from scratch, we pre-train the model with ERM using the recommended hyper-parameters in Koh et al.
(2020), and ﬁne-tune with Fish.
For this dataset, we ﬁnd that Fish performs the best when starting from a pretrained model that has not yet converged, achieving much higher accuracy than the ERM model.
we provide an ablation study on this in Appendix E.
Evaluation: Average accuracy.
We evaluate the average accuracy of this binary classiﬁcation task.
Following Koh et al.
, we show the mean and standard deviation of results over 10 random seeds runs.
The number of random seeds required here is greater than other W ILDS datasets due to the large variance observed in results.
Note that these random seeds are not only applied during the ﬁne-tuning stage, but also to the pretrained models to ensure a fair comparison.
Results: Following the practice in W ILDS , we adopt DenseNet-121’s
(Huang et al., 2017) architecture for models trained on this dataset.
See results in Table 7.
The results show that Fish signiﬁcantly outperforms all baselines – its test accuracy surpasses the best performing baseline by 6% .
Also note that for all other baselines, there is a large gap between validation and test accuracy ( 11% ∼ 27% ).
This is because W ILDS chose the hospital that is the most difﬁcult to generalize to as the test split to make the task more challenging.
Surprisingly, as we can observe in Table 7, the discrepancy between test and validation accuracy of Fish is quite small ( 3% ).
The fact that it is able to achieve a similar level of accuracy on the worst-performing domain further demonstrates that Fish does not rely on domain-speciﬁc information, and instead makes predictions using the invariant features across domains.
To demonstrate that Fish’s strong performance on C AMELYON 17’s selected test set is not merely coincidental [Table] we randomly chose 3 ways of assigning the 5 domains as train/test/val (3/1/1 domains) splits.
See results collected over 10 random seeds in Table 8.
Notably, on these different splits of C AMELYON 17 Fish still on average outperforms ERM (with only 2 exceptions), demonstrating that Fish can generally achieve good performance on the Camelyon dataset. [Table] [Table]
[Table]
C. FM
[Table]
W-W ILDS
Task
[Table]
Infrastructure classiﬁcation (62 classes).
Domains [Table] 80 (16 years x 5 regions)
Similar to C AMELYON 17-W ILDS , since the number of domains is large, we sample N = 5 domains for each inner-loop.
To speed up computation, we also use a pretrained ERM model and ﬁne-tune with Fish; different from Appendix C.2, we ﬁnd the best-performing models are acquired when using converged pretrained models (see details in Appendix E).
Evaluation: Average & worst-region accuracies .
Following W ILDS , the average accuracy evaluates the model’s ability to generalize over years, and the worst-region accuracy measures the model’s performance across regions under a time shift.
We report results using 3 random seeds.
Results:
Following Koh et al.
(2020), we use a DenseNet-121 pretrained on ImageNet for this dataset.
Results in Table 9 show that Fish has the highest worst-region accuracy on both test and validation sets.
It ranks second in terms of average accuracy, right after ERM.
Again, Fish’s performance is notably stable with the smallest standard deviation across all metrics compared to baselines.
C. C IVIL C OMMENTS -W ILDS
Task: Toxicity detection in online comments (2 classes).
Domains: 8 demographic identities.
The C IVIL C OMMENTS -W ILDS contains 450,000 comments collected from online articles, each annotated for toxicity and the mentioning of demographic identities.
Again, we use ERM pre-trained model to speed up computation, and sample N = 5 domains for each inner-loop.
Evaluation: Worst-group accuracy .
To study the bias of annotating comments that mentions demographic groups as toxic, the W ILDS benchmark proposes to evaluate the model’s performance by doing the following: 1) Further separate each of the 8 demographic identities into 2 groups by toxicity – for example, separate black into black, toxic and black, not toxic ; 2) measure the accuracies of these 8 × 2 = 16 groups and use the lowest accuracy as the ﬁnal evaluation of the model.
This metric is equivalent to computing the sensitivity and speciﬁcity of the classiﬁer on each demographic identity, and reporting the worse of the two metrics over all domains.
Good performance on the group with the worst accuracy implies that the model does not tend to use demographic identity as an indicator of toxic comments.
Again, following Koh et al.
we report results of 3 random seed runs.
Results: We compare results to the baselines used in the W ILDS benchmark over 3 random seed runs in Table 10.
All models are trained using BERT (Devlin et al., 2018).
The results show that Fish outperforms the best baseline by 4% and 7% on the test and validation set’s worst-group accuracy respectively, and is competitive in terms of average accuracy with ERM (within standard deviation).
The strong performance of Fish on worst-group accuracy suggests that the model relies the least on demographic identity as an indicator of toxic comments compared to other baselines.
ERM, on the other hand, has the highest average accuracy and the lowest worst-group accuracy.
This indicates that it achieves good average performance by leveraging the spurious correlation between toxic comments and the mention of certain demographic groups.
Note that different from all other datasets in W ILDS that focus on pure domain generalization (i.e, no overlap between domains in train and test splits), C IVIL C OMMENTS -W
ILDS is a subpopulation shift problem, where the domains in test are a subpopulation of the domains in train.
As a result, the baseline models used in W ILDS for this dataset are different from the methods used in all other datasets, and are tailored to avoiding systematic failure on data from minority subpopulations.
Fish works well in this setting too without any changes or special sampling strategies (such as ∗ and + in Table 10).
This further demonstrates the good performance of our algorithm on different domain generalization scenarios.
[Equation]
Task: Animal species (186 classes).
Domains : 324 camera locations.
The dataset consists of over 200,000 photos of animal in the wild, using stationary cameras across 324 locations.
Classifying animal species from these heat or motion-activated photos is especially challenging: methods can easily rely on the background information of photos from the same camera setup.
Fish models are pretrained with ERM till convergence, and for each inner loop we sample from N = 10 domains.
Evaluation: Macro F1 score .
Across the 186 class labels, we report average accuracy and both weighted and macro F1 scores (F1-w and F1-m, respectively, in Table 11).
We run 3 random seeds for each model.
Results: All models reported in Table 7 are trained using a ResNet-50.
We ﬁnd Fish to outperform baselines on both test accuracy and weighted F1, with a 1% improvement on both metrics over the best performing model (ERM).
However, this comes at the cost of lower macro F1 score, where Fish performs 1% worse than ERM models that we trained and 3% than the ERM reported in W ILDS .
This suggests that Fish is less good at classifying rarer species, however the overall accuracy on the test dataset is improved.
Although Fish did not outperform the ERM baseline on the primary evaluation metric proposed in
Koh et al.
, we found the improvement of Fish in both accuracy and weighted F1 to be robust across a range of hyperparameters.
See more details on this in Appendix D.
C. A MAZON -W ILDS Task: Sentiment analysis (5 classes
[Table] reviewers.
The dataset contains 1.4 million customer reviews on Amazon from 7,676 customers, and the task is to predict the score (1-5 stars) given the review.
Similarly, we pretrained the model with ERM till convergence, and due to the large number of domains ( S = 5008 in train)
we sample N =reviewers for each inner loop.
Evaluation: 10th percentile accuracy .
Reporting the accuracy of the 10th percentile reviewer helps us assess whether the model performance is consistent across different reviewers.
The results in
Table 12 are reported over 3 random seeds.
Results:
The model is trained using BERT (Devlin et al., 2018) backbone.
While Fish has lower average accuracy compared to ERM, its 10th percentile accuracy matches that of ERM, outperforming all other baselines.
[Table] [Table]
D H YPERPARAMETERS
In Table 13 we list the hyperparameters we used to train ERM.
The same hyperparameters were used for producing ERM baseline results and as pretrained models for Fish.
In val.
metric we report the metric on validation set that is used for model selection
[Table]
and in cut-off we specify when to stop training when using ERM to generate pretrained models.
[Table]
In Table 14 we list out the hyperparameters we used to train Fish.
Note that we train Fish using the same model, batch size, val metric and optimizer as ERM – these are not listed in Table 14 to avoid repetitions.
Weight decay is always set as
[Table]
E A BLATION S TUDIES ON P RE - TRAINED M ODELS
In this section we perform ablation study on the convergence of pretrained ERM models.
Note that the pretraining is done on the same domain generalization dataset that Fish is trained on later, not on
ImageNet.
We study the performance of Fish with the following three conﬁgurations of pretrained
ERM models:
[Table]
Model is trained on 10% of the data (epoch 1); 2.
Model is trained on 50% of the data (epoch 1); 3.
Model at convergence.
By comparing the results between these three settings [Table] we demonstrate how the level of convergence affects the Fish’s training performance.
See results in Table 15.
Note that P OVERTY is excluded here because the dataset is small enough that we are able to train Fish from scratch.
[Table]
We see that C IVIL C OMMENTS sustain good performance using pretrained models at different convergence levels.
FM O W
and I W ILD C AM on the other hand seems to have strong preference towards converged model
[Table]
and the results worsen as the amount of data seen during training goes down.
C AMELYON 17 achieves the best performance when only 10% of data is seen, and the test accuracy decreases while training with models with higher level of convergence.
F A BLATION S TUDIES ON HYPERPARAMETERS
α and (cid:15)
We study the effect of changing Fish’s inner loop learning rate α and outer loop learning rate (cid:15) .
To make the comparisons more meaningful, we keep α · (cid:15) constant while changing their respective values.
See results in Figure 8.
Meta steps N
For most of the datasets we studied (all apart from C AMELYON 17 where T = 3 ) we sample a N -sized subset of all T domains available for training (see Table 14 for T of each dataset).
Here we study when N = 5 , 10 , 20 .
[Table]
In general altering these hyperparameters don’t have a huge impact on the model’s performance, however it does seem thet when N = 20 the performance on some datasets (P OVERTY and C IVILC OMMENTS ) degrade slightly. [Figure] [Figure]
G T RACKING GRADIENT INNER PRODUCT
In Figure 9, we demonstrate the progression of inter-domain gradient inner products during training using different objectives.
We train both Fish (blue) and ERM (yellow) untill convergence while recording the normalized gradient inner products (i.e. cosine similarity) between minibatches from different domains used in each inner-loop.
The gradient inner products are computed both before (dotted) and after (solid) the model’s update.
To ensure a fair comparison, we use the exact same sequence of data for Fish and ERM (see Appendix I for more details).
Inevitably, the gradient inner product trends differently for each dataset since the data, types of domain splits and the choice of architecture are very different.
In fact, the plot for C D S PRITES -N and
P OVERTY are signiﬁcantly different from others, with a dip in gradient inner product at the beginning of training – this is because these are the two datasets that we train from scratch.
On all other datasets, the gradient inner products are recorded when ﬁne-tuning with Fish.
Despite their differences, there are some important commonalities between these plots: if we compare the pre-update (dotted) to post-update (solid) curves, we can see that ERM updates often result in the decrease of gradient inner product, while for Fish it can either increase signiﬁcantly (Figure 9c and Figure 9e) or at least stay at the same level (Figure 9a, Figure 9b, Figure 9d and Figure 9f).
As a result of this, we can see that the post-update gradient inner product of Fish is always greater than
ERM at convergence.
The observations here shows that Fish is effective in increasing/maintaining the level of inter-domain gradient inner product and sheds some lights on its efﬁciency on the datasets we studied
. H S UMMARY OF D OMAIN B ED
D OMAIN B ED is a testbed for domain generalization that implements consistent experimental protocols across SOTA methods to ensure fair comparison.
It contains 7 popular domain generalization datasets, including Colored MNIST (Arjovsky et al., 2019), Rotated MNIST (Ghifary et al., 2015),
VLCS (Fang et al., 2013), PACS (Li et al., 2017), OfﬁceHome (Venkateswara et al., 2017), Terra
Incognita (Beery et al., 2018) and DomainNet (Peng et al., 2019), and offers comparison to a variety of SOTA domain generalization methods, including IRM (Arjovsky et al., 2019), Group DRO (Hu et al., 2018; Sagawa et al., 2019), Mixup (Yan et al., 2020), MLDG (Li et al., 2018a), Coral (Sun and
Saenko, 2016), MMD (Li et al., 2018b), DANN (Ganin et al., 2016) and CDANN (Li et al., 2018). I A LGORITHM FOR TRACKING GRADIENT INNER PRODUCT To make sure that the gradients we record for ERM and Fish are comparable, we use the same sequence of S -minibatches to train both algorithms. See Algorithm 6 for details.
J T -SNE V ISUALISATION OF L EARNED R EPRESENTATIONS
In this section we visualize the representation learned by Fish and ERM with t-SNE (van der Maaten and Hinton, 2008) for PACS, VLCS and OfﬁceHome datasets.
See Figure 10.
We can see that Fish and ERM are both capable of forming distinctive label-clusters on PACS (ﬁrst row) and VLCS (second row), however with ERM we can observe within each label cluster that sub-clusters of domains are forming.
This is particularly the case for the the red and yellow cluster in
Figure 10a and cyan cluster in Figure 10c.
The domain clustering phenomena is not observed for
Fish.
On the other hand, for the OfﬁceHome dataset where Fish outperforms ERM by more than 1% , we clearly see that Fish exhibits better label clustering performance than ERM.
1: Initialize Fish θ f ← θ , ERM θ e ← θ 2: for i = 1 , 2 , · · · do 3: //Get
all minibatches 4: for D n ∈ {D 1 , D 2 , · · · , D N } do 5:
Sample batch d n ∼ D n 6: end for 7: //GradInnerProd
before update 8: ¯¯
g
Fb =
GIP ( { d 1 , d 2 , · · · , d N } , θ f ) 9: ¯¯
g Eb =
GIP ( { d 1 , d 2 , · · · , d N } , θ e ) 10: //Fish
training 11: ˜ θ ← θ f 12: for d n ∈ { d 1 , d 2 , · · · , d N } do 13: g n = ∂l ( d n ; ˜ θ ) /∂
˜ θ 14:
Update ˜ θ ← ˜ θ − αg n 15: end for 16: θ f ← θ f + (cid:15) (˜ θ − θ f ) 17: //Rearrange minibatches 18: d = shuffle ( concat ( d 1 , d 2 , · · · , d N )) 19: { ˜ d 1 , ˜ d 2 , · · · , ˜ d N } = split ( d ) 20: //ERM
training 21: for ˜ d n ∈ { ˜ d 1 , ˜ d 2 , · · · , ˜ d N } do 22: g n = ∂l ( ˜ d n ; θ e ) /∂θ e 23:
Update θ e ← θ e − αg n 24: end for 25: //GradInnerProd after update 26: ¯ ¯ g Fa =
GIP ( { d 1 , d 2 , · · · , d N } , θ f ) 27: ¯¯
g Ea = GIP ( { d 1 , d 2 , · · · ,
[Figure]