Towards Robust Out-of-Distribution Generalization Bounds via Sharpness

Thank you for the feedback and recognition of our contributions. Below are our responses.

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

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.

Q2: Have you considered the impact of different hyperparameters on the performance of your method? Can give some deeper discussion on this.

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.

Thank you for your valuable feedback and constructive suggestions. We have carefully reviewed and revised the proofs, as well as corrected any typos in the paper. Please find our responses to your comments below:

Main weaknesses

W1: The constants in Theorem 3.1 don't seem to be fully correct.

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.

W2: Theorem 3.4 makes sense as a statement, but is very hard to read and contains plenty of minor errors.

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.

Minor comments

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.

[Typos]

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.

Questions

Q1: 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?

We have added proof of Proposition C.5 and corrected the error constant.

Q2: Lemma D.5

We have thoroughly revised the proof.

Q3: What is going on in Equation (64)?

We derived the first inequality by the chain rule. We made the choice of loss functions explicitly in Appendix B.1, please refer to the changes.

Q4: Could you fix the "definitions" in corollary 3.5?

We have refined our statement and introduced a supplementary assumption to Corollary 3.5, stipulating that the model size ($m$) is considerably greater than the data dimension ($d$). This new assumption is widely used in overparameterized regimes within which we have re-proven Lemma D.1. Revised Lemma D.1 has allowed us to present more precise and refined results for Corollary 3.5 in the revised version of our paper.

Overall, we do appreciate your recognition of our novelty as well as for providing thorough and thoughtful reviews that improve the rigor of our manuscript. In response to your insightful comments and constructive feedback, we have revised the proofs thoroughly. We are grateful for the opportunity to improve our manuscript with your guidance.

Thank you for the feedback and suggestions. Below are our responses.

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).

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.

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].

References

[1] Cha, Junbum, et al. "Swad: Domain generalization by seeking flat minima." Advances in Neural Information Processing Systems 34 (2021): 22405-22418

[2] Kwon, Jungmin, et al. "Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks." International Conference on Machine Learning. PMLR, 2021.

[3] Andriushchenko, Maksym, et al. "A modern look at the relationship between sharpness and generalization."

[4] Petzka, Henning, et al. "Relative flatness and generalization." Advances in neural information processing systems 34 (2021): 18420-18432.

[5] Redko, I., Morvant, E., Habrard, A., Sebban, M., & Bennani, Y. (2020). A survey on domain adaptation theory: learning bounds and theoretical guarantees. arXiv preprint arXiv:2004.11829.

Thank you for the feedback and suggestions. Below are our responses.

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.

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.

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.

References

[1] Xu, H., & Mannor, S. (2012). Robustness and generalization. Machine learning, 86, 391-423.

The authors have carefully answered my questions. Additional experiments are also added to show its empirical effectiveness. I think this paper shows some insight into OOD generalization. The theoretical analyses are extensive. Therefore, I would like to raise my score to 8, but my confidence remains at 3.