Paper: Towards Robust Out-of-Distribution Generalization Bounds via Sharpness (ICLR-2024)

Reviewer comment (Writing):
The author can give a clear explanation of ε(S) when describing definition of robustness (Def 2.2).

Author response:
Thanks for the nice suggestion. We have added the description 
> This definition captures the robustness of the model in terms of the input. Within each partitioned set
$C_i$, the loss difference between any sample z belonging to $C_i$ and training sample $\boldsymbol{s} \in C_i$ will be
upper bounded by a robustness constant $\epsilon(S)$. Hence, the number of partition $K$ and robustness constant $\epsilon(S)$ define robustness for learning models.

below the Definition 2.2 in Section 2.2. Please refer to the highlighted revision in our new version.

Reviewer comment (Experiments):
Have you considered the impact of different hyperparameters on the performance of your method? Can give some deeper discussion on this.

Author response:
Thanks for your nice suggestion. We conducted an ablation study on the spurious feature synthetic experiment with different numbers of partitions $K$ (Figure 7 in Appendix A.2). The results show that the bounds with partitioning ($K > 1$) can better capture the OOD generalization error changes along the distribution shifts.

Reviewer comment (Theory):
The constants in Theorem 3.1 do not seem to be fully correct. The proof relies on a combination of Lemma 1 and Lemma C.3 from the arXiv version of [1], both relying on excluding a low probability event. Combining both results requires a union bound, so all the $\delta$s should be changed to $\delta/2$, resulting in different constants throughout. The same issue appears in the proof of Proposition C.5.

Author response:
Thank you for your careful review. We have corrected the constant in Theorem 3.1 as well as Proposition C.5. Following your suggestion, we have added the proof of Proposition C.5 and corrected the notions of VC-dimension by pseudo-dimension.

Reviewer comment (Theory):
The combination of results in Theorem 3.1 and its cited lemmas should be done explicitly, as the constants are wrong again due to not applying the union bound. It is also unclear if the authors are actually using Theorem 3.2 rather than a simple generalization bound in terms of VC dimension.

Author response:
Thank you for your careful review. We have corrected the constant in Theorem 3.1 as well as Proposition C.5. Following your suggestion, we have added the proof of Proposition C.5 and corrected the notions of VC-dimension by pseudo-dimension.

Reviewer comment (Theory):
Theorem 3.4, while making sense as a statement, is very hard to read and contains numerous minor errors. For example, the uniqueness of $z_i(A) = \argmax...$ is questionable; it seems the intended meaning was $z_i \in \argmax...$. The definition of $\mathcal{M}_i$ as "the set of global minima in $C_i$" is ambiguous, as it actually represents the argmin of the objective function over $C_i$. This is not clearly written.

Author response:
We carefully reviewed all the proofs. Our revisions for Theorem 3.4 and Corollary 3.5 can be summarized below:
- We have re-proven the Lemma D.1 by replacing the assumption $L_{x^*}>0$ by $m=Poly(d)$.
- We have corrected the numbers in Lemma D.5.
- We refine the definition of $\mathcal{M}_i$ and $z_i(A)$.
- Revised the setting and assumptions in Appendix B.

Reviewer comment (Writing):
There is inconsistency in the definition style used for $\mathcal{M}_i$ and for $z_i(A)$. The overall statements are written vaguely using words instead of mathematics, increasing the difficulty of understanding the intent.

Author response:
We carefully reviewed all the proofs. Our revisions for Theorem 3.4 and Corollary 3.5 can be summarized below:
- We have re-proven the Lemma D.1 by replacing the assumption $L_{x^*}>0$ by $m=Poly(d)$.
- We have corrected the numbers in Lemma D.5.
- We refine the definition of $\mathcal{M}_i$ and $z_i(A)$.
- Revised the setting and assumptions in Appendix B.

Reviewer comment (Writing):
Some transitions in the text are sloppily worded. For instance, on page 14, 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" contains the unclear term "rich regime."

Author response:
For point 2, the term "rich regime" refers to a non-kernel regime in which training overparameterized networks through gradient descent leads to implicit biases that diverge from RKHS norms. We have supplemented this term with a previously omitted reference. For the remaining points, we have conducted the necessary revisions. Please consult our recently updated version for details.

Reviewer comment (Writing):
There are many grammar, spelling, and punctuation issues that harm understanding. Examples include: in the abstract, "sharpness of learned minimum influences" should be "the sharpness of learned minima influences"; on page 2, "as shown in \cite... that when the loss lanscape..." should have a comma instead of "that."

Author response:
Thank you for your careful review of our manuscript and for taking the time to point out the typographical errors in the original submission. The revised version of the manuscript has been uploaded with corrections highlighted in red.

Reviewer comment (Experiments):
The 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.

Author response:
**Q1: Empirical evidence from existing work demonstrates that sharpness is limited in the context of small datasets like CIFA10 or SVHN**

[1] has demonstrated that Stochastic Weight Averaging towards flat minima can improve OOD generalization on five standard DG benchmarks, namely PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet. In addition, we have observed a **consistent phenomenon on larger-scale datasets** such as PACS, and WILDS-Camelyon17. Please refer to Figure 5,6 in Appendix A in our revision.

-------

**Q2: [3] shows a negative correlation between sharpness and out-of-distribution (OOD) generalization suggesting that sharper minima might actually generalize better, contradicting the authors' proposed theory.**

The distinction in the definition of sharpness can lead to disparate observations and interpretations. It is crucial to note that our measure (Definition 3.3) is different from the one presented in [3] where its sharpness is defined by 
$$S_{avg}^\rho(w, c):=  E_{S \sim P_m, \delta \sim(0, \rho^2 \text{diag}(c^2)} L_{S}(w+\delta)-L_{S}(w) ~~~~(1)$$
which captures the loss changes within a neighbor area. In contrast, our sharpness (following [4]) integrates the norm of weights and the trace of the Hessian matrix for each layer:
$$\kappa(w, S, A):= \langle w, w \rangle \text{tr}(H_{S, A}(w)). ~~~(2)$$
Our theoretical results derived the dependence of sharpness and robustness upon Eqn(2). Both experiments in [1] and our paper have validated the positive correlation between flatness and OOD generalization.

-------

**Q3: It is essential for the authors to conduct more extensive experiments on larger and more intricate datasets using complex models.**

Thanks for the nice suggestion. To better evaluate the correlation between our sharpness and Out-Of-Distribution (OOD) generalization, we conducted additional experiments on PACS with ResNet50, and on Wilds-Camelyon17 with DenseNet121 according to your suggestion (Figure 5 Appendix A in our revision).

Reviewer comment (Theory):
Numerous 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.

Author response:
**Q4: I anticipate seeing the theoretical results related to Adaptive Sharpness and reparametrization-invariant sharpness.**

We have adopted the "relative sharpness" defined in [4] to account for **reparameterization-invariant sharpness** that alters certain flatness measures while leaving generalization unaffected. Moreover, [4] has established the rigorous link between sharpness and generalization. Works on adaptive sharpness such as ASAM [2] and [3] examine reparametrization-invariant sharpness from a different perspective. The differences we have shown in Q2.

Reviewer comment (Theory):
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?

Author response:
**Q5: Do the theoretical results still apply to transfer learning scenarios?**

Yes, we believe our results are applicable to transfer learning contexts, specifically domain adaptation. The framework of our bound is commonly shared in domain adaptation [5].

Reviewer comment (Experiments):
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?

Author response:
**W1-(a): Why conducting partition on all training data is still not clear to me.**

According to the original definition in [1], the (whole) sample space $\mathcal{Z}$ is partitioned into $K$ disjoint sets, not just on training data such that all unknown samples from OOD may fall into the same set as the training data. An example of formulation can be found in Example 1 in [1] where 

> Example 1 Fix $\gamma>0$ and put $K=2 N (\gamma / 2, X, ||\cdot||_2 )$. If $\mathcal{A}s$ has a margin $\gamma$, then $\mathcal{Z}$ can be partitioned into $K$ disjoint sets,  such that if $s_j$ and $z \in \mathcal{Z}$ belong to a same $C_i$, then $|l(\mathcal{A}s, s_j)-l(\mathcal{A}s, z)|=0$.

This example suggests that according to the covering number, sample space $X$ can be partitioned into $N (\gamma / 2, X, ||\cdot||_2 )$ subsets such that each subset has a diameter less or equal to $\gamma$.


**W1-(b): Could it be possible to conduct an ablation study to justify that such kind of partition is indeed helpful for OOD generalization?**

Yes, we conducted an ablation study on the spurious feature synthetic experiment with different numbers of partitions $K$ (Figure 7 in Appendix A.2). The results show that the bounds with partitioning ($K > 1$) can better capture the OOD generalization error changes along the distribution shifts.

Reviewer comment (Evaluation):
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?

Author response:
**W2: The results in Figure 4 are interesting. However, some intuitive explanations are missing.**

Thanks for the nice suggestions, we added the following intuitive explanation on Page 9. Please refer to our highlighted revision.

> As shown in Figure 4(d), when $p_{\text{maj}} = 0.5$, there is no distributional shift between training and test data due to no spurious features correlated to training data. Thus, the training and test distributions align closer and closer when $p_{\text{maj}}<0.5$ and increase, resulting in an initial decrease in the test error for the worst-case group.
However, as $p_{\text{maj}} > 0.5$ and deviates from 0.5, introducing the spurious features, a shift in the distribution occurs. This deviation is likely to impact the worst-case group differently, leading to an increase in the test error.

Reviewer comment (Experiments):
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.

Author response:
**W3: Experimental results on larger datasets such as DomainBed and WILDS.**

We have added the new experimental results on PACS and WILDS-Camelyon17. To extend the sharpness measure to complicated architectures, such as DenseNet-121, we compute the Hessian matrix from the last FC layer. Even though, the result is still consistent with what we report in Figure 3. The phenomenon that **sharper minima can lead to poor OOD generalization performance** can also be observed on large-scale datasets.