PG R A D : L EARNING P RINCIPAL G RADIENTS F OR D OMAIN G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Machine learning models fail to perform when facing out-of-distribution (OOD) domains, a challenging task known as domain generalization (DG).
In this work, we develop a novel DG training strategy, we call PGrad , to learn a robust gradient direction, improving models’ generalization ability on unseen domains.
The proposed gradient aggregates the principal directions of a sampled roll-out optimization trajectory that measures the training dynamics across all training domains.
PGrad ’s gradient design forces the DG training to ignore domain-dependent noise signals and updates all training domains with a robust direction covering main components of parameter dynamics.
We further improve PGrad via bijection-based computational reﬁnement and directional plus length-based calibrations.
Our theoretical proof connects PGrad to the spectral analysis of Hessian in training neural networks.
Experiments on DomainBed and WILD S benchmarks demonstrate that our approach effectively enables robust DG optimization and leads to smoothly decreased loss curves.
Empirically, PGrad achieves competitive results across seven datasets, demonstrating its efﬁcacy across both synthetic and real-world distributional shifts.
I NTRODUCTION
Deep neural networks have shown remarkable generalization ability on test data following the same distribution as their training data.
Yet, high-capacity models are incentivized to exploit any correlation in the training data that will lead to more accurate predictions.
As a result, these models risk becoming overly reliant on “domain-speciﬁc” correlations that may harm model performance on test cases from out-of-distribution (OOD).
A typical example is a camel-and-cows classiﬁcation task (Beery et al., 2018; Shi et al., 2021), where camel pictures in training are almost always shown in a desert environment while cow pictures mostly have green grassland backgrounds.
A typical machine learning model trained on this dataset will perform worse than random guessing on those test pictures with cows in deserts or camels in pastures.
The network has learned to use the background texture as one deciding factor when we want it to learn to recognize animal shapes.
Unfortunately, the model overﬁts to speciﬁc traps that are highly predictive of some training domains but fail on
OOD target domains.
Recent domain generalization (DG) research efforts deal with such a challenge.
They are concerned with how to learn a machine learning model that can generalize to an unseen test distribution when given multiple different but related training domains.
Recent literature covers a wide spectrum of DG methods, including invariant representation learning, meta-learning, data augmentation, ensemble learning, and gradient manipulation (more details in Section 2.4).
Despite the large body of recent DG literature, the authors of (Gulrajani & LopezPaz, 2021) showed that empirical risk minimization (ERM) provides a competitive baseline on many real-world DG benchmarks.
ERM does not explicitly address distributional shifts during training.
Instead, ERM calculates the gradient from each training domain and updates a model with the average gradient.
However, one caveat of ERM is its average gradient-based model update will preserve domain-speciﬁc noise during optimization.
This observation motivates the core design of our method.
We propose a novel training strategy that learns a robust gradient direction for DG optimization, and we call it PGrad .
PGrad samples an optimization trajectory in high dimensional parameter space by updating the current model sequentially across training domains.
It then constructs a
Θ !"
Θ #" Θ $" Θ " 𝐷 !
𝐷 " 𝐷 # Θ " Θ !"
Θ #" Θ $"
Principal Vectors’ Construction
Direction & Length Calibration Θ " Θ !"
Θ #" Θ $" Θ " Θ !"
Θ #" Θ $"
Noise Suppression 𝛻
Eval: ℒ(Θ ∗ ;)
𝑇 $ Θ,𝐷 !
,𝐷 # ,𝐷 $ , , ·· ·· Θ "(! =
[Figure] [Figure]
Train: local coordinate system to explain the parameter variations in the trajectory.
Via singular value decomposition (SVD), we derive an aggregated vector that covers the main components of parameter dynamics and use it as a new gradient direction to update the target model.
This novel vector - that we name the ”principal gradient” - reduces domain-speciﬁc noise in the DG model update and prevents the model from overﬁtting to particular training domains.
To decrease the computational complexity of SVD, we construct a bijection between the parameter space and a low-dimensional space through transpose mapping.
Hence, the computational complexity of the PGrad relates to the number of sampled training domains and does not depend on the size of our model parameters.
This paper makes the following contributions: (1) PGrad places no explicit assumption on either the joint or the marginal distributions.
(2) PGrad is model-agnostic and is scalable to various model architecture since its computational cost only relates to the number of training domains.
(3) We theoretically show the connection between PGrad and Hessian approximation, and also prove that
PGrad beneﬁts the training efﬁciency via learning a gradient in a smaller subspace constructed from learning trajectory.
(4) Our empirical results demonstrate the competitive performance of PGrad across seven datasets covering both synthetic and real-world distributional shifts.
M ETHOD
Domain generalization (Wang et al., 2021; Zhou et al., 2021) assumes no access to instances from future unseen domains.
In domain generalization, we are given a set of training domains
D tr = { D i } ni =1 and test domains T te =
{ T j } mj =1 .
Each domain D i (or T j ) is associated with a joint distribution P D i X×Y (or P T j X×Y ), where X represents the input space and Y is the output space.
Moreover, each training domain D i is characterized by a set of i.i.d samples { x ik , y ik } .
For any two different domains sampled from either D tr or T te , their joint distribution varies P D
i X×Y (cid:54) =
P D j X×Y , and most importantly, P D i X×Y (cid:54) =
P T j X×Y .
We consider the prediction task from the input x ∈ X to the output y ∈ Y .
Provided with a model family whose parameter space is Θ ⊂ R d and the loss function L : Θ × ( X × Y )
→ R + , the goal is to ﬁnd an optimal Θ ∗ te on test domains so that:
[Equation]
In DG setup, any prior about T te , such as inputs or outputs, are unavailable in the training phase.
Despite not considering domain discrepancies from training to testing, ERM is still a competitive method for domain generalization tasks (Gulrajani & Lopez-Paz, 2021).
ERM naively groups data from all training domains D tr together and obtains its optimal parameter Θ ∗ tr via the following to approximate Θ
[Equation] [Equation]
In the rest of the paper, we omit the subscript in Θ tr and use Θ for simplicity (during DG training, only training domains D tr will be available for model learning).
When optimizing with ERM on DG across multiple training domains, the update of Θ follows:
[Equation]
where ∇ Θ
[Equation]
L D i = ∇ E ( x , y ) ∼P
Di X×Y L [Θ t , ( x , y )] calculates the gradient of the loss on domain D i with respect to the current parameter Θ t and γ is the learning rate.
The gradient determines the learning path of a model.
When using ERM in DG setting, each step of model updates uses an average gradient and may introduce and preserve domain-speciﬁc noise.
For instance, if one training domain includes the trapping signals like cows always in pastures and camels always in deserts (as mentioned earlier).
When investigating across multiple training domains, we envision such domain-speciﬁc noise signals will not be the main components of parameter variations across all domains.
This motivates us to design PGrad as follows.
2. PG RAD : P RINCIPAL G RADIENT BASED M ODEL U PDATES
We extend ERM with a robust gradient estimation that we call PGrad .
We visualize an overview in
Figure 1 to better explain how it works.
Brieﬂy speaking, given the current model parameter vector, we sample a trajectory of parameters by training sequentially on each training domain.
Next, we build a local principal coordinate system based on parameters obtained from the sampled trajectory.
The chosen gradient direction is then built as a linear combination of the orthogonal axes of the local principal coordinates.
Our design also forces the learned gradient to ﬁlter out domain-speciﬁc noise and follows a direction that maximally beneﬁts all training domains D tr ; we refer to this extracted direction as the principal gradient.
In the following, we cover details of the trajectory sampling, local coordinate construction, direction and length calibration, plus the noise suppression for our principal gradient design.
Trajectory Sampling.
Denote the current parameter vector as Θ t .
We ﬁrst sample a trajectory S through the parameter space Θ by sequentially optimizing the model on each training domain:
[Equation]
We refer to the process of choosing an order of training domain to optimize as trajectory sampling .
Different ordering arrangements of training domains will generate different trajectories.
Principal Coordinate system Construction.
Now we have a sampled trajectory S
=
{ Θ ti } ni =0 ∈
R ( n +1)
× d , that were derived from Θ t .
Note: the inclusion of the starting location Θ t 0 as part of the trajectory is necessary; see the proof in Appendix (A.3).
Then we construct a local principal coordinate system to explain the variations in S .
We are looking for orthogonal and unit axes V =
[ v T z ] nz =0 ∈ R ( n +1)
× d to maximally capture the variations of the trajectory.
Each v z ∈ R
d is a unit vector of size d , the same dimension as the parameters Θ t .
[Equation]
The above objective is the classic principal component analysis formulation and can be solved with singular value decomposition (a.k.a. SVD).
Eq.
(5)) has the closed-form solution as follows (the revised computational complexity is n ):
[Equation]
Here λ z , v z denote the z -th largest eigenvalue and its corresponding eigenvector.
ˆ S denotes the centered trajectory by removing the mean from S .
In the above Eq.
, the computational bottleneck lies in the SVD, whose computational complexity comes at O ( d 3 ) due to ˆ S T ˆ S ∈ R d × d .
d denotes the size of the parameter vector and is fairly large for most state-of-the-art (SOTA) deep learning architectures.
This prohibits the computation of the eigenvalues and eigenvectors from Eq.
(6) for
SOTA deep learning models.
Hence, we reﬁne and construct a bijection as follows to lower the computational complexity (to O (( n + 1) 3 ) ):
[Equation]
Eq.
(7) indicates that if λ z , e z are the z -th largest eigenvalue and corresponding eigenvector of ˆ
S ˆ S T , the z -th largest eigenvalue and corresponding eigenvector of ˆ S T
ˆ S are λ
z , ˆ S T e z (i.e., v z = ˆ S T e z ).
This property introduces a bijection from eigenvectors of ˆ S ˆ S T ∈ R ( n +1) × ( n +1) to those of ˆ S T
ˆ S ∈ R d × d .
Since n - the number of training domains - is much smaller than d , calculating eigen-decomposition of ˆ S ˆ S T ∈ R ( n +1) × ( n +1) is therefore much cheaper.
Directional Calibration.
With the derived orthogonal axes V
= [ v T z ] nz =0 from Eq.
(7), now we construct a local principal coordinate system with each axis aligning with one eigenvector v z .
These principal coordinate axes V are ordered based on the magnitude of the eigenvalues.
This means that v i explains more variations of the sampled trajectory S than v j if i < j , and they are all unit vectors.
In addition, these vectors are unoriented, which means either positive or negative multiple of an eigenvector still falls into the eigenvector set.
Now our main goal is to get a robust gradient direction by aggregating information from V .
First we calibrate the directions of each eigenvectors so that they point to the directions that can improve the
DG prediction accuracy.
Learning an ideal direction is impossible without a reference.
The choice of the reference is ﬂexible, as long as it is pointing to a direction that climbs up the loss surface.
We want the reference to guide the principal gradient in the right direction for gradient descent based algorithms.
For simplicity, we use the difference between the starting point Θ t 0 and the end point Θ tn of the trajectory S as our reference ∇ r
=
Θ t 0 − Θ tn .
So for each coordinate axis, we revise its direction so that the resulting vector w z is positively related to the reference ∇ r in terms of the inner product:
[Equation]
Constructing Principal Gradient.
The relative importance of each w z is conveyed in the corresponding eigenvalue λ z .
Larger λ
z implies higher variance when projecting the trajectory S along w z direction.
We weight each axis with their eigenvalues, and aggregate them together into a weighted sum.
This gives us the principal gradient vector being calculated as follows:
[Equation]
There exists other possible aggregation besides Eq.
For instance, another choice of weight could be λ z / ||
λ
|| 1 or simply λ z since the eigenvalue of a semi-positive deﬁnite matrix is non-negative.
Gradient normalization has been widely recommended for improving training stability (You et al., 2017; Yu et al., 2017).
Our design in Eq.
(9) automatically achieves L 2 normalization, because:
[Equation]
Length Calibration.
As training updates continue, a ﬁxed length gradient operator may become too rigid, causing ﬂuctuations in the loss.
We, therefore, propose to calibrate the norm of ∇ p with a reference, for achieving adaptive length tuning.
Speciﬁcally, we propose to multiply the aggregated gradient from Eq.
(9) with the L 2 norm of ∇ r :
[Equation]
With this length calibration via ||∇ r || 2 , the norm of the proposed gradient is constrained by the multiplier, and is automatically tuned during the training process.
Noise Suppression.
Most w z axes correspond to small eigenvalues and may relate to just domainspeciﬁc noise signals.
They may help the accuracy of a speciﬁc training domain D i , but mostly hurt the overall accuracy on D tr .
We therefore deﬁne the principal gradient as follows and show how to use it to solve DG model training via gradient descent optimization (where k is a hyperparaemter):
[Equation]
2. T HEORETICAL A NALYSIS
In Appendix (A.5.1), we prove that 1 n ˆ S T ˆ S in Eq.
(6) provides us with the mean of all training domains’ domain-speciﬁc Fisher information matrix (FIM).
Since FIM is the negative of Hessian under mild conditions, PGrad essentially performs spectrum analysis on the approximated Hessian matrix.
Moreover, in Appendix (A.5.2), we show that PGrad improves the training efﬁciency of neural networks by recovering a subspace from the original over-parameterized space Θ .
This subspace is built from the top eigenvectors of the approximated Hessian.
We visualize the evolution of the eigenvalue distributions in Figure 8.
Our theoretical analysis connects to the machine learning literature that performs spectrum analysis of Hessian matrix (Gur-Ari et al., 2018) and connects its top subspace spanned by the eigenvectors to training efﬁciency and generalization in neural networks (Wei & Schwab, 2019; Ghorbani et al., 2019; Pascanu et al., 2014).
For example, (Hochreiter & Schmidhuber, 1997) shows that small eigenvalues in the Hessian spectrum are indicators of ﬂat directions.
Another work (Sagun et al., 2017) empirically demonstrated that the spectrum of the Hessian contains both a bulk component with many small eigenvalues and a few top components of much more signiﬁcant positive eigenvalues.
Later, (Gur-Ari et al., 2018) pointed out that the gradient of neural networks quickly converges to the top subspace of the Hessian.
2. V ARIATIONS OF PG
RAD
There exist many ways to construct a sampled trajectory, creating multiple variations of PGrad .
• PGrad-F : The vanilla trajectory sampling method will sample a trajectory of length n + 1 by sequentially visiting each D i in a ﬁxed order.
See appendix for the results of the rigid variation.
• PGrad : We can randomly shufﬂe the domain order in the training, and then perform to sample a trajectory.
This random order based strategy is used as the default version of PGrad .
• PGrad-B : We can split each training batch into B smaller batches and construct a long sampled trajectory that is with length n ∗ B + 1 .
• PGrad-BMix : Our method is model and data agnostic.
Therefore it is complementary and can combine with many other DG strategies.
As one example, we combine the random order based PGrad-B and MixUp (Zhang et al., 2017) into PGrad-BMix in our empirical analysis.
In PGrad and PGrad-F , the principal gradient’s trajectory covers all training domains D tr exactly once (per domain).
There are two possible limitations.
(1) If the number of training domains n is tiny, a length- ( n + 1) trajectory will not provide enough information to achieve robustness.
In the extreme case of n = 1 , we will only be able to get one axis w z , that goes back to ERM.
The current design can only eliminate discrepancies between different domains.
Notable intra-domain variations also exist because empirically approximating the expected loss may include noise due to data sampling, plus batch-based optimization may induce additional bias.
Based on this intuition, we propose a new design for sampling a trajectory by evenly splitting { x ik , y ik } from a training domain
D i into B small batches.
This new strategy allows us to obtain nB pseudo training domains.
Such a design brings in two beneﬁts: (1) We can sample a longer trajectory S , as the length changes from n to nB .
(2) Our design splits each domain’s data into B batches and treats each batch as if they come from different training domains.
By learning the principal gradient from these nB pseudo domains, we also address the intra-domain noise issue.
We name this new design PGrad-B .
Appendix (A.1) includes a ﬁgure comparing vanilla trajectory sampling with this extended trajectory sampling.
2. C ONNECTING TO R ELATED W ORKS
We can categorize existing DG methods into four broad groups.
Invariant element learning.
Learning invariant mechanisms shared across training domains provides a promising path toward DG.
Recent literature has equipped various deep learning components
- especially representation modules - with the invariant property to achieve DG (Li et al., 2018d;e;
Muandet et al., 2013).
The central idea is to minimize the distance or maximize the similarity between representation distributions P ( f ( X ) | D ) across training domains so that prediction is based on statistically indistinguishable representations.
Adversarial methods (Li et al., 2018d) and moment matching (Peng et al., 2019; Zellinger et al., 2017) are two promising approaches for distributional alignment.
A recent line of work explores the connection between invariance and causality.
IRM (Arjovsky et al., 2019) learns an invariant linear classiﬁer that is simultaneously optimal for all training domains.
Under the linear case and some constraints, the invariance of the classiﬁer induces causality.
Ahuja et al.
further extend IRM by posing it as ﬁnding the Nash equilibrium (Ahuja et al., 2020) and adding information bottleneck constraints to seek theoretical explanations (Ahuja et al., 2021).
However, later works (Kamath et al., 2021) show that even when capturing the correct invariances, IRM still tends to learn a suboptimal predictor.
Compared to this stream of works, our method places no assumption on either the marginal or joint distribution.
Instead, the PGrad explores the promising gradient direction and is model and data-agnostic.
Optimization methods.
One line of optimization-based DG works is those related to the Group
Distributionally robust optimization (a.k.a DRO) (Sagawa et al., 2019).
Group DRO aims to tackle domain generalization by minimizing the worst-case training loss when considering all training distributions (rather than the average loss).
The second set of optimization DG methods is optimizationbased meta-learning.
Optimization-based meta-learning uses bilevel optimization for DG by achieving properties like global inter-class alignment (Dou et al., 2019) or local intra-class distinguishability (Li et al., 2018b).
One recent work (Li et al., 2019) synthesizes virtual training and testing domains to imitate the episodic training for few-shot learning.
Gradient manipulation.
Gradient directions drive the updates of neural networks throughout training and are vital elements of generalization.
In DG, the main goal is to learn a gradient direction that beneﬁts all training domains (plus unseen domains).
Gradient surgery (Mansilla et al., 2021) proposes to use the sign function as a signal to measure the element-wise gradient alignment.
Similarly, the authors of (Chattopadhyay et al., 2020) presented And-mask, to learn a binary gradient mask to zero out those gradient components that have inconsistent signs across training domains.
Sandmask (Shahtalebi et al., 2021) added a tanh function into mask generation to measure the gradient consistency.
They extend And-mask by promoting gradient agreement.
Fish (Shi et al., 2021) and Fishr (Rame et al., 2021) are two recent DG works motivated by gradient matching.
They require the parallel calculation of the domain gradient from every training domain w.r.t a current parameter vector.
Fish maximizes the inner product of domain-level gradients; Fishr uses the variance of the gradient covariance as a regularizer to align the per-domains’ Hessian matrix.
Our method PGrad differs gradient matching by learning a robust gradient direction.
Besides, our method efﬁciently approximates the Hessian with training domains’ Fisher information matrix.
Appendix (A.4) includes a detailed analysis comparing parallel versus sequential domain-level training.
Furthermore, we adapt PGrad with parallel training, and compare it against PGrad with sequential training and ERM to justify our analysis, see visualizations in Figure 6.
We then show that gradient alignment is not necessary a sufﬁcient indicator of the generalization ability in Figure 7.
Others.
Besides the categories above, there exist other recently adopted to conquer domain generalization.
Data augmentation (Xu et al., 2021; Zhang et al., 2021; 2017; Volpi et al., 2018), which generates new training samples or representations from training domains to prevent overﬁtting.
Data augmentation can facilitate a target model with desirable properties such as linearity via Mixup (Zhang et al., 2017) or object focus (Xu et al., 2021).
Other strategies, like contrastive learning (Kim et al., 2021), representation disentangling (Piratla et al., 2020), and semi-supervised learning (Li et al., 2021), have also been developed for the DG challenge.
E XPERIMENTS
We conduct empirical experiments to answer the following questions: Q1.
Does PGrad successfully handle both synthetic and real-life distributional shifts?
Q2.
Can PGrad handle various architectures (ResNet and DenseNet), data types (scene and satellite images), and tasks (classiﬁcation and regression)?
Q3.
Compared to existing baselines, does PGrad enable smooth decreasing loss curves and generate smoother parameter trajectories?
Q4.
Can PGrad act as a practical complementary approach to combine with other DG strategies? 2 Q5.
How do bottom eigenvectors in the roll-out trajectories affect the model’s training dynamic and generalization ability?
[Table] [Table]
3. D OMAIN B ED B ENCHMARK
Setup and baselines.
The DomainBed benchmark (Gulrajani & Lopez-Paz, 2021) is a popular suite designed for rigorous comparisons of domain generalization methods.
DomainBed datasets focus on distribution shifts induced by synthetic transformations, We conduct extensive experiments on it to compare it with SOTA methods.
The testbed of domain generalization implements consistent experimental protocols for various datasets, and we use ﬁve datasets from DomainBed (excluding two MNIST-related datasets) in our experiments.
See data details in Table 1.
DomainBed offers a diverse set of algorithms for comparison.
Following the categories we summarized in Section 2.4, we compare with invariant element learning works: IRM (Arjovsky et al., 2019),
MMD (Li et al., 2018d), DANN (Ganin et al., 2016), and CORAL (Sun & Saenko, 2016).
Among optimization methods, we use GroupDRO (Sagawa et al., 2019) and MLDG (Li et al., 2018b).
The most closely related works are those based on gradient manipulation, and we compare with Fish (Shi et al., 2021) and Fishr (Rame et al., 2021).
Of the representation augmentation methods, we pick two popular works: MixUp (Zhang et al., 2017) and ARM (Zhang et al., 2021).
DomainNet’s additional model parameters in the ﬁnal classiﬁcation layer lead to memory constraints on our hardware at the default batch size of 32 .
Therefore, we use lower batch size 24 .
For our method variation PGrad-B , we set B = 3 for all datasets except using B = 2 for DomainNet.
We default to Adam (Kingma & Ba, 2017) as the optimizer to roll-out a trajectory.
All experiments use the DomainBed default architecture, where we ﬁnetune a pretrained ResNet50 (He et al., 2016).
Results analysis.
We aggregate results on each dataset by taking the average prediction accuracy on all domains, and the results are summarized in Table 2.
The per-domain prediction accuracy on each dataset is available in Appendix (A.7).
We summarize our observations: 1).
ERM remains a strong baseline between all methods, and gradient alignment methods provide promising results compared to other categories.
PGrad ranks ﬁrst out of 11 methods based on average accuracy.
Concretely, PGrad consistently outperforms
ERM on all datasets and gains 1 . 8% improvement on VLCS, 2 . 8% on OfﬁceHome, 2 . 9% on TerraIncognita, and no improvement on DomainNet.
Our variation PGrad-B outperform PGrad on almost all datasets except VLCS (where it is similar to PGrad ).
This observation showcases that intra-domain noise suppression can beneﬁt OOD generalization.
A longer trajectory enables
PGrad to learn more robust principal directions.
The combined variation PGrad-BMix outperforms MixUp across all datasets.
On average (see last column of Table 2), PGrad-BMix is the best performing strategy.
This observation indicates our method can be effectively combined with other
DG categories to improve generalization further.
[Table]
[Figure]
Tuning k for noise suppression.
As we pointed out in Section 2.1, we achieve domain-speciﬁc noise suppression by only aggregating coordinate axes { w z } kz =0 when learning the principal gradient ∇ p .
To investigate the effect of k , we run experiments with different values of k for both PGrad and PGrad-B .
The analysis results on PACS dataset are collected in Table 3.
Note that for default version of PGrad , the maximum number of training domains is n = 3 , therefore, the length of the
PGrad trajectory is upper bounded by 4 .
Table 3 shows that the generalization accuracy initially improves and then drops as k increases.
If we use k = n + 1 (as k = 4 for PGrad ), domain-speciﬁc noise is included and aggregated from principal coordinate W and the performance decreases compared with k = 3 .
The same pattern can also be observed in PGrad-B (note: the length of the trajectory is upper bounded by nB +1 = 10 ).
Training loss curve analysis.
Learning and explaining a model update’s behavior is an important step toward understanding its generalization ability.
To answer Q3 , we look for insights by visualizing domain-wise training losses as updates proceed.
To prevent randomness, we plot average results together with the standard deviation over 9 runs.
The results for ERM and PGrad are visualized in Figure 2.
Compared to ERM, our method PGrad has smoother decreasing losses as training proceeds.
Besides, all training domains beneﬁt from each update in PGrad .
On the other hand,
ERM’s decreasing loss on one domain can come with increases on other domains, especially late in training.
We hypothesize this is because domain-speciﬁc noise takes a more dominant role as training progresses in ERM.
PGrad can effectively suppress domain-speciﬁc noise and optimize all training losses in unison without signiﬁcant conﬂict.
Moreover, the loss variances across training domains are stably minimized, achieving a similar effect as V-REx (Krueger et al., 2021) without an explicit variance penalty.
In Appendix (A.2), we visualize four training trajectories trained with
PGrad and ERM.
ERM trajectories proceed over-optimistically at the beginning and turn sharply in late training.
PGrad moves cautiously for every step and consistently towards one direction.
[Table] [Table]
3. WILD S B ENCHMARK
WILD S
(Koh et al., 2021) is a curated benchmark of 10 datasets covering real-life distribution shifts in the wild such as poverty mapping and land use classiﬁcation.
We apply our method to its two vision applications.
Our goal is to investigate the scalability of PGrad under different model architectures and metrics, and more importantly, its performance when facing real-world
OOD shifts.
We conduct experiments on WILD S to explore both Q1 and Q2 .
For each dataset, we use the recommended metrics and model architecture.
A summary of the dataset, the metrics
[Table]
and the model architectures are provided in Table 4.
(I)
The P OVERTY M AP dataset is collected for poverty evaluation, where the input x ∈ X is an 8-channel multispectral satellite image, the output y ∈ Y is the real-value asset wealth index.
The dataset includes satellite images from 23 different countries and covers both urban and rural areas for each country.
We use 13 countries as training domains, pick 5 other countries for model selection, and use the remaining 5 countries for test purpose.
We calculate the Pearson correlation r between the predicted index and the groundtruth, and report the average r across 5 test domains and two different areas.
(II)
The objective of the FM O W dataset is to categorize land use based on RGB satellite images spanningyears and 5 geographical regions.
The training domains contain images from the ﬁrst 11 years, with the middle 3 years as validation domains and the last 2 years as test domains.
We report the average region accuracy on both validation and test domains to evaluate our method under the geographical distributional shift challenge.
The training details and the hyperparameters we used can be found in Appendix (A.9).
We compare ours with SOTA methods including IRM (Arjovsky et al., 2019),
Coral (Sun & Saenko, 2016), and Fish (Shi et al., 2021).
We repeat each experiment with three random seeds and report both recommended metrics and their standard deviations on each dataset.
In Table 5, PGrad achieves state-of-the-art results on both datasets, and the variation PGrad-B further improves the performance.
On P OVERTY M AP , PGrad demonstrates a strong correlation between the predicted wealth index and the ground truth by achieving the highest Pearson coefﬁcient on both validation and testing domains.
Its low standard deviation indicates PGrad is stable across different random seeds.
PGrad-B achieves better domain generalization by enabling more extended trajectory sampling.
Similarly, on the FM O W data, PGrad improves over baseline IRM and Coral, and is on par with Fish.
PGrad-B achieves the best accuracy.
These experimental results demonstrate that the proposed methods are effective across different model architectures and can successfully handle the real-life distributional shift.
C ONCLUSION
In this paper, we made an effort toward domain generalization by proposing a new training strategy.
Our method learns robust gradient directions by analyzing the parameter path created by sequential optimization on training domains.
The learned gradient direction, which we call the principal gradient, aggregates and explains the variation of the sampled optimization trajectory.
Principal gradients improve domain generalization by reducing the impact of gradient directions speciﬁc to individual training domains.
Our design also enables convenient gradient normalization and re-calibration for smooth multi-domain training.
Empirically, PGrad demonstrates state-of-the-art performance across two DG benchmark suites, covering both synthetic and real-world distribution shifts.
A S UPPLEMENTARY FOR PGrad : Learning robust gradients for domain generalization .
This appendix includes the following contents:
• We include ﬁgures describing the different trajectory sampling methods in A.1.
• To understand how PGrad changes the optimization path of model’s parameters, we show the TSNE projection of parameters during training in A.2.
In Figure 4, we visualize the projection for different datasets and test domains.
• The inclusion of the starting point during trajectory sampling is important; we analyze the reason for this effect in A.3.
• PGrad differs from related work in that it adopts a sequential training strategy when sampling a trajectory.
We include a detailed discussion comparing the parallel training used in gradient matching methods like Fish with our sequential training strategy in A.4.
We empirically demonstrate that gradient alignment is not necessarily a sufﬁcient indicator of the generalization ability, see Figure 7.
• To justify the analysis between PGrad under sequential training and parallel training, We design experiments comparing PGrad-B (Sequential),
PGrad-B (Parallel), and ERM.
The observations are consistent with the intuitive analysis.
Using ERM as baseline, PGrad-B (Sequential) enlarges clean direction and PGrad-B (Parallel) enlarges noise.
See Figure 6 in A. •
We show that our method has two major theoretical contributions.
First, instead of being driven by gradient alignment, PGrad learns a gradient ﬂow from the tangent space spanned by the eigenvectors of the Hessian matrix - see A.5.1 for the detailed analysis.
Second, in A.5.2, we show that PGrad has a deep connection to the efﬁcient training of neural networks by projecting the high-dimensional parameter space to some low-dimensional subspace.
To further understand the evolution of eigenvalue distributions over time, we visualize their average log values over 1k training step intervals in Figure 8.
• A critical question is how bottom eigenvectors affect the training dynamic and generalization ability.
We design experiments showing that those bottom eigenvectors span the normal subspace perpendicular to the tangent space of the loss landscape.
Updating the model with any directions lie within the normal subspace will make no changes to the training loss but hurt the generalization ability.
See A.6 for details. • Finally, in A.7, A.8, and A.9, we show the per-domain prediction accuracy of our PGrad variants and the experimental details on both benchmarks.
A. I LLUSTRATIVE EXPLANATION OF THE PROPOSED TRAJECTORY SAMPLING METHODS
We compare the three proposed trajectory sampling methods in Figure 3.
As vanilla trajectory sampling only considers inter-domain variations, the long trajectory sampling variant splits perbatch to smaller batches to eliminate intra-domain discrepancies from data collection, data sampling, etc.
[Figure]
A. R EAL O PTIMIZATION T RAJECTORY V ISUALIZATION
We visualize the optimization paths for both our method PGrad and the ERM baseline with TSNE projection.
Concretely, we save the model’s parameters into the memory buffer after everytraining steps and project them to the xy -plane after ﬁnishing 5,000 training steps.
The trajectories for different datasets and test domains are shown in the following Figure 4.
ERM moves over-conﬁdently with a large step size at the beginning, the turns sharply in late training.
Our method PGrad ‘thinks fast’ but ‘moves cautiously’.
It samples roll-out optimization trajectory [Figure] 1⋆ [Figure]
2⋆ Θ 3⋆ and aggregates the principal directions.
The training curves are consistently smooth even late in training.
A. T RAJECTORY S AMPLING
The inclusion of the starting point Θ t 0 in the trajectory sampling is important, otherwise the learned principal gradient will skip the gradient information from the ﬁrst training domain from each update.
We show the derivation in the following.
To simplify the notation, we substitute ∇ Θ ti − 1 L D i with ∇ Θ ti and assume we have only two training domains.
Sampled Trajectory With Starting Point Θ t
[Equation]
Trajectory Centering
[Equation]
[Equation]
Sampled Trajectory W/O Starting Point Θ t
S = {
[Equation]
Trajectory Centering
[Equation]
After trajectory centering, the gradient information from the ﬁrst training domain ∇ Θ t 1 will be skipped if trajectories are sampled without the starting point Θ t 0 .
To learn from all training domains, we use the left policy for sampling.
A. C OMPARISON BETWEEN THE PARALLEL TRAINING AND SEQUENTIAL TRAINING
In this subsection, we detailed the comparison between our method PGrad with the other two gradient-based methods: Fish and Fishr.
Both Fish and Fishr are inspired by gradient alignment.
They made efforts to align the gradients from different domains and adding the alignment as a penalty to the loss function.
In their vanilla implementation, both Fish and Fishr calculate the perdomain gradient w.r.t the current parameter Θ t 0 .
We can name the training paradigm as ‘parallel training’.
As contrast, our method learns a robust gradient direction from the optimization trajectory sampled by sequential training.
We explain why parallel training is not a proper choice when learning the principal direction in PGrad .
Principal directions learn dominant changing directions.
If we apply parallel training in PGrad , the centering of the trajectory will suppress the shared pattern and reinforce the domain-speciﬁc noises, see Figure 5(b).
Instead, the sequential training keeps enlarging shared gradient patterns with each of the multi-step updating.
We visualize and compare the two cases using Figure 5.
Moreover, the sequential training is more efﬁcient comparing with parallel training.
In Figure 5, we intuitively explain that sequential training will reinforce the learning of a clean direction and parallel training will signiﬁcantly suppress it.
We now design experiments to justify our hypothesis.
We adapt PGrad-B by learning the principal direction through parallel training.
We ﬁx all random mechanisms and compare PGrad-B under sequential training, PGrad-B under parallel training, and ERM.
The test accuracy as functions of the training epoch is visualized in
Figure. In the left panel, we use C, P, R as training domains. In the right panel, we use A, P, R as training domains. We attribute the performance drop of the PGrad-B (parallel) to the enlarged noise component and suppressed clean direction.
The observations are consistent with above analysis.
Another interesting observation is that curves from PGrad-B are smoother compared to ERM.
[Figure]
We highlight the novelty of our method by showing that PGrad achieves domain generalization without constraining the gradient alignment.
Speciﬁcally, we deﬁne a gradient alignment measurement as the mean of the cosine gradient similarity across all training domain pairs:
[Equation]
In Figure 7, We visualize the test domain accuracy and training domains gradient alignment as functions of the training epoch.
The ﬁgure implies the gradient alignment is not a sufﬁcient indicator of the generalization ability.
The test accuracy of the PGrad is lower bounded by Fish, but Fish aligns training domains’ gradient better than PGrad .
Secondly, the right ﬁgure shows the smoothness of the gradient alignment curves has a positive correlation with the test accuracy.
Starting from 3,000 epochs, the alignment curve of the PGrad becomes smooth, and the model achieves higher prediction accuracy in the test domain.
The starting phase of the Fish reveals the same pattern.
A. T HEORETICAL ANALYSIS
The basic motivation of PGrad is to learn a gradient ﬂow combines the top eigenvectors of the Hessian which are approximated by the average of the Fisher information matrix calculated sequentially on each training domain.
A.5. C ONNECTION TO H ESSIAN AND F ISHER I NFORMATION M ATRIX .
Hessian matrix, which is widely used for analyzing the model’s training behaviour, is deﬁned as the second order derivative of the function.
Similar to NTK (Jacot et al., 2018), we assume the loss function L is a functional acting on parameters Θ ,
[Equation]
We have Taylor expansion around parameter Θ :
[Equation]
where H i,j = ∂ 2 L ∂ Θ i ∂ Θ j .
The Hessian matrix H contains local geometric properties of the loss landscape around Θ .
The calculation of the second-order gradient is impractical, especially for modern neural network architectures.
Under certain mild regularity conditions and equipped with log-likelihood as loss function (Schervish, 2012), we can approximate H with the outer product of the gradient.
Speciﬁcally,
[Equation]
(16) where I is also known as Fisher information matrix (FIM).
We explain how our method PGrad automatically approximates and aggregates the eigenvalues of the Hessian matrix by following the proposed training procedures.
We sample a trajectory as
S =
{ Θ 0 , Θ 1 , · · · , Θ n } ∈ R ( n +1) × d .
In the following, we show the trajectory centering operation is equivalent to taking the average of the training domains’ Hessian approximations.
Lemma A.1 The centered trajectory ˆ S is the linear transformation of the domain-speciﬁc gradient, whose columns can be interpreted as the domain-wise gradient vector starting from the same initialization.
The shared initialization is the trajectory center of S .
The ﬁrst half of the Lemma is easy to show.
Any vector within a convex hull can be recovered by the linear combination of its edges.
See Appendix A.3 for the derivation on a simple case.
It implies that the centered trajectory ˆ S contains as rich information as training domains’ gradient matrix.
For the sampled trajectory S , the center is calculated as S o =
[Equation]
n + 1 .
We can re-interpret the sampled trajectory with local coordinate centered with S o .
Since the update step is small enough, we have:
[Equation]
The new gradients { ˆ ∇ i } n i =0 shares pseudo initialization S o .
We then proceed to show the covariance of the centered trajectory gives us the mean of the training domains’ FIM.
PGrad uses the training domains’ average FIM to approximate the real expected FIM.
[Equation]
where I i , H i are domain-speciﬁc FIM and Hessian matrix, respectively.
The approximation is a covariance matrix, which is positive semi-deﬁnite (SPD) and has the eigenvalue-eigenvector pairs { ( λ z , v z ) } nz =0 with λ 0 > λ 1 > · · · > λ n .
The eigenvalue λ z is the curvature of the loss in direction of v z in the neighborhood of S o .
The training behaviour of the neural network is determined by the distribution of the eigenvalues.
Speciﬁcally, ﬁrst-order optimization methods slow down signiﬁcantly when { λ z } nz =0 are highly spread out (Bottou & Bousquet, 2007; Ghorbani et al., 2019).
The property inspires us to zero out insigniﬁcant directions and use the directions with the large curvature to derive our gradient direction ∇ p .
As a contrast, Fishr (Rame et al., 2021) uses the current parameter value as the initialization and parallelly approximates the per-domain Hessian matrix with its gradient variance.
It deﬁnes the
Fishr regularization as the square of the Euclidean distance between gradient variance matrices to bring closer domain-wise Hessian.
Its goal is to align the second order gradient of the training domains.
The high computational cost from parallel training constrains them to operate only on the last classiﬁcation layer in practice.
Comparing with other gradient manipulation works which emphasize alignment or matching ,
PGrad uses the average of the FIMs to approximate the underlining Hessian matrix under the DG setup.
The operation can also reduce the noise variance by 1 /n , where n is the size of the training domains.
The property also provides an explanation of why PGrad-B achieves better generalization ability compared with PGrad .
Second, instead of matching eigenvalues of per-domain Hessian, we learn a robust gradient ﬂow which is the combination of the eigenvectors reﬂecting the dominant changes and zero out the remaining directions, which are empirically proved to be generalization toxic, see Table 3.
Our approximation of the Hessian matrix and introduced bijection allows us to efﬁciently operate on the high-dimensional parameter space .
A.5. C ONNECTION TO THE LEARNING BEHAVIOUR OF THE NEURAL NETWORKS
Modern SOTA neural networks are usually over-parameterized (Allen-Zhu et al., 2019; Zou & Gu, 2019).
Improving the training efﬁciency of high-dimensional neural networks is an active research direction (Li et al., 2018a).
Recent works (Gur-Ari et al., 2018; Gressmann et al., 2020) have shown that deep neural networks can be optimized in some subspace of much smaller dimensionality than their native parameter space, we show how PGrad connects to the line of the work.
It was proved by previous work (Li et al., 2018c; Gur-Ari et al., 2018; Gressmann et al., 2020) that the functional induced by the neural network varies most along only some speciﬁc directions.
We, therefore, focus on recovering the low-dimensional subspace where the loss function L varies the most on average, and project the native parameter space to the subspace.
We formulate the discussion in the following.
Given a direction v , the directional derivative of the loss L at Θ is deﬁned as:
[Equation]
and we can measure the expected scale (or length) of the directional derivative as:
[Equation]
However, the distribution of the Θ is unavailable, we turn to utilize the parameters from different training domains for empirical approximation.
[Equation]
We demonstrate that learning the principal gradient direction enables us to ﬁnd a low-rank updating space which is noise resistant by showing the following lemma.
Lemma A.2 Suppose we are searching for a k dimensional linear projection M ∈ R k × d of the original parameter space Θ ∈ R d , such that it keeps largest directional derivatives with respect to the loss L .
If the eigenvalue-eigenvector pair of the outer product matrix E Θ [ ∇ Θ L (Θ) ⊗ ∇ Θ L (Θ)] are { λ z , v z } nz =0 with λ 0 > λ 1 · · · > λ n .
We have:
[Equation]
To understand how eigenvalue changes during training, we visualize the relative contribution of each eigenvalue by normalizing with their summation and then plotting the log of their average over 1k training epochs in Figure 8.
In our experiments, we learn ∇ p by aggregating eigenvectors corresponding to the top-4 eigenvalues.
The ﬁgure indicates that the contribution from the smallest eigenvalues keeps decreasing as training continues.
Our method PGrad can effectively distinguish domain-speciﬁc noise and update with shared patterns by suppressing the information from small eigenvalues.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
The length of the trajectories is nB + 1 = 10 .
We plot the distribution of the top-9 eigenvalues since the smallest one is out-of-precision.
We calculate the contribution of each component by normalization with their sum.
To smooth the results, we take the log value of the average across per-thousand training epochs.
The ∇ p is learned by aggregating the top-4 eigenvectors.
A. W HAT ARE THE TAIL EIGENVECTORS ?
To answer the last question we proposed in Section 3, we conduct the analysis to show the effect of the bottom eigenvectors on the training dynamics. Ablation studies in Table 3 indicate that including bottom eigenvalues into principal gradient will hurt the generalization ability. We add new experiments to clarify that the bottom eigenvectors are noise signals of ‘special’ properties.
Concretely, we design three different strategies to update the model with PGrad :
• Always from bottom eigenvectors.
• Start from the top eigenvectors and then switch to the bottom vectors in the middle.
• Always from the top eigenvectors.
The training losses keep being constant for case (1) and case (2) after the switching, even when we set the step size to be meaningfully large.
The training loss keeps getting decreased for the setup (3).
These results imply that those bottom components span the subspace perpendicular to the tangent space of the loss landscape.
They do not hurt the training loss but are not helpful for generalization.
Similar to the original setup, we calibrate the direction and length for optimization purpose.
We show the training loss curves in Figure 9.
A. E XPERIMENTAL S ETUP D ETAILS FOR D OMAIN B ED
Following the training protocol described in DomainBed, we run experiments on each domain with a random mechanism to reduce the bias of hyperparameter selection.
The DomainBed experiments (Gulrajani & Lopez-Paz, 2021) select the best model from random samples of the hyperparameter search space, and repeat this search process 3 times.
Speciﬁcally, for each domain, we selectrandom combinations of hyperparameters from a relatively narrow range and repeat the searchtimes to report conﬁdence intervals.
For a dataset with n domains, we train n models.
T te , held out for testing, and use the rest as training domains D tr .
This design leads to a total of 6 n experiments per dataset for each method.
[Figure] Lacking access to data from the test distribution during training makes it hard to perform model selection for DG compared to other supervised learning tasks.
In the main [Figure] we follow the popular setup where we sample 20% of the data from each training domain and group them as a validation set [Figure]
The validation accuracy will be used as an indicator of the optimal model.
We show the prediction accuracy of PGrad and the extensions on each domain in Tables 7-11.
A. H YPERPARAMETER SEARCH
We following a similar training and evaluation protocol to the DomainBed experiments (Gulrajani & Lopez-Paz, 2021).
We list our hyperparameter search space in Table 6. [Table]
We design a narrow search space to prevent undersampling and potential performance degeneration from a smaller number of random trials.
We use a ﬁxed batch size of 32, which ensures there are adequate samples to provide precise gradient directions when sampling with our long and highentropy method.
For example, there are still 11 samples in each small batch if B = 3 .
We also found the outer learning rate γ = 0 .
1 works consistently well on each dataset.
DomainNet’s additional model parameters in the ﬁnal classiﬁcation layer lead to memory constraints on our hardware at the default batch size of 32 .
Therefore, we use a lower batch sizes 24 , and increase the training step to 6 , 000 .
A. E XPERIMENTAL D ETAILS ON WILD S B ENCHMARK
A summary of the two vision datasets in the W ILDS benchmark is shown in Table 4.
For the P OVERTY M AP task, the inputs are 8-channel satellite images, therefore, we tuned the ﬁrst convolutional layer of the ResNet18 (He et al., 2016) to accommodate a multispectral input.
To sample a trajectory of length ˆ n +1 in high dimensional parameter space Θ , we randomly sample ˆ n training domains and sequentially update the parameter starting from Θ t on each of them.
We use the Adam optimizer for trajectory construction and set the default learning rate to be 1 e − 3 without weight decay.
We set the outer learning rate γ = 0 .
1 / ˆ n , which adjusts the step size based on the number of sampled training domains.
We train the model with the proposed PGrad for 200 epochs and select the optimal model based on the average Pearson coefﬁcient on validation domains.
For the land use classiﬁcation task FM O W, we follow the exact same training and evaluation protocol applied in Fish (Shi et al., 2021).
We ﬁnd a good starting point by updating the ERM objective with an Adam optimizer for 24,000 iterations with a learning rate of 1 e − 4 .
After pretraining, we proceed with our proposed PGrad and keep tuning the model for 10 epochs with outer learning rate γ = 0 .
01 / ˆ n .
After training is completed, we report the worst regional accuracy on the test domains.
This measurement is designed by W ILDS (Koh et al., 2021) to test models’ generalization ability when facing both time and regional distribution shift.
The numerical results for both datasets are available in Table
[Table] [Table]
[Table] [Table]
[Table] [Table]