Evolution of Neural Tangent Kernels under Benign and Adversarial Training
Abstract
Two key challenges facing modern deep learning are mitigating deep networks’ vulnerability to adversarial attacks and understanding deep learning’s generalization capabilities.
Towards the first issue, many defense strategies have been developed, with the most common being Adversarial Training
(AT).
Towards the second challenge, one of the dominant theories that has emerged is the Neural Tangent
Kernel (NTK) – a characterization of neural network behavior in the infinite-width limit.
In this limit, the kernel is frozen, and the underlying feature map is fixed.
In finite widths, however, there is evidence that feature learning happens at the earlier stages of the training (kernel learning) before a second phase where the kernel remains fixed (lazy training).
While prior work has aimed at studying adversarial vulnerability through the lens of the frozen infinite-width NTK, there is no work that studies the adversarial robustness of the empirical/finite NTK during training.
In this work, we perform an empirical study of the evolution of the empirical
NTK under standard and adversarial training, aiming to disambiguate the effect of adversarial training on kernel learning and lazy training.
We find under adversarial training, the empirical NTK rapidly converges to a different kernel (and feature map) than standard training.
This new kernel provides adversarial robustness, even when non-robust training is performed on top of it.
Furthermore, we find that adversarial training on top of a fixed kernel can yield a classifier with 76 . 1% robust accuracy under PGD attacks with ε = 4 / 255 on CIFAR-10.
Introduction
Modern deep learning, while effective in tackling clean and curated datasets, is often very brittle to domain and distribution shifts [76].
Perhaps the most notorious failure mode of deep learning to domain shifts is adversarial examples [73]: images with small, bounded perturbations which consistently fool state-of-the-art classifiers.
While work has been dedicated to mitigating [25, 18, 51], explaining [24, 39, 37], and harnessing [66, 54, 67] this peculiar behavior, the problem still remains largely open, with the state-of-the-art robust classifiers still falling far behind in benign accuracy compared to standard non-robust networks [18].
Indeed, tackling this idiosyncrasy of deep learning seems almost impossible given that the behavior of deep models under benign data and training is still largely unexplained [13, 62, 63], let alone adversarial training.
The community has tried developed numerous theories on better interpreting [77, 27, 43] and understanding deep learning [74, 5, 38, 64, 28, 14, 31, 30].
One emerging theory is the notion of deep learning as kernel learners [48, 21, 1].
That is, deep networks trained with gradient descent learn a kernel whose feature map is embedded in the tangent space of the network outputs with respect to its parameters.
After a brief phase of kernel learning, neural networks behave similarly to lazy learners , i.e., linear and slow in this neural tangent feature map.
This two-stage theory of deep learning has been observed in vision networks such as residual networks [32] used in practice and has been the subject of many recent empirical [71, 48, 21, 7] and theoretical works [29, 35, 2].
While these reports have verified this theory for benign training, no work has looked at this property at the intersection of adversarial training.
To this end, here, we perform the first empirical study of how the empirical/finite neural tangent kernel evolves over the course of adversarial training.
In section 2, we present our experimental setup designed to isolate the effect of adversarial training on the two stage of deep learning under the kernel learner theory: 1) kernel learning and 2) linear fitting.
By using this framework, we study and show the following:
Similar to standard training, adversarial training results in a distinct kernel which quickly converges in the first few epochs of training (section 3) 2.
Adversarial robustness can be inherited from the kernel made from adversarial training , even when the second stage classifier has no access to adversarial examples during training (section 4).
Adversarial training is effective on top of the learned kernels given by standard training (in the frozen feature regime), providing a testbed where many of the fixed-feature assumptions present in theoretical works on adversarial training are met, while still providing high practical performance (section 7).
Eigenvectors in the learned NTK of networks with adversarial training contained visually interpretable features, while the initial NTK and benign training NTK do not (section 8).
1.1 Background and Related Works
Neural Tangent Kernel.
Explaining the empirical success of deep learning models from a theoretical standpoint is an exciting and active line of research that is still in its infancy.
One of the most popular tools for understanding neural network behavior is the neural tangent kernel (NTK)
Under this theory training deep networks with gradient descent in the infinite-width limit corresponds to training training a kernel-based classifier, with the corresponding kernel being the NTK defined as k NTK ( x, x ′ ) =
E θ ∼ p
( θ ) [ ∇ θ f ( x ) T ∇ θ f ( x ′ )] , for parameter distribution p ( θ ) and network architecture f θ .
Analogously, training an infinite width Bayesian neural network corresponds to Gaussian process regression using the NNGP kernel [55, 23, 45, 52, 56, 50].
In both these infinite-width limits, the underlying feature space is fixed, determined entirely by the architecture and parameter distribution.
The frozen nature of the kernel stands in contradiction to the ability of deep models to learn useful features and representations.
Indeed, it has been shown theoretically, under different limiting assumptions, or with finite-widths, that feature learning does occur, with a time-evolving and datadependent kernel, k t ( x, x ′ )
=
∇ θ t f θ t ( x ) T ∇ θ t f θ t ( x ′ ) , with time-dependent network parameters θ t [78, 29, 35, 2, 58, 59].
This learned NTK aligns itself with dataset labels to accelerate training and improve generalization.
Empirical evidence supports this, with works finding that this meta-kernel evolves quickly within the early stages of network training before stabilizing [48, 21, 7, 71, 40, 8].
In this new setting, there exists two stages in deep network training: 1) kernel learning and 2) linear fitting .
In the first stage, the kernel rapidly evolves to align with the dataset’s features and labels, while in the second phase, the kernel changes only minimally and the network behaves linearly in the
NTK feature map ϕ ( x ) = ∇ θ ′ f θ ′ ( x ) (cid:12)(cid:12)(cid:12) θ ′ =
θ t , a regime sometimes referred to as lazy training [17, 46].
Adversarial Examples and adversarial training.
Adversarial examples present one of the most infamous failure modes of deep learning [73].
These are images within bounded perturbations of naturally occurring images that consistently fool deep networks over a wide array of architectures [53].
There is much literature dedicated to designing techniques for securing networks against these attacks[25, 51, 18, 26, 44, 80, 61, 16, 42].
In this work, we focus on adversarial training with iterated projected gradient descent (PGD), which seeks to minimize
L rob =
E x,y ∼ p ( x,y ) (cid:104) max x ′ ∈
B ϵ ( x ) L standard ( x ′ , y ) (cid:105) .
That is, minimizing the worst-case loss of samples in a small neighborhood around training examples, where the inner maximization is computed during training using iterated PGD [25, 51].
From a theoretical standpoint, much work has been devoted to explaining the causes of adversarial vulnerability and the limitations of adversarial robustness.
Popular theories include the notion of there existing “robust and non-robust features" present in data [37, 76], adversarial examples being an outcome of high dimensional geometry [24, 39, 69], an outcome of stochastic gradient descent (SGD) [70, 3], and many more [14].
Based on the NTK theory, there have been works studying the presence of adversarial examples under the NTK and NNGP kernels [22, 9, 15, 60, 10].
Recently, NTK theory has been used to generate attack methods for deep networks [75, 79].
Adversarial examples have been shown to arise in simplified linear classification settings [68, 37], which readily transfer to the
NTK settings as the NTK classifier is linear in its underlying feature representation.
However, these reports on the adversarial vulnerability of NTK and NNGP kernels have been focused solely on the infinite-width limit of these kernels, with no literature on the adversarial robustness of the learned and data-dependent NTK that is present in practical networks.
Experimental setup
Definitions and problem setup.
We first define the necessary objects and terminologies for our experimental setup.
Our scheme follows closely that of [21]; however, we consider the additional dimension of adversarial robustness in addition to benign accuracy.
First, we define the parameterdependent empirical NTK: k ENTK,θ ( x, x ′ ) =
∇ θ f θ ( x ) T ∇ θ f θ ( x ′ ) .
In the case where θ = θ ( t ) , where θ ( t ) refers the parameters of the network after training for t epochs, we use the shorthand k t =
k ENTK,θ ( t ) .
Next, we define three training dynamics.
Standard dynamics as f standard ,θ = f θ ( t ) , that is, network behavior with no modifications.
Linearized (also referred to as lazy) dynamics around epoch
t : f lin ,θ,t ( x ) =
f θ t ( x ) + ( θ − θ t ) T ∇ θ ′ f
θ ′ ( x ) (cid:12)(cid:12)(cid:12) θ ′ =
θ t .
I.e., linearized training dynamics corresponds to performing a firstorder Taylor expansion of the network parameters around a point θ t .
We refer to f θ t as the parent network and t as the spawn epoch.
Note that in linearized dynamics, the underlying feature map ϕ ( x ) = ∇ θ ′ f θ ′ ( x ) (cid:12)(cid:12)(cid:12) θ ′ =
θ t , and the corresponding kernel is fixed.
Fort et al.
[21] studied this regime and showed that in practice, deep networks training with SGD undergo two stages in training, in which the first phase was chaotic with a rapid changing kernel until some relatively early epoch t , and the second stage behaves like linear training about θ t .
Centered
linear training (or centered for short): where we subtract the parent network’s output [34, 47, 49]: f centered ,θ,t ( x ) =
( θ − θ t ) T ∇ θ ′ f
θ ′ ( x ) (cid:12)(cid:12)(cid:12) θ ′ =
θ t .
This corresponds to linear training, with the zeroth order term removed.
Now the output is strictly dependent on the difference between θ and θ t , and the contribution of the first t epochs is only through how it modifies the feature map.
Studying this setting lets us isolate the properties of the learned kernel k t , without worrying about its contribution from earlier.
Efficient implementation of both these linearized versions is done using forward mode differentiation in the JAX library [11].
Experimental design We follow closely that of [21], where we train networks in two distance stages with an added dimension of adversarial training.
Additional details are available in ?? .
Stage 1: Standard dynamics.
We train Resnet18s on CIFAR-10 or CIFAR-100 [41] for t epochs either using benign data (i.e. no data modification), or adversarial training, for 0 ≤ t ≤ 100 .
Stage 2: Linearized dynamics.
Following stage 1, we take the parent network from stage 1 and train for an additional 100 epochs with either linearized or centered linearized training dynamics.
Note that for centered training, after stage 1, the networks will output entirely zeros, as the zeroth order term has no effect, so the classifier does not have a “warm start."
In this stage, we train using benign data, with no adversarial training.
Additionally, we freeze the batchnorm running mean and standard deviation parameters after stage 1, as allowing these to change would implicitly change the feature map ϕ ( x ) = ∇ θ ′ f θ ′ ( x ) (cid:12)(cid:12)(cid:12) θ ′ =
θ t .
For all experiments in the main text, we use the standard ε = 4 /adversarial radius under a L ∞ norm, but we verify that the results hold for ε = 8 / 255 in sectionwith additional results for ε = 8 / 255 in the appendix.
Evolution of the NTK under standard and adversarial training
First, we look at the evolution of the kernel under different training conditions.
To do this, we calculate the empirical NTK on a random subset of 500 class-balanced training points on CIFAR10 and CIFAR-100 for resnets trained for 100 epochs SGD under standard dynamics with either benign training or adversarial training.
We also consider a third scenario where we perform benign training for 50 epochs then change to adversarial training for the remaining 50 epochs so we have an additional control point of the effect of adversarial training.
For networks with multiple outputs, the kernel matrix is a rank-four tensor in R C × C × N × N with C being the class count and N being the dataset size.
The calculation for these entries is given by k c,c ′ ( x, x ′ )
= ∇ θ f cθ ( x ) T ∇ θ f c ′ θ ( x ′ ) , and the resulting subclass kernel matrix K c,c ′ ∈ R N × N .
Unless otherwise stated, for the statistics we present, we calculate the trace kernel, the average of the diagonal elements of the class-specific matrix: ¯ K = 1 C (cid:80)
Cc =1 K c,c .
Firstly, we look at the kernel distance between neighboring epochs.
This kernel distance is given by:
S ( K 1 , K 2 ) = 1 −
Tr ( K T 1 K 2 ) ||
K
||
F || K
|| F , which equals to
iff K 1 = K 2 .
The kernel velocity is given by dSdt , which we approximate as a finite difference between neighboring epochs (See fig. 1a).
We also consider the kernel distance from the current epoch to the kernel at the end of training (fig. 1b).
From these metrics we see that both for benign and adversarial training, the kernel converges withinepochs to close to the final kernel, in accordance with the results in Fort et al.
This suggests that after these few epochs, the underlying feature set is fixed, and the remainder of the training is spent performing linear classification on this feature set.
Surprisingly, adversarial training also quickly converges, albeit to a different kernel whose underlying feature set is more robust (as we will see in later sections).
In the third setting, we observe a small spike in kernel velocity at epoch 50 when we swap from standard to adversarial training.
The change is small compared to the initial rapid kernel evolution, suggesting that the standard training kernel is more similar to the adversarial kernel than it is to the initial NTK.
Likewise, when we plot the kernel distance to the final kernel, both models nearly converge after epoch 40, while in the third setting the model stabilizes at the standard training kernel before changing to the adversarial kernel after 50 epochs.
The third metric we compute is the effective rank of the kernel matrix [65].
This measures the dispersion of the matrix over its eigenvectors, and is given by: erank(K) =
exp ( − (cid:80) N i =1 p i log p i ) ,
With p i
=
λ
i (cid:80) Ni =1 λ i with λ i being the eigenvalues of the kernel matrix K .
This value is bounded between 1 when the matrix has one dominant direction and N when all eigenvalues are equal.
Previous work has shown that deeper networks are biased towards low effective rank (of the conjugate kernel) at initialization [36].
Alternatively, the effective rank could be interpreted as the complexity of the dataset under the kernel, and existing work has shown that robust classifiers require more model capacity [14, 68], although it is unclear the precise relationship between the effective rank notion of complexity and those given in prior work relating to adversarial robustness.
As a fourth metric, we consider the mean kernel specialization.
As discussed earlier, for multi-classed outputs, the full NTK is a 4 D tensor in R C ×
C × N × N with C class-specific kernels corresponding to the diagonal entries of the first two dimensions.
Over training, it has been observed that these individual class-specific kernels specialize in aligning themselves with their classes [71].
The kernel specialization, defined as: KSM
[Equation]
where A
[Equation]
S ( K c,c , y c ′ y Tc ′ ) , i.e. the cosine similarity of the kernel matrix and the one-hot class labels.
Intuitively, this compares how aligned class c specific kernel is to the labels of class c ′ compared to other class kernels.
This quantity is bounded between 0 and C and higher diagonal entries (when c = c ′ ), indicate higher specialization.
We define the mean kernel specialization as C − 1 Σ Cc =1 KSM ( c, c ) , i.e. the average of diagonal entries of the KSM matrix.
We calculate this for only CIFAR-10, as computing all 100 class-specific kernels for CIFAR-100 is too costly.
From fig.
1d, where we see that adversarial training is associated with a lower mean kernel specialization, we take away that adversarial training promotes features that are more broadly shared between classes, although this needs further investigations which will be the focus of our continued effort.
Performance of Linearized and Centered Training
Next, we look at the performance of linearized and centered training, where we vary the spawn epoch at which we begin stage 2 training, i.e. the epoch t described in section 2.
We then plot the benign and adversarial performance of the classifier after stage 2 training, in comparison with the performance at the end of stage 1 training in fig.
2 for CIFAR-10 and CIFAR-100.
We show the results when either stage 1 is performed with benign training or adversarial training.
Additionally, for the case of linearized training, we plot the relative magnitude of the zeroth order and first-order terms in the linearization dynamics equation.
Specifically, let f
=
f
θ t ( x ) and ∆ f = ( θ − θ t ) T ∇ θ
[Equation]
f
[Equation]
( x
[Equation]
We plot the relative magnitude of the two components, given by ∆ f T ∆ f f T 0 f 0 , averaged over the test set.
The most surprising observation is that centered training gains significant robustness over standard and linearized training dynamics, both in adversarial and benign training.
Because centered training does include the zeroth order term in its predictions, all the robustness given by centered training is inherited entirely through the learned NTK , and not through adversarial training (as we perform benign training in stage 2).
This lends credence to the “robust feature" hypothesis given in Ilyas et al.
[37], however, now what matters is not necessarily what features are present in the data, but what features the networks learns in its kernel.
On the final adversarial kernel, centered training brings about a 3 .
7% adversarial accuracy gain over standard adversarial training on CIFAR-10, going from 51 .
77 ± 0 .
80% to 55 .
46 ± 1 .
Similar improvements are also seen on CIFAR-100 with an improvement from 24 .
42 ± 0 .
14% to 27 .
35 ± 0 .
At around epoch 10, the centered network reaches 62 .
09 ± 1 .
76% robust accuracy on CIFAR-10, an over 10% performance gain.
Equally surprising is that even under benign training, the network learns a robust kernel over time, with centered training reaching 41 .
41 ± 9 .
75% after epoch 100 for CIFAR-10 and 14 .
16 ± 1 .for CIFAR-100.
Despite this, the neural network under standard dynamics still sees no adversarial accuracy, suggesting that the contribution from the robust learned kernel is minor compared to the contribution of the non-robust kernels, which are present early in training.
Roughly speaking, the final network outputs could be written as the integral over time of instantaneous kernel-dependent functions [19].
We can verify that the contribution of the kernel at later epochs is small by the relative magnitude plots in fig.
2c, which shows that the first-order additional term of linearized dynamics becomes very minor by epoch 30, meaning that the network has essentially fully fit the dataset, at which point the learned kernel reaches it maximum robustness (as seen by the centered plots).
We observe that the learned kernel only gains robustness over time, but its robustness stagnates as kernel velocity slows down.
Linearized training sees a peak in the relative magnitude of first to zeroth order components around epoch 10-15, coinciding with the small bump in robustness for the benign training kernel under linearized training, and for the adversarial training kernel, a peak in accuracy in benign and adversarial settings for linearized and centered training.
We conjecture that this bump in relative magnitude correlates with the stage transition between kernel learning and linear fitting, where the kernel is well aligned with the data, but the network has not yet fully fit the training data, causing the first-order component to contain the majority of the data fit.
Adversarial Transferability of Linearized Dynamics
We next study the differences between standard and linearized dynamics from the perspective of adversarial transferability, specifically, how well do adversarial examples generated with one dynamics transfer to models with other dynamics?
From the experiments described in section 2, we have a set of classifiers with standard, linearized, and centered dynamics: { f stan ,t } , { f lin ,t } , { f cen ,t } for 0 ≤ t ≤ 100 .
For each of these classifiers, we generated adversarial examples on the test set with L ∞ PGD (details provided in ?? ) , { ˆ X stan ,t } , { ˆ X lin ,t } , { ˆ X cen ,t } .
To see how these adversarial examples compare to those generated by a network with standard dynamics, we then evaluated the accuracy of f stan ,t =100 on these adversarial examples:
Acc ( · ) ,t ′ stan ,t =100 =
E x,y ∼ ˆ X ( · ) ,t ′
[ I ( y = f stan ,t =100 ( x ))] .
We then define the adversarial transferability as: τ ( · ) ,t ′ =
Acc
[Equation]
, where
Acc b stan ,t =100 is the accuracy of the standard dynamics classifier with a benign test set.
This metric measures the relative effectiveness of adversarial examples generated by another network to adversarial examples generated by the network itself.
Values of τ = 1 mean that the adversarial examples perform just as well as adversarial examples generated by the classifier itself, while τ =would imply that adversarial examples do not transfer whatsoever.
We plot these for CIFAR-10 and CIFAR-100 trained Resnet-18s with either benign or adversarial training in fig.
Comparing the plots of benign training and adversarial training, we see that in both cases, the adversarial examples of the standard and linearized dynamics transfer similarly well to the full trained network.
In contrast, the adversarial examples made by centered training are far less effective, suggesting that the feature set used by the kernel does not necessarily correspond to the
[Table]
same feature set used by a fully trained network.
For standard training, the adversarial transferability plots rise much more quickly than for adversarial training, suggesting that features from early kernels are used more in the final model than for adversarial training, as seen by the slower initial rise of the standard and linearized training in fig.
3b and fig.
3d.
As seen earlier, centered training on earlier
NTK kernels provides little robust accuracy, suggesting that some of the lack of robustness of fully trained networks could be due to the contribution of these early kernels.
For standard training, the adversarial transferability of fully trained networks rises sharply in the first 15 epochs, then slowly decays over time.
This suggests that the adversarial behavior of the fully trained network is closer to some of the earlier kernels than to the final kernel.
However, in both cases, adversarial examples generated by the initial NTK do not transfer at all to fully trained networks, suggesting that the initialization NTK analysis of adversarial examples in some theoretical work [22, 9] may not be accurately describe the behavior of adversarial examples in practical networks.
SGD vs Linearized Dynamics at low learning rates
In section 4, we showed that adversarial robustness could arise from linearized and centered training, in which the underlying feature set is fixed.
In this section, we investigate whether freezing the kernel is necessary for robustness gains.
To investigate this, we trained networks with adversarial training with standard dynamics until epoch 10.
We then performed centered training on top of the resulting kernel.
We chose ten epochs as it corresponds to the large jump in robust accuracy in fig.
The centered network serves as the control when the feature set is frozen.
Then, we varied stage 2 training in two ways.
First, we trained a network using “centered standard" dynamics, where f cen, stan, θ ( x ) =
f θ ( x ) − f θ t ( x ) , that is, we subtract the network from its initialization’s output, and train with standard gradient descent.
We use learning rates of either η = 0 . 0001 or η = 0 .
01 , corresponding to low and high learning rates.
Previously, we had argued that freezing the batchnorm running mean and standard deviation is necessary so that the feature map ∇ θ ′ f θ ′ ( x ) (cid:12)(cid:12)(cid:12)
θ ′ =
θ t does not change.
Here, we question this assumption and either keep the batchnorm parameters frozen or let them change.
Table 1 shows the resulting benign and adversarial accuracies on CIFAR-10, as well as the kernel distance between the final network’s kernel and the kernel after stage 1.
Additional results for different spawn epochs and for CIFAR-100 are reported in ?? . We see that in the frozen batchnorm setting, both SGD variants lose robust accuracy.
SGD does not freeze the kernel, we also see that the kernel changes too (as seen from the kernel distance), with SGD with a larger learning rate seeing a larger change.
Letting the batchnorm parameters change results in low adversarial and benign accuracy for all training models, with correspondingly higher kernel distances.
This interpretation of freezing batchnorm parameters to preserve the NTK provides a different perspective as to why more sophisticated batchnorm schemes improve performance in few-shot learning [12].
Adversarial Training on a Frozen NTK
In fig.
1, we observed that adversarial robustness could be inherited from a kernel in stage 1 of training, even if stage 2 was performed on benign data.
We now investigate whether we can gain robustness by performing adversarial training in stage 2 on top of a frozen kernel.
Much theoretical analysis of adversarial robustness has been in the frozen kernel regime or linear classification [22, 9, 60, 37, 68], as analysis of neural network training simplifies when the underlying feature set does not change.
For this experiment, we ran adversarial training using centered dynamics for three kernels.
First, we used the initialization kernel, that is, K t =0 .
Then we used the post-training kernels for networks
[Table]
trained either using benign training or adversarial.
As a baseline, we compare the accuracies achieved by adversarial training under standard dynamics.
These are shown in table 2, with
The initialization kernel fails to learn the dataset under adversarial training, achieving near-random benign accuracy and no meaningful adversarial accuracy.
Adversarial training sees a very limited gain in robust accuracy, improving only 2% , and most surprisingly, the benign training kernel sees a dramatic robustness gain, achieving 76 .
15% , surpassing the robustness of all other methods described in this paper.
Similar results are reported in ?? for CIFAR-100, with the adversarial kernel seeing little gain and the benign kernel surpassing the robustness of all other methods.
While we do not have a clear explanation as to why we see such a dramatic improvement, we do have a few speculative hypotheses.
Firstly, it has been argued that fitting global structure before fitting local structure can improve robust accuracy [16, 61].
Under this view, we can consider stage 1 training as learning global structure by choosing task predictive features and stage 2 as fitting local structure by performing adversarial training.
Secondly, separating the two tasks can reduce some of the instability caused by adversarial training [26, 72], as only the linear weights changed in stage 2.
From a theoretical perspective, adversarial training in linear classification is far easier to analyze, and thus this regime has been the interest of several papers [68, 22, 9].
However, in practice, these analyses have not been very practically useful as neural network behavior needs to deviate far from the linear ideal in order to achieve high standard accuracy [47].
As we have shown here, linearized training can achieve high benign and robust accuracy, provided that we use the learned NTK as opposed to the initial one.
This linearized regime can thus serve as a middle ground where many of the analysis tools of adversarial training are applicable while still retaining high practical performance.
Visualizing NTK eigenvectors
In the lazy dynamics regime, the resulting classifier is linear in its features given by ϕ ( x ) = ∇ θ ′ f θ ′ ( x ) (cid:12)(cid:12)(cid:12)
θ ′ =
θ t .
Furthermore, the effective classifier weights ( θ − θ t ) lie in the span of the space of the features of the training set: J =
[ ϕ ( x 1 ) , ϕ ( x 2 ) ...
ϕ
( x N )] T ∈ R N × P , with P the number of parameters.
The top eigenvectors of the NTK K =
JJ T correspond to the dominant features of the dataset in the NTK feature space.
While much work has looked at the properties of the NTK eigenvectors/value at initialization to understand neural network generalization or how the eigenspace evolves over training, here we visualize these eigenvectors directly.
It has been shown that not only are adversarial examples generated by robustly trained networks more visually interpretable but also that they make reasonable generative models by directly maximizing the class logits log p ( y
= c | f θ ( x )) w.r.t to the inputs x [20].
By visualizing the eigenvectors of the network directly, we can disambiguate whether this property is caused by the linear fit on top of the eigenvectors or directly from the feature space caused by adversarial training.
To visualize these eigenvectors, we describe the following algorithm.
First, we calculate the feature space representation of the top eigenvectors.
This corresponds to the first columns of V in the singular value decomposition of the training set features: J = U Σ V T .
Typically P > 10 6 and
N ∼ 10 5 , so calculating this SVD directly is prohibitive.
Instead, we note that Σ V T =
U T J , so the i th eigenvector v i ∝ J T u i (because Σ is a rectangular diagonal matrix), where u i is the i th eigenvector of K .
We can calculate K =
U Λ U T since N << P .
Calculating J directly is too expensive, so instead, we note that J T u i =
∂∂θ ( u Ti f θ ( X )) , i.e., we take derivatives of a weighted sum of the network on training samples, with respect to the parameters.
Now with the feature space representation of eigenvectors v i , we visualize them by optimizing x i to optimizing cos( v i , ϕ ( x i )) , giving us 2 N images, with the factor 2 arising from whether we maximize or minimizing the cosine similarity.
As this algorithm describes how to visualize eigenvectors for only scalar functions f , we chose the dog class of the neural network as the chosen function to visualize, however, in practice, we found that eigenvectors for other classes looked very similar.
As before, computing K is too expensive for the full dataset, so we pick a subset of 500 images from the training set to calculate the NTK.
The visualizations of the top 4 eigenvectors of the kernels generated for the initial NTK, the kernel after 100 epochs of benign training, and the kernel after 100 epochs of adversarial training are visualized in fig.
The initial NTK has visually meaningless eigenvectors which look indistinguishable from pure noise.
Meanwhile, the eigenvectors after benign training seem to distinguish the edges of CIFARclasses but still are largely uninterpretable.
Finally, the adversarial kernel results in the eigenvectors that clearly belong to different classes.
This observed difference between NTK eigenvectors of the adversarial and benign NTKs provides support to the “feature purification" mechanism described in
Allen-Zhu and Li [3], in which they show that adversarially trained networks provably learn robust features.
Higher ε regime
In the previous sections, we showed various properties of benign and adversarial finite NTKs under the ε = 4 / 255 adversarial budget.
However, it is known that the behavior of adversarial training can vary greatly depending on ε .
In this section, we verify that the results we observed for ε = 4 /hold for ε = 8 / 255 .
We repeat all experiments with ε = 8 / 255 .
fig.
5 shows the resulting kernel evolution plots for CIFAR-10, and the resulting performance of standard, linearized, and centered training dynamics.
Consistent with our previous observations, we notice that centered dynamics see an increase of robust accuracy over standard dynamics ( 49 .
15 ± 3 .
49% vs. 35 .
80 ± 0 .
76% ) and that 32 . 29 ± 8 .robust accuracy can be achieved from centered training on the benign kernel.
We observe the robust overfitting phenomenon, where the best adversarial accuracy occurs earlier in training.
Comparing fig.
1 to fig.
5, we observe that there is a slightly larger kernel distance between the benign kernel and the adversarial one with larger ε values and also that the adversarial kernel takes longer to converge with larger ε .
This motivates future work to look at how different ε values affect the stability and convergence of the finite NTK.
Additionally, in ?? and ? , we show that the results in section 6 and section 7 also hold for ε = 8 / 255 .
Discussions, limitations and Conclusions
The experiments in this paper showed that adversarial training results in a distinct and measurable change in the NTK compared to standard training.
While under both training regimes, the NTK converges quickly in the first few epochs of training, the adversarial NTK appears to have a high effective rank and lower mean kernel specialization compared to the standard training NTK.
Visually, the top eigenvectors in the adversarial NTK correspond to much more visually interpretable images compared to the standard NTK and the initialization NTK.
Furthermore, we showed that classifiers built on top of the adversarial Kernel retain significant amounts of adversarial accuracy, despite being trained without ever seeing adversarial examples.
The initial NTK results in classifiers with no adversarial robustness, and surprisingly, the standard training
NTK sees a dramatic increase in adversarial robustness over time.
We additionally demonstrated that a frozen feature set is necessary for retaining the full robustness aspect of these kernels, as either allowing the batchnorm parameters to vary or using SGD as opposed to linearized training in the second stage results in a drop in robust accuracy.
Our results strongly motivate further research utilizing the adversarial properties of linearized training we observed as a defense mechanism for adversarial attacks.
Furthermore, we note that linearized training costs approximately double the cost of benign training, however, adversarial training is significantly more expensive.
As shown in fig.
2, one can perform “kernel early-stopping,” where we switch to linearized training on clean data so that the remainder of training can be done without adversarial training.
This observation certainly motivates a more efficient way to gain adversarial robustness in practice, which future work can investigate further.
Limitations.
We performed our analysis on robustness to an L ∞ PGD adversary, which is an uncertified attack but was designed to fool networks under standard dynamics and is unproven against networks under centered or linearized dynamics.
More advanced state-of-the-art adversarial defense mechanisms were not studied in this paper, nor was the optimal tuning for hyperparameters for adversarial or linearized training.
Furthermore, while we presented many interesting empirical insights about adversarial training on kernels, there is a great opportunity for developing the theory of the observed interplay of adversarial training and NTK evolution.
Previous research had to impose limiting assumptions, e.g., fixed kernels on NTKs due to the nonlinear nature of neural nets [68, 22, 9].
We hope that our study of adversarial training in the frozen kernel regime with learned kernels could bridge the gap between these two extremes by having simple linear dynamics but still matching or exceeding the performance of standard networks.
As shown in section 7, the performance of adversarial training in the frozen feature regime can significantly exceed that of standard dynamics.
Moreover, based on our observations, we believe that centered and linearized dynamics can be used as a defense against adversarial examples.
Acknowledgments
This research has been funded in part by the Office of Naval Research Grant Number Grant N0001418-1-2830, DSTA Singapore, and the J. P. Morgan AI Research program.