PGrad: Learning Principal Gradients For Domain Generalization
Authors: Zhe Wang
Jake Grigsby
Yanjun Qi
Venue: ICLR 2023 poster (2023)

Abstract: 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 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 refinement 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 WILDS 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 efficacy across both synthetic and real-world distributional shifts.

Published as a conference paper at ICLR 2023
PGRAD LEARNING PRINCIPAL GRADIENTS FOR Do- MAIN GENERALIZATION
Zhe Wang, Jake Grigsby, Yanjun Qi Department of Computer Science University of Virginia Charlottesvilll, VA 22903, USA {zw6sg, jeg6dn, yq2h sq2h]evirginia.ee
ABSTRACT
Machine learning models fail to perform when facing out-of-distributio (OOD) domains, challenging task known as domain generalization (DG). this work, we develop novel DG training strategy, we call PGrad. to learn robust gradient direction, improving models' generalization ability on unseen domains. The pro- posed gradient aggregates the principal directions of sampled roll-out optimiza- tion trajectory that measures the training dynamics across all training domains. PGrad is gradient design forces the DG training to ignore domain-dependent noise signals and updates all training domains with a robust direction cover- ing main components of parameter dynamics. We further improve PGrad via bijection-basse computational refinement and directional plus length-base cali- brations. Our theoretical proof connects PGrad the spectral analysis of Hes- sian in training neural networks. Experiments on DomainBed and WILDS bench- marks demonstrate that our approach effectively enables robust DG optimization and leads to smoothly decreased loss curves. Empirically, PGrad achieves com- petitive results across seven datasets, demonstrating its efficacy across both syn- thetic and real-world distributional shifts.
INTRODUCTION
Deep neural networks have shown remarkable generalization ability on test data following the same distribution their training data. Yet, high-capacity models are incentivized to exploit any cor- relation the training data that will lead to more accurate predictions. As a result, these models risk becoming overly reliant on "'domain-specific" correlations that may harm model performance on test cases from -ittt--distriuttion (OOD). A typical example is a camel samel-and-cows classification task (Beery et al. 2018 Shi al.) 2021), where camel pictures training are almost always shown in desert environment while cow pictures mostly have green grassland backgrounds. A typical ma- chine 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 overfits to specific 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 chal- lenge. They concerned with how to learn 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 learn- ing, meta-learning data augmentation, ensemble learning, and gradient manipulation (more details in Section 2.4p. Despite the large body of recent DG literature, the authors of (Gulrajani & Lopez- Paz 2021 showed that empirical risk minimization (ERM) provides a competitive baseline on many weal-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 av- erage gradient. However, one caveat of ERM is its average gradient-based model update will pre- serve domain-speciff noise during optimization. This observation motivates the core design of our method.
In the rest of this paper, use the terms "domain" and "distribution" interchangeably.

---

Published as a conference paper at ICLR 2023
Train:
O,D Dgii
D3
Bf -1
Ri
Evai:
8F
4(8%)
Principel Veciors' Construction
Diractin & Longth Calibration
Noisa Suppression
Figure/ 1: Overview of our PGrad training strategy. With a current parameter O', we first obtain rollout trajectory Of Of + O Og by sequentially optimizing across all training domains Drr WE Then PGrad updates O by extracting the principal gradient direction V, of the trajectory. target model's generalization is evaluated on unseen (OOD) domains T
We propose novel training strategy that learns 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 local coordinate system to explain the parameter variations the trajectory. Via singular value decomposition (SVD), we derive 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 name the "principal gradient" reduces domain-specific noise in the DG model update and prevents the model from overfitting to particular training domains. To decrease the computational complexity of SVD, we construct a bijection between the parameter space and 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 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 benefits the training efficiency via learning gradient in a smaller subspace constructed from learning trajectory. (4) Our empirical results demonstrate the competitive performance PGrad across seven datasets covering both synthetic and real-world distributional shifts.
2 METHOD
Domain generalization (Wang et al., [021] Zhou et al. 2021) assumes no access to instances from future unseen domains. domain generalization, we are given set of training domains D.r = {D.) and test domains Te. = (T, (m Each domain D; (or T;) is associated with joint distribution PKxy (or PKxy), where X' represents the input space and ) is the output space. Moreover, each training domain D; is characterized by set of i.i.d samples {x] yk For any two different domains sampled from either D,, or Tee their joint distribution varies PKxy # P PXxyy and most importantly, Prxxy # P
We consider the prediction task from the input x X to the output y e e. Provided with a model family whose parameter space is θ dR and the loss function L O x (X x y)) R,, the goal to find an optimal O on test domains so that:
0 arg min EI,~TI (x.))~ (x 410, (x. ece
In DG setup, any prior about Ti. 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-Pa7 2021). ERM naively groups data
2

---

Published as a conference paper at ICLR 2023
from all training domains Dr together and obtains its optimal parameter Or via the following to approximate Oze
i arg min ED,-D,, IE, (x.y) (xy)) C (x,
og
In the rest of the paper, we omit the subscript in ①r and use θ for simplicity (during DG training, only training domains Drr will be available for model learning).
When optimizing with ERM DG across multiple training domains, the update of O follows:
θ = = -0'-2 Cer
(3)
where VerLD, = VLE (x Pxxy
A (x, y)] calculates the gradient of the loss on domain D; with respect to the current parameter Oc and is the learning rate.
The gradient determines the learning path of model. When using ERM in DG setting, each step of model updates uses an average gradient and may introduce and preserve domain-specific 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 do- mains, envision such domain-specific noise signals will not be the main components of parameter variations across all domains. This motivates us to design P PGrad as follows.
2.1 PGRAD: PRINCIPAL GRADIENT BASED MODEL UPDATES
We extend ERM with robust gradient estimation that we call PGrad. We visualize an overview in Figure [ to better explain how it works. Briefly speaking, given the current model parameter vector, we sample a trajectory of parameters by training sequentially on each training domain. Next, we build 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 filter out domain-specific noise and follows direction that maximally benefits all training domains Der; 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.
Denote the current parameter vector as O We first sample a trajectory S through the parameter space θ by sequentially optimizing the model each training domain:
Trajectory Sampling.
e @' 0 0i -
aVe; ,.D, i (1....,,)) ,n]
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 sampled trajectory S (0
R(*+1jxd that were derived from O'. Note: the inclusion of the starting location O as part of the trajectory is necessary; see the proof in Appendix (A.3).
Then we construct local principal coordinate system to explain the variations in S. We are looking for orthogonal and unit axes V = | E (x++ddx to maximally capture the variations of the trajectory. Each Ug E Idd a unit vector of size d, the same dimension the parameters et
max Variance((So..), s.l. VTV =Ia- 0,
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)):
>e. U. = SVD. 1 -TS
S S),
Here Ag v. denote the z-th largest eigenvalue and its corresponding eigenvector. å denotes the cen- tered trajectory by removing the mean from S. In the above Eq. (6)), the computational bottleneck lies in the SVD, whose computational complexity comes at ((d3) due to ST Š E Raxd d denotes
3

---

Now our main goal is to get 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 reference. The choice of the reference flexible, 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 0f and the end point Published as a conference paper at ICLR 2023
the size of the parameter vector and is fairly large for most state-of-the (SOTA) deep learning architectures. This prohibits the computation of the eigenvalues and eigenvectors from Eq. (6) for SOTA deep learning models. Hence, we refine and construct bijection as follows to lower the computational complexity (to O((n + 1)9)):
S d2e= = STSST =
U2 ST e,
Eq. (7) indicates that if Y are the z-th largest eigenvalue and corresponding eigenvector of SST the z-th largest cigenvalue and corresponding cigenvector of S S are X2 ST ez (i.e., Ug = ez) This property introduces a bijection from eigenvectors of S3" E [(=++)x(n+1) to those of S E Rdxd Since n the number of training domains is much smaller than d, calculating eigen- degeppecompooitionn of ST E (==++))((n=1) is therefore much cheaper.
Directional Calibration. With the derived orthogonal axes V [v]] from Eq. 7 now construct a local principal coordinate system with each axis aligning with one eigenvector v.. These principal coordinate axes V are ordered based on the magnitude of the eigenvalues. This means that v; explains more variations of the sampled trajectory S than v, if < 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.
O,, of the trajectory S as reference V, O O So for each coordinate axis, revise its direction so that the resulting vector w, is positively related to the reference V, in terms of the inner product:
if (v,, 2,) -1, otherwise.
1,
W² U2, 1-=1-1,
(8)
Constructing Principal Gradient. The relative importance of each w, is conveyed in the cor- responding eigenvalue Az. Larger A implies higher variance when projecting the trajectory S along w, 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:
wT A = [Ao A1. An]
There exists other possible aggregation besides Eq. (9) For instance, another choice of weight could be X=/ A or simply A since the eigenvalue of semi-positive definite matrix is non-negative. Gradient normalization has been widely recommended for improving training stability (You al. Yu et al. 2017). Our design in Eq. (9) automatically achieves L2 normalization, because:
2017
1 1
(10)
Length Calibration. As training updates continue, a fixed length gradient operator may become too rigid, causing fluctuations in the loss. We, therefore, propose to calibrate the norm of Vp with a reference, for achieving adaptive length tuning. Specifically, we propose to multiply the aggregated gradient from Eq. (9) with the L2 norm of V,:
n
A 1 ||||
w.
(11)
With this length calibration via |r multiplier, and is automatically tuned during the training process.
1|2. the norm of the proposed gradient is constrained by the
Noise Suppression. Most w axes correspond to small eigenvalues and may relate to just domain- specific noise signals. They may help the accuracy of specific training domain D., but mostly hurt
4

---

Published as a conference paper at ICLR 2023
the overall accuracy Dry We therefore define the principal gradient as follows and show how to use to solve DG model training via gradient descent optimization (where is a hyperparaemter).
V | 2
22, 0 = ~VV.
(12)
2.2 THEORETICAL ANALYSIS
In Appendix (A.5.1). J we prove that ST in Eq. (6) provides us with the mean of all training 71
domains' domain-specific Fisher information matrix (FIM). Since FIM 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 efficiency of neural networks by recovering subspace from the original over-parameterized space O This subspace is built from the top eigenvectors of the approximated Hessian. We visualize the evolution of the eigenvalue distributions in Figure
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 training efficiency and generalization in neural networks (Wei & Schwab) [0119] Ghorbani al., 2019 Pascanu et al. 2014). For example, (Hochreiter & Schmidhuber 1997) shows that small eigenvalues in the Hessian spectrum are indicators of flat directions. Another work (Sagun et al. 2017) empirically demonstrated that the spectrum of the Hessian contains both bulk component with many small eigenvalues and few top components of much more significant positive eigenval- ues. Later, (Gur-Ari et al. 2018) pointed out that the gradient of neural networks quickly converges to the top subspace of the Hessian.
2.3
VARIATIONS OF PGRAD
There exist many ways to construct sampled trajectory, creating multiple variations of PGrad.
PGrad-F: : The vanilla trajectory sampling method will sample a trajectory of length n + by sequentially visiting each D, in fixed order. See appendix for the results of the rigid variation.
PGrad: We can randomly shuffle the domain order in the training, and then perform to sample 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 long sampled trajectory that is with length n B +
PGrad-BMix: Our method is model and data agnostic. Therefore it 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-BM
our empirical analysis.
In PGrad and PGrad-F, the principal gradient's ' trajectory covers all training domains Der exactly once (per domain). There are two possible limitations. (1) If the number of training domains 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 get one axis wx, that goes back ERM. (2) 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 xxf.y from training domain D, into B small batches. This new strategy allows us to obtain nB pseudo training domains. Such design brings two benefits: (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 if they come from different training domains. By learning the principal gradient from these nB pseudo domains, also address the intra-domain noise issue. We name this new design PGrad-B. Appendix (A.1) includes a figure comparing vanilla trajectory sampling with this extended trajectory sampling.
2.4 CONNECTING TO RELATED WORKS
We can categorize existing DG methods into four broad groups.
Invariant element learning, Learning invariant mechanisms shared across training domains pro- vides promising path toward DG. Recent literature has equipped various deep learning components
5

---

Published as a conference paper at ICLR 2023
especially representation modules with the invariant property to achieve DG ( et al. 2018d ej Muandet et al. 2013). The central idea is to minimize the distance or maximize the similarity be- tween representation distributions ((f(X)] across training domains that prediction is based
on statistically indistinguishable representations. Adversarial methods (Li et al., (018d) and mo- ment matching (Peng et al., 2019] Zellinger al., 2017) two promising approaches for distribu- tional alignment. A recent line of work explores the connection between invariance and causality. IRM (Arjovsky et al., 2019) learns an invariant linear classifier that is simultaneously optimal for all training domains. Under the linear case and some constraints, the invariance of the classifier in- duces causality. Ahuja et al. further extend IRM by posing it as finding the Nash equilibrium (Ahuja 2020) and adding information bottleneck constraints to seek theoretical explanations (Ahuja et al. 2021 However, later works (Kamath et al., 202T) show that even when capturing the cor- rect 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.
et al.
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 dis- tributions (rather than the average loss). The second set of optimization DG methods is optimization- based meta-learning. Optimization-basee ing properties like global inter-class alignment (Dou et al. 2019) or local intra-class distinguisha- bility (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.
Optimization methods. One line of optimization-
peeanlle uses bilevel optimization for DG by achiev-
Gradient manipulation. Gradient directions drive the updates of neural networks throughout train- ing and are vital elements of generalization. In DG, the main goal is to learn a gradient direction that benefits all training domains (plus unseen domains). Gradient surgery (Mansilla al., 202I) pro- poses 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. Sand- mask (Shahtalebi et al., (02T) added 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. 202T) 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 current parameter vector. Fish maximizes the inner product of domain-level gradients; Fishr uses the variance of the gradient covariance as regularizer to align the per-domains' Hessian matrix. Our method PGrad differs gradient matching by learning a robust gradient direction. Besides, our method efficiently approximates the Hessian with training domains' Fisher information matrix. Ap- pendix (A.4) includes 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 ( We then show that gradient alignment is not necessary a sufficient indicator of the generalization ability in Figure 7
Others. Besides the categories above, there exist other recently adopted to conquer domain gen- eralization. Data augmentation (Xu et al.] 2021: Zhang et al.) 2021 2017: Volpi et af. 2018). which generates new training samples representations from training domains prevent over- fitting. Data augmentation can facilitate a target model with desirable properties such as linearity via Mixup (Zhang et al., (017)) or object focus (Xu al. (202T). Other strategies, like contrastive learning (Kim et al. 20221) representation disentangling (Piratla et al.) 2020). and semi-supervised learning (Li ct al. 2021), have also been developed for the DG challenge.
3 EXPERIMENTS
We conduct empirical experiments to answer the following questions: QI. Does PGrad success- fully handle both synthetic and real-life distributional shifts? Q2. Can PGrad handle various ar- chitectures (ResNet and DenseNet), data types (scene and satellite images), and tasks (classification 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 comple- mentary approach to combine with other DG strategies? [ Q5. How do bottom eigenvectors in the roll-out trajectories affect the model's training dynamic and generalization ability?
Note: we leave hyperparameter tuning details and some ablation analysis results in Appendix A.7 A.9

---

Published as a conference paper at ICLR 2023
Table 1: A summary on DOMAINBED dataset, metrics, and architectures we used.
Dataset
of Images
Domains
# of Classes
PAC' et 2017
9,991
Artpaint, Cartonn Sketches, Photo
VLCS
Fang al 2013
10,729
PASCAL VOC 2007, LabelMe, Caltech, Sun
OFFICEHOME Venkatesware el 2017
15,588
Art. Cliputt, Product, Real-World Location #100, #38, #43, #46
65
TERRAINCOGNITA
ar 2015
24,788
DOMAINNET
PPnn of a. 2019
586,575
Clipart, Infograph. Painting, Quickdraw, Real, Sketch
345
Table 2: Test accuracy (%) on five datasets from the DomainBed benchmark. We group 20% data from each training domain to construct validation set for model selection.
Categories
Algorithms
VLCS
PACS
OfficeHome
Terralic
DomainNet
Arg
Baseline
ERM
77.5
0.4
85.5 0.2
66.5 0.3
46.1
1.8
40.9 + 0.1
63.3
IRM
78.5 + 0.5
83.5 + 0.8
64.3 + 2.2
47.6 + 0.8
33.9 + 2.8
61.6
MMD
77.5 0.9
84.6
66.3 0.1
42.2 = 1.6
23.4 + 9.5
58.8 -48
Invariant
DANN
1..6
83.6 :: 0.4
65.9 0.6
46.7 = 0.5
38.3
4 0.1
62.6 a7
CORAL
78.8 1.8:000 0.6
86.2 = 0.3
68.7 + 0.3
47.6 1.0
41.5 0.1
64.5 41.2
Optimization
GroupDRO
0.6
84.4
66.0 + 0.7
43.2 1.1
33.3
60.7 -2.6
MLDG
0.4
84,9 1.0
66.8
47.7 0.9
41.2 0.1
63.6 .03
Augmentation
MixUp
0..
84.6 0.6
68.1 + 0.3
47.9 + 0.8
19.2++0.1
63.4
ARM
0.3
85.1
64.8 0.3
45.5 0.3
35.5 - 0.2
61.7 -16
Fish
77.8
85.5
68.6 -
= 0.2
63.9 406
Fishr
77.8 1: 0.1
85.5 0.4
67.8 4: 0.1
47.4 +: 1.6
41.7 +: 0.0
64.0
Gradient Manipulation
PGrad
79.3 41
85.1 1
69.3 0.1 +2.8
49.0
0.3 +2.9
41.0 1 0.1 -01
64.7 -1.4
PGrad-B
78.9 : 0.3 +14
87.0
: 0.1 -18
69.6
.3.1
49.4
+3.3
41.4 0.1 -05
65.3 -2,0
PGrad-BMLL BMix
78.9 1: 0.2
86.2 1 0.4 .077
69.8 0.1
50.7 1 0.6
42.6 1 0.2
65.7
3.1 DOMAINBED BENCHMARK
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 compare it with SOTA methods. The testbed of domain generalization implements consistent experimental protocols for various datasets, and we use five datasets from DomainBed (excluding two MNIST-related datasets) in our experiments. See data details in Table !
DomainBed offers a diverse set of algorithms for comparison. Following the categories we summa- rized in Section 2.4 compare with invariant element learning works: IRM (Arjovsky et al., 2019), MMD (Li et al. 20 2018d), DANN (Ganin et al. 20166), and CORAL (Sun & Saenko (016). Among use GroupDRO (Sagawa al. 2019) and MLDG (Li al. (018b). The most closely related works are those based on gradient manipulation, and we compare with Fish (Shi al. 2 and Fishr (Rame et al.] 2021). Of the representation augmentation methods, we pick two popular works: MixUp (Zhang et al. (017) and ARM (Zhang et al. 202T). DomainNet's 's additional model parameters in the final classification layer lead to memory constraints on hardware the default batch size of 32. Therefore, we use lower batch size 24. For our method variation P PGrad-B for DomainNet. We default to Adam (Kingma)
optimization methods,
we set B for all datasets except using & Ba 2017) as the optimizer to roll-out trajectory. All experiments use the DomainBed default architecture, where we finetune 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
each dataset is available in Appendix (A.7).
We summarize our observations: 1). ERM remains a strong baseline between all methods, and gra- dient alignment methods provide promising results compared to other categories. 2). PGrad ranks first 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 OfficeHome, 2.9% on Ter- ralncognita, and no improvement on DomainNet. 3) 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 benefit OOD generalization. A longer trajectory enables PGrad to learn more robust principal directions. 4) The combined variation PGrad-BMix outper- forms MixUp across all datasets. On average (see last column of Table 2), PGrad-BMa is the best performing strategy. This observation indicates our method can be effectively combined with other DG categories to improve generalization further.
7

---

Published as a conference paper at ICLR 2023
Table 3: Analysis the effect of varying k. The experiments are performed on PACS dataset. We highlight first and second best results.
Method
Algorithms
A
C
Avg
k =
98.0+0.2
07.
76,8+0.4
13
83.9
PGrad
k 97.8+0.0
87.5+0.3
78.2:0..8
74.0
84.4
k 97.8++.0
87.8+0.4
78.410.6
77.2+11
85.3
k
97.4+0..1
87.6+0.3
79.1+1.0
76.3+112
85.1
k =
97.5+0.1
89.1+0.8
80.3+0.6
77.5+0.4
86.1
k 97.7+0.2
88.5+1.0
79.9+1.1
79.2+0.7
86.4
PGrad-B
k 98.010.2
89.9
80.010.6
80.1
87.0
k
97.6++.33
88.2++0.8
81.1+
1..0
86.5
(a) ERM: Tec V
(b) ERM: Tic L
(c) ERM: Tee C
(d) ERM: Tee = 8
(e) PCrad: Te: (f) PGrad: The
(g) PGrad: Tre = C
(h) PGrad: Tie S
Figure 2: Visualzing domain-wise training losses on VLCSS. Curves are average over 9 runs, surrounded by =o the standard deviation. For comparison, the loss curves start from 1,000 epochs.
Tuning k for noise suppression. As we pointed out in Section 2.1 noise suppression by only aggregating coordinate axes {w, 1 ===0 when learning the principal gradi- ent Vp. To investigate the effect k., we run experiments with different values of & for both PGr ad and PGrad-B. The analysis results PACS dataset are collected in Table Note that for default version of PGrad, the maximum number of training domains is = 3, therefore, the length of the PGrad trajectory is upper bounded by 4.
we achieve domain-speciff
Table 3 shows that the generalization accuracy initially improves and then drops as increases. we use k: = n 1 (as k = 4 for PGrad), domain-specifie noise is included and aggregated from principal coordinate W and the performance decreases compared with 3. The same pattern can also be observed in PGrad-B B (note: the length of the trajectory is upper bounded by nB + 10).
Training loss curve analysis. Learning and explaining model update's behavior is important step toward understanding its generalization ability. To answer Q3, we look for insights by visualiz- ing 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 Figure Compared ERM, our method PGrad has smoother decreasing losses as training proceeds. Besides, all training domains benefit 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 training progresses in ERM. PGrad can effectively suppress domain-specific noise and optimize all training losses in unison without significant conflict. Moreover, the loss variances across training domains are stably minimized, achieving a similar effect as V-REx (Krueger et al. 2021) without explicit variance penalty. In Appendix (A.2), visualize four training trajectories trained with PGrad and ERM. ERM trajectories proceed over-optimistically at the beginning and turn sharply in late training. PGrad d moves cautiously for every step and consistently towards direction.
noise takes more dominant role as
8

---

Published as a conference paper at ICLR 2023
Table 4: A summary on WILDs dataset, metrics, and architectures we used.
Train
Test
Dataset
Domain Types
Input
Output
Domains
Domains
Domains
Metric
Arch.
MAP
Urban/Rural (2)
Countries (23).
Satchlic
Images
Asset (jeal valued)
(f)
ReiNei-18
FMoW
Regions (5)
Time (16).
Satellite
Images
Land Use (62 classes)
11
3
2
Aug Region Acc.
DenseNet. 121
Table 5: Results on WILDs benchmark. Top two results are highlighted.
FMOW
POVERTYMAP
Method
Val. Accuracy (%)
Test Accuracy (%)
Val. Pearson
Test Pearson
IRM
57.2.10.01
50.9110.2
0.804
0.78
Coral
56.71:0.06
50,5
0.8010.04
0.77
Fish
57.31:.01
51,8::012
0.8010.01
0.80#0.01
PGrad
57.5.10.15
51,9110,10
0.81
0.8010.01
PGrad-B
57,9110088
52.1:10.09
0.82.18044
0--2-15e-3
3.2 WILDS BENCHMARK
WILDs (Koh al. 2021) is curated benchmark of 10 datasets covering real-life distribution shifts in the wild such as poverty mapping and land use classification. We apply our method to its two vision applications. Our goal is 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
WILDs to explore both Q1 and Q2.
For each dataset, we use the recommended metrics and model architecture. A summary of the dataset, the metrics, and the model architectures are provided in Table 4 (I) The POVERTYMAP dataset collected for poverty evaluation, where the input x E X' an 8-channel multispectral satellite image, the output y E >/ 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 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 across test domains and two different areas. (II) The objective of the FMoW dataset to categorize land use based on RGB satellite images spanning 16 years and 5 geographical regions. The training domains contain images from the first 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 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 metries 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 POVERTYMAP. PGrad demonstrates strong correlation between the predicted wealth index and the ground truth by achieving the highest Pearson coeffi- cient on both validation and testing domains. Its low standard deviation indicates PGrad is stable across different random seeds. PGrad-B B achieves better domain generalization by enabling more extended trajectory sampling. Similarly, on the FMoW data, PGrad improves over baseline IRM and Coral, and is on par with Fish. PGrad-B achieves the best accuracy. These experimental re- sults demonstrate that the proposed methods are effective across different model architectures and can successfully handle the real-life distributional shift.
CONCLUSION
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 gradi- ent, aggregates and explains the variation of the sampled optimization trajectory. Principal gradients improve domain generalization by reducing the impact of gradient directions specific to individ- ual training domains. Our design also enables convenient gradient normalization and re-calibration for smooth multi domti- training. Empirically, PGrad demonstrates state-of-the-as performance across two DG benchmark suites, covering both synthetic and real-world distribution shifts.
We follow the exact same training protocols as in Fish el. with ERM suboptimal starting and then proceed with PGrad.
2021 On FMoW, we pretrain the model

---

Published as a conference paper at ICLR 2023
REFERENCES
Kartik Abuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. In ICML. pp. 145-155. PMLR, 2020.
Kartik Ahuja, Fthan Caballero, Dinghuai Zhang, ChnnstCpitttpheee Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. NeurIPS, 34, 2021.
Zeyuan Allen-Zhu Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameter ized neural networks, going beyond two layers. NewrIPS, 32, 2019.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 0tt:::9977.02000 002893, 2019.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In ECCV, pp- 456-473, 2018.
Léon Bottou and Olivier Bousquet. The tradeoffs of large scale learning. NeurIPS, 20, 2007.
Prithvijit Chattopadhyay, Yogesh Balaji, and Judy Hoffman. Learning to balance specificity and invariance for and out of domain generalization. ECCV, pp. 301-318. Springer, 2020.
Qi Dou, Daniel C. Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. In NeurIPS, 2019.
Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: the utilization of multiple datasets and web images for softening bias. In ICCV, pp. 1657-1664, 2013.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural net- works. The journal of machine learning research, 17(1):2096-2030, 2016.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. ICML, pp. 2232-2241 PMLR, 2019.
Frithjof Gressmann, 7ach Raton-Rosen, and Carlo Luschi. Improving neural network training in low dimensional random bases. NeurlPS, 13:12140-1215) 2020.
Ishaan Gulrajani and David Lopez-Paz. search of lost domain generalization. In ICLR, 2021.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.0477 2018
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. NeurIPS, 31, 2018.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk mini- mization capture invariance? In AISTATS, pp. 4069-4077 PMLR, 2021.
Daehee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jaekoo Lee. Selfreg: Self- supervised contrastive regularization for domain generalization. In ICCV, pp. 9619-9628, Oc- tober 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wile wild distribution shifts. In ICMI., pp. 5637-5664. PMLR, 2021.
10

---

Published as a conference paper at ICLR 2023
David Krueger, Ethan Caballero, Joem-Henrik Jacobsen, Amy Zhang. Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex). In ICML pp. 5815-5826. 1MLR,20211
Bo Li, Yezhen Wang, Shanghang Zhang, Dongsheng Li, Kurt Keutzer, Trevor Darrell, and Han Zhao. Learning invariant representations and risks for semi-supervised domain adaptation. In CVPR. pp. 1104-1113, 2021.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. arXiv preprint arXiv: 1804. 08838, 2018a.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In CVPR, pp. 5542-5550, 2017.
Da Li, Yongxin Yang, Yi-Zhe Song. and Timothy M Hospedales. Learning to generalize: Meta- learning for domain generalization. AAAI, 2018b.
Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy M Hospedales. Episodic training for domain generalization. In ICCV, pp. 1446-1455, 2019.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. NeurIPS, 31, 2018c.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex Kot. Domain generalization with adversarial feature learning. In CVPR, pp. $400-54001 2018d.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In ECCV, pp. 624- 639, 2018e.
Lucas Mansilla, Rodrigo Echeveste, Diego H Milone, and Enzo Ferrante. Domain generalization via gradient surgery. ICCV, pp. 6630-6638, 2021.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In JCML, pp. 10-18. PMLR, 2013.
Razvan Pascanu, Yann N Dauphin, Surya Ganguli, and Yoshua Bengio. On the saddle point problem for non-convex optimization. arXiv preprint arXiv:1405. 4604,, 2014.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV, pp. 1406-1415, 2019.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Efficient domain generalization via common-specific low-rank decomposition. In ICML, pp. 7728 -7388. PMLR, 2020.
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. arXiv preprint arXiv:2109.0200 2021.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust
neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint arXiv: 1et::1911...1 2019.
Levent Sagun, Utku Evei, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv:1706.04455 2017.
Mark J Schervish. Theory of statistics. Springer Science & Business Media, 2012.
Soroosh Shahtalebi, Jean-Christophe Gagnon-Audet, Touraj Laleh, Mojtaba Faramarzi, Kartik Ahuja, and Trina Rish. Sand-mask: An enhanced gradient masking strategy for the discovery of invariances domain generalization. arXiv preprint arXiv:2106.00 022666 2021.
Yuge Shi, Jeffrey Seely, Philip H. S. Torr, N. Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv: 104.0993, 2021.

---

Published as a conference paper at ICLR 2023
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In ECCV, pp. 443-4444555 450. Springer, 2016.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In CVPR, pp. 5018-5027, 2017.
Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing to unseen domains via adversarial data augmentation. NeurIPS, 31, 2018.
Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Wenjun Zeng, and Tao Qin. Generalizing to unseen domains: A survey on domain generalization. arXiv preprint arXiv:2103.03099, 2021.
Mingwei Wei and David J Schwab. How noise affects the hessian spectrum in overparameterized neural networks. arXiv preprint 1rXiv:1910.00195, 2019.
Qinwei Xu, Ruipeng Zhang, Ya Zhang. Yanfeng Wang. and Qi Tian. A fourier-based framework for domain generalization. In CVPR, pp. 14383-1439) June 2021.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 100.003888,88 2017.
Adams Wei Yu, Lei Huang, Qihang Lin, Ruslan Salakhutdinov, and Jaime Carbonell. Block- normalized gradient method: An empirical study for training deep neural network. arXiv preprint arXiv:170700422. 2017.
Werner Zellinger. Thomas Grubinger. Edwin Lughofer, Thomas Natschläger, and Susanne Saminger-Plas Central moment discrepancy (emd) for domain-invarian representation learn- ing. arXiv preprint arXiv:/702.088// 1, 2017.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Pat Paz. mixup: Beyond empirical risk minimization. arXiv preprint 171.0:1710 2017.
M. Zhang, H. Marklund, N. Dhawan, A. Gupta, S. Levine, and C. Finn. Adaptive risk minimization: Learning adapt domain shift. In NeurlPS, 2021.
Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. arXiv e-prints, pp. arXiv-2103, 2021.
Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural networks. NewrIPS, 32, 2019.
12

---

Published as a conference paper at ICLR 2023
A SUPPLEMENTARY FOR PGrad: Learning robust gradients for domain
generalization.
This appendix includes the following contents:
We include figures describing the different trajectory sampling methods in A.1
To understand how PGrad changes the optimization path of model's parameters, we show TSNE projection of parameters during training in A.2 In Figure 4 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 A.3
PGrad differs from related work in that it adopts sequential training strategy when sampling a trajectory. We include detailed discussion comparing the parallel training used in gradient matching methods like Fish with our sequential training strategy in A..4. We empirically demon- strate that gradient alignment is not necessarily sufficient indicator the generalization ability, see Figure
To justify the analysis between PGrad under sequential training and parallel training, We design experiments comparing PGr Pdrad-B (Sequential), P 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.4
We show that our method has major theoretical contributions. First, instead of being driven by gradient alignment, PGrad learns a gradient flow from the tangent space spanned by the eigen- vectors of the Hessian matrix see A.5.I for the detailed analysis. Second, in A..52 we show that PGrad has deep connection to the efficient training of neural networks by projecting the high-dimensional parameter space 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 S
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 show the per-domain prediction accuracy of our PGrad variants and the experimental details on both benchmarks.
A.1 ILLUSTRATIVE EXPLANATION OF THE PROPOSED TRAJECTORY SAMPLING METHODS
As vanilla trajectory sampling only considers inter-domain variations, the long trajectory sampling variant splits per- batch to smaller batches to eliminate intra-domain discrepancies from data collection, data sampling, etc.
We compare the three proposed trajectory sampling methods in Figure 3
13

---

Published as a conference paper at ICLR 2023
Detaset: VLCS. Test Domainss: (0)
Dataset: VLCS, Test Domains:
-10
Dataset: Terralncognita. Test Domeins: [0]
Jataset: Terralncagnita, Test Domains:
Kone
110
Figure 4: Two-dimensional projection of the parameters' trajectories on different datasets. We use ResNet50 the backbone and apply TSNE for projection. Both PGrad and ERM start from sim- ilar random initialization. Increasing path thickness represents the later training phase. To reinforce the visual effect, smooth the curve within window of size 8.
0 D
+ D
Ø
O
0'
O
P13
D
0
D3
+
0
O'
O
Figure 3: Comparison three trajectory sampling methods. Assume the number of training domains 3. The top-left box shows the fixed-order trajectory sampling, which is PGrad-F The top- right box shows the random-order trajectory sampling, the default method PGrad. The bottom box represents PGrad-B, version of the long trajectory sampling with B 2.
A.2 REAL OPTIMIZATION TRAJECTORY VISUALIZATION
We visualize the optimization paths for both our method PGrad and the ERM baseline with TSNE projection. Concretely, we save the model's 's parameters into the memory buffer after every 100 training steps and project them the æyy-lane after finishing 5,000 training steps. The trajectories different datasets and test domains are shown
following Figure
ERM moves over-confidently with a large step size at the beginning, the turns sharply in late training. Our method PGrad 'thinks fast' but 'moves cautiously". samples roll-out optimization trajectory
14

---

Published as a conference paper at ICLR 2023
9
-0
(a): Sequential training will reinforce robust direction
Parallel training suppress robust direction
Figure 5: Comparison of sequential training and parallel training to learn principal gradient. Stars represent domain-spectio optimal minimum, @ is the starting point, (o forms a trajectory, the ellipse captures the current principle coordinate chart.
and aggregates the principal directions. The training curves are consistently smooth even late in training.
A.3 TRAJECTORY SAMPLING
The inclusion of the starting point 0j in the trajectory sampling is important, otherwise the learned principal gradient will skip the gradient information from the first training domain from each update. We show the derivation in the following. To simplify the notation, substitute Ver_, Lp. with Ve; and assume we have only two training domains.
Sampled Trajectory With Starting Point O
Sampled Trajectory W/O Starting Point Of
S ( (0% Ve;) (0) Ve: Va;1)
S 10} (0) Veg)}
Trajectory Centering
Trajectory Centering
$ e + Veg
Vos + Ves
-V - -2
8 = 1 Voz -- V
3
3
After trajectory centering, the gradient information from the first training domain Vay will be skipped if trajectories are sampled without the starting point O To learn from all training domains, we use the left policy for sampling.
A.4 COMPARISON 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 inspired by gradient alignment. They made efforts to align the gradients from different domains and adding the aligmment as penalty to the loss function. In their vanilla implementation, both Fish and Fishr calculate the per- domain gradient w.r.t the current parameter Oß. We can name the training paradigm as 'parallel training'. As contrast, our method learns a robust gradient direction from the optimization trajec- tory 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-specifie noises, see Figure 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 Moreover, the sequential training is more efficient comparing with parallel training.
In Figure 5, we intuitively explain that sequential training will reinforce the learning of clean direction and parallel training will significantly suppress it. We now design experiments to justify our hypothesis. We adapt PGrad-B B by learning the principal direction through parallel training. We fix all random mechanisms and compare PGrad-B under sequential training, PGrad-B under parallel training, and ERM. The test accuracy functions of the training epoch is visualized in Figure. In the left panel, we use C, P. R as training domains. the right panel, we use A, P, R as training domains. We attribute the performance drop of the PGrad-3 (parallel) the enlarged noise
15

---

Published as a conference paper at ICLR 2023
component and suppressed clean direction. The observations consistent with above analysis.
Another interesting observation is that curves from PGrad-B are smoother compared to ERM.
Ta)Test Accuracy, OfficeHome: C, P. A
(b)Test Accuracy, OfficeHome: A, P. C
0.5
0.5
0.4
0.3
PGrad-B Sequential
0.2
PGrad with Sequential
PGrad w Parallel
PGrad 8 with Porallel
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 6: Comparison of PGrad- (sequential), PGrad-B (parallel), and ERM on OfficeHome dataset.
We highlight the novelty of our method by showing that PGrad achieves domain generalization without constraining the gradient alignment. Specifically, define gradient alignment measure- ment as the mean of the cosine gradient similarity across all training domain pairs:
< <V()), V(D,)= >
1/(5) Σ VIJ IV(D)III)(
(13)
In Figure 7 We visualize the test domain accuracy and training domains gradient alignment func- tions of the training epoch. The figure implies the gradient alignment is not a sufficient 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 figure shows the smooth- ness of the gradient alignment curves has positive correlation with the test accuracy. Starting from 3,000 epochs, the alignment curve 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) PACS: P.S
OfficHome: P,
0.6
0.8
0.4
PGrad Align PGrad
hign_PGrad
Align_PGrad
0.4
Act. Fish Acsh
0.2
Acc Fish Align
Align Fish
0.2
0.0
0.0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Figure 7: Visualization test domain accuracy and training domain gradient alignment.
A.5 THEORETICAL ANALYSIS
The basic motivation of PGrad is to learn a gradient flow combines the top cigenvectors of the Hes- sian which are approximated by the average of the Fisher information matrix calculated sequentially on each training domain.
A. CONNECTION TO HESSIAN AND FISHER INFORMATION MATRIX.
Hessian matrix, which is widely used for analyzing the model's training behaviour, is defined as the second order derivative of the function. Similar to NTK (Jacot et al. 2018), assume the loss
16

---

Published as a conference paper at ICLR 2023
function & is functional acting parameters θ
((0) = (x,
x Σ [ (x),
(14)
We have Taylor expansion around parameter O:
[(e') = [(0) + (0'- () - - O(16'' -
(15)
0
where H4j = landscape around θ
00,00; The Hessian matrix H contains local geometric properties of the loss
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 loss function (Schervish 2012), we can approximate H. with the outer product of the gradient. Specifi- cally.
02C
I VeL & VoC
-
(16)
202
where Z 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 trajectory as S (0o, O,,... On. } E [R(++)xdd In the following, show the trajectory centering operation 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 -pnain--peccfic gradi- ent, whose columns be interpreted the domain-wise gradient vector starting from the same initialization. The shared initialization is the trajectory center of S.
The first half of the Lemma is easy to show. Any vector within convex hull can be recovered by the linear combination of its edges. See Appendix for the derivation on a simple case. It implies that the centered trajectory contains as rich information as training domains' gradient matrix. For the
sampled trajectory S, the center is calculated as S, Σn= θ
We can re-interpret the sampled trajectory with local coordinate centered with S.. Since the update step is small enough, we have:
n +1
= | - -S, θ -S, ,, S.l = (.. v..... V.l
(17)
The new gradients ( 1 shares pseudo initialization S..
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.
STS
1 Vo. V
V
+ V..
V1,
V.
(18)
7
n
n Σ
Σ Hi
(19)
n
where I, H4 are domain-speciff FIM and Hessian matrix, respectively. The approximation is covariance matrix, which is positive semi-definies (SPD) and has the eigenvalue-eigenvecoor pairs {(Az, vz) with Ag > A₁ >.... > An. The eigenvalue A is the curvature of the loss in direction of v. in the neighborhood of S,. The training behaviour of the neural network is determined by the distribution of the eigenvalues. Specifically, first-order optimization methods slow down sig- highly spread out (Bottou & Bousquet 2007 Ghorbani et al.) 2019). The property inspires us to zero out insignificant directions and use the directions with the large curvature derive gradient direction V,.
nificantly when {>r)
17

---

Published as a conference paper at ICLR 2023
As contrast, Fishr (Rame et al. 2021 uses the current parameter value as the initialization and parallelly approximates the per-domail Hessian matrix with its gradient variance. It defines the Fishr regularization 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 operate only the last classification layer in practice.
Comparing with other gradient manipulation works which emphasize alignment or matching, PGrad uses the average the FIMs approximate the underlining Hessian matrix under the DG setup. The operation can also reduce the noise variance by 1/n, where is the size of the training domains. The property also provides an explanation why PGrad-B achieves better generalization ability compared with PGrad. Second, instead of matching eigenvalues of per-domain Hessian, learn a robust gradient flow which is the combination the cigenvectors reflecting 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 efficiently operate the nigh-dimensional parameter space.
A.5.2 CONNECTION TO THE LEARNING BEНIVIIOR OF THE NEURAL NETWORKS
Modern SOTA neural networks are usually over-parameterized (Allen-Zhu et al. 2019 Zou & Gu 2019). Improving the training efficiency of high-dimensional neural networks is an active research direction (i al. 2018a). Recent works (Gur-Ari al. 2018 Gressmann 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 the line of the work.
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 specific directions. We, therefore, focus on recovering the oiv-dimensional subspace where the loss function L varies the most 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 C at θ is defined as:
( [Ve v,
(20)
and measure the expected scale (or length) of the directional derivative as:
Ee (e)) = Eslv" VeL(e) & o² Ve[(e)] = v Ee (e(()) & VeL(e)]u,
(21)
However, the distribution of the is unavailable, we turn to utilize the parameters from different training domains for empirical approximation.
Ee EeloC. = v² 2 V V 1 1 1 5" So.
(22)
We demonstrate that learning the principal gradient direction enables us to find 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 E R of the original parameter space O E Rd, such that it keeps largest directional derivatives with respect to the loss L. the
pair of the outer product matrix Ee (V)(() @ VeL(O))
are ()z, U2) in with do > X1 An- We have:
Span M1. .M&} = Span(vo. U1.... vR
(23)
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 g In our experiments, we learn V, by aggregating eigenvectors corresponding to the top-4 eigenvalues. The figure indicates that the contribution from the smallest eigenvalues keeps decreasing training continues. Our method PGrad can effectively distinguish domain-specific noise and update with shared patterns by suppressing the information from small eigenvalues.
18

---

Published as a conference paper at ICLR 2023
Epochs G-1000
Epoch: 1600-2000
Epoch: 2000-5000
Epoch: 300C-w1000
Epoch: 4000-55000
-10
-10
11
10
-30
-22
12
12
-22
-22
(a) PGrad-B OfficeHome, The A
Epoch: O-1000
Epoch: 1000
Epoch: 2000-3000
Epoch: 300C-4000
Epoch: 4000-5000
+2
+2
-20
-10
-22
12
(b) PGrad-3 3 OfficeHome, Tha C
Figure 8: The changing of the eigenvalue distribution with our proposed PGrad-B. where 3. The length of the trajectories is nB + = 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 Vp learned by aggregating the top-4 eigenvectors.
A.6 WHIAT ARE TAIL EIGENVECTORS?
To answer the last question we proposed in Section 3 we conduct the analysis show the effect of the bottom eigenvectors on the training dynamics. Ablation studies in Table Jindictee that including bottom eigenvalues into principal gradient will hurt the generalization ability. We add new exper- iments 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 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 Figure
A.7 EXPERIMENTAL SETUP DETAILS FOR DOMAINBED
Following the training protocol described in DomainBed, we run experiments on each domain with random mechanism to reduce the bias of hyperparameter selection. The DomainBed experiments (Gulrajani & Lopez-Paz Paz 202I> select the best model from random samples of the hyperparameter search space, and repeat this search process 3 times. Specifically, for each domain, we select 2 random combinations of hyperparameters from a relatively narrow range and repeat the search times report confidence intervals. For dataset with domains, train models. Tre. held out for testing, and use the rest as training domains D.,. This design leads to a total of 6n experiments per dataset for each method.
19

---

Published as a conference paper at ICLR 2023
(a) Start with bottom eigenvectors
(b) Switch to bottom eigenvectors
[c) No bottom eigenvectors
Training domain:
Training domain:
4.0
4.0
Training domain:
4.0
Training domain:
Training domein: R
Training domain: R
3.5
3.5
3.5
3.0
3.0
2.5
2.5
2.5
2.0
2.0
2.0
1.5
1.5
1.0
1.0
1.0
domain:
Trairing domain: P
0.5
Trairing
0.0
0.0
0.0
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Figure 9: In the left figure, we learn principal gradient always from bottom eigenvectors; the middle figure starts with top eigenvectors and switches to bottom ones after 200 epoch; the right figure always uses top cigenvectors. The vertical line indicates when the intervention happened. The y-axis is the training loss, the z-axis is the training epoch.
Lacking access to data from the test distribution during training makes it hard to perform model se- lection for DG compared to other supervised learning tasks. In the main, we follow the popular setup where we sample 20% of the data from each training domain and group them as validation set. The validation accuracy will be used accuracy of PGrad and the extensions on each domain in Tables 7
indicator of the optimal model. We show the prediction
A.8 HYPERPARAMETER SEARCH
We following a similar training and evaluation protocol to the Domain experiments (Gulrajani
& LLopzz-Paz 202I). We list our hyperparameter search space in Table 6
Table 6: Sample space for hyperparameters.
parrameparr
Default value
Random distribution
Inner learning rate
5e
5e
Outer learning rate
0.1
0.1
Batch size
32
32 if not DomainNet else 24 RandomChoicee(() 0.1, 0.5)) 10Unilomm( -4)
ResNet dropout
0
Weight decay
0
Training step
5,000
5,000 if not DomainNet else 6,000
We design narrow search space to prevent undersampling and potential performance degeneration from a smaller number of random trials. We use fixed batch size of 32, which ensures there adequate samples to provide precise gradient directions when sampling with our long and high- entropy method. For example, there are still 11 samples in each small batch if B 3. We also found works consistently well on each dataset. DomainNet's additional
the outer learning rate 7 model parameters in the final classification layer lead to memory constraints on our hardware the default batch size 32. Therefore, we use lower batch sizes 24, and increase the training step to 6, 000.
A.9 EXPERIMENTAL DETAILS ON WILDS BENCHMARK
A summary of the two vision datasets in the WILDS benchmark is shown in Table For the POVER- TYMAP task, the inputs are 8-channel satellite images, therefore, we tuned the first convolutional layer of the ResNet18 (He et al. 2016) accommodate multispectral input. To sample trajectory of length h++ in high dimensional parameter space θ we randomly sample ń training domains and sequentially update the parameter starting from O' each of them. We the Adam optimizer for trajectory construction and set the default learning rate to le 1e-3 3 without weight decay. We set the outer learning rate 7 0.1 / which adjusts the step size based the number of sampled training
20

---

Published as a conference paper at ICLR 2023
domains. We train the model with the proposed PGrad for 200 epochs and select the optimal model based on the average Pearson coefficient on validation domains.
For the land use classification task FMOW, we follow the exact same training and evaluation proto- col applied in Fish (Shi et al., (021). We find a good starting point by updating the ERM objective with an Adam optimizer for 24,000 iterations with learning rate of Le 4. After pretraining, we proceed with our proposed PGrad and keep tuning the model for epochs with outer learning rate 0.01/á. /n. After training is completed, we report the worst regional accuracy on the test domains. This measurement designed by WILDS (Koh al. (021) test models' generalization ability when facing both time and regional distribution shift. The numerical results for both datasets are available in T
Table 7: Per-domain prediction accuracy (%) on VLCS dataset
Algorithm
C
L
V
Avg
PGrad-F
98.7 98.7 1 0.2
64.4 = 0.9
73.6 0.2
76.91004
78.4
PCrad
98.3 1:3::00 0.2
64.4
7.
79.9:1.77
79.3
PGrad-B
98.7 + 0.3
+3.91.11
0..6
78.5+0.6
78.9
PGrad-BMix x
99.1 + 0.1
63.8 + 0.7
73.5
+3.5+0..5
79.0 + 0.5
78.9
Table Per-domain prediction accuracy (% on PACS dataset
Algorithm
A
C
S
Avg
PGrad-F
88,0 0.4
79.2 = 0.3
98.2 + 0.2
76.6
1.S
85.5
PGrad
87.6 0.3
79.1 =
0.1
1.2
85.1
PGrad-B
0.2
80.0
9.2
80.1
0.9
87.0
PGrad BMix
0.3
78.9 = 97.7 + 0.3
1.0
86.2
Table 9: Per-domain prediction accuracy (%) on OfficeHome dataset
Algorithm
A
C
P
R
Avg
PGrad-F
64.4 0.4
54.7 0.3
77.0
0.3
78.6
0.2
68.7
PGrad
56.0
77.4 0.2
0.3
PGrad-B
65.2
5..9
0.5
0.3
PGrad-BMix
0.2
55.4
78.0
80.0 0.4
Table 10: Per-domain prediction accuracy (%) on TerraIncognita dataset
Algorithm
L100
L38
L43
I.66
Avg
PGrad-F
$2.0 1.3
42,0
57.2 + 0.9
43.2
2.1
48.6
PGrad
51.2 0.8
43.4
60.0 0.6
41.3
0.8
49.0
PGrad-B
2.1
43.5
0.5
41.9
0.3
PGrad-BK BMix
61.2 2.2
45.7 =
58.2 2.2
37.9
0.7
50.7
Table 11: Per-domain prediction accuracy (%) on DomainNet dataset
Algorithm
clip
info
paint
quick
real
sketch
Avg
PGrad
57.0 =
18.2 0.2
48.4
1.1
60.9 + 0.1
48.8 0.1
41.0
PGrad-3
57.2 =
0.2
48.3 =
13.1 0.1
61.2
49.9 +
41.4
FGrad-BMix
59.4 - 0.1
19.8 0.1
49.1 = 0.4
13.7 0.1
61.8
0.11
42.5
21


Reviews:
[{"review_id": "v-QSkeTF071", "reviewer": "Reviewer_ap7X", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper develops a new optimization method for domain generalization. The idea of the proposed PGrad approach is to use a robust gradient direction for parameter update. The robust direction is estimated in 2 steps. First, the rollout of the weights trajectory is collected by sequentially applying training updates in each of the training domains. Then, the principal components of the collected weight space point cloud are estimated and the final update direction is constructed as a combination of the leading principal directions with the downplayed or truncated contribution of the tail directions. The intuition behind this approach is that the low-variance directions are likely to correspond to domain-specific noise and suppression of this noise presumably leads to robust training updates and better generalization on unobserved test domains.\n\nThe authors conduct an extensive set of experiments demonstrating the effectiveness of the approach on large scale domain generalization benchmarks: DomainBed and WILDS. The method delivers improvement over recent domain generalization methods including closely related approaches (Fish and Fishr) that are based on cross-domain gradient penalties.\n", "summary_of_the_review": "The authors develop a simple and effective method for domain generalization which improves \nOOD generalization on large-scale benchmarks. \n", "strength_and_weaknesses": "Strengths:\n* The paper proposes a simple and effective method for domain generalization.\n* The evaluation and ablations of the proposed approach are comprehensive.\n* The authors thoroughly discuss the connection and differences between the proposed approach and the recent related work. I find the argument about the benefit of the sequential gradient computation approach compared to parallel gradients particularly interesting (Appendix A.4).\n\nWeaknesses:\n\nI did not find any critical weaknesses. Below are some questions that came to my mind when reading the paper. These are not meant to be taken as a criticism of the work done by the authors but rather as open questions and potential directions for future work.\n\nAfter reading the paper, it is clear that the method delivers improvements on domain generalization benchmarks. However, the mechanisms behind the success of the method and theoretical justification are not completely clear. The authors provide several insightful yet not very concrete and rather intuition-level explanations (such as noise suppression and connection between trajectory of covariance matrix and Hessian of the loss), but some questions remain unanswered:\n\n* Do the bottom eigenvectors actually correspond to “domain-specific” noise? Is there any way to validate this claim directly (theoretically or in experiment)?\n\n\n* The authors point out that, compared to approaches based on representation alignment or invariance, PGrad does not place any explicit assumptions on the data distributions across domains and the type of distribution shift. On the one hand it is a nice property as the method can be applied to any domain generalization problem in principle, on the other hand it tells us that we do not really know in what scenarios the method is expected to work well and more importantly when it will fail. Clearly, no single method can be successful in all possible domain generalization settings. Figuring out assumptions on the data generating process and the model under which the method is guaranteed to succeed is an important open question.\n\n* What can be said about the method from the optimization perspective? Does PGrad have convergence guarantees?\n\n* How does the wall-clock time of PGrad training compare to that of SGD training? How significant is the computational overhead of PGrad?\n\n* I find the intuitive explanation of the benefit of the sequential training vs parallel training (Figure 5) very interesting. Is it possible to develop this argument in more precise theoretical statements or design experiments highlighting the effect?\n\nI realize that answering most of these questions requires conducting separate investigations, but I would still encourage authors to include a discussion of these issues mentioned above in the revised version of the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "syAcP5ZzPv7", "reviewer": "Reviewer_jbQd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper addresses the problem of how to train a single model to handle predictions in multiple domains (domain generalization), more specifically how to learn a model that will generalize well across the different domains.  The authors argue that one deficiency of current approaches is that models often fail to lear nt the distinction between shared common features and domain-specific noise, hindering their generalization.  They define the principal gradient (or PGrad), which is derived as the subspace of the gradients which capture the largest directions of variation (akin to PCA), and argue that updating the model by steps in the direction of the principal gradient will allow for greater generalization capabilities.  They show that different variations on PGrad achieve strong results on DomainBed benchmark datasets, and competitive results on WILDS. They also establish with experimental evidence that variation in the training loss is much smoother.  \n\n", "summary_of_the_review": "The main technical contributions of this paper are to adapt ideas from principal component analysis to the estimation of gradient directions that are common among disparate domains.  While simple on the surface, the authors show clearly that this is a very clever and elegant solution to domain generalization.\n\n", "strength_and_weaknesses": "### Strengths\n- The paper is written very directly.  The intuition behind the idea of finding principal gradient directions is well explained, and the connection to the estimation of the Fisher Information Matrix is helpful.  Though I will say that Appendix 5.1 is a bit brisk in its explanation. \n- The clever trick in equation (7) is compactly explained.  \n- Section 2 of the paper establishes links to similar ideas in invariant feature learning and gradient regularization.\n- Building off of DomainBed makes the method immediately more reproducible and straightforwardly comparable to other domain generalization methods.\n\n### Weaknesses\n- In section 2, the authors state that \n> When using ERM in DG setting, each step of model updates uses an average gradient and may introduce and preserve domain-specific noise. For instance, if one training domain includes the trapping signals like cows always in pastures and camels always in deserts (as mentioned earlier).\n\nGulrajani and Lopez-Paz managed to use ERM to achieve great performance on the DomainBed datasets despite being open to this risk.  Given their strong ERM baseline results, which are unable to distinguish domain specific noise from common features, do you think it is a real risk or simply a possible one?\n\n- In section 2.1 just above equation (6), the authors claim “the revised computational complexity is $n$”.\nIt is unclear what it meant here. The complexity of computing the SVD of the matrix $\\hat{\\mathbb{S}}^{T}\\hat{\\mathbb{S}}$ is cubic in the parameter vector size.  Even after introducing the bijection allowing this SVD to be performed on an n x n matrix, the computational complexity is still cubic in $n$.  Could the authors clarify here?\n- In section 2.1, equation (9) is justified by scaling the eigenvectors $w_z$ by the ratio of $\\frac{\\lambda_z}{\\mathbf{\\lambda}}$.  This is justified under the rubric of normalizing the gradient to reduce training instability, but it is unclear why equation (11) needs to further scale by accounting for the length of $\\Delta_r$.  Could the authors provide a footnote or some more justification for this?  Why not simply adjust the learning rate later in training?\n-  For equation (12), this exposition can be simplified.  It might help to establish a link between the degree to which each individual $\\theta_i$ is optimal (estimated itself via SGD, but not specified IIRC) and the degree to which PGrad can mitigate errors in the directions of any group of $\\theta_i$\n- Section 3 begins with a series of questions that the authors aim to probe in their experiments.  Question 2 is about whether Pgram will work with different architectures or data types.  I feel like Q2 is almost bound to be answered affirmatively.  How could it not? Any parameterized model will yield gradient estimates, and so the principal gradient can always be computed.  As the authors claim in section 2.4, their model is model and data agnostic.  \n\nA question I think would be more impactful, is to see is how well the principal gradient performs in the presence of noisy data that yield high variance gradient estimates?  Could the principal gradient be combined with other variance reduction techniques to yield more stable updates in the presence of such noisy (and variable) gradient estimates?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2wHwVbcejX", "reviewer": "Reviewer_Fpjy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper develops a new optimization algorithm named PGrad for domain generalization. The proposed method attempts to filter out domain-specific noise and provide an update direction that maximally benefits all training domains. The experimental results show improved domain generalization performance on multiple datasets.", "summary_of_the_review": "This paper is easy to follow and has shown some promising results. However, it lacks important theoretical or experimental analysis, and ablation studies.", "strength_and_weaknesses": "Strength\n\nThe technical part of this paper is easy to follow, and the empirical results have shown the potential of the proposed method.\n\nWeaknesses\n1. PGrad is motivated by a few intuitions, which are not supported or verified by any theoretical or experimental analysis. For instance, it's unclear why PGrad \"forces the learned gradient to filter out domain-specific noise\nand follows a direction that maximally benefits all training domains\", where the domain-specific noise is not well-defined to begin with.\n2. PGrad incorporates multiple techniques, including trajectory sampling and its SVD decomposition, directional calibration, L2 normalization, length calibration, and noise suppression. Without a proper ablation study, it's unclear how each technique contributes to the overall improvement. In particular, it is possible that simple techniques such as sequential optimization on each training domain and gradient normalization are responsible for most of the improvement, in which case the motivation would be even less convincing.\n3. The computational overhead of trajectory sampling seems to be significant, but is not discussed in the paper.\n4. While Table 2 include 11 methods for comparison, the experiments on the WILDS dataset (Table 5) only include three methods.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster