S2O: Enhancing Adversarial Training with Second-Order Statistics of Weights

Abstract

Adversarial training has emerged as a highly effective way to improve the robustness of deep neural networks (DNNs). It is typically conceptualized as a min-max optimization problem over model weights and adversarial perturbations, where the weights are optimized using gradient descent methods, such as SGD. In this paper, we propose a novel approach by treating model weights as random variables, which paves the way for enhancing adversarial training through \textbf{S}econd-Order \textbf{S}tatistics \textbf{O}ptimization (S$^2$O) over model weights. We challenge and relax a prevalent, yet often unrealistic, assumption in prior PAC-Bayesian frameworks: the statistical independence of weights. From this relaxation, we derive an improved PAC-Bayesian robust generalization bound. Our theoretical developments suggest that optimizing the second-order statistics of weights can substantially tighten this bound. We complement this theoretical insight by conducting an extensive set of experiments that demonstrate that S$^2$O not only enhances the robustness and generalization of neural networks when used in isolation, but also seamlessly augments other state-of-the-art adversarial training techniques. The code is available at https://github.com/Alexkael/S2O.

Figures (4)

• Figure 1: Illustration of the theoretical framework: perturbation bound with consideration of correlated weights. Under this framework, a standard generalization bound is extended to a robust generalization bound with weight correlation matrix, and further to a bound with weight correlation matrix estimated over both clean data and adversarial data.
• Figure 2: Illustration of the optimization framework: \ref{['thm:advbound']} and \ref{['thm:advbound2']} show the influence of the second-order statistics of weights, over clean and adversarial data, on the robust generalization performance. In \ref{['sec:fnorm']}, we demonstrate that these statistics can be approximately optimized using the Frobenius norm of the correlation matrix. To simplify this optimization process, we employ the Laplace approximation in \ref{['sec:laplace']} and propose our adversarial training methodology, S$^2$O, in \ref{['sec:optimize']}.
• Figure 3: We sample 10000 9-dimensional correlation matrices and demonstrate (a)$\|\mathbf{R}_{l,\mathbf{x}}\|_F^2$, $\|\mathbf{R}_{l,\mathbf{x}'}\|_F^2$ w.r.t $\Lambda^c_{l,\max}$ or $\Lambda^r_{l,\max}$; (b)$\|\mathbf{R}_{l,\mathbf{x}}\|_F^2$, $\|\mathbf{R}_{l,\mathbf{x}'}\|_F^2$ w.r.t $\Lambda_{l,\min}^{k_l} \Lambda_{l,\max}^{h^2-k_l}$; (c)$\|\mathbf{R}_{l,\mathbf{x}}\|_F^2$, $\|\mathbf{R}_{l,\mathbf{x}'}\|_F^2$ w.r.t $\Lambda^c_{l,\max}$ or $\Lambda^r_{l,\max}$; (d)$\|\mathbf{R}_{l,\mathbf{x}}\|_F^2$, $\|\mathbf{R}_{l,\mathbf{x}'}\|_F^2$ w.r.t $\Lambda_{l,\min}^{k_l} \Lambda_{l,\max}^{h^2-k_l}$. Note that the horizontal axis represents both $\|\mathbf{R}_{l,\mathbf{x}}\|_F^2$ and $\|\mathbf{R}_{l,\mathbf{x}'}\|_F^2$.
• Figure 4: We train PreAct ResNet18 with AT and AT+S$^2$O on CIFAR-10, and show the results of partial weights. The results estimated from adversarial data are presented in the top figures, while the bottom figures show the results estimated from clean data. (a) shows the normalized spectral norm of $\mathbf{R}^c$, ${\mathbf{R}}^r$, and the determinant of $\mathbf{R}$, with sampling estimation and Laplace approximation respectively. (b) and (c) demonstrate the absolute correlation matrix of partial weights, for AT and AT+S$^2$O respectively.

Theorems & Definitions (23)

• Remark 1
• Remark 2
• Lemma 2.1: neyshabur2017pac
• Theorem 2.2: Standard PAC-Bayesian generalization bound, neyshabur2017pac
• Theorem 2.3: Robust generalization bound, farnia2018generalizable
• Theorem 2.4: Robust generalization bound, xiao2023pac
• Definition 3.1: Correlation matrix
• Theorem 3.2: Robust generalization bound with correlation matrices
• proof
• Remark 3