Towards Robust Out-of-Distribution Generalization Bounds via Sharpness
Authors: Yingtian Zou
Kenji Kawaguchi
Yingnan Liu
Jiashuo Liu
Mong-Li Lee
Wynne Hsu
Venue: ICLR 2024 spotlight (2024)
Keywords: Out-of-Distribution generalization
Sharpness
Robustness
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: In this paper, we propose a sharpness-based OOD generalization bound by considering the model robustness.

Abstract: Generalizing to out-of-distribution (OOD) data or unseen domain, termed OOD generalization, still lacks appropriate theoretical guarantees. Canonical OOD bounds focus on different distance measurements between source and target domains but fail to consider the optimization property of the learned model. As empirically shown in recent work, sharpness of learned minimum influences OOD generalization. To bridge this gap between optimization and OOD generalization, we study the effect of sharpness on how a model tolerates data change in domain shift which is usually captured by "robustness" in generalization. In this paper, we give a rigorous connection between sharpness and robustness, which gives better OOD guarantees for robust algorithms. It also provides a theoretical backing for "flat minima leads to better OOD generalization". Overall, we propose a sharpness-based OOD generalization bound by taking robustness into consideration, resulting in a tighter bound than non-robust guarantees. Our findings are supported by the experiments on a ridge regression model, as well as the experiments on deep learning classification tasks.

Published as a conference paper at ICLR 2024
TOWARDS ROBUST OUT-OF-DISTRIBUTION GENERALIZATION BOUNDS VIA SHARPNESS
Yingtian Zou, Kenji Kawaguchi, Yingnan Liu, Jiashuo Liu", Mong-Li Lee, Wynne Hsu
School of Computing, National University of Singapore
Institute of Data Science, National University of Singapore
"Department of Computer Science & Technology, Tsinghua University, China
ABSTRACT
Generalizing to aut-of-diitribution (OOD) data or unseen domain, termed OOD generalization, still lacks appropriate theoretical guarantees. Canonical OOD bounds focus on different distance measurements between source and largel do- mains but fail to consider the optimization property of the learned model. As empirically shown in recent work, the sharpness of learned minima influences OOD generalization. To bridge this gap between optimization and OOD general- ization, we study the effect of sharpness on how a model tolerates data change in domain shift which is usually captured by "robustness" in generalization. In this paper, we give a rigorous connection between sharpness and robustness, which gives better OOD guarantees for robust algorithms. It also provides theoreti- cal backing for "flat minima leads to better OOD generalization". Overall, we propose a sharpness- based OOD generalization bound by taking robustness into consideration, resulting in a tighter bound than non-robust guarantees. Our find- ings are supported by the experiments on a ridge regression model, as well as the experiments on deep learning classification tasks.
INTRODUCTION
Machine learning systems are typically trained on a given distribution of data and achieve good performance on new, unseen data that follows the same distribution as the training data. Out-of- Distribution (OOD) generalization requires machine learning systems trained in the source domain to generalize to unseen data or target domains with different distributions from the source domain. myriad of algorithms (Sun & Saenko, 2016; Arjovsky et al., 2019; Sagawa et al., 2019; Koyama & Yamaguchi, 2020; Pezeshki al., 2021; Ahuja et al., 2021) aim to learn the invariant components along the distribution shifting. Optimization-base Duchi & Namkoong, 2018; Liu et al., 2021; Rame et al., 2022) focus on maximizing robustness by optimizing for worst-case eтror over uncertainty distribution set. While these methods are sophisticated, they do not always perform better than Empirical Risk Minimization (ERM) when evaluated across different datasets (Gulrajani & Lopez-Paz, 2021; Wiles al., 2022). This raises the question of how to understand the OOD generalization of algorithms and which criteria should be used select models that are provably better (Gulrajani & Lopez-Paz, 2021). These questions highlight the need for more theoretical research in the field of OOD generalization (Ye et al., 2021).
methods such as (El Ghaoui & Lebret, 1997;
To characterize the generalization gap between the source domain and the target domain, a canonical method (Blitzer et al., 2007) from domain adaptation theory decouples this gap into an In-Distribution (ID) generalization and a hypothesis-speciiff Out-of-Doe Dissrribtiti (OOD) distance. However, this distance is based on the notion of VC-dimension (Kifer et al., 2004), resulting in a loose bound due to the large size of the hypothesis class in the modern overparameterized neural networks. Subsequent works improve the bound based on Rademacher Complexity (Du et al., 2017), whereas Germain et al. (2016) improves the bound based on PAC-Bayes. Unlike the present paper, these works did not consider algorithmic robustness, which has natural interpretation and advantages for distribution shifts In this work, we consider algorithmic robustness to derive the OOD generalization bound. The key idea is to partition the input space into K -verlapppg subspaces such that the error difference in the model's performance between any pair of points in each subspace is bounded by some constant e. Within each subspace, any distributional shift is considered subtle for the robust

---

Published as a conference paper at ICLR 2024
model thus leading to less impact on OOD generalization. Figure illustrates this with the two distributions where the target domain has a distributional shift from the source domain. Compared to existing non-robust OOD generalization bounds Zhao et al. (2018),
generalization error does not depend on hypothesis size, which is more reliable in the overparameterized regime. Our goal is to measure the generalizability of model by considering how it is robust to this shift and achieves tighter bound than existing works.
Although robustness captures the tolerance dis- tributional shift, it is intractable to compute robust- ness constant due to the inaccessibiiity of target distribution. The robustness definition in Xu & Mannor (2012) indicates that the loss landscape induced by the model's parameters is closely tied its robustness. To gain deeper understanding of robustness, we further study the learned model from an optimization perspective. As shown (Lyu et al., 2022; Petzka et al., 2021), when the loss landscape is "flat", there is a good general- ization, which is also observed in OOD settings (Izmailov et al., 2018; Cha et al., 2021). How- ever, the relationship between robustness and this geometric property of the loss landscape, termed Sharpness, remains an open question. In this pa- per, we establish provable dependence between robustness and sharpness for ReLU random neural network classes. It allows us replace robustness constant with the sharpness of learned model which is only computed from the training dataset that addresses the problem mentioned above. Our result of the interplay between robustness and sharpness can be applied to both ID and OOD generalization bounds. We also show example to generalize our result beyond our assumption and validate it empirically.
Dicretine different robust areas
JySTy> 6
(87)
Contimuous Density
Discrete Density
Figure I:
An example of a target distribution (red) directly translated from source distribu- tion (blue). The ID density reflects the marginal distribution. Unlike existing works (left), di- vide the distributions into disjoint partitions as a small change in distribution for robust model is negligible (right). The sharpness of the model will decide the tolerance of change thus affecting the partitions. If two sub-distributions S, T have small shifts such that they fall into the same par- tition (red grid), their distance measure d'(S, T) by considering robustness will zero.
Our main contributions can be summarized as follows:
We proposed a new framework for Out-of-distributiinn Out-of-domain generalization bounds. In this framework, leads to tighter upper bounds generally.
robustness capture the tolerance of distribution shift which
We reveal the underlying connection between the robustness and sharpness of the loss landscape and use this connection to enrich our robust OOD bounds under one-hidden layer ReLU NNs. This is the first optimization-based bound in generalization.
We studied two cases in ridge regression and classification which support and generalize our main theorem well. All the experimental results corroborate our findings well.
PRELIMINARY
Notations We use [n] to denote the integers set [i) We use || to denote the morm (Euclidean norm). vector form, w, denotes the i-th instance while the wy the j-th element of the vector w and we use |w] for the element-win wise absolute value of vector w. We use n, d for the training set size and input dimension. O is the Big-O notation.
2.1 PROBLEM FORMULATION
Consider a source domain and a target domain of the OOD generalization problem where we use Dg, Dr to represent the source and target distribution respectively. Let each D be the probability
2

---

Published as a conference paper at ICLR 2024
measure of sample x from sample space 2 == X x x with &' E RR. In the source domain, have training set S { 1= Vi E [n] ži ~ Ds while the target is to learn a model f E F with S and parameters 0 θ where f @ x X' + R generalizes well. Given loss function f: R x R R... which is for short, the expected risk over the source distribution Ds will be
(f(fe) - EE-DD [f)(=]] = DE
= 2, [((())]]
We use (z)z) as the shorthand. The OOD generalization is to measure between target domain expected risk Lr(So) and the source domain empirical risk (fo) which involves two parts: (1) In-domain generalization eπrоr gap between empirical risk and expected risk Cs(fe) in the source domain. (2) -it-o---Dmmann distance between source and target domains. model-agnostic example in Zhao et al. (2018) gave the following uniform bound:
Proposition 2.1 (Zhao et al. Zhao et al. (2018) Theorem 2 & 3.2). With hypothesis class F and pseudo dimension Pdim(F)= d', unlabeled empirical datasets from source and target distribution Di and Dr with size each, then with probability at least for all f jJF.
Cr(f) - Es(f) 1+=222=0 dFAF (Dr:Ds) +0 (Ve/n)
where dFAF(Dr; = (SuDAAAAAToveerrourrr
and @ is the XOR PD [bslAl] P PD, LAi]
operator. Specific form of O (FF/)) is defined in Appendix C.6.
2.2 ALGORITHMIC ROBUSTNESS
Definition 2.2 (Robustness, Xu & Mannor (2012)). A learning model fo on training set S' (K, (-))- robust, for K E if Z can be partitioned into K disjoint sets, denoted by {Ci) such that Vs S we have
8, C2 Vi
[ = (()()) fo(=)] c(S).
This definition captures the robustness of the model in terms of the input. Within each partitioned set C₁ the loss difference between any sample z belonging to C; and training sample E C; will be upper bounded by the robustness constant (S)) The generalization result given by Xu & Mannor (2012) provides framework to bound the empirical risk with algorithmic robustness which has been stated in Appendix C. Based on this framework, we are able to reframe the existing OOD generalization theory.
3 MAIN RESULTS
In this section, we propose new Out-of-Distribttion (OOD) generalization bound for robust algorithms that have not been extensively studied yet. We then compare our result to the existing domain shift bound in Proposition 2.1 and discuss its implications for OOD and domain generalization problems by considering algorithmic robustness. To further explain the introduced robustness, we connect it to the sharpness of the minimum (a widely concerned geometric property in optimization) by showing a rigorous dependence between robustness and sharpness. This interplay will give us a better understanding of the OOD generalization problem, and meanwhile, provide more information on the final generalization bound. Detailed assumptions are clarified in Appendix B.I.
3.1 ROBUST OOD GENERALIZATION BOUND
The main concern in OOD generalization is to measure the domain shift of learned model. However, existing methods fail to consider the intrinsic property of the model, such as robustness. Definition 2.2 gives us a new robustness measurement to the model trained on dataset S where the training set S is collection of i.idd. data pair (x, sampled from source distribution Dş with size The measurement provides an intuition that if a test sample from the target domain is similar to a specific group of training samples, their losses will be similar as well. In other words, the model's robustness is a reflection of its ability to generalize to unseen data. Our first theorem shows that by utilizing the "robustness" measurement, we can more effectively handle domain shifts by setting a tolerance
3

---

Published as a conference paper at ICLR 2024
range for distribution changes. This robustness measurement, therefore, provides useful tool for addressing OOD generalization.
Theorem 3.1. Let Dr be the empirical distribution of size n. drawn from Dr. Assume that the loss & is upper bounded by M. With probability least - ô (over the choice of the samples S), for every fe trained on S satisfying (K, (rS)))- we have
2K In 222221)
(r(So) S Ls(So) + Md DT) 2e(S) + 3M
where the total variation distance d(e,K) yo.yeee for discrete empirical distributions is defined by:
K Σ
(s(
(D()
Vi [il # ( (S,
(2)
n
72
and ni ns(S), nz (Dr) are the number of samples from and Dr that fall into the set C. respectively, Remark. The result can be decomposed into in-domain generalization and out-domain distance (fo) Ls (ffo (please refer to Lemma C.1). Both of them depend on robustness c(S).
See proof in Appendix C. The last three terms on the RHS of (1) are distribution distance, robustness constant, and error term, respectively. Unlike traditional distribution measures, partition the sample space and the distributional shift separately in the K sub-groups instead measuring point-wisely. We argue that the di (S, Dr) can be zero measure if all small changes happen within the same partition where 2D illustrative case is shown in Figure 1. Under the circumstances, our distribution distance term will be significantly smaller than Proposition 2.1 as the target distribution essentially perturbation of the source distribution. As a robust OOD generalization measure, our bound characterizes how robust the learned model is to negligible distributional perturbations. To prevent a bound that expands excessively, we also propose an alternate solution tailored for non-robust algorithms (K oo) as follows.
Corollary 3.2. If + co, the domain shift bound (fr( - Cs(ff)) can be replaced to the distribution distance in Proposition 2.1 where
1 Cr(fe) - Ls(fe)) S - ( +
1
(3)
where the pseudo dimension Pdim(FF d' The proof is in Appendix C.I. As dictated in Theorem 3.1, when K + co, the use infinite number of
partitions on the data distribution leads meaningless robustness. However, Corollary C.7 suggests that our bound can be replaced by (Ds in the limit of infinite K. This avoids computing a vacuous bound for non-robust algorithms. In summary, Theorem 3.1 presents a novel approach for quantifying distributional shifts by incorporating the concept of robustness. Our framework is particularly beneficial when a robust algorithm is able to adapt to local shifts in the distribution. Additionally, data-dependent result remains valid and useful in the overparameterized regime, since K does not depend on the model size.
3.2 SHARPNESS AND ROBUSTNESS
Clearly, robustness is inherently tied to the optimization properties of model, particularly the curvature of the loss landscape. One direct approach to characterize this geometric curvature, referred to "sharpness," involves analyzing the Hessian matrix (Foret et al., 2020; Cohen al., 2021). Recent research (Petzka et al., 2021) has shown that the concept of "relative flatness", the sharpness in this paper, has strong correlation with model generalization. However, the impact of relative flatness on OOD generalization remains uncertain, even within the convex setting. To address this problem, we aim to investigate the interplay between robustness and sharpness. With the following definition of sharpness, we endeavor to establish OOD generalization bound rooted in optimization principles.
Definition 3.3 (Sharpness, Petzka et al. (2021)). For twice differentiable loss function L(w) Eoes (8). w E IR" with a sample set S, the sharpness is defined by
k(w, S, A) (w, w) tr (..)(w))
(4)
where Hs.A is the Hessian matrix of loss L(w) w.r.t. u with hypothesis set A and input set S.
4

---

Published as a conference paper at ICLR 2024
As per Definition 3.3, sharpness is characterized by the sum of all the eigenvalues of the Hessian matrix, scaled by the parameter norm, Each eigenvalue of the Hessian reflects the rate of change of the loss derivative in the corresponding eigenspace. Therefore the smaller value of indicates flatter minimum. In Cha et al. (2021), they suggest that flatter minima will improve the OOD generalization, but fail to deliver an elaborate analysis of the Hessian matrix. In this section, random ReLU Neural Networks parameterized by 0 = (fa:)emm) w) where w = [W],.., Wm] is the trainable parameter. Let A = [a....., aml, the whole function class defined as
begin with the
f(w, A, x)
Vd wao (ær, a:) : we E IR, a, ~ Unif((s-'((V))) i E [m]
(5)
where o(-) is the ReLU activation function and are random vectors uniformly distributed on n-dim hypersphere whose surface is a manifold. We then define any convex loss function ([S(1, A,;),, v) R + RR, The corresponding empirical minimizer in the source domain will be: v = arg minl - f(f(w, A, x.).). With w. we are interested in loss geometry over the sample domain (z)) for short). Intuitively. flatter minimum on the loss landscape is expected
to be more robust to varying input. Suppose the sample space 2 can be partitioned into K disjoint the loss difference upper bounded by c(S, A). Given S, we have
sets. For each set C, i E [K],
c(S, A) A max sup (( (z)))
BE[K] x,''EC,
As an alternative form of robustness, the A( in (6) captures the "maximum" loss difference between any two samples in each partition and depends on the convexity and smoothness of the loss function in the input domain. Given training set S and any initialization wo. the robustness c(S, A) of a learned model fe will be determined. explicitly reflects the smoothness of the loss function in each (pre-) prretppariiioned set. Nevertheless, its connection the sharpness of the loss function in parameter space still remains unclear. In order to address this gap, we establish a connection between sharpness and robustness in Theorem 3.4. Notably, this interplay holds implications not only for OOD but also for in-distribution generalization.
Theorem 3.4. Assume for any A, the loss function (z) w.r. sample z satisfies the L- Hessian Lipschitz continuity (refer to Definition B.2) within every set Ci Vi [K]. Let (A) = arg maxzec,ns (z). Define M; to be the set of global minima in C. sup- pose (A) Mi such that for some Pi(L) > 0, [|z((A) surely. then let Pmax = (L), E []] || R(d) and n' S .n,E N wap
2 (A)ll I po(L) almost
1.
2 क over {a, ) Unif(Sd-1(Vi)) have p min arccos (R(d)=t). V.Tlia)
Pmax (L)²
(nr n' +0 (n't
k(② S, + 4pmax (L)
c(S, A)
2L
3
Remark. Given the training set S, we can estimate factor ñ that n' < ît by comparing the maximum Hessian norm w.r.t. žy to the sum of all the Hessian norms over {z, [z])e]] Note that the smoothness condition only applies
every partitioned set (locally) where it is much weaker than the global requirement for the loss function to be satisfied We also discuss the difference between our results and Petzka et al. (2021) in Appendix F. The chosen family of loss functions that applied our theorem can be found in Appendix B.I.
Corollary 3.5. Let wmin be the minimum value of |(). Suppose Vo Unif(Sd-1 (Vd)) and (( A, x)x)) is bounded by [Mr Md). If m Poly(d), > 2, Pmax (L) < (wi2 (d, (2)/ taking expectation over all x, E S.j E [n] and all a; E ~ (ail(sd='(Vd))))) E [m]) we have
7pmax (L)²
Es,A (c(S, A)] S ES,A
(, A,))
(8)
6I. Ga) 0 is the minimum eigenvalue and G
where (d, m) =E, Ba--uiit( ( (Vd)) Amiu ET (,a)GG)
is product constant of Gegenbauer pooyennmmerrplyyommiil (definition be founded in Appendix B).
5

---

Published as a conference paper at ICLR 2024
See proof in Appendix D. From Theorem 3.4 we see, the robustness constant c(S, .A) is (pointwise) upper bounded by the sharpness of the learned model, as measured by the quantity A à S. A), and the parameter Pmax (L). It should be noted that the parameter Pmax(L) depends on the partition, and as the number of partitions increases, the region Pimax (.) of the input domain becomes smaller, thereby making sharpness the dominant term in reflecting the model's robustness within the small local area. In Corollary 3.5, we show a stronger connection when the partition satisfies some conditions. Overall, this bound states that the larger the sharpness of the model k(t0, S, A), the larger the upper bound on the robustness parameter c(S, A). This result aligns with the intuition that a sharper model is more prone to overfitting the training domain and is less robust in the unseen domain. While the dependency is not exact, it still can be regarded as an alternative approach that avoids the explicit computation of the intractable robustness term. By substituting this upper bound for eS, A) into Theorem 3.1, we derive a sharpness-based OOD generalization bound. This implies that the OOD generalization error will have a high probability of being small if the learned model is flat enough. Unlike existing works, our generalization bound provides more information about how optimization property influences performance when generalizing to OOD data. It bridges the gap between robustness and sharpness which also be generalized to non-OOD learning problems. Moreover, we provide a better theoretical grounding for an empirical observation that a flat minimum improves domain generalization (Cha et al., 2021) by pinpointing clear dependence on sharpness.
3.3
CASE STUDY
To better demonstrate the relationship between sharpness and robustness, we provide two specific examples: (1) linear ridge regression; (2) two-layer diagonal neural networks for classification.
Example 3.6. In ridge regression models, c(S, has a reverse relationship to the regularization parameter B. B 1, the more probably flatter minimum 4 and less sensitivity 4 the learned model could be. Following the previous notation, we have 36g > such that c(S, A) S Ō, S) + δa where og has a smaller order than x(0, for large d (proof refer to Appendix E.
As suggested in Ali et al. (2019), let's consider a generic response model y/0. (X0_,o²) where X E Rkaxd d >> n. The least-square empirical minimizer of the ridge regression problem will be: 0
i = arg min
|| - y + 2 101 = (X + y
2n
Let S be the training set. It's trivial to get the sharpness of a quadratic loss function where
( S) |²1 tr(x X/n 3I.) = |101 tr(M)
(10)
It's obvious that both of the above two equations depend on the same matrix M - BIn For fixed training samples X, we have k(@, S) O(8-1) in the limit of 3. Then it's clear that a higher penalty 3 leads to flatter minimum. This intuition is rigorously proven in Appendix F. According to Theorem 3. and Theorem 3.4, a flatter minimum probably associates with lower robustness constant (SS, A). Thus it enjoys a lower OOD generalization error gap. In ridge regression, this phenomenon can be reflected by the regularization coefficient B. Therefore, in general, the larger is, the lower the sharpness x(O () and variance are. As a consequence, larger 3 learns a more robust model resulting in lower OOD generalization error gap. This idea is later verified in the distributional shift experiments, shown as Figure 2.
We consider a classification problem using a 2-layer diagonal linear network with exp-loss. The robustness c(S, has similar relationship in Theorem 3.4. Given training set S, after iterations t > T., _2² > 0, A( & 02 sup:2T, k))
Example 3.7.
In addition to the regression and linear models, we have obtained a similar relationship for 2-layer diagonal linear networks, which are commonly used in the kernel and rich regimes as well as in intermediate settings (Moroshko et al., 2020). Example 3.7 demonstrates that the relationship also holds true when the model is well-trained, even exp-loss does not satisfy the PL condition. By extending our theorems to these more complex frameworks, we go beyond our initial assumptions and offer insights into broader applications. Later experiments our statements. However, we still need a unified theorem for general function classes with fewer assumptions.
non-lineat NN also support

---

Published as a conference paper at ICLR 2024
4 RELATED WORK
Despite various methods (Sun & Sacnko, 2016; Sagawa et al., 2019; Shi et al., 2021;
Shaficezadch Abadeh et al., 2015; Li et al., 2018; Cha et al., 2021; Du et al., 2020; Zhang et al., 2022) that have been proposed to overcome the poor generalization brought by unknown distribution shifts, the underlying principles and theories still remain underexplored. As pointed out in Redko et al. (2020); Miller et al. (2021), different tasks that address distributional shifts, such as domain adaptation, OOD, and domain generalization, are collectively referred as "transfer transductive learning" and share similar generalization theories. In general, the desired generalization bound (ID) generalization error and Out-of- www.ooiitrbuttiooooomeaeo (OOD) distance. Since Blitzer et al. (2007) establish a VC-dimension-buseee framework to estimate the domain shift gap by a divergence term, many following works make the effort to improve this term in the following decades, such as Discrepancy (Mansour et al., 2009). Wasserstein measurement (Courty et al., 2017; Shen et al., 2018), Integral Probability Metrics (IPM) (Zhang et al., 2019b; Ye et al., 2021) and B-divergence (Germain et al., 2016). Among them, new generalization tools like PAC-Bayes, Rademacher Complexity, and Stability are also applied. However, few of them discuss how the sharpness reacts to data distributional shifts.
will be split into In-Distribution/Domme
Beyond this canonical framework, Ye et al. (2021) reformulate the OOD generalization problem and provide generalization bound using the concepts of "variation" and "informativeness: The causal framework proposed in Peters et al. (2017); Mojas-Carullaa et al. (2018) focuses on the impact of interventions robust optimization over test distributions. However, none of these frameworks consider the optimization process of a model and how it affects OOD generalization. Inspired by previous investigation on the effect of sharpness on ID generalization (Lyu al.. 2022; Petzka et al., 2021), recent work in Cha et al. (2021) found that flatter minima can also improve OOD generalization. Nevertheless, they lack a sufficient theoretical foundation for the relationship between the "sharpness" of a model and OOD generalization, but end with union bound of Blitzer et al. (2007)'s result. In this paper, we aim to provide a rigorous examination of this relationship.
EXPERIMENTS
In light of space constraints, we present only portion of our experimental results to support the validity of our theorems and findings. For comprehensive results, please refer to the Appendix
5.1 RIDGE REGRESSION IN DISTRIBUTIONAL SHIFTING
Following Duchi & Namkoong (2021), we investigated the ridge regression distributional shift. We randomly generate 88 E R in spherical space, and data from the following generating process:
X / NO 1), y = X0%- To simulate distributional shift, we randomly generate a perpendicular unit vector Of to 0% Let 00 0j be the basis vectors, then shifted ground-truth will be computed from the basis by 0; = 0/ cos(a) + 0 sin(a). For the source domain, we use Oj as our training distribution. We randomly sample 50 data points and train a linear classifier with a gradient descent of 3000 iterations. By minimizing the objective function in (9), we can get the empirical optimum Ò. Then we gradually shift the distribution by increasing a to get different target domains. Along distribution shifting, the test loss É(O, ya) will increase. As shown in Figure 2, the test loss will culminate in around rads due the maximum distribution shifting. Comparing different levels of regularization, we found that the larger L2-penalty / brings lower OOD generalization error which is shown as darker purple lines. This plot bears out our intuition in the previous section. As stated in the aforementioned case, the sharpness of ridge regression should inversely depend on B. Correspondingly, we compute sharpness using the definition equation (4) by averaging ten different results. For each trial, regressor is shown in the legend of Figure 2. As we can see, larger 3 leads to less sharpness.
the same training and test data for every B. The sharpness of each ridge
5.2 SHARPER MINIMUM HURTS OOD GENERALIZATION
In our results, we proved that the upper bound of OOD generalization error involves the sharpness of the trained model. Here we empirically verified our theoretical insight. We follow the experiment setting in BemainBed (Gulrajani & Lopez-Paz, 2021). To easily compute the sharpness, we choose
7

---

Published as a conference paper at ICLR 2024
8-0.0
1027
96.5
Envo
6-
6-1.0
70.56
Test Accuracy
96.0
Env2
B=1.5
Env3
8=2.0
85.76
50
95.5
95.0
2
94.5
0
0.0
0.5
1.0
1.5 2.0
25
3.0
180
200
220
240
200
200
300
320
Angles (rad)
Sharpness
Figure 2: OOD test losses increase along dis- tributional shifting. The X-axis is the shifting angle a and the Y-axis the test loss of the model which trained distribution a O. Lines are average test losses and shadows are variances of 10 trials. Larger regularization (darker color) causes a lower increase in test loss but smaller sharpness.
Figure 3: The relationship between out-of- domain test accuracy and model sharpness RotatedMNIST dataset. Here we show 4 dif- ferent OOD environments: 15,, 30" .45%, ,, 60 rotation as the OOD test set respectively. Each marker denotes minimum of an algorithm with a specific seed. The marker style means the models trained in the same environment.
the 4-layer MLP RotatedMNIST dataset where RotatedMNISTT is rotation of MNIST handwritten digit dataset (LeCun, 1998) with different angles ranging from [0 15", , -00 75%]. codebase, each environment refers selecting domain (a specific rotation angle) as the test domain/OOD test dataset while training on all other domains. After getting the trained model of cach environment, we compute the sharpness using all domain training sets based on the implementation of Petzka et al. (2021). To this end, we plot the performances of Empirical Minimization Risk (ERM), SWAD (Cha et al., 2021), Mixup (Yan al., 2020) and GroupDRO (Sagawa al., 2019) with 6 seeds of each. Then we measure the sharpness of all these minima. Figure shows the relationship between model sharpness and out-of-domain accuracy. The tendency is clear that flat minima give better OOD performances. In general, different environments can not be plotted together due to different training sets. However, we found the middle 4 environments are similar tasks and thus plot them together for clearer trend. In addition, different algorithms lead to different feature scales which may affect the scale of the sharpness. To address this, we align their scales when putting them together. For more individual results, please refer to Figure in the appendix.
In this
5.3 COMPARISON GENERALIZATION BOUNDS
To analyze our generalization bounds, we follow the toy example experiments in Sagawa et al. (2020). In this experiment, the distribution shift terms and generalization error terms can be explicitly com- puted. Furthermore, their synthetic experiment considers the spurious correlation across distribution shifts which now general formulation OOD generalization (Wald et al., 2021; Aubin et al., 2021; Yao et al., 2022). Consider data [Neone Tspul that consist of two features: feature and spurious feature. The features are generated from the following rule:
Teore
y 1 Ny Ta)
Tspo a N (al,
where E 11 } is the label, and a E (-1, is the spurious attribute. Data with forms the majority group of size Rmaj- and data with y = -a forms minority group of size Re Total number of training points n Tmaj + Romn- The spurious correlation probability, Pmai probability of in training data. In testing, we always have Pmaj 0.5. The metric, worst-group error Sagawa al. (2019) is defined
defines the
Errwg (w)
Ez (wo (x, (F,9))] sE[4]
where fu-1 is the 0 loss in binary classification. Here compare the robustness of our proposed OOD generalization bound and the baseline in Proposition 2.1. We also give the comparison to other baselines, like PAC-Bayes DA bound in the Appendix G.
8

---

Published as a conference paper at ICLR 2024
S Error
Bascline Ours
Generalizztion Bound
Beselie Ours
Test Error
average
DF Model Size (b)
Model Size
Model Size (c)
average worst group
Distribttion Disterce
Test Error
Bassline
cetton
Soseline Ours
Ours
Correlation Probability
Correlation Probability
Correlation Probability
(d)
(c)
(l)
Figure 4: Spurious feature synthetic experiment. Each dot represents trained model. The dash curves are the smoothed function fit by the test data points. The baseline is Proposition 2.1. (d), the generalization error of the logistic regression models with increasing the model size/correlation probability. (b): concentration error term in domain shift bound. (e): comparison of distribution distance bounds. (c),(f): comparisons of generalization bounds. Note that model size 500 is the overparameterized regime. The further the correlation probability is from 0.5, the greater the distributional shift is.
Along model size We plot the generalization error of the random feature logistic regression along the model size increases in Figure 4(a). In this experiment, we follow the hyperparameter setup of Sagawa et al. (2020) by setting the number of points n = 500, data dimension 2d = 200 with 100 on each feature. majority fraction Pmaj 0.9 and noises 2 1, 0 100. The worst-group eтror turns out to be nearly the same as the model size increases. However, in Figure 4(b), the error term in domain shift bound Proposition 2. (A) will keep increasing when the model size is increasing. In contrast, our domain shift bound order /K independent of the model size which addresses the limitation of their bound. We follow Kawaguchi et al. (2022) to compute K in an inverse image the e-covering in a randomly projected space (see details in appendix). We set the same value K 1, 000 in experiment. Different from the baseline, K is data dependent and leads to a constant concentration error term along with model size increases. Analogously. our OOD generalization bound will not explode model size increases (shown in Figure 4(c)).
Along distribution shift In addition, we are interested in characterizing OOD generalization when test distribution shifts from train distribution by varying the correlation probability Pmaj during data generation. As shown in Figure 4(d), when Pnaj = 0.5, there is no distributional shift between training and data due to no spurious features correlated training data. Thus, the training and test distributions align closer and closer when Pmaj < 0.5 and increase, resulting in an initial decrease in test error for the worst-case group. However, as Pmaj > 0.5 and deviates from 0.5, introducing the spurious features, shift in the distribution occurs. This deviation is likely to impact the worst-case group differently, leading to an increase in the test error. As displayed in Figure 4(c) and Figure 4(1), our distribution distance and generalization bound can capture the distribution shifts but are tighter than the baseline.
CONCLUSION
In this paper, we provide a more interpretable and informative theory to understand Out-of- Distribution (OOD) generalization Based on the notion of robustness, we propose robust OOD bound that effectively captures the algorithmic robustness in the presence of shifting data distribu- tions. In addition, our in-depth analysis of the relationship between robustness and sharpness further illustrates that sharpness has a negative impact on generalization. Overall, our results advance the understanding of OOD generalization and the principles that govern it.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENT
This research/project is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: ALSC-GC-2011-001-2AA (AISG Award No: AISG2-TC- (023-010-SGIL) and the Singapore Ministry of Education Academic Research Fund Tier 1 (Award No: TI 251RES2207). We would appreciate the contributions of Fusheng Liu, Ph.D. student at NUS.
REFERENCES
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean- -hrishhist Gagnon- Augeo,-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. Advances in Neural Information Processing Systems, 34:3438-3450, 2021.
Alnur Ali, J Zico Kolter, and Ryan J Tibshirani. A contimuous-time view of early stopping for least squares regression. In The 22nd international conference on artificial intelligence and statistics, pp. 1370-1378. PMLR, 2019.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint (907:229000089 2019.
Benjamin Aubin, Agnieszka Stowik, Martin Arjovsky, Leon Bottou, and David Lopez-Paz. Linear unit-tests for invariance discovery. arXiv preprint arXiv:2102.10865 2021.
Peter Bandi, Osear Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, al. From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge. IEEE Transactions on Medical Imaging, 2018.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79:151-175, 2010.
John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman. Learning bounds for domain adaptation. Advances neural information processing systems, 20, 2007.
Jean-Christophe Bourin, Yun- Lee, and Minghua Lin. Positive matrices partitiooned into a small number of hermitian blocks. Linear Algebra and its Applications, 438(5):2591-2595 2013.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405-22418, 2021.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.0066 00060, 2021.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamoniy. Joint distribution optimal transportation for domain adaptation. Advances in Neural Information Processing Systems, 2017.
Simon S Du, Jayanth Koushik, Aarti Singh, and Barnabás Póczos. Hypothesis transfer learning via transformation functions. Advances in neural information processing systems, 30, 2017.
Yingjun Du, Xiantong Zhen, Ling Shao, and Cees GM Snoek. Metanorm: Learning to normalize few-shot batches across domains. In International Conference on Learning Representations, 2020.
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally 2018
robust optimization. arXiv preprint
John C Duchi and Hongscok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. The Annals of Statistics, 4938))1378-1406 2021.
Laurent El Ghaoui and Hervé Lebret. Robust solutions to least-squares problems with uncertain data.
SIAM Journal on mairix analysis and applications, 1997.
10

---

Published as a conference paper at ICLR 2024
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXi::2010.01422, 2020.
Pascal Germain, Amaury Habrard, François Laviolette, and Emilie Morvant. A pac-bayesian approach for domain adaptation with specialization to linear classifiers. In International conference on machine learning. PMLR, 2013.
Pascal Germain, Amaury Habrard, François Laviolette, and Emilie Morvant. A new pac-bayesian perspective domain adaptation. In International conference on machine learning. PMLR, 2016.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Av- eraging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.085 3054077 2018.
Kenji Kawaguchi, Zhun Deng, Kyle Luh, and Jiaoyang Huang. Robustness implies generalization via data-dependent generalization bounds. In International Conference on Machine Learning. PMLR, 2022.
Daniel Kifer, Shai Ben-David, and Johannes Gehrke. Detecting change in data streams. In VLDB, volume 4, pp. 180-191. Toronto, Canada, 2004.
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: benchmark of in-the-wild distribution shifts. In International Conference Machine Learning, pp. 5637-5664. PMLR, 2021.
Masanori Koyama and Shoichiro Yamaguchi. Out-of-distribution generalization with maximal invariant predictor. 2020.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/nmist// 1998.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference computer vision, pp- 5542-5550, 2017.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning generalize: Meta-learning for domain generalization. In Proceedings of the AAAI conference on artificial intelligence, 2018.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In International Conference on Machine Learning, pp. 6804-6814, PMLR, 2021.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. In Advances in Neural Information Processing Systems, 2022.
Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms. arXiv preprint arXiv:0902.3430 2009.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75 (4):667-766, 2022.
John Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang, Yair Carmon, and Ludwig Schmidt. Accuracy the line: the strong correlation between distri--iitiiiitionn and in in-distribution generalization. International Conference on Machine Learning. PMLR, 2021.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018
11

---

Published as a conference paper at ICLR 2024
Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Advances in neural information processing systems, 33:22182-221921 2020.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms. The MIT Press, 2017.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. Advances in Neural Information Processing Systems, 34: 18400-1432, 2021.
Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron C Courville, Doina Precup, and Guil- laume Lajoie. Gradient starvation: A learning proclivity in neural networks. Advances in Neural Information Processing Systems, 34: 125- 2021.
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out- of-distribution generalization. In International Conference on Machine Learning, pp. 18347-18377. PMLR, 2022.
levgen Redko, Emilie Morvant, Amaury Habrard, Marc Sebhan, and Younès Bennani. sur- vey on domain adaptation theory: learning bounds and theoretical guarantees. arXiv preprint 0rXiv:2004.11829 2020.
Herbert Robbins. A remark on stirling's formula. The American mathematical monthly, 62(1):26-29, 1955.
Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 2018.
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 generalization. 2019.
arXiv preprint
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning. PMLR, 2020.
Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. Advances in Neural Information Processing Systems, 2015.
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv:2104.09937, 2021.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450. Springer, 2016.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domail generalization. Advances in neural information processing systems, 2021.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre-Alvise Rebuffi, Ira Ktena, Krishnamurthy Dj Dvijotham, and Ali Taylan Cemgil. A fine-grain gnninnedd analysis distribution shift. In International Conference on Learning Representations, 2022.
Huan Xu and Shie Mannor. Robustness and generalization. Machine learning, 86(3):391-423, 2012.
Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training, arXiv preprint arXiv:2001.0067 2020.
Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation. arXiv preprint arXiv:2201.00299 2022.
12

---

Published as a conference paper at ICLR 2024
Haotian Ye, Chuanlong Xie, Tianle Cai, Ruichen Li, Zhenguo Li, and Liwei Wang. Towards theoretical framework outtof-distribution generalization. Advances in Neural Information Processing Systems, 2021.
Xiao Zhang, Yaodong Yu, Lingxiao Wang, and Quanquan Gu. Learning one-hidden-layer relu networks via gradient descent. In The 22nd international conference on artificial intelligence and statistics, pp. 1524 1534. PMLR, 2019a.
Xingxuan Zhang, Linjun Zhou, Renzhe Xu, Peng Cui, Zheyan Shen, and Haoxin Liu. Towards unsupervised domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 4910-4920, 2022.
Yuchen Zhang, Tianle Liu, Mingsheng Long, and Michael Jordan. Bridging theory and algorithm for domain adaptation. In International Conference Machine Learning. PMLR, 2019b.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costeira, and Geoffrey Gordon. Adversarial multiple source domain adaptation. Advances in neural information processing systems, 31, 2018
Kai Zhong, Zhao Song, Prateek Jain, Peter Bartlett, and Inderjit Dhillon. Recovery guarantees for onc-hidden-layer neural networks. In International conference machine learning, pp. 4140-4149. PMLR, 2017.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL EXPERIMENTS
A.1 SHARPNESS V.S. OOD GENERALIZATION ON PACS AND WILDS- CLLDD-CAMELYON17
64
1.0
Envo
Env2
Test Accuracy
Test Accuracy
0.8
0.7
8.7
63
0.4-
645
Sharpness
Sharpness
Figure 5: The relationship between out-of- distribution (OOD) test accuracy the test environment and model sharpness (of last FC layer) on the Wilds-Camelyon17 dataset. Each marker denotes a model trained using ERM with different seed and hyperparameters.
Figure 6: The relationship between OOD test accuracy and model sharpness the PACS dataset. Each marker denotes a model trained using ERM with different seeds and hyperpa- rameters. The marker style shows the out-of- distribution test environment of the model.
To evaluate our theorem more deeply, we examine the relationship between our defined sharpness and OOD generalization
larger-scale real-world datasets, Wilds-Camelyon17 Bandi et al. (2018); Koh et al. (2021) and PACS Li et al. (2017). Wilds-Camelyon17 1 dataset includes 455,954 tumor and normal tissue slide images from five hospitals (environments). One of the hospitals is assigned as the test environment by the dataset publisher. Distribution shift arises from variations patient population, slide staining, and image acquisition. PACS dataset contains 9,991 images of 7 objects in 4 visual styles (environments). art painting, cartoon, photo, and sketch. Following the common setup in Gulrajani Lopez-Paz (2021), each environment used as a test environment in turn. We follow the practice in Petzka et al. (2021) to compute the sharpness using the Hessian matrix from the last Fully-Connectee test the sharpness of 18 ERM models trained with different random seeds and hyperparameters. Figure 5 shows the result. For the PACS dataset, we run 60 ERM models with different random seeds and hyperparameters for each test environment. To get clearer correlation, align the points from 4 environments by their mean performance. Figure 6 shows the result. From the two figures, we can observe clear correlation between sharpness and mut-of-distributiin (OOD) accuracy. Sharpness tends to hurt the OOD performance of the model. The result is consistent with what we report in Figure It shows that the correlation between sharpness and OOD accuracy can also be observed on large-scale datasets.
(FC) layer each model. For the Filds-Camelyon17 dataset,
B NOTATIONS AND DEFINITIONS
Notations We use [n] denote the integers set {i) represents f2-norm ||-1|2 for short. Without loss of generality, which is denoted as (0(z)) and we use n. d for training set size and input dimension. Note that we generally follow the notations in the original papers.
use (f, x)) y) for the loss function of model data pair (æ, 8).
Ls, Lr: expected risk of the source domain and target domain, respectively. The corre- sponding empirical version will be Es, Cr
{C K partitions on sample space and C₁ denotes each partitioned set.
Ds, Dy: distributions of source and target domain. Their sampled dataset will be denoted Ds, Dr accordingly.
14

---

Published as a conference paper at ICLR 2024
0: In our setting, 0 (w, {a]}") denotes the model parameters where is the trainable parameters and {a, ) are the random features (we also use = [a_,.... aml for short notation in many places). v0 is the minimizer of empirical loss.
M: upper bound of the loss function.
{( } [.]] training data of size n.
Definition B.1 (Robustness, Xu & Mannor (2012)). A learning algorithm A on training set S is
(K, eK,)(())-oobust, for K N, can partitioned into K disjoint sets, denoted by {C* ) such that for all E S E Z we have
Vs, Cr (K),Ie(ss (As, 2)1 c(S,
Definition B.2 (Hessian Lipschitz continuous). For twice differentiable function f IR" - R, has I.-Lipschitz continuous Hessian for domain æ; y are vectors in C2 if
f(o) - V f(x) < Lilly
where L; > 0 depends on input domain C, and || || La norm. Then for all K domains U have
Cs let
L == max [K]] be the uniform Lipschitz constant,
[() - v² f())) S Lilly y æ s L æl1 vi [ (x, E EC
which is uniformly bounded with
Lemma B.3 (Hessian Lipschitz Lemma). If f is twice differentiable and has L-Lipschitz continuous Hessian, then
f(y) - f(z) (Vf(x)) x) (V'
-Ff(=)(-2))((1 - z)) 1 - fly x/1²
Gegenbauer Polynomials
We briefly define Gegenbauer polynomials here whose details can be found in Appendix of Mei & Montanari (2022). First, we denote (r-1(r) {x E Rd ||>||| = =r} as the uniform spherical distribution with the radius on d manifold. Let Ta be the probability measure on $d-1 and and the inner product in functional space L2 d), paa denoted as (:)) and !|
[ (3--1(Va)
(1
f(z)g()))al dær).
For any function a E [.2([-d, d). ra), where Ta is the distribution (ær, y) /v/vd d (ax, y ~ (VT). the orthogonal basis } forms the Gegenbauer polynomial of degree t(t 0), its spherical harmonics coefficients (d) can be expressed as:
>a.e(o) - (Vdr)ra(t))
then the Gegenhauer generating function holds in 1 (I-vd, Vdl. sense
(x) Σ \a.lo)Naa, (Vddz)
where Na, is the normalized factor depending the norm of input.
B. ASSUMPTIONS
We discuss and list all assumptions used in our theorems. The purposes are offer clarity regarding the specific assumptions required for each theorem and ensure that the assumptions made in theorems are well-founded and reasonable, reinforcing the validity and reliability of results.
15

---

Published as a conference paper at ICLR 2024
OOD Generalization
(Setting): Given a full sample space 2, souree and target distributions are two different measures over this whole sample domain Z. The purpose is study the robust algorithms in the OOD generalization setting.
(Assumptions): For any sample vz Z, the loss function is bounded (0(2) [0, M). This assumption generally follows the original paper Xu & Mannor (2012). While is possible to relax this assumption and derive improved bounds, primary objective is to formulate framework for robust OOD generalization and establish a clear connection with the optimization properties of the model.
Robustness and Sharpness
(Setting): In order to give fine-gras fine-grained analysis, we follow the common choice where two-layer ReLU Neural Network function class widely analyzed most literature, i.e. Neural Tangent Kernel, non-kernel (rich) regime Moroshko al. (2020) and random feature models. Among them, we select the following random feature models as our function class:
f(w, A)
wao (æ, a;) wy E RR a, Unif(s" (Vd)), i [m]
where m is the hidden size and A [a,,..., am contains random vectors uniformly distributed on n-dim hypersphere whose surface is n manifold. a(a x) (a x x)l{{a x} denotes the ReLU activation functiin and I is the indicator function. w [w',, .... Wm] is the trainable parameter. We
choose the common loss functions: (1) Homogeneity in regression; (2) (Binary) Cross-Entropy Loss;
(3) Negative Log Likelihood (NLL) loss;
(Assumptions):
(i) Let C., i E [K] be any set from whole partitions U C, we assume Vz E C., the loss function (=) satisfies T.-Hessian Lipschitz for all E [K] (See details in Definition B.2). Note that we only require this assumption to hold within each partition, instead of holding globally. In general, the smoothness and convexity condition actually equivalent to locally convex which is a weak assumption for most function classes.
(ii) Consider an optimization problem each partition C, E [K]. Let one of training points (A) e SNC, be the initial point and x (A) E M, is the corresponding nearest local minima where M; is the local minima set of partition C₁. For some P.(L) > 0, we assume ||x, x [A)]| S
p. (L)/L holds a.s.. It ensures the hessian norm || (2; (A)) has the lower bound. Similar conditions and estimations can be found in Zhang et al. (2019a).
(iii) To simplify the computation of probability, assume & a. x obeys rotationally invariant distribution.
(iv) For Corollary 3.5, make additional assumptions that loss function E() satisfied bounded condition where the second derivative Yz Unif(Sd-1 (vV),, (oPd)))(@P(((()) (f A, x). v)/ with re- spect its argument f(wØ A, æ) should be bounded by [M1, M2]. Note, we consider the case where the data Væ (ini(()=I((Vd)) while m = Poly(d) is to ensure the positive definiteness of a E ddd almost surely.
C PROOF TO DOMAIN SHIFT
Lemma C.I. Let DT be the empirical distribution of size n drawn frownfrrom DT. The loss & is upper bounded by M. With probability least 1 - 8 (over the choice of the samples), for every fo trained on S, we have
2K In + 2l1(1
Cr(JO) S Ls(So) +
Dr) +cS) 2M
(11)
where
k Σ
nc(S) ny (T(DT)
Vi (eln),n,ss) == #(x (Sniin)),,
(S,
(12)
16

---

Published as a conference paper at ICLR 2024
and n4(S), are the number of samples from S and Dr that fall into the set C₁ respectively.
Proof. In the following generalization statement, we use {(So, z) to denote the enror obtained with input z and hypothesis function Se for better illustration. By definition we have,
Cr(So) Ls(fo) =- Eer ((S. x') (ELwD,L(Je.> z).
(13)
Then we make the K partitions for source distribution Ds. Let n, be the size of collec- tion set of points æ fall into the partition C₁ where is the n..d.muutinomial random vari- able with (ps (CC))( We use parallel notation for target distribution Dr with S ~ (pe (C1), (CCK))) Since
EgNDe ((fo,
Σ
EgnDs ( 2/) C.)).(C.)
i=1
(14)
Σ ED (l(fo. (e e C.),,(C.)
and thus we have
(r(fe) Ls(fe) = Σ [(((ffo,2))) C.)p(CC)) E(((f,,))) (C.)
=E(e[Jo z')1z' (c)),(C.)
(E(((Jo, 2')1=' E c.)) (p.(
(15)
[B(((fo, 2')1=' e C.) E(((fo.)) C.)lp.(C.)
K
(E(l(fo. 2')2z' E CJ) (p: - P. +
If we sample empirical distribution 8, Dr of size n each drawn from Ds and Dr respectively. (T.,..., rg) are the i.idd random variables belongs to C,, We use the parallel notation ni for target distribution.
K Σ
(,()))
nT(DT)
di, (S.
(16)
n
Further, we have
(m(DT)
ni(S)
(B(((Jo,'))<' E CC). ( pc())) Σ
(((((Jo,'')|=' EC:))
n
n
(p(CC)) ni(Dr)
[E(l(Jo, z')1z' e EC:)) (0))(()))) (p.(C.) (S(
(B(E(Jo: 2')12' E Ci))
n
K
PC. n((Dr)
K
(S)
< M
+M
(C(C,)
(17)
With Breteganolle-Huun Hullee---bb---e Huber-Carol inequality we have
K
((())
2K In + ln(1/8)
Σ
(C:)
S
(18)
n
n
17

---

Published as a conference paper at ICLR 2024
To integrate these two inequalities, we have
K
Σ B[L[foo 2')12' E ())
K (B(((Jo,>))] = CC)) nx(Dr) n4(5)
2K In + (1(1/))
+ 2M
2K In + (1
S M) 2M
n
(19)
In summary with probability 1 1 we have
2K In + (1((1/))
Cr(so) S (s( + (), c()) 2M
(20)
which completes the proof.
With the result the domain (distribution) shift and the relationship between sharpness and robustness, move forward to the final OOD generalization error bound. First, we state the context of ID robustness bound in Xu & Mannor (2012) as follows.
Lemma C.2 (Xu et al.Xu & Mannor (2012)). Assume that for all h E H. and z E Z, the loss is upper bounded by M i.e., f(h, 2) < M. If the learning algorithm A is (K, (())(())oobbss. then for any > O, with probability at least
over an iid draw of n samples S = (zi) it holds that:
2K 1n + 2l1(1/))
E, (As, 2) =
Σ i=1
21s c(S) M
As the conclusive results, we briefly prove the following result by summarizing Lemma C.2 and Lemma C.I.
Theorem C.3 (Restatement of Theorem 3.1). Let Dr be the empirical distribution of size drawn from Dr. The loss E upper bounded by M. With probability at least (over the choice of the samples), for every So trained S, have
2K In 2 21n(2/8)
Lx (0) < Ls(0) (d 2c(S) + 3M
(21)
Proof. Firstly, with Lemma C. and probability as least
2. we have
2K In + 2ln(2
Cr(fo) s Ls(fe) + Mdde. Dr) + 2M
+ c(S)
Secondly, with Lemma C.2 (Xu & Mannor (2012) Theorem 3) and probability as least 1 2. we have
2K In 2 2ln(288)
(s(()) Ls (S)) My
n
By taking the union bound, we conclude our final result that with probability at least
2K In 2 + (2/66 22(S) + Md, D7))
(r(SO) s Ls(Jo)+ + 3M
(22)
n
Here e(S) is the robustness constant that replace with any sharpness measure.
18

---

Published as a conference paper at ICLR 2024
C.I PROOF TO COROLLARY 3.2
Definition C.4. dFAF (Dr; Ds) :- 2sUPA(fIEAPA Pro.(A()) PRD, (A())) and FAF is defined
FAF = (f(x) @ [
where @ is the XOR operator e.g. [(f + f(z)).
Lemma C.5 (Lemma 2 Zhao et (2018)). If = d', then VCdim(FAFF) S 2d'.
Proposition C.6 (Zhao (2018)). Let F be hypothesis class with pseudo dimension Pdlm(
d' Ds is the empirical distributions generated with n i...d. samples from source domain, and Dr is the empirical distribution on the target domain generated from samples without labels, then with probability at least for all f E we have:
2d log
log 3
Cr(1) S [s(f)+ +
72
2n
2d' (n(2n) + in
(23)
draF (Dri +
n
(Empirical div Ernor)
where 8x E- + Cr( is the total error of best hypothesis f over source and target domain.
Proof.
With Lemma 4 (Zhao et al., 2018). we have
Cr(f) Es(f) 2 ddsar E
where
=- inf Cs(f) Cr(f). PeF
Lemma (Zhao et al., 2018), which actually Lemma in (Ben-David et al., 2010), shows following results
drAF Ds): < drAF (Dr:Ds) + +41
VCdim(FAF))
3
As suggested in Zhao et al. (2018), VCdim(FAF) is at most 2d', Further, with Theorem 2 (Ben- David et al., 2010), we have at probability at least 1
1 Cr(f) s Cs(f) + dFAF (Dr:Ds) + +
VCdim(FAF)) In(mn) + In
+
n
(24)
1 $ [s(f)+ 7 (DriDs) + +41
2d (n(2n) In 2 +
n
Using in-domain generalization error Lemma 11.6 (Mohri et al., 2018), with probability at least 1 the result is
2d log
log 5
Cs(J) < F.)(s) My
+ M
n
2n
Note in Zhao et al. (2018), the M for the normalized regression loss. Combine them all, conclude the proof.
Corollary C.7. If - 0o, M = 1, domain shift bound (Lr(Jo) Ls(so) will be reduced to (Empirical div Error) in Proposition C.6 where
[r(fo) Ls(fo)) s ddsr Dr) < (Empirical diy Error)
(25)
19

---

Published as a conference paper at ICLR 2024
Proof. According Theorem C.3, we have
K
Cr(fo) Ls(fo) - Σ B(l(fo: 2'1Z'E - C.)p,(C:)
(26)
If K + oo, to, let's define a domain that U == = C.. The equation (26) will be
- I ((Jo, (((22) - / ((Jo: 22)x
[ JwED
f 2EDS
(27)
(SS: (2)
ELWD, ED ()..2))
In this case, we have,
(CT(Jo) Ls(So))
[B [(o: z') EZzDS
PPDD (e(f(0, 2')0, > t) dt - -Prp, (((f(0,z),)) > t) dtl
/ PPDD (e(f(0,x'),
Pros (L(f(0, x) t) dt (M
(28)
I sup sup PrD, (e(f(0. 1) - PRds (4(f(0, > t) ((0,)FF
sup [Pror (A())) Pros (A()))
A(f)EAFar
2 drsr(Ds;DDp) Dr) < (Empirical div error)
where AFAF represents a learning algorithm under the hypothesis FAF (f(x) (cc): F}. which completes the proof.
D SHARPNESS AND ROBUSTNESS
Lemma D.1 (positive definiteness of Hessian). Let Wmin be the minimum value of [u) and x² arg ( ((f(x, A,) For any A = (a1...., am), a, ( (m) denote (dd m) = Amin E a,a; Ga) > be the minimum eigenvalue, where Gij = Exe X2 (o) (( a,) /Vd/v) the polynomial product constant. If m Poly(d). the hessian H(x") can be lower bound by
12 winn M), Ta-
E
x Unit(( ((a))) H(x))
2
(29)
d
Proof. As suggested Lemma D.6 of (Zhong et al., 2017), we have a similar result to bound the local positive definiteness of Hessian. By previous definition, the Hessian w.r.t. æ has a following partial order
H(x*)= D2(x""."" ,y²)
w,dga,a d' daa (aTTa""'''
M
a o' (0,Tx'))) (a)
(30)
d
1=1
2 min
Mr
ΣΣ
úu d'aa, x (1)))")
d
20

---

Published as a conference paper at ICLR 2024
For the ReLU activation function, further have aa T x²) 2 d'(aaaaa x²)
(31)
We extend the E (2 vdl, Ta) (where Ta is the distribution of (0), 12) Vd) by Gegenhauer polynomials that
a (x) Σ Na
(32)
Let A = (a1...., am) E Roxd We assume Vz E Unif(S)-((Va)) Lemma C.7 in (Mei & Montanari, 2022), suggests that
U tna x') (0 A.j((m]
mm
(33)
which shows matrix U is a positive definite matrix. Similarly, taking the expectation x", terms RHS of (30) bracket can rewritten as
m
m
E Uni ((V2)
a. o' (a (aaa"="'''o x'))
(34)
a:a Re B (a x² [ x
Besides, we have the following property Gegenbauer polynomials,
1. For v. S
(0) +), Q ((v,))) 12($a-1(Va)))))
Q6 ((2, ((x, v))-
Na
2. For v, E (Vd)
1 Naa
((r,))) NA.
Σ ( (x) ( (g).
where spherical harmonics (Y15 hSiSNa, forms an orthonormal basis which gives the following results
Ex. (a x0 [V
Na, (1a, = Gú <0.
t=
(35)
Hence, we have
m m
Σ
a.a Eey
x² a /Vd)o(aa x² /Va)) = Σ a Ga + (ar/))
(36)
1=1
Since m Poly (d), and (aheem) are i..d, then rank (Im (G)) rank (AAT) d. Let
id, m) = Eg Amin (Imh a.a; G.) > 0 we have
22 (d, m) m
Eg. H(x))) x
la
(37)
21

---

Published as a conference paper at ICLR 2024
Lemma D.2. Let Uk=1 Ck be the whole domain, the notion of (e, rorbbustnsss is described by
c(S, A)
max C.CUS C& =.,
sup (L,,A(z) - (z))).
Define M, be the set of global minima in Cs. where
M; - (z(A)(((A) = min (z))
zEC'
suppose for some maximum fraining loss point
z.(A)
max L0.A(z) - Lo L (A)) =EC;NS
there 3zg (A) where
2) - arg min ||2- 2((1))
*EM,
such that (A) (A)] S e almost surely hold for any A E Unif(S"-I (vd) and any A, (a,A(2) is I--Hessian Lipschitz continuous. Then the ((,, A) can be bounded by
2t( (( 4p:(L) 3
<(S,A) S max
SE KE[K] 2L
Proof. Let z E S be collection of (a. y) from the training set S and z/ denote any collection from the set C. We define local minima set M,, (which is the global minima of C:). Assume that for some maximum point (A) E max (z), there exists 2, (A) M, almost surely for all A Uni(())-I(V))) such that
arg min f (= (I =11 s.t. I|z(AA) =11 Pt(L) zEM,
z7 (A)
(38)
T
By definition, A) can be rewritten as
e(S, A)
max sup <E[K]
(()) - (z/))
iE[K] 2. (0,1(2) (2))(2")
sup
(39)
max (z))))((A)) (z)(((((A)). =E[K]
According to Lemma B.3, we have
A(S,A) max l (zz(A)) - (2,
iE[K]
(4)
S max (V Veo. (2) 2, -2 =A))
{E|K]
L
A
(A)),, =( =A)) + - (12 213
(V (A) =f(A)), zi(A) =i()) -1 (A))
kE[K]
\2 (A)) (A) - (A)I2
max [
T
(A) - 2ズ (A)l (Cauchy-Schwarr,
(40)
where (i) support by the fact (z²) = 0. With Lipschitz continuous Hessian we have
(2 (A)) (A) 2 (A) ( (z,
(41)
22

---

Published as a conference paper at ICLR 2024
Overall, have
1 (x))) (I)) ==(ai) - (((((A)) =A)1
e(S, A) S max
KE[K] 2
$ max
(A)=(A) =+(ail)) | - 25(A)1?
kE[K] 2
+ 1. 1 (12) - 2 -1) (A)l}
(42)
+ Pi(L)²
Pc(L)²
$ max
+
K 2
L
6L2
P=(L)²
(D;(L)3
( +
max
=E[K] 2L2
3L
which completes the proof.
Lemma D.3 (Lemma 2.1 Bourin et al. (2013)). For every matrix in M partitioned into blocks, we have decomposition
A
X
A
0
0 0 O V
U
U
+
X
0
B
for some unitaries U, U.V E Mm+mm
Lemma D.4. Then, given an arbitrary partitioned positive semi-definite matrix,
X
|||l||
|B||
X
B
for all symmetric norms.
Proof. In lemma D.3 have
A X
X
A
0
0
U
U +
V
for some unitaries U, V E MM+m- The result then follows from the simple fact that symmetric norms
are non-decreasing functions of the singular values where f f= ||
ll M + R, we have
X
(r
ss S
U +
g
x B
Lemma D.5. For (Va and æ are some vector Rdd with norm |||||| R(dd) d. we have
2 arccos
VR(d)
22d 4 RTR(d
P((e, a)22 [|a|3) 2 min
exp
(43)
4d
Proof.
We can replace the unit vector of a with by
P((x, a)² 2 ||a|l²) = P((x e)² 2 Similarly, we can replace æ by unit vector such that
(44)
P((x, e) 2 1) P
(s,c)² 2 R(d)
1
(45)
Solving e), Ridj. we get
e)² cos² 1000= R(d)
o arccos +
(46)
/ RR(d)
23

---

Published as a conference paper at ICLR 2024
In this case, the probability will converge to R(d) increases. As is known to us, surface area of
gd-1 equals
Ad = , 2/2
(47)
T
An area Ca of the spherical cap equals
A L'"
2n (d 11/2
Aa-1 sin 0rd00
sin
odo.
(48)
r d-1
(2(n-1)) where I' (n - 1) 24 11 VT.
When d = almost surely we have
P((x ee)
1) =1
(49)
2) When d = 2, we have a ~ S where S2 is a circle r = 1 and the probability is the angle between 8, how much the vectors span within the circle where
2 2 rd
2
(8, e) R(d)
(50)
π
π
3) When d 2 the probability equals
P ( ((A,e)
1
A (r)
A
cap (r)
(51)
VR(d)
2 Aa
2 Aa
where (r))(r) is the remaining area of cutting the hyperspherical caps in half of the sphere,
2z 2d 1)/2
Acap (r)
sin 2ind-2
r
d-1
2
2x (d- 1)/2
sin 0d0
(-=1)
(52)
2m (d- 11/2
(- 008
0os Φ)
r (d-1)
2
2n (d- 1)/2
COS Φ
3-a) If d is even then r (4) = (5 1) so
T r( ((-1) (d - 2F
VT
2-2
(53)
r (4) 2 1)
2d-2
d-2
2
Robbins' bounds (Robbins, 1955) imply that for any positive integer d
4d
T
exp
(- 8d Bd-11
<
2d
d
4d
exp
1 8d + 1 (1
(-
(54)
So we have
2A (r)
21 (1)
2d-1
d d-2
-1
Ad
(
d-1
π
d_2
2
cos Φ
2d-1
2)/2
cos Φ
(55)
2d-2 exp
Id=3
2d -4
Cos Φ.
TT exp
-Td=g
24

---

Published as a conference paper at ICLR 2024
So the probability will be
V2d 4
1
(s, (s, 2 R(d)
exp
cos Φ.
(56)
VT
4d -
Suppose R(d) = kd, k > 1, we have
2d -
2
lim d0000
(57)
Vakden txkdexp/ "dd-
3-b) Similarly, if d is odd, then r (-1) [==3 1233) 80
2d-2 ( 1=3 12
2d-2
d-3 3
(58)
r (d - 2
(d - 2) 2)
d-3
2
If d 3 then
(0,ee) R(d)
2 π
P
2
cos Φ
(59)
2/7
VR(d)
If d > 3 then Robbins' bounds imply that
(
2d-4
-1
n (dl
3)/2
1
2r (5)
d 2
d-3
d-2
exp
4d - 11
(60)
2
Thus, the probability will be at least
d-2
1
(s, es,e)²
>
exp
cos Φ.
(61)
R(d)
Vn(d - 3)
4d - 11
To simplify the result, compare the minimum probability that for Vd 2
1
d
P (8, e)² 2
exp
cos Φ
R(d)
VT(d -3)
4d
11
d-2
1
exp
v -3)
4d
11
(62)
d-2
1
exp
cos Φ
V 3)R(d)
4d
11
1 V2d 4
1
exp
[R(dd
4d -
Overall, we have Vdl N4.
2 arccos
1
VRta)
2d 4
1
9 ) }
P (s. 2) 2 R(d)
min
1
cxp
(63)
R(d(
4d
D.1 PROOF TO THEOREM 3.4
Proof. Let (ai)ieid) (.dd Unir(Sd-((Vd)). We consider the random ReLU NN function class
to be
1 Σ wao (x: a.)) w E RR, á [m]]
Frd) s(w, Aw,
Va
where A [a1,..., am] E RRdx The empirical minimizer of the source domain is
i min weRd
Σ
{(fww, A, A, = 1 Σ (zz))
(64)
ni 2.ES
25

---

Published as a conference paper at ICLR 2024
Then with the chain rule, the first derivative of any input at ti will be
Vel de (f A, v) dl A, x,
do (Ax)
0f(ū, A)
o(Az)
dx
(x. v)
Vd a 1 {a. 2 0}
(65)
D Da v) d'
vd
where short notation (xx y) denotes the first order directional derivative of f(ú, A, A. d'Az) Rmxd d the Jacobian matrix w.r.t. input æ. Apparently, the second order derivative is represented as D3 (z, v). thus the Hessian will be
A,1 (3(() d
ΣΣ a 2 20,aaa 20}
(66)
()( d' (Ax)) A o' dAxa)
d
Similarly, we have
(e A, (ar, v)
(sgn(y))² 2 (z
(67)
where syn(y) is the sign function. holds under our choice of the family of loss functions.
1. Homogeneity in regression, ie. LI, MSE, MAE, Huber Loss, we have [2] 8)l (z. v)l:
2. (Binary) Cross-Entropy Loss:
(e( = 0² Σ exp(x)) exp(xx) 10y² = 0;
3. Negative Log Likelihood (NLL) loss: D3 v)
Besides, as convex loss function, D3 v) 2 0. Hence, the range of D2 y) will be [0, (1, v)].
To combine with robustness, we denote (ar; Rd+1, Therefore, the Hessian of will be
(' 8ydn
(x) =
(68)
('EL(((.AA.)..))
(A A, x)
dyee
With Lemma D.4, the spectral norm of Hessian z will be bounded by
[=)|| ( A,*). + (,
(69)
The first term in (69) can be further bounded by
ww T l (Az)o' (Ax)) [| DJ (x. (Ax) "r of (Ax) || < Dy (x. v)
(70)
= D3 (r. 1| 1|2 ||o' (la (Az) (Am)
where the convexity of loss functions Vx, y, D2 y) 2 0 supports the last equation. The right term has the facts that
||o'(Am) (Aa)) 11o' (Ax)o' (Ax)) T ||
(Aa)))'(AAr)]]
(71)
26

---

Published as a conference paper at ICLR 2024
In summary, we have the following inequality:
(1 [ T| +11
|H(2)]| (z( y)
I (z( y)
[110 (o (Ax)) ) +
(72)
m (d ||
D3(a, y)
0 +1 la 1 20
In Lemma D.2, it depends on some z; (2i E ESnC, that
Pc(L)²
[H(()) + Ap:(L)
c(S, A) I max (E[K] 2L
Pc(
|1b|1
Api(L)
max kE[K]
212
(2,,
| [1aj x. 20 +1
+
3
(73)
Pmax (T.)
D3( ykk
m
4 4pmax (.) Ōz
| Σ
||yل| 11 [{ 2b 2 0)
2L2
d
where Pmax (L) max (L) the o O (J, Fh. yk 11012 + ||y]] [la]1"14 [ a 2 0} d/m)
smaller order term compared first term, since taken we find maximum (æt,, yn) Ô E {Ci)eek]. Because æk, y/ E S is of the training
d. Last equality, the maximum can be
sample le e [m]. there must exist n'' [0, n] that
m
(x 10001
Σ
|a 2
[a] wR 20 0}
(74)
n'
m (2 mm
m
|| n
||a 1| az 2 0}.
Recall that the sharpness of parameter à is defined by
16( <( S.
== ||00| tr (((H).
= 1
Dj (zzz U,) (o(A);)o(Azz;)
n2
(75)
n D (ar, Mj)
m
12 1
Σ (a x. 21 aa; Φj 2 0]
7
d
Let the ; a. ~~ D(E) and the expectation of E(g: > 0) = qi where D(E) is some rotationally invariant distribution, i.e. uniform normal distribution. Under this circumstance, the sample mean of & still obeys the same family distribution as D(EE). Thus, we have
4i
4
P
EI (Es 2 03
Σ ==1
|a
1I 2 0}
P
Σ
E,
|| |ay11
(76)
j=1
IP((a
x)² > ||a| 22) Exp(x)
With Lemma D.5, we have at least probability at
P((a T x)² > [|a]|?) = min f
(m(d)-t),
/2d 4
arccos
(77)
RxR(d)
27

---

Published as a conference paper at ICLR 2024
the following inequality holds,
e(S, A) z Pmaxx (L)
4pmax (L.)
((
+ 0r.
212
3
(78)
Pmax (L)²
(n''
n' +0 (mm
4 k(10, S, A) 3 P
21,
m
D.2 PROOF TO COROLLARY 3.5
Corollary D.6 (Restatement Corollary 3.5). Let Wmin be the minimum value of |aol. Sup-
pose Vo ~ Uni((st-((Va)) and [ A, x), f2 is bounded by [M1, My). If m >
Preax(1) < (w Mí (d, ()d/(2d) for any A (a1,..., am), a; ~ Unif(S--I(Vd)itt E [m]. taking expectation over all 2j E Unif(Sa-1 (Vd) in S, we have
7 Pmmax (L)² (n'n(10, (, S. + MM)
ES,A [c(S,A)] & Es.A
(79)
6L2
where ö(d, m) = Eai Eie a.a. Ga)) > is the minimum eigenvalue and Ga is product
constant of Gegenbauer polynomials
G = Σ
(a.(o) ( (far, Vv)//
I=0
Proof. In our main theorem, with some probability, we have the following relation
E(S, A) S Punax(
( S, + 4Pmax(()) + dx
21,2
So, we are concerned about the relation between the k(w S, A) second term. If (;()) < x(t0, S, A). we may say the RHS is dominated by sharpness term x(w0, S, A) well as the main effect is taken by the sharpness. As suggested Lemma D.1, we have
w min MIōA(m, d)
B Unitsai ((Va) H(x))
la
(80)
d
where x' is the global minimum over the whole set. As defined in (38), the following condition holds true
3z7 (A) e Mr ||x -, =A)ll S Pc(L)
(81)
L
and with Hessian Lipschitz, the relation almost surely for arbitrary x that (x;ca))l((=((AA) - (A))lI (A) =AAil S pc(L)
2 min M1ō m))
Ia
2d
(82)
< EA 2 (x=")l
Ex / | [2 (A))l
Obviously, [a))())) > 2P.(L). Following Lemma A.2 of Zhang et al. (2019a), for 2 (A), (z(A) E C, have similar result that
Fmin H (=))) 2 5min
(minn(((2)))))) (H(z, [[H(((A)) - H(2 (A)) (t(
(83)
where Huin denotes the minimum singular value. With Lemma D.2, we know that
(L)
4p;(L)
ES,AC(S, A) S E5S. max
[HH((
(84)
vE(K)
2L3
3
28

---

Published as a conference paper at ICLR 2024
We also have another condition that
@Pe(f(E0 A,)),
1D (z. v)
E [Mi M2), vx, y-
(85)
0ʃ²
Combine all these results, we finally have
Pc(L)²
4
Es,Ac(S, & Es,A max iE(K] 2L
1+
(2(())))
3
(86)
Pmax (
D2 (aR yk.)
I Es.A
|
||,1| 1| 2 1 {aj 2 20} + M2
6L2
d
==1
Recall the definition of n(( S, A) in the main theorem that
Es Es,A S, A) < (2)= ||
1
Dz 81)
(a x.
) )²1 {a, az 2 0] (87)
n
d
Look at the second sum, we have
m
Ei (a,
aa) {az my 2 2 Σ Bay Ra; |a: 10:121100511 cos cO)T a aj 2 0]
(88)
Σ Bmy Bai |a: 11 æj 0} coos (8).
Suppose x and a are i..dd from Unif((d-1(11)) let u = a). we have well-known result that
En (x, 2
||||=| |2 all cos (3)
(89)
EEcoo= d'
a E Rdd
Therefore, in (88).
Σ En, Fe. |a:11311 x 2 0) dc) Σ Exe Ba, 1ail æz 2 0}
(90)
and we have (based on proof of main theorem), Es,
Pauax (L)²
n
D&
( Ui)
m
||00
Σ
Re E. ||a,1111((/a x, 2 0) Mi
612
7
d
(91)
1=1
Ppmax (I)2
< ES.A
(1'n(10,SAA) S, + MM)
6L2
E CASE STUDY
To better illustrate our theorems, we here give two different cases for clearly picturing intuition. The first case the very basic model, ridge regression. As is known to us, ridge regression provides straightforward way (by punishing the & norm of the weights) to reduce the "variance" of the model in order to avoid overfitting. In this case, this mechanism is equivalent to reducing the model's sharpness.
Example F.I. In ridge regression models, the robustness constant has a reverse relationship to regularization parameter B where 3 1, the more probably flatter minimum k. and less sensitivity of the learned model could be. Follow the previous notation that e(S, A) denotes the robustness and x(,,S) is the sharpness on training set S. then we have
<>0, 0, (0, e(S, A) S cx(O, S) 5a where da is much smaller order than x(θ,S) S).
29

---

Published as a conference paper at ICLR 2024
E.1 RIDGE REGRESSION
We consider a generic response model as stated in Ali et al. (2019).
y 00.- (X0., o2)
ridge regression minimization problem is defined by
1 min - |xe yl1 マ 101
3
X E Raxd n < d.
(92)
2n
The least-square solution of ridge regression is
Ô - x x y.
(93)
With minimizer 0, we now focus on its geometry w.r.t. æ. Let S {z}; (X, y) be the training set, (2, E, p) be a measure space. Consider the bounded sample set Z such that
3M > 0, p(2) < to0.
(94)
The can be partitioned into K disjoint sets [Cilie(k]. By definition, we have robustness defined by each partition C₁,
Vz, E C. |(0 2) - E(θ, z') - <(S,
(95)
For this convex function (0, 2), we have the following upper bound in the whole sample domain
<(S) max sup
[(0, 2) E(0,, =')
EE E. EC,
max sup E(Ō, z) ((, 2 C.)
(96)
kE(K] ZEC;
< sup ((0 2,) - f(0,z') x²)
2,E
where the =² 2" are the global minimum point in C₁ and whole domain Z = U C. respectively. zy is a training data point that has the maximum loss difference from the optimum. Specifically, it as well as the augmented form of Ô can be expressed as
[1 La. v]
0 10,,... -1]_, Rd+1
Then the loss difference can be rewritten as
l,_ (z) = 0 x)² = (ŌT - v) + H(6g (z))is P.S.D matrix.
It is convex function with regards such that
((S) < sup fa_ (z₂) La_ (z")
EYES
sup Vlg, (x*) x²) + 1 [2( -x') H((g(>>))))) =2)
2
x, ES
(97)
sup [z) =T)T((((
"1(6__ (z)) 2²)
2,es
sup ( (z*)) 2012
2,e3
where the second equality is supported by convexity and the third equality is due to &g. (z,) 0. Further, with Lemma D.4, we have
80 (#s)) dyow
0
(4 (z*))]I
(o't,, (xx)) diedy
T
< =000T +
(98)
30

---

Published as a conference paper at ICLR 2024
In (97), can also bound the norm of the imput difference by
|| z ==11 < ||ag x x=112 + (Us - 3²²
(99)
and then for simplicity, assume
| ==1 - || 1| we have
1 2
(1000 +1 (1|ae) ===12 (y5 1²)
((S) S sup
[10|] + 1 (||=,|1 1| 2 + ||00 11 + (z) 5²)
sup 2, ES 2
(100)
$ sup
1
[6]
(1|=y)1 |12 + ||>0|1 O(d)
ES
2
S sup e
0
|| 1| + (())
where 0 2 |,, 11 = O(d²) is the dominate term for large d. Now, let's look at the relation to
sharpness. By definition,
X + (11) n
x(0, S) = = tr (Ho(e(0 sy) ||0
(101)
Since
X X
X
X
tr
+ 31 = tr
+ tr (B1) = tr
+ B =
| 112 1
(102)
n
n
so we have
k( x(0, 5) = =10112 n2
Σ x + 0
(103)
As is known to us, the "variance" of ridge estimator Ali et al. (2019) can be defined by
Var(@) a tr (0OT) = =tr ((x*x + n311 -1 x yy x (x + n.31) -
(104)
Note that T is PSD with rānk( 1, thus it has only eigenvalue x(θĀT) = =1:011 0.
Var(O) tr (00) | 0(8-2)
(105)
By definition, the covariance matrix Elyy is diagonal matrix with entries of o², Averagely,
have
tr (00T) o² tr (x X n311 -1 X x (XT x +n31) "]
σ² tr
X
X +11
-2
(106)
XxX X/n)
(Ax( X/n)+
where X and xx=xx + B1)
are simultaneously diagonalizable and commutable. Therefore,
the greater 3 is, the smaller tr (00T) is.
Conclusions From our above analysis, we have the following conditions.
Upper bound of robustness E(S) < sup Es
10 (d).
Sharpness expression x(,,)) = = Emy ||, + 3)
31

---

Published as a conference paper at ICLR 2024
Variance expression Var(0) ||0]2.
1) First, let's discuss how 3 influences sharpness.
x(0,5) 5) = |101| (H)(((),S))
= y (xT + n3I1) -2 x
E EJ
|0, + B
(107)
n
= -((3-1).
As dictated in above equation, the k(O, = 0 where sharpness holds an inverse relationship to 3.
2) Now, it's trivial to get the relationship between robustness and sharpness by combining the first two points. Because z Vj E in supermum one of the training samples there exists a constant c < n and r0 = ((d)) such that
c(S)
[1011112
ōs
n
c Σ (10111123)) n
c3) c311011
(108)
cx(6, ō..
This relation is consistent with Theorem 3.4 where the robustness is upper bounded by n' times sharpness S). Besides, the relation is simpler here, where robustness only depends sharpness without other coefficients before x(0 S).
3) Finally, the relation between /3 and robustness will be
X X/n) Ej |x
((S) S C
[Var(0)12e,113]) + ō
Co
(x X/n) + 8)²
²² X/n) X/X)
Σg Xx(XTXXmm X/n)
(109)
+ öz
= 3)
where the e(S) somehow is the order of ((3-2) 2) in the limit of B. It's clear to us that the greater 3 is, the less sensitive (more robust) model we can get. practical, we show that "over-robust" may hurt the model's performance (we show the detail empirically in Figure 7).
E.2 2-LAYER DIAGONAL LINEAR NETWORK CLASSIFICATION
However, our main theorem assumes the loss function satisfying Dolyak-Lojasiewice (PL) condition. To extend our result more general case, here study 2-layer diagonal linear network classification problem whose loss is exponential based and not satisfied the PL condition.
We consider classification problem using 2-layer diagonal linear nerwork with exp-loss. The robustness (S, A) has similar relationship in Theorem 3.4. Given training set S, after iterations 1. To, 302 > <S,, A) C2 supr2T, ())
Example E.2.
Given a training set 5 = (X, y) X [1, .... LnI. X E Rdx=, A depth-2 diagonal linear networks specified by:
with parameters u [u., u TR 2d
f(u, x) (u² n2 2
We consider exponential loss where C(t) = 1/2 E exp(- O(t)y:) and y; E {-1,1), 1). J. has the same tail hehavior and thus similar asymptotic properties as the logistic or cross-entropy loss. WLOG,
32

---

Published as a conference paper at ICLR 2024
assume Vi E [n] such that z; y, Suggest by Moroshko et al. (2020), have
((t) 2a sinh 4x S. r(s)ds)
where r(t) n e -X (t)) and ||x(t]]1= = C(t). Note that
d0(t) 4 -=0 4a X exp (-x 0()) A(1)x(())
(110)
di
where A(t) diag (4 +
Part i, sharpness. First derivative and Hessian can be obtained by
Vol(t) (Xr(tt)
Ho(((t)) = a(Xr(1)
(111)
80(() Σ [
With the definition of sharpness, get
x(O(t),S) = |10(t)]2 (r(H)(((()))
|10(t)]1 tr
(112)
1|0(t)] tr
Σ (1)xx X;
||((t)]?
Σ r+(0))
Part ii, robustness. Now, let's use the same discussion of the robustness constant in the previous case. Follow the previous definition of c(S, A). after some iteration number T, is defined by
cSS (S)A))= = sup |nr,(t) - (t)l
(113)
iE(m],+2T.
where (t) is the (denormalized) point-wise loss of æ, and r" (t) denotes the minimum loss of point x*. There exists n' n, we have
c(S, A) sup n' (1)(())l)-)) (t)) s sup n' ((])((t)) 12T,
(114)
12T,
Let [mmia|] min |<:|| the above equation
A(S x
sup (n' ||x(t)]] 2mnn||
@mial ?|?
127.
n'
Ømin ||
sup 12
(115)
n
[miin
||
127,
sup
1=1
Part iii, connection. Compare the last part of (112) and (115), we found that robustness and sharpness depend on the same term. Further, we can say that for any step f, [|@(t)] will have the upper bound.
From Lemma of Moroshko al. (2020), have the following inequality.
|0)()) s 2a sinh
7(1)
(116)
22
where C(I) exp(-7(1)). Then, can bound the [10(1))² via:
C < [10(1)) 410(1)133 = 4do sinh?
2
7(0)
< 00
(117)
33

---

Published as a conference paper at ICLR 2024
Note that C₁ > 0 In summary, have
n'
((, Σ
sup (()(t))) S C2 sup s()(t), 127,
(118)
a 127,
where C2 millmmia
0 is constant.
Part iv, sanity check- asymptotic. Asymptotically, as + oc, L(t) + 0 while []0(t) will be explode. So if sharpness n: ((O(t), X) co, then it will fail to imply robustmess.
(0((),
|0
Σ (1)))112
[((())1²
(t) =mxxl
(119)
[ [((((((((())))))) Fmax 1
Kmax S)
Let Kmax (f()) be a upper bound of (t)( X) at any time step t. The dynamics will drimax (0)
a max ()))))))(11
dt
max ?? tr((x] (10(()))2(())) + (2000))))200+0+0000000)00000
Pmaxx 1| () (X0)::/))))XXX000000000000000000
Fmax
txx(Xx
- (1000))((*+()) A(1)Xr(1)
(120)
As 0o, have C(t) O, thus it is easy converge that
lim drmma ()( - (1000 ((()(()))) =I| || ((0((()
= 0
dt
(121)
As can see, the dynamics of the derivative of Kimax (0(t), X) decreasing as -- oo which means the sharpness (()(f), X) will be upper bounded by converged number k(O(t + 00), X)
Coo oc. sharpness will not explode.
COMPARISON TO FEATURE ROBUSTNESS
F.1 THEIR RESULTS
Definition F.1 (Feature robustness Petzka et al. Petzka et al. (2021)), Let f: ( x y - denote a loss function, and 8 two positive (small) real numbers, S C X x y a finite sample set, and A E R matrix. A model (xx (+b (x)) with o(X) C Rm is called ((6, S, e)eeatuee robust, if ( (AA)) S e for all 0 S < ó. More generally, for probability distribution A on perturbation matrices in R. we define
(f, EANA (T)
and call the model ((6, A), ff-auuee robust on average over A, if (f,, s for 0
<
Theorem F.2 (Theorem 5 Petzka et al. Petzka et al. (2021)), Consider a model f(z w) g(w((())
as above, loss function and sample set S, and let Om Rmxm denote the set of orthogonal matrices. Let be a positive (small) real number and w = w E Bdxm denote parameters at local minimum of the empirical risk on sample set S. If the labels satisfy that (ØSA (x+)) y (o (x+)) for all (Ii, 3i) E S and all |||]| < 1, then o fx, w) is ((5, S, Om) e robust
average Om for
82 R (w) O
2m
34

---

Published as a conference paper at ICLR 2024
Proof sketch
1. With assumption that v(osA(I;)) = Us for all (2a, Ui) S and all ||]]] Σ 1, feature perturbation around
S, a.) + Eemp (w, S) Ee + A.w.A
2. Since Taylor expansion for local minimum w w will only remains second order term, thus
a²
(w + awA, (wm
Σ (G.4))H.(w..d($))))
2 8,1=1
3. With basic algebra, one can easily get
8²
d
EA~Om (S(()S,aA)
s
Σ
Banom
(we He,t (urA)² + +0
8,1=1
8
(w,, wk)
Tr (T((H,,) O
2m
8,t=1
5²
2m
F.2 COMPARISON OUR RESULTS
We first summarize the commonalities and differences between their results and ours:
Both of us consider the robustness of the model. But they define the feature robustness while study the loss robustness Xu & Mannor (2012) which has been studied for many years.
They consider non-standard generalization gap by decomposing it into representativeness and the expected deviation of the loss around the sample points. But we strive to integrate sharpness into the general generalization guarantees.
For point 1, their defined feature robustness trivially depends on the sharpness. Because the sharpness (the curvature information) is just defined by the robust perturbation areas around the desired point. From step 2 in the above proof sketch we can see, the hessian w.r.t. w is exactly the second expansion of perturbed expected risk. So we think this definition provides less information about the optimization landscape. In contrast, we consider the loss robustness for two reasons: 1) it is easy to get in practice without finding the orthogonal matrices O.. first. 2) we highlight dependence on the data manifold.
For point 2, we try to integrate this optimization property (sharpness) into the standard generalization frameworks in order to get clearer interplay. Unlike feature robustness, the robustness defined by loss function will be easier analyzed in generalization tools, because it's hard and vague to define the "feature" in general. Besides, our result will also benefit the data-dependent bounds Xu & Mannor (2012); Kawaguchi et al. (2022).
G EXPERIMENTS AND DETAILS
G.1 1 RIDGE REGRESSION WITH DISTRIBUTIONAL SHIFTING
As we stated before, we followed the Duchi & Namkoong (2021) to investigate the ridge regression on distributional shift. We randomly generate 0g E IRd in spherical space, and data from
X (0, y X
(122)
To simulate distributional shift, we randomly generate perpendicular unit vector 0 to 0ü- Let 0. 6ū be the basis vectors, then shifted ground-truth will be computed from the basis by 0r = 0% cos(a) + + 0.00f-sin sin((a) For the source domain, use 0j our training distribution. We randomly
35

---

Published as a conference paper at ICLR 2024
sample 50 data points and train linear classifier with gradient descent of 3000 iterations. Starting from a = 0, we gradually increase the a to generate different distributional shifts.
that larger penalty suffers from lower loss increasing when 3 ranges from 0 to 2. Since we consider a cycling shift of label space, 180 corresponds to the maximum shift thus leading the highest loss increase. According larger 3 means a flatter minimum and more robustness, resulting in a better OOD generalization. This experiment verifies our theoretical analysis. However, it is important to note that large coefficient of &2 regularization will hurt the performance. As shown in Figure (right panel), the curse of underfitting (indicated by brown colors) appears when
From the left panel in Figure 7
analysis of ridge regression,
Maximum Shift
Maximum Shift
8-0.0
#-0.0
p-10
6-0.
8
4+2.
P-10
0:3.0
Poor
3-4.0
fitting
Loss
50
100
250 Angles (degree)
150
200
300
350
100
150 Angles (degree)
20c
250
500
350
Figure 7: OOD test losses are increasing along distributional shifting. The X-axis is the shifting angle a and the Y-axis is the test loss of the model which is trained on distribution a Left: The larger regularization 3 causes lower increase in test loss (Darker purple lines have lower test losses). Right: Too large & penalty coefficient 3 brings poor fitting and thus fails to generalize on both ID and OOD datasets (orange lines). Best viewed in colors.
G.2 ADDITIONAL RESULTS SHARPNESS
Here we show more experimental results on RotatedMNIST which is a rotation of MNIST handwritten digit dataset with different angles ranging from 0", 15 300, .45 60% 75 (6 domaais/distributtions)) For each environment, we select a domain as a test domain and train the model on other domains. Then OOD generalization test accuracy will be reported on the domain. In our experiments, we run each algorithm with 12 different seeds, thus getting 12 trained models of different minima. Then we compute the sharpness (see Algorithm 1) for all these models and then plot them Figure 8. For algorithms, we choose Empirical Risk Minimization (ERM), and Stochastic Weight Averaging (SWA) as the plain OOD generalization algorithm which shown the first column. robust optimization algorithms, we choose Group Distributional Robust Optimization Sagawa et al. (2019). We also choose CORALSun & Saenko (2016) as a multi-source domain generalization algorithm. Among these different types of out-of-domain generalization algorithms, affect the test accuracy on the OOD dataset.
conclude that sharpness will
Experimental configurations are listed in Table 1. choose the last model as the test model. To ease the computational burden, we choose the 3-layer MLP to compute the sharpness.
For each model, the 5000 iterations and
Algorithms ERM(SWA)
Optimizer
1r
WD
batch size
MLP size
cta
MMD
Adam
0.001
64
265*3
DRO
Adam
0.001
0
64
265*3
0.01
CORAL
Adam
0.001
0
64
265*3
Table 1: Hyperparameters we use for different DG algorithms the experiments.
From Figure 8 we can see, the sharpness has an inverse relationship with the out-of-domain general- ization performance for every model in each individual environment. To make clear, plot similar
36

---

Published as a conference paper at ICLR 2024
tasks from environment to 4 as the last row. Thus, verified our Theorem 3.1. Note that all algorithms have different feature scales. One may need to normalize the results of different algorithms when plotting them together.
clearer tendency in all algorithms.
Algorithm 1 Pseudocode of model sharpness computation
Require: feature layer (x)), training loss
Ensure: Sharpness S
Get Jacobian matrix w.r.t. feature layer J v(([(z),)) for each gradient vector J in do aJ
Compute Hessian w.r.t. element j of f(x) by 88( end for
We store Hessian in the variable H Initialize sharpness for i in feature layer.shape(0] do for j in feature layer.shaee(O] do Retrieve the hessian value of j element via H[; Sharpness Sij + Trace(h) * (x)
= S + --j end for end for
G.3 COMPARE OUR ROBUST BOUND TO OTHER OOD GENERALIZATION BOUNDS
G.3.1 COMPUTATION OUR BOUNDS
First, we follow Kawaguchi et al. (2022) to compute the K in an inverse image the e- covering in randomly projected space. The main idea is to partition input space in projected space with transformation matrix Å. The specific steps will be (1) To generate random matrix A, we i.i.d. sample each entry from the Uniform Distribution 4( 1). (2) Each row of the random matrix A E R 3xd is then normalized so that Ar E [0, 1]3, i.e. A,j = A.j/ Σ Aig (3) After generating random matrix A, we use the c-covering of the space = Ax to define the pre-partition {Ć: {K
G.3.2 COMPUTATION OF PAC-BAYES BOUND
We follow the definition to compute expected dis, in Germain et al. (2013) where
Definition G.I. Let H be hypothesis class. For any marginal distributions Ds and Dr over X, any distribution p on H. the domain disagreement disp (Ds, Dr) between Ds and Dr is defined by,
disp (Ds, Dr)
def
ht.h RDr (h, R.S (h, h')]
Since the dis, (Ds, Dr) is defined as the expected distance, we can compute its empirical version according to their theoretical upper bound follows.
Proposition G.2 (Germain al. Germain et al. (2013) Theorem 3). For any distributions Ds and Dr over X, any set of hypothesis H, any prior distribution over H, any 8 (0, 1), and any real number a > O, with probability at least d over the choice of S x T (Ds Dr)" forr every on H, we have
2a dis T) + 2KL(p][=) In
+ -
disp (Ds, Dr) <
20
With (Ss,(S,T), we can then compute the final generalization bound by the following inequality
Vp on H. Rp, (Gp) Rp, (G,²) S Rps (Gp) + dis, (Ds, Dr) + RD, (Gps Gpz) + RDS (Gp: Gpz)
37

---

Published as a conference paper at ICLR 2024
ERM
DRO
CORAL
0.3
045
042
042
03.0
Out-of- -ooDain Generalization Accuracy
302
6015
104
15.
275
Dre3
Und
Sharpness
Figure 8: Domain generalization test accuracy on RotatedMNIST.. From top to the bottom: environ- ment 2, 3, 4, 5 with angles: [0 1,",,0 30" 455 ,,00 75%] and plot together. Each column shows the 12 runs of each algorithm.
38

---

Published as a conference paper at ICLR 2024
with PT argmin Rp, (Gp) is the best target posterior, and RD (Gps Gpz) Eho RD h').
Note that ignore the expected errors the best hypothesis by assuming the RD (Gps Gpz) O.
We apply the same operation in 8 of Proposition C.6 as well.
G.3.3 COMPARISONS
In this section, we add some additional experiments on comparing to other baselines, i.e. PAC-Bayes bounds Germain al. (2013). As shown in the first row of Figure 9, robust framework has smaller distribution distance in the bound compared to the two baselines when increasing the model size. In the second row, we have similar results in final generalization bounds. From the third and fourth rows we can see, our bound is tighter than baselines when suffering distributional shifts.
39

---

Published as a conference paper at ICLR 2024
Dassttern Dintence
18
Model Size
Model Sip (b)
Model Sire
(a)
Bayas
Emandine de Dar
Model Size
Model Size
Model Size
(d)
(e)
(f)
aneaae top
BBC Boysa
Test
Correlation Probsbility g
Correlation Probebility
Correlation Probebility
(h)
wxx oop
Dur
Error
Beaeline TBC Boyen
Correlation Probability
Correlation Probability
15
Correlation Probability
16
(k)
(I)
Figure 9: Additional spurious feature synthetic experiment. Each dot represents a trained model. The dash curves the smoothed function fit by the test data points. The baseline is Proposition 2. (g) the generalization error of the logistic regression models with increasing the model size/correlation probability. (b), (e): distribution distances for PAC-Bayes Germain al. (2013) along model size increases. (h), (k): distribution distances for PAC-Bayes Germain et al. (2013) along
model distribution shift. (c), (f), (i), (I): comparisons among baselines Proposition C.6, Germain et al. (2013) and ours. Note that model size > 500 is the overparameterized regime. The further the correlation probability is from 0.5, the greater the distributional shift is.
40


Reviews:
[{"review_id": "zYkEL0UV6q", "reviewer": "Reviewer_2MLh", "summary": "This paper studies the relationship between robustness and sharpness in improving OOD generalization performance. The authors provided rigorous theoretical analyses through the PAC learning framework to give an intuitive insight into why minimizing sharpness can lead to better OOD performance. This paper is the first work to reveal the benefits of flatness to OOD generalization. Moreover, based on existing theoretical results, a titer error bound is proposed. Empirically, the authors choose two simple datasets RotateMNIST and a self-generated dataset to validate the findings of their theory. This is a pure theoretical paper which is a solid contribution to related fields.", "questions": "Please see the weaknesses for detials.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The theoretical derivation is rigorously demonstrated. It is solid work by providing support for encouraging flat minima on OOD generalization problems.\n- All the results are very intuitive and confirm many assumptions of my own.\n- This paper is well-written and easy to follow.", "weaknesses": "- Why conducting partition on all training data is still not clear to me. It is reasonable that a small distribution shift in the same sub-group could lead to noise and should be ignored. But how such a strategy is formulated is unknown. Besides, could it be possible to conduct an ablation study to justify that such kind of partition is indeed helpful for OOD generalization?\n- The results in Figure 4 are interesting. However, some intuitive explanations are missing. For example, why did the test error of the worst-case group first decrease along the correlation probability and then increase?\n- It would be much more helpful if the authors could provide the experimental results on larger datasets such as DomainBed and WILDS. I understand theoretical papers could still have gaps in realistic application, but it would be helpful for us to find problems and look for possible directions.", "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": "", "extra_scores": {}}, {"review_id": "dK30DSxQid", "reviewer": "Reviewer_m8jH", "summary": "The authors introduce an out-of-distribution (OOD) generalization bound based on Sharpness that incorporates robustness, leading to a more precise bound compared to non-robust guarantees. The given results are validated through experiments conducted on both a ridge regression model and some deep learning classification tasks.", "questions": "Furthermore, transfer learning has become the standard choice for vision and language tasks. However, there is limited understanding of sharpness in this context. Do the theoretical results still apply to transfer learning scenarios?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "-\tThe given bound seem to be tighter than existing non-robust guarantees.\n-\tThe experiments are well conducted on ridge regression and simple classification tasks using MLP.", "weaknesses": "-\tThe datasets being used are overly simplistic, making it difficult to ensure the validity of theoretical outcomes. Empirical evidence from [1], [2], [3], [4], [5] demonstrate that sharpness is limited in the context of small datasets like CIFA10 or SVHN. Comprehensive evaluations involving transformers, ConvNet, CLIP, and BERT  trained on larger datasets such as ImageNet and MNLI reveal a weak correlation between sharpness and generalization. Surprisingly, a consistent negative correlation between sharpness and out-of-distribution (OOD) generalization suggests that sharper minima might actually generalize better, contradicting the authors' proposed theory. Moreover, experiments in [6] indicate that the choice of sharpness definition heavily depends on the specific dataset. To substantiate their theoretical claims, it is essential for the authors to conduct more extensive experiments on larger and more intricate datasets using complex models.\n\n-\tNumerous studies propose that flatter minima might offer better generalization for both standard and out-of-distribution (OOD) data such as [7] and [8]. However, traditional sharpness definitions often fall short in capturing generalization accurately, potentially because they are not consistent under reparametrizations that do not change the model itself. To tackle these issues, Adaptive Sharpness and reparametrization-invariant sharpness have been developed. I anticipate seeing the theoretical results related to Adaptive Sharpness and reparametrization-invariant sharpness.\n\n[1] Kwon, J., Kim, J., Park, H., and Choi, I. K. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. ICML, 2021.\n[2] Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2016.\n[3] Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. ICLR, 2020.\n[4] Dziugaite, G. K. and Roy, D. Entropy-sgd optimizes the prior of a pac-bayes bound: Generalization properties of entropy-sgd and data-dependent priors. In ICML, pp. 1377–1386. PMLR, 2018.\n[5] Bisla, D., Wang, J., and Choromanska, A. Low-pass filtering sgd for recovering flat optima in the deep learning optimization landscape. AISTATS, 2022.\n[6] Andriushchenko, Maksym, et al. \"A modern look at the relationship between sharpness and generalization.\" arXiv preprint arXiv:2302.07011 (2023).\n[7] Xing, C., Arpit, D., Tsirigotis, C., and Bengio, Y. A walk with sgd. arXiv preprint arXiv:1802.08770, 2018.\n[8] Zhou, P., Feng, J., Ma, C., Xiong, C., Hoi, S. C. H., et al. Towards theoretically understanding why SGD generalizes better than Adam in deep learning. NeurIPS, 2020.", "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": "", "extra_scores": {}}, {"review_id": "HiGNODjRRi", "reviewer": "Reviewer_q8BG", "summary": "This paper provides generalization bounds for Out of Distribution Generalization based on the sharpness of the loss function minimum attained. One of the main ideas is to partition the input space according to the suggestion in [1]: it is assumed that the loss function doesn't vary too much inside each \"relevant\" subset of a partition of space (this is loosely related to the flatness of the minima) where \"relevant\" means close to an observed sample. Based on this assumption, the authors are able to bound (in Theorem 3.1)  the error as a sum of: \n\n1 the empirical error in the source domain\n2 the relaxed TV distance between the distributions in the partition considered\n3 The robustness of the algorithm/loss (how much the loss can vary for points in a given partition element close to a sample\n4 some function class capacity term. \n\nThe proof is a straightforward manipulation of the definitions.\n\nThe, corollary 3.2 attempts to show that the limiting behavior of Theorem 3.1 is the same as that of the main result of [1] which relies on the traditional TV distance. \n\nIn the next results, the authors rely on some ideas from [2] as well as brute force calculations to relate the concept of robustness of [1] and this work to that of the flatness of the local minima in the case of a kernel regression problem equivalent to a two layer neural network where only the last layer is trained.  The flatness and the norm of the parameter vectors are then claimed to be related in the case of ridge regression, and synthetic data experiments demonstrate superior OOD performance is associated with stronger genralization. \n\n\n\n\n\n\n\n[1] Robustness and generalization. Huan Xu and Shie Mannor. Machine Learning. 2012.\n\n[2] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, Mario Boley, \"Relative Flatness and Generalization\", NeurIPS 2021.\n\n\n\n\n\n\n\n\n\n\n\n\n\n=====================Post rebuttal=================\n\n\n\n\nPost-Rebuttal\n\nI thank the authors for the rebuttal and for substantially improving the quality of the paper by fixing many of the errors I mentioned in the proofs. \nI am not an expert in this particular field, so it is difficult for me to vouch for the correctness of all of the work, especially given that the first version had a very large number of errors that made reading impossible and necessitated a very deep revision. To the best of my understanding, the revision makes sense at a high level, but there are still minor errors, including both in parts of the paper I had not originally commented on, and parts of the paper that were rewritten for this revision (details below). **However, since those remaining errors appear minor and reviewers 2MLh and b2N6 are very positive about the paper and the authors have made an effort already, I am increasing my score from 5 to six.  In addition, since I have not found any fatal or serious errors in the revision, I am increasing my soundness rating to fair (from poor).**\n\nHowever, I **very strongly recommend** the authors take at least two weeks of **full time** work to re-read all of the paper and **fix all the errors** and improve the proofs. To the best of my understanding, this is still necessary.  I don’t mean just fixing the errors that I caught, I mean fixing the ones that none of the reviewers caught.", "questions": "Could you write down a proof of  Proposition C.5 clarifying if you are using a standard VC bound or Theorem 3.2 from arxiv [1], and fixing the constants from the union bound? \n\n\nDo you stand by your proof of Lemma D.5? It seems sloppy to the point of being completely incomprehensible, though I believe a similar result probably holds. Is equation (46) an equation or an approximation ? It seems that the probability (LHS) should be equal to $2\\phi/\\pi$ and not involvea $cos$ at all. For equation (47), it seems that you are calculating the bound the wrong way around: equation (47) holds if the inequality is reversed inside the probability statement. In the next substantial revision you upload, could you also include a precise restatement and citation for \"Robbin's bounds\".\n\n\nWhat is going on in Equation (64)? The statement of Theorem 3.4 appears to be agnostic w.r.t. the choice of loss function, but the proof says \"ohh we can verify that this holds for common loss functions such as L1, MSE, MAE etc.'. How do you get the first equality in (64), and does it hold for all loss functions? Please rewrite the proof accordingly. \n\nCould you fix the \"definitions\" in corollary 3.5? When taken at face value, I maintain that $L_{x^*}=0$: there is something missing there. Also, although I guess that $\\hat{w}_min$ is the minimum absolute value of any entry of $\\hat{w}$ this is not clear from the definition since the notation $|\\cdot|$ to mean element wise absolute value is non standard. \n\n\n \nPlease rewrite all the proofs much more carefully and upload a revision. It may help to ask another author to read the proofs.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "Assuming they are correct, the results are certainly very interesting and promise to be the basis for further research. There is quite a lot of non trivial work in this paper, and many proofs require complicated calculations and take inspiration from a variety of existing work.", "weaknesses": "The appendix and many parts of the main paper very severely lack polishing to the extent that communication is difficult. Many of the proofs look more like a first draft than a serious submission. In particular, I can't vouch for the correctness of any results past Theorem 3.1 (which is not 100 percent correct either. In addition to poor language and presentation, the organization of the paper in terms of which assumptions are made is not good. The paper just needs more calm work and rewording: the authors of the appendix appear in a rush. \n\n\nHere are more details: \n\nFirst, although this is a minor point, it is worth mentioning that the constants in Theorem 3.1 don't seem to be fully correct. Indeed, the proof relies on a combination of Lemma 1 and Lemma C.3 from the arXiv version of [1], with both theorems relying on excluding a low probability event. Since combining both results requires a union bound, all the $\\delta$s should be changed to $\\delta/2$, resulting in different constants throughout. The same issued shows up in the proof of Proposition C.5. Note also that the result cited from [1] in Lemma C.3. is not in the published version, but in the arXiv version only. There are no Theorems 2 and 3.2 in the NeurIPS version, THis is Theorem 2 from NeurIPS and Theorem 3.2 from arXiv (\nHowever, the same result can be lifter from [3] instead.) In addition, the combination of both results should be done explicitly since the constants are wrong again due to not applying the union bound. Also, I am really not sure the authors are actually using Theorem 3.2 rather than a simple generalization bound in terms of VC dimension. \n\n\nTheorem 3.4 makes sense as a statement (though I couldn't check the full proof), but is very hard to read and contains plenty of minor errors. Why is $z_i(A)=argmax...$ unique? It seems the authors mean to say that $z_i\\in argmax....$. When defining $\\mathcal{M}_i$ as \"the set of gobal minima in $C_i$\", the authors actually mean that $\\mathcal{M}_i$ as a set is the argmin of the objective function over $C_i$ (in particular, the elements are \"global\" minima inside the partition element $C_i$, but not necessarily global minima for the whole problem). This is absolutely not clearly written. It is also not clear why a different definition style is used for $|mathcal{M}_i$ and for $z-i(A)$ in that case Every statement in this paper is written vaguely using words instead of maths, adding to the burden of trying to make sense of what is attempted here. \n\n\n\n\n\n\n\n==================Minor comments============\n\n\nWhat is Theorem C.2 a restatement of? It seems it is Theorem 3.1 but it is not written (only the title \"restate\"). \n\nThe wording of transitions  is quite sloppy. For instance, in page 14 we have the sentence \"In order to give a fine-grained analysis, we follow the common choice where a two-layer NN function class is widely analysed in most literature, i.e. Neural Tangent Kernel, rich regime, random feature models.\" What do the authors mean by \"rich regime\"? \n\nNote also at the top of page 14, the integral is not properly written as I believe the authors actually meant to write $\\int f(x)g(y) \\tau_d(d(x,y))$ instead of $\\int f(x)g(y) \\tau_d(dx)$  (otherwise the quantity is just equal to 1). \n \n\nThe statement of Corollary 3.2 is not really mathematically well defined. What the authors mean is that the bound from Theorem 3.1 converges to the RHS of equation (3), but what is written is just that the OOD generalization error is bounded by it. The \"proof\" is not very clean either. \n\nThe second sentence of Theorem 3.1 should contain a word such as \"assume that\". \n\nThe description of the case where $K=1$ above corollary 3.2 seems quite dodgy as it implies that the loss is approximately constant in the whole of input space! \n\n\nNote that even the definition of one of the main concepts of this paper, i.e. robustness as in Definition 2.2, is not very well introduced: there \"for all\" which relates to $s,z$ is duplicated, and there is a \"for all i$ after the $forall $s,z\\in C_i$ (which depends on $i$). This makes this much harder to read than necessary. Note that the definition in [1] is much cleaner. It should be copied more faithfully. \n\n\n\n====================Typos====================\n\nThere are plenty of issues with the language, not only in terms of improper use of grammar to an extent that harms understanding, but also in terms of polishing, spelling and punctuation. Below is a highly non exhaustive list. \n\n\nIn the abstract, fourth line \"sharpness of learned minimum influences\" should be \"the sharpness of learned minima influences\" \n\nPage 2: to the left of the figure, \"as shown in \\cite... that when the loss lanscape...\": a coma would be better than \"that\"\n\nEvery theorem which is lifted from a reference involves the text of the name of the authors written twice (once manually and once via the \\cite function). This is probably due to copy and paste from another template with a different citation style without checking through it again. (prop. 2.1, definition 3.3\n\nPage 21, \"under this case\" should be \"in this case\" \n\nEquation (44), there is an equal sign with three bars instead of 2 where I think the authors actually mean equality in the strict sense. \n\nLine 3 of Section 3.1 \" $(K,\\epsilon)$-robust\" should be \"..robustness...\"\n\nIn page 5 (theorem 3.4), there is a \"w.p. a\" which probably should just be \"w.p.\"\n\n\nThere are plenty of \"w.r.t\" and \"i.i.d\" (missing period) throughout the text. \n\n\nIn appendix D, note the use of \"where definition of ....please refer to the corollary 3.5\" \n\nIn page 18 at the beginning of the proof of Lemma D.1, although I agree that equation (29) makes sense, the sentence which introduces it speaks of the smallest eigenvalue whilst the equation itself shows the equivalent statement expressed as an PSD inequality between matrices. \n\nJust below equation (27), it is unclear what the authors mean by saying that \"$\\mathcal{A}_{\\mathcal{F}\\Delta \\mathcal{F}}$ represents the algorithm\" WHat algorithm? \n\n\n====================References===========================\n\n[1] Robustness and generalization. Huan Xu and Shie Mannor. Machine Learning. 2012.\n\n[2] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, Mario Boley, \"Relative Flatness and Generalization\", NeurIPS 2021. \n\n[3] Learning Bounds for Domain Adaptation. John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman. NeurIPS 2007", "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": "", "extra_scores": {}}, {"review_id": "JStDFBYjZD", "reviewer": "Reviewer_b2N6", "summary": "This paper gives a tight upper bound of out-of-distribution (OOD) by taking robustness into account and discusses the connection between robustness and sharpness of the loss landscape. The author proves the robust generalization bound by partitioning the sample space and measuring the distributional shift of each sub-group. As to the connection between robustness and sharpness, the author proves a theorem stating that sharper models give a larger upper bound on robustness, which aligns with the intuition that sharper models are more prone to overfitting, which hurts OOD generalization.", "questions": "1. The author can give a clear explanation of ε(S) when describing definition of robustness (Def 2.2).\n2. Have you considered the impact of different hyperparameters on the performance of your method? Can give some deeper discussion on this.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Using the partitioning method to achieve a tighter bound is interesting and also leads to a good result.\n2. The paper is well-written and easy to understand, with clear explanations of key concepts and terminology.\n3. The new distance metric better tolerates distributional shift compared with previous work.", "weaknesses": "The assumptions and constraints are a bit hard to meet and also may hurt the generalizability of the results.(See my questions)", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)