Understanding Adversarially Robust Generalization via Weight-Curvature Index

We thank the reviewer again for appreciating our work's contribution. Below, we clarify your questions. Please do not hesitate to follow up if there are any remaining questions.

The method you use to generate adversarial examples in Section B.2 raises some concerns. Why did you choose not to resort to state-of-the-art techniques, such as AutoAttack?

The concern regarding our choice of adversarial example generation methods in Section B.2 is appreciated. The primary reason for not initially employing state-of-the-art techniques like AutoAttack was the significant computational cost associated with these methods. Specifically, AutoAttack requires approximately 40 minutes per epoch on CIFAR-10 for both training and testing, resulting in an 80-fold increase in total testing time compared to our selected method. Given the scale and scope of our experiments, this was computationally prohibitive for the initial analysis.

Nevertheless, we've conducted an additional set of experiments using AutoAttack. These experiments confirmed that our proposed method effectively controls robust overfitting even under this rigorous evaluation. While accuracy decreased by approximately $5\%$ on CIFAR-10, the overall robustness trends and generalization indicators remained consistent, demonstrating the reliability of our approach. We believe these results highlight the scalability of our method while ensuring robustness across a variety of adversarial attack strategies. Moving forward, we recognize the value of employing state-of-the-art attacks like AutoAttack for benchmarking and aim to incorporate such evaluations more extensively as computational resources permit.

Please also comment on the validity of the assumptions.

Our thereotical results rely on the following assumptions:

1. For any $\theta'$ sampled from the posterior distribution $\mathcal{Q}$, the robust loss function can be approximated around $\theta$ using a second-order Taylor expansion:

$$\mathcal{L}\_\mathcal{S}(\theta', x + \delta(\theta'), y) \approx \mathcal{L}\_\mathcal{S}(\theta, x + \delta(\theta), y) + \tfrac{1}{2} \Delta \theta^\top \nabla^2\_{\theta} \mathcal{L}\_\mathcal{S}(\theta, x + \delta(\theta), y) \Delta \theta,$$

where $\Delta \theta = \theta' - \theta$ and $\delta(\theta)$ denotes the worst-case perturbation that maximizes the robust loss (we explicitly express the dependence of $\delta$ on $\theta$ to avoid confusions).

2. $\theta$ is considered to be at or near a stationary point of $\mathcal{L}\_\mathcal{S}(\theta, x + \delta(\theta), y)$, since models (considered for adversarilly robust generalization) are typically learned using first-order gradient optimization algorithms, such as PGD-based adversarial training. Consequently, we can discard the first-order term in the above approximated second-order Taylor expansion of the loss. This setting is widely used in theoretical analyses of generalization and robustness (see [1-2], for example).

[1] Stephan, Mandt, Matthew D. Hoffman, and David M. Blei. "Stochastic gradient descent as approximate bayesian inference." Journal of Machine Learning Research 18.134 (2017): 1-35.

[2] Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

We thank the reviewer again for the detailed review and constructive feedback. If you agree that we've managed to address all issues, please consider raising the score. But if there are any remaining concerns, please do not hesitate to follow up - we would be more than happy to address your further questions.

There are some references missing, such as [1-3], which would help the reader place the work into context.

Thank you for pointing out these references! We agree that incorporating these works will help place our contributions into a broader context and clarify their relationship to existing research. These references will be added in the revised manuscript, and their relevance will be explicitly discussed in the related work and conclusion section, to strengthen the contextualization of our work and emphasize its contributions to the field.

Below, we provide a brief summary of these references with connection to our work. The role of the Hessian trace in generalization has been extensively studied, such as in [1], where it was shown that trace minimization correlates with improved generalization across tasks. WCI extends this concept by integrating the trace and Frobenius norm of weights, creating a robust proxy for adversarial training contexts. Regularization techniques targeting the Fisher Information Matrix trace were discussed in [2], highlighting the importance of early-phase curvature control. These insights complement our findings that robust overfitting is mitigated when WCI regularization is incorporated during training. PAC-Bayesian bounds that incorporate Hessians, as explored in [3], provide a theoretical framework supporting our derivations of WCI. Specifically, the inclusion of curvature measures aligns with the PAC-Bayesian methodology, ensuring tight bounds on robust generalization errors.

[1] Jastrzebski, Stanislaw, et al. "Catastrophic fisher explosion: Early phase fisher matrix impacts generalization." International Conference on Machine Learning. PMLR, 2021.

[2] Golatkar, Aditya Sharad, Alessandro Achille, and Stefano Soatto. "Time matters in regularizing deep networks: Weight decay and data augmentation affect early learning dynamics, matter little near convergence." Advances in Neural Information Processing Systems 32 (2019).

[3] Ju, Haotian, Dongyue Li, and Hongyang R. Zhang. "Robust fine-tuning of deep neural networks with hessian-based generalization guarantees." International Conference on Machine Learning. PMLR, 2022.

I recommend at least testing this with muliple seeds, doing proper correlation analysis, and also comparing the difference in correlation against individual components using suitable statistical significance tests. At the very least, multiple seeds ought to be run, a confidence interval ought to be computed.

Following the reviewer's suggestions, we have run additional experiments and summarized the results in the following table, which consolidates the statistical significance of the correlations between WCI and both robust error gap and robust loss gap across multiple random seeds. These results reinforce the validity of WCI as a reliable indicator of robust generalization:

Key Observations:

1. Across all seeds and both gap types, WCI demonstrates strong and consistent correlations with robust error gap and robust loss gap, as evidenced by high correlation coefficients and statistically significant p-values (adjusted using Holm-Bonferroni and Benjamini-Yekutieli corrections).
2. The $R^2$-values and confidence intervals of the WCI coefficient indicate the predictive reliability of WCI across different seeds and metrics.

We commit to including this table in the revised manuscript, along with updated figures showing these trends. These additions will provide a more comprehensive visualization of our findings and further substantiate our claims regarding WCI’s effectiveness as a robust generalization indicator.

We thank the reviewer again for the detailed feedback. Your questions and comments touch on critical aspects of the theoretical framework. Below, we clarify the assumptions, derivations, and implications of our theoretical results presented in Section 2. We believe that we've clarified all your concerns, so please reconsider the score if this is the case. Otherwise, we would be more than happy to address any of your remaining questions.

The main concern is that the theoretical results are not convincing, or underdescribed at least. To properly consider $l(\theta,x+\delta,y)$, a supremum term on perturbations should be included. Therefore the proofs are needed to be checked whether they are correct under the supremum or considering the fact that $\delta$ depends on $l,(x,y),\epsilon,\theta$ also - e.g. PAC Bayes bound, second-order differentiability and else.

We've clarified the notation $\delta$ in our General Response 1 and explained why it does not affect the validity of our theoretical results and conclusions in our General Response 2.

When to define an upper bound in theorem 2.5, $\sigma_k^2$ should be controlled to minimize the bound. If not, the upper bound does not hold since it needs an equality condition for an arithmetic–geometric mean inequality. I think the Frobenius norm of layer weights $W_k$ and the trace of Hessian $H_k$ is not controllable during training, where their control is crucial to the control of $\sigma_k^2$. Hence, I conjecture that the fluctuation of WCI rises in learning curves since the upper bound does not holds all the time. This hinder the reason that WCI works.

To derive Theorem 2.5, we set $\sigma_k^2 = \sqrt{\|\mathbf{W}_k\|_F^2 / (2\lambda \mathrm{Tr}(\mathbf{H}_k))}$. As reviewer correctly recognized, such a choice enables us to achieve the equality condition for an arithemic-geometric mean inequality, thereby minimizing the generalization bound. This is desirable, since we want the upper bound to be tight when determining $\sigma_k$. However, we want to clarify a potential misunderstanding that layer weights and its Hessian matrix observed during training do not necessarily need to correspond to how we set $\sigma_k$ in deriving our PAC-Bayesian robust generalization bound and the WCI metric.

We further explain why this is the case. When we fix the prior variances, the KL term is proportional to the squared Frobenius norm of parameters. However, since we introduced the special prior, we can arbitrarily change the prior variance after training and control the KL term. Similar to [6-8], one has to make sure that $\mathcal{Q}$ and $\mathcal{P}$ have the same support, which is the case when they are from the same parametric family of distributions, such as Gaussian or Laplace. Recall that PAC-Bayesian Robust Generalization bound holds for any potential prior $\mathcal{P}$ and posterior $\mathcal{Q}$. Specifying particular priors and posteriors does not violate the bound but only affects the tightness. In other words, there is no restriction on the value of $\sigma_k^2$. We fix $\sigma_k^2$ just for a tighter bound.

In the original experiments where overfitting occurred, Frobenius norm of layer weights and the trace of Hessian matrices are indeed not controllable during training, which we think is the main reason for overfitting. In our extended algorithm, we control the learning rate by monitoring the value of WCI to achieve the effect of controlling these two items. Since our main purpose is to illustrate the correlation between WCI itself and the generation of gap, we did not focus on directly controlling these two items. The fluctuation of WCI needs to be understood, and perhaps Frobenius norm of layer weights $W_k$ and the trace of Hessian $H_k$ are worth further consideration.

In Figure 4, the authors explained that the Frobenius norm of the weight matrix is the most critical factor during the initial training stages. However, to me, it seems that the trace term dominantly influence to WCI since they share the whole shape. This implies that considering WCI is not basically different from considering the trace of the Hessian. Combined with above weaknesses, we may conclude that considering WCI is not basically different from simply using SAM with an adversarial risk.

The WCI combines both the Frobenius norm of the weights and the trace of the Hessian, reflecting both the magnitude of the weights and the curvature. The WCI provides a more holistic view by accounting for both aspects. While the trace may dominate at certain points, the weight norm plays a crucial role in other stages or scenarios.

Differentiation from SAM: Stochastic Weight Averaging (SWA) and Sharpness-Aware Minimization (SAM) focus on flattening the loss landscape. Flatness-based measures tend to exhibit poor correlations with the robust generalization gap (see [1]). The WCI provides a specific metric that combines weight magnitude and curvature, which may offer different insights, and we have confirmed that WCI is strongly correlated with gap generation. [2] suggested that sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. This calls for the search for other explanations for the generalization of over-parameterized neural networks.

[1] Kim, Hoki, et al. "Fantastic robustness measures: the secrets of robust generalization." Advances in Neural Information Processing Systems 36 (2024).

[2] Wen, Kaiyue, Zhiyuan Li, and Tengyu Ma. "Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization." Advances in Neural Information Processing Systems 36 (2024).

The authors illustrated that the prior works have shown that reducing the classifier variability term also decreases the perturbation discrepancy term. Does it also hold when to consider an adversarial risk, not a standard risk?

Yes, these findings still hold for adversarial risk. The relationship between reducing the classifier variability term and the perturbation discrepancy term remains valid even in the adversarial setting. This is because the core derivation relies on the same mathematical framework. If we apply the same Taylor expansion and approximation to the perturbation discrepancy term under adversarial risk, the coefficients in the resulting expressions may change due to the inclusion of adversarial perturbations. However, the structure of the results, including the conclusions about reducing classifier variability and its impact on perturbation discrepancy, remains unchanged. The adversarial perturbation simply modifies the coefficients, but it does not alter the underlying dependency between the terms. Thus, our conclusions about the robustness and generalization properties hold for both standard and adversarial risks.

In the proof of lemma 2.4, how can you consider $\mathcal{Q}=N(\theta,\mathbf{I})$? I’m quite confusing that is it possible to fix the mean as a certain point of $\theta$?

As we clarified in general respose 2, PAC-Bayesian Robust Generalization bound holds for any potential prior $\mathcal{P}$ and posterior $\mathcal{Q}$. Specifying particular priors and posteriors does not violate the bound but only affects the tightness. In other words, there is no restriction on the value of $\sigma_k^2$. We fix $\sigma_k^2$ just for a tighter bound. Fixing the mean of $\mathcal{Q}$ to $\theta$ is not a restriction. It is a deliberate choice to analyze the uncertainty around $\theta$. The PAC-Bayesian framework remains valid for any $\mathcal{Q}$, and this choice ensures the bound reflects the generalization of the learned parameters. Other forms of $\mathcal{Q}$ could be used, but $\mathcal{Q}=N(\theta,\mathbf{I})$ is particularly useful for theoretical analysis due to its simplicity.

We thank the reviewer for recognizing the usefulness of our derived PAC-Bayesian bounds on our theoretical results. We believe that there are misunderstandings about the soundness of our results. If you agree that we've managed to address all issues, please consider raising the score. But if there are any remaining concerns, please do not hesitate to follow up - we would be more than happy to address your further questions.

The authors consider the Hessian of the loss function $\ell(\theta , x+\delta ,y)$ in Lines 204-205.

We've clarified the notation $\delta$ in our General Response 1 and demonstrated why it will not affect the validity of our theoretical results in our General Response 2.

Can the authors clarify how they compute the Hessian of the adversarial loss function $\ell(\theta , x+\delta ,y)$?

In the definition of WCI in Equation (5), I wonder how the authors compute $\text{Tr}(H_k)$.

Below, we answer the reviewer's questions regarding the computation of the Hessian matrix $\mathbf{H}$. Let $\mathbf{H}\_k$ be the Hessian matrix for the $k$-th layer of the neural network $h\_\theta$. A short answer is that we empirically compute the Hessian matrix and its trace by employing projected gradient descent (PGD) to find an approximated perturbation to the worst-case (i.e., $\delta_{\mathrm{pgd}}$) and fix the perturbation as $\delta$ in the loss function when computing the Hessian.

Theoretically, $\mathbf{H}\_k$ (see Equation 10 in the appendix) is defined as:

$$\mathbf{H}\_k = \nabla^2\_{\theta} \ \mathbb{E}\_{(x, y) \in \mathcal{S}}\big[ \ell(\theta, x + \delta(\theta,x,y), y) \big ] = \sum\_{i,j} \frac{\partial^2 \mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta,x,y), y)}{\partial W\_k[i,j]^2},$$

where $\delta(\theta, x, y)$ denotes the worst-case perturbation (its dependence on $\theta$ and $(x,y)$ is clarified in our General Response 1) and $W_k$ denotes the weights of the $k$-th layer of the neural network. For simplicity, we write $\delta(\theta) = \delta(\theta,x,y)$ in the following discussions. The trace of the Hessian matrices are introduced in Lemma 2.4 to simplify the Classifier Variability term in Equation 1 of our work.

In our numerical analysis and experiments, however, the Hessian computation deviates from the idealized theoretical formulation, due to the challenges of optimizing over both $\delta$ and $\theta$ simultaneously. Thus, we resort to approximation and optimization tools, such as PGD to approximate the worst-case perturbation $\delta$, to compute the Hessian matrix and its trace. Specifically, our implementation proceeds in two distinct stages: (a) use Projected Gradient Descent (PGD) to obtain $\delta_{\mathrm{pgd}}$ that serves as a good approximation of the worst-case perturbation $\delta$ with respect to the corresponding $\ell_p$-norm ball, and (b) use auto differentiation to compute the Hessian by fixing the input perturbation as $\delta_{\mathrm{pgd}}(\theta)$ (i.e., the Hessian is computed with respect to $\mathcal{L}\_{\mathcal{S}}(\theta, x + \delta\_{\mathrm{pgd}}(\theta), y)$.

There may exist a gap between the theoretic Hessian definition (Equation 10) and the empirically computed Hessian matrix used in our experiments. However, we believe the empirical Hessian matrix can largely preserve the Hessian structure with respect to the (worst-case) robust loss $\mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta), y)$ and will not affect the conclusion of our numerical analysis and experiments. In fact, PGD-based adversarial training is still a widely adopted method for adversarial robustness, thus using $\delta_{\mathrm{pgd}}$ is a reasonable choice for understanding robust generalization of adversarially-trained neural networks.

Lemma 2.4 is based on an unspecified assumption that "the empirical loss $L_S(\theta ,x+\delta ,y)$ can be approximated using second-order Taylor expansion". Can the authors clarify the precise mathematical statement of the assumption?

Below, we lay out the precise assumption of the second-order loss approximation that we used in Lemma 2.4. Let $\mathcal{L}\_{\mathcal{S}}(\theta, x+\delta(\theta), y)$ be the robust loss with $\delta$ being the worst-case perturabtion (clarified in our General Response 2). For any $\theta'$ sampled from the posterior distribution $\mathcal{Q}$, we assume that the empirical robust loss at $\theta'$ can be approximated by: $\mathcal{L}\_\mathcal{S}(\theta', x + \delta(\theta'), y) \approx \mathcal{L}\_\mathcal{S}(\theta, x + \delta(\theta), y) + \tfrac{1}{2} \Delta \theta^\top \nabla^2\_{\theta} \mathcal{L}\_\mathcal{S}(\theta, x + \delta(\theta), y) \Delta \theta.$ Note that we keep the Taylor expansion terms up to the second-order terms. Here, the first-order term can be discarded because $\theta$ is considered to be at or near a stationary point of the robust loss function, which is a common setting considered in prior literature on deep learning generalization (see [1-2], for example).

[1] Stephan, Mandt, Matthew D. Hoffman, and David M. Blei. "Stochastic gradient descent as approximate bayesian inference." Journal of Machine Learning Research 18.134 (2017): 1-35.

[2] Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

First, standard ReLU activation functions are not twice-differentiable and so the Hessian analysis does not apply to them. Also, $L_S$ is supposed to be the adversarial loss, which means the authors assume that the worst-case loss $\max_{\Vert\delta\Vert\le \epsilon } \ell(\theta, x+\delta , y)$ has a bounded Hessian. Therefore, the flatness of the adversarial loss seems a very strong assumption which very likely would not work even using twice-differentiable activation functions.

We would like to clarify that our analysis does not rely on the assumption that the adversarial loss $\ell(\theta, x+\delta(\theta), y)$ has a bounded Hessian or that the flatness of the adversarial loss holds universally.

While ReLU activations are not twice-differentiable (due to the kink at zero), this does not invalidate the Hessian analysis. The Hessian is computed over the weights (parameters) of differentiable layers (e.g., linear or convolutional layers). The non-differentiability of ReLU affects the loss landscape indirectly but does not preclude meaningful second-order approximations in regions where the network behaves smoothly. In practice, deep networks often operate in regions where the activation paths induced by ReLU are stable. In these regions, the overall loss landscape can still be locally approximated using second-order methods.

The second-order Taylor expansion assumption for the adversarial loss is a simplification that may not strictly hold in all cases, particularly for networks with ReLU activations or sharp adversarial loss surfaces. Despite the non-differentiability of ReLU, second-order approximations have been shown empirically to align with network behaviors in both adversarial and standard training scenarios. This includes works that analyze sharpness, curvature, and generalization bounds (see Equation 8 in [3]). In practice, it has been shown to capture meaningful characteristics of the loss landscape in regions where the network behaves smoothly.

[3] Wu, Dongxian, Shu-Tao Xia, and Yisen Wang. "Adversarial weight perturbation helps robust generalization." Advances in neural information processing systems 33 (2020): 2958-2969.

Also, when I checked the proof of this lemma, I could not understand Lines 804-807. How does the "second-order Taylor series expansion around $\theta$" yield Equation (8)?

Equation 8 directly follows by expanding the loss difference using the second-order Taylor expansion at $\theta$ and leveraging the properties of Gaussian-distributed perturbations to bound the expected loss difference. Below, we explain the detailed derivations steps:

• Step 1. Because of the second-order loss approximation, we obtain $\mathcal{L}\_\mathcal{S}(\theta', x + \delta(\theta'), y)] - \mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta), y)\approx \frac{1}{2}\Delta\theta^\top \nabla^2\_{\theta} \big[ \mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta), y) \big ] \Delta\theta,$ where $\Delta = \theta' - \theta$ follows $\mathcal{N}(0,\Sigma)$ because we set the posterior $\mathcal{Q}$ as $\mathcal{Q} = \mathcal{N}(\theta, \Sigma)$.

• Step 2. By taking expectation over $\theta'\sim\mathcal{Q}$ with respect to the above, we derive $\mathbb{E}\_{\theta'\sim\mathcal{Q}}\mathcal{L}\_\mathcal{S}(\theta', x + \delta(\theta'), y)] - \mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta), y)\approx \frac{1}{2}\mathbb{E}\_{\Delta\theta\sim\mathcal{N}(0,\Sigma)}\Delta\theta^\top \nabla^2\_{\theta} \big[ \mathcal{L}\_{\mathcal{S}}(\theta, x + \delta(\theta), y) \big ] \Delta\theta.$

We will revise the proof of Lemma 2.4 in the revised manuscript to ensure clarity and rigor to eliminate any typos, ambiguities or confusion. Specifically, we will carefully review all mathematical expressions, assumptions, and steps to ensure they are accurate and align with the intended derivation shown above. Your feedback is invaluable, and we are committed to providing a clear and precise explanation of Lemma 2.4 in the final version.

Thank you once again for your constructive comments. We have implemented several key modifications in the manuscript to address your feedback:

Definition 2.1 (Adversarial Risk and $\delta$): Modified in Section 2, Lines 119-124, to clarify the dynamic optimization of $\delta$ for each input example.

Lemma 2.4 (Hessian Calculation): Detailed computational methods now included in Appendix A.3, Lines 932-937, emphasizing our approach to second-order loss approximation.

Empirical Loss Approximation: Expanded discussion in Section 2.2, Lines 220-225, about the conditions and justification for using the Taylor expansion.

We believe these revisions address the concerns you raised and enhance the manuscript’s technical depth. We look forward to your thoughts on these updates and any further suggestions you might have.

We appreicate the reviewer's follow up comments and apologize for any confusion caused by our previous clarification. In the following, we clarify why the theoretical Hessian matrix can be estimated by first finding an approximated solution $\delta\_{\mathrm{pgd}}$ to the inner maximization problem using PGD then computing the second-order derivative with respect to the loss function at $(x+\delta\_{\mathrm{pgd}}, y)$. In particular, we will explain both from the general perspective based on the Danskin's Theorem and using the specific example provided by the reviewer.

Consider $\ell\_\infty$ perturbations bounded by $\epsilon$ and the corresponding inner maximization problem with respect to $(x,y)$. According to the Danskin's Theorem, we can obtain:

$$\nabla\_{\theta} \bigg[ \max\_{||\delta||\_\infty \leq \epsilon} \ell(\theta, x+\delta, y) \bigg] = \nabla\_{\theta} \big[ \ell(\theta, x+\delta^*, y) \big]\approx \nabla\_{\theta} \big[ \ell(\theta, x+\delta\_{\mathrm{pgd}}, y) \big],$$

where $\delta^* = {\mathrm{argmax}}\_{||\delta||\_\infty\leq \epsilon} \: \ell(\theta, x+\delta, y)$ denotes the worst-case perturbation and $\delta\_{\mathrm{pgd}}$ denotes the perturbation returned by PGD attack. As the reviewer mentioned, the above equation characterizes the theoretical basis of adversarial training. Taking the derivative with respect to $\theta$ for both sides of the above equation, we immediately have:

$$\nabla\_{\theta}^2 \bigg[ \max\_{\|\delta\|\_\infty \leq \epsilon} \ell(\theta, x+\delta, y) \bigg] = \nabla\_{\theta}^2 \big[ \ell(\theta, x+\delta^*, y) \big] \approx \nabla\_{\theta}^2 \big[ \ell(\theta, x+\delta\_{\mathrm{pgd}}, y) \big].$$

This suggests that to estimate the theoretical Hessian matrix with the worst-case perturbation $\delta^*$, one can first find an approximated solution $\delta_{\mathrm{pgd}}$ using PGD attack then compute the Hessian matrix with respect to the loss function at $(x+\delta_{\mathrm{pgd}}, y)$. We want to emphasize that both $\delta^*$ and $\delta_{\mathrm{pgd}}$ should not be understood as fixed constant vectors, since their values essentially depend on the model parameter $\theta$ as the input.

Below, we further illustrate why this is the case using the example provided by the reviewer. Consider the following optimization problem:

$$\max\_{\delta} L(\theta, \delta), \quad \text{where} \quad L(\theta, \delta) = \delta^{\top}\theta - \frac{1}{2}||\delta||\_2^2.$$

Theoretically, we know the optimal solution is $\delta^* = \theta$ and we immediately have:

$$\nabla^2\_\theta \bigg[\max\_\delta L(\theta, \delta)\bigg] = \nabla^2\_\theta \ L(\theta, \delta^*) = \nabla\_\theta^2 \left( \frac{1}{2} ||\theta||\_2^2 \right) = \mathbf{I},$$

suggesting the theoretical Hessian is an identity matrix. To simulate PGD attack used for approximating the worst-case, we implement a simple gradient ascent algorithm to solve the maximization problem $\max_\delta L(\theta, \delta)$, then compute the Hessian matrix with respect to the $L(\theta, \hat\delta)$ with $\hat\delta$ denoting the gradient descent returned solution. We verify that $\nabla^2_{\theta} L(\theta, \hat\delta)$ is very close to an identity matrix (instead of a zero matrix), which confirms the validity of our numerical computation of the Hessian matrix. We provide Python code of this implementation for the reviewer to check:

import torch

def counterexample_test():
    # Define the counterexample function L(delta, theta)
    # L(delta, theta) = delta^T * theta - 0.5 * ||delta||_2^2
    def loss_function(delta, theta):
        return torch.dot(delta, theta) - 0.5 * torch.norm(delta, p=2) ** 2

    # Compute the optimal perturbation delta^* as a function of theta
    def compute_delta_star(theta, max_iters=10, lr=0.9):
        delta = torch.zeros_like(theta, requires_grad=True)  # Initialize delta
        for _ in range(max_iters):
            loss = loss_function(delta, theta)  # Calculate the current loss
            grad = torch.autograd.grad(loss, delta, create_graph=True)[0]  # Calculate the gradient of loss with respect to delta
            delta = delta + lr * grad  # Gradient ascent update delta
        return delta

    # Calculate g(theta) = max_delta L(delta, theta)
    def max_function(theta):
        delta_star = compute_delta_star(theta) 
        return loss_function(delta_star, theta)

    # Initialize theta
    theta = torch.rand(5, requires_grad=True)

    # Compute the Hessian using numerical methods
    hessian = torch.autograd.functional.hessian(max_function, theta)

    print(hessian)

if __name__ == "__main__":
    counterexample_test()

Even if we lower the value of max_iters or lr, the result will never be $\mathbf{0}$. It only affects the size of the error between the numerical calculation and the actual value.

We believe the above example and explanations should clarify the reviewer's concern about the validity of our Hessian computation regarding the adversarial loss. For the sake of completeness, we explain how the Hessian matrix is computed in our experiments. In particular, we calculate the hessian following the lines 311 to 333 in the code repository of the "Robust Overfitting" paper (Rice et al., 2020).

robust_output = model(normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit)))
robust_loss = criterion(robust_output, y)

Then in line 334 we add a function to compute WCI as well as hessian:

WCI, WCI_norm, WCI_trace = funcWCI(model, robust_loss)

The specific content of this function is as follows:

def funcWCI(model, loss, epsilon=1e-6):
    sum_WCI = 0
    sum_norm = 0
    sum_trace = 0

    grads = torch.autograd.grad(loss, [p for p in model.parameters() if p.requires_grad], create_graph=True)
    
    for param, grad in zip(model.parameters(), grads):
        if grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > epsilon:
            grad_clipped = torch.clamp(grad, -1e6, 1e6)
            grad2 = grad_clipped.pow(2) + epsilon
            trace = torch.autograd.grad(grad2.sum(), param, retain_graph=True)[0]
            if trace is not None and torch.isfinite(trace).all():
                hessian_trace = trace.sum()
                layer_frobenius = torch.norm(param, p='fro') ** 2
                WCI = (layer_frobenius * hessian_trace + epsilon).sqrt()
                
                # Collecting layer-specific details
                if param.requires_grad:
                    if torch.isfinite(WCI):
                        sum_WCI += WCI
                    sum_norm += layer_frobenius
                    sum_trace += hessian_trace

    return sum_WCI, sum_norm, sum_trace

Thanks for pointing out the additional reference and the challenges for Hessian computation, particularly in cases where the adversarial maximizer is not unique. We recognize the validity of the concern that a single adversarial example can yield a subgradient rather than a true gradient, as highlighted in the referenced work. The proposed approach of selecting a minimum norm convex combination from the subdifferential indeed ensures safe descent and could be a robust strategy for handling subgradient challenges. Extending such approaches to the Hessian level, including the use of a bordered Hessian, is an interesting direction for future work.

Nevertheless, we want to emphasize that the numerical Hessian matrix computation used in our paper is valid, as we clarified in our new response to Reviewer msia. In paricular, we explain why such a computation is reasonable assuming the Danskin's Theorem is applicable (i.e., the optimal solution to the inner-maximization problem is unique) and validate numerically using a gradient ascent approximation with respect to a specific example $\max_{\delta} \left(\delta^{\top}\theta - \frac{1}{2}||\delta||_2^2\right)$. As the authors of https://openreview.net/forum?id=I3HCE7Ro78H said, PGD and DDi seem to behave similarly in the later stages of training. So here, we use PGD as the basis, given its wide adoption in adversarial training and its practical efficiency. Our approach aligns with standard practices and has been shown to work robustly in the settings we analyze.

While we appreciate the suggestion to consider bordered Hessians, which could provide additional theoretical guarantees in addressing subgradient issues, we believe our current method is accurate and efficient for the scope of our study. Future work could further explore these ideas to generalize the analysis to settings where non-uniqueness is a significant concern. Thanks again for raising this important point and for the helpful reference. It has been a valuable addition to this discussion.

An important observation is that for the above PAC-Bayesian generalization bound to hold, there are no assumption imposed on the loss function $\ell$, thus it directly applies to the adversarial loss (with respect to the worst-case perturbation $\delta$). In fact, the only required assumption is that the prior distribution $\mathcal{P}$ is data independent such that the Fubini's theorem can be applied to exchange the intergration with $\theta$ and the sample (see Equation 2.2 in [1]). Therefore, when deriving robust generalization bound, we can simply replace the standard loss in the above inequality by their robust counterpart. More specificaly, denote by $\ell\_{\mathrm{rob}}(\theta, x, y) = \underset{\delta' \in \mathcal{B}\_{\epsilon}(0)}{\max} \ell(\theta, x+\delta', y)$ the individual robust loss with the worst-case perturabtion, then we have $\mathbb{E}\_{\theta\sim\mathcal{Q}} \mathbb{E}\_{(x,y)\sim \mathcal{D}}[\ell_{\mathrm{rob}}(\theta, x , y)] \leq \mathbb{E}\_{\theta\sim\mathcal{Q}} \bigg[\frac{1}{|\mathcal{S}|}\sum\_{(x,y)\in\mathcal{S}}\ell\_{\mathrm{rob}}(\theta, x , y)\bigg] + \frac{1}{\lambda} \mathrm{KL}(\mathcal{Q} || \mathcal{P}) + \text{const.},$ This lays the foundation for Lemma 2.2 in our paper, which establishes a generic robust generalization bound that holds for any classifer $h_\theta$ (not necessary neural network), posterior distribution $\mathcal{Q}$ and data-independent prior $\mathcal{P}$.

To extract meaningful insights from the generic robust generalization bound, however, we need to set $\mathcal{P}$ and $\mathcal{Q}$ properly such that the generalization bound is tight and can be simplified into analytical terms. These are exactly what we have done in Lemma 2.3, Lemma 2.4 and Theorem 2.5, where we set $\mathcal{P}$ as $\mathcal{N}(0, \sigma^2 \Sigma)$ and $\mathcal{Q}$ as $\mathcal{N}(\theta, \Sigma)$ with covariance matrix $\Sigma = \text{diag}(\sigma_1^2 \mathbf{I}_1, \sigma_2^2 \mathbf{I}_2, \dots, \sigma_K^2 \mathbf{I}_K)$ to simplify the terms (Lemma 2.3 and Lemma 2.4), and then choose the variance for the $k$-th layer $\sigma_k^2 = \sqrt{\|\mathbf{W}_k\|_F^2 / (2\lambda \mathrm{Tr}(\mathbf{H}_k))}$ to minimize the bound (Theorem 2.5). Selecting Gaussian prior and posterior aligns well with existing literature that employs the PAC-Bayesian framework (see [2,3]). Such choices of $\mathcal{P}$ (which is clearly data-dependent) and $\mathcal{Q}$ do not violate the soundness of our theoretical results, again because of the generic nature of PAC-Bayesian framework.

[1] Pierre Alquier et al., "User-friendly introduction to pac-bayes bounds". Foundations and Trends in Machine Learning, 17(2):174–303, 2024.

[2] Mbacke, Sokhna Diarra, Florence Clerc, and Pascal Germain. "PAC-Bayesian generalization bounds for adversarial generative models." International Conference on Machine Learning. PMLR, 2023.

[3] Jin, Gaojie, et al. "Enhancing adversarial training with second-order statistics of weights." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.