Faithful and Efficient Explanations for Neural Networks via Neural
  Tangent Kernel Surrogate Models
A recent trend in explainable AI research has focused on surrogate modeling,
where neural networks are approximated as simpler ML algorithms such as kernel
machines. A second trend has been to utilize kernel functions in various
explain-by-example or data attribution tasks. In this work, we combine these
two trends to analyze approximate empirical neural tangent kernels (eNTK) for
data attribution. Approximation is critical for eNTK analysis due to the high
computational cost to compute the eNTK. We define new approximate eNTK and
perform novel analysis on how well the resulting kernel machine surrogate
models correlate with the underlying neural network. We introduce two new
random projection variants of approximate eNTK which allow users to tune the
time and memory complexity of their calculation. We conclude that kernel
machines using approximate neural tangent kernel as the kernel function are
effective surrogate models, with the introduced trace NTK the most consistent
performer. Open source software allowing users to efficiently calculate kernel
functions in the PyTorch framework is available
(https://github.com/pnnl/projection\_ntk).
Andrew Engel, Zhichao Wang, Natalie S. Frank, Ioana Dumitriu, Sutanay Choudhury, Anand Sarwate, Tony Chiang
# Robust Explanations for Deep Neural Networks via Pseudo Neural Tangent Kernel Surrogate Models

###### Abstract

One of the ways recent progress has been made on explainable AI has been via explain-by-example strategies, specifically, through data attribution tasks. The feature spaces used to attribute decisions to training data, however, have not been compared against one another as to whether they form a truly representative surrogate model of the neural network (NN). Here, we demonstrate the efficacy of surrogate linear feature spaces to neural networks through two means: (1) we establish that a normalized psuedo neural tangent kernel (pNTK) is more correlated to the neural network decision functions than embedding based and influence based alternatives in both computer vision and large language model architectures; (2) we show that the attributions created from the normalized \(\mathrm{pNTK}\) more accurately select perturbed training data in a data poisoning attribution task than these alternatives. Based on these observations, we conclude that kernel linear models are effective surrogate models across multiple classification architectures and that pNTK-based kernels are the most appropriate surrogate feature space of all kernels studied.

## 1 Introduction

Explainability remains a critical open problem for deep neural networks and the applications built upon them [23]. Explain-by-example [54; 35] has emerged as a major technique for generating explanations; this technique works by attributing decisions to specific training datapoints then visualizing the highest-attributed training data as exemplars of model behavior [21]. Kernel functions measure the similarity between individual data points via an inner product in a given feature space [2]. If a space can be found where higher-similarity datapoints are more influential on the decision making process, then we can use the kernel functions for the data attribution task. Thus, kernel-based approaches have become a natural choice for data attribution tasks and a wide range of feature spaces have been studied in literature [21; 1; 38; 27]; however, each feature space represents a different approximation of the neural network, so their comparative evaluation is a major topic of study [16]. In this work, we propose a methodology that achieves a higher-quality of approximation over the state-of-the-art methods and subsequently improves the performance for data attribution-based explanation-by-example tasks.

We measure the quality of a kernel-based feature representation by assessing how well a kernel generalized linear model (kGLM) [18] approximates the softmax probabilities of the original neural network using a rank-correlation score [38]. A faithful reconstruction of the softmax probabilities
confirms the kernel model's ability to serve as an accurate surrogate model to the neural network, so we can use the learned weighted attributions from our analytically simpler kGLM as the attributions for the complex neural network. Recent research on the neural tangent kernel (NTK) [20] shows that kernel models computed from finite-width architectures achieve comparable performance to the neural network itself [27; 49]; however, computing the full NTK is computationally intensive [33]. Approximations of the full NTK such as TRAK [38] seek to alleviate this problem by projecting the feature space of the NTK down to a smaller dimension at the cost of requiring many model instances to achieve a high-fidelity surrogate model. In contrast, we utilize an approximation referred to as the \(\mathrm{pNTK}\)[31] that allows us to analyze a single trained model instance with a single kernel [27]. Considering that it is common to use a single model instance with pre-trained weights (e.g., foundation models [7]), our approach brings in dual benefits: 1) we can explain the single model instance rather than the class of models that share the same architecture. Given the known variation between model instances [32], it is possible that different instances need different explanations. 2) Kernels for the most widely used foundation model instances could be pre-computed and disseminated to the community for analysis.

#### Contributions

We make three major contributions in this work:

1. We establish that kGLM based on a normalized pseudo Neural Tangent Kernel (pNTK) provides an approximation of deep neural networks including both computer vision and large language models.
2. We use the correlation between the kernel surrogate model and the neural network to evaluate different kernel functions as a feature space and demonstrate the superiority of pNTK-based kGLM approach.
3. We demonstrate that a pNTK-based kGLM offers higher performance for explanation-by-example tasks. We design an adversarial learning experiment and demonstrate that our approach can retrieve poisoned data samples with higher accuracy than competing methods.

#### Related Work

Surrogate models for explaining model behavior.Recent work in explainable AI has focused on determining when neural networks are exactly equivalent to other common ML algorithms [25; 5; 45], including kernel machines [20; 12]. It has been shown that infinitely wide neural networks are equivalent to a kernel machine trained with feature space chosen as the neural tangent kernel [20]. These infinitely wide models, however, do not replicate the feature learning behavior seen in finite-width networks [9; 55; 50]. Subsequently, researchers turned to investigating properties of finite-width models with NTK computed at various checkpoints and/or after training. This framework was used to explore inductive biases [34], feature learning [43], learning dynamics [14; 4], and adversarial robustness [47; 28]. Critically, kernel support vector machines (via the empirical NTK calculated after training) were also shown to achieve the same test accuracy as the underlying NN [4; 27; 49]. Our work extends this last thread by evaluating whether kGLMs approximate the underlying neural network itself.

Kernels for attribution.Kernel functions with different feature spaces have been proposed to explain the behavior of a neural network. Influence functions were introduced to perform data attribution [21], but require computing and inverting the Hessian, which is computationally expensive. As such, a scalable approximation (TraceIn [42; 1]) has since been developed. Other works have investigated data attribution via the "data counterfactual" task which predicts the behavior of a class of model when a subset of training data is removed from the dataset as a measure of the importance of those training data to the model [19]. Park et al. [38] (hereafter Trak ) showed that many instances of models could be combined with a scalable NTK approximation to perform the data counterfactual task. We report performance comparison with Trak [38] and TraceIn [1] as a choice of kernel space and demonstrate the superiority of the \(\mathrm{pNTK}\)-based approach in our single model instance regime.

Evaluating kernel attribution.In this paper, we use two evaluation strategies. The first focuses on evaluating the accuracy of the surrogate model, while the second evaluates its performance on data-attribution tasks. For the former, our primary evaluation methodology uses rank-based correlation,
which was suggested in contemporaneous work for evaluating data counterfactual models [38]. For the data attribution task, we follow the methodology in [46] to evaluate the model via precision and recall in tracing decisions on poisoned test data back to poisoned training data. Previous work on kernel evaluation on data attribution did not use a surrogate model-based approach; instead, they relied upon whether the attributions trace to training data of the correct class [16]. We believe that the surrogate model framework provides a superior approach because the surrogate model approximates the neural network behavior, e.g., predicting both the same correct or incorrect class as the neural network point for point.

## 2 Background

Notations.We use boldface text to denote vectors (lower case) and matrices (upper case). The Hadamard (element-wise) product is written as \(\odot\), and unless indicated otherwise, the norm \(\left\|\cdot\right\|\) is the Euclidean norm.

Neural networks.We consider the supervised classification problem where a neural network denoted as \(F(\mathbf{x}\,;\mathbf{\theta})\) parameterized by the vector \(\mathbf{\theta}\) is learned via iterative first-order optimization to minimize the cross entropy loss function between a target label vector and softmax probability vector evaluated on a set of data inputs \(\mathbf{x}\) and labels \(\mathbf{z}\subseteq\mathbb{R}^{C}\). We represent the input as the matrix \(\mathbf{X}=[\mathbf{x}_{1},\dots,\mathbf{x}_{N}]\in\mathbb{R}^{N\times p}\), with \(p\) being the original dimensionality of inputs, and N is the number of training samples. We denote the \(c\)-th scalar output (i.e., the \(c\)-th output neuron) of the network by \(F^{c}\). The output vector of the neural network (i.e., the final activations) are fed to the softmax function \(\sigma\):

\[\sigma(F(\mathbf{x}\,;\mathbf{\theta}))=\frac{\exp(F^{c}(\mathbf{x}\,;\mathbf{\theta}))}{\sum_ {c=1}^{C}\exp(F^{c}(\mathbf{x}\,;\mathbf{\theta}))}.\]

The output of the softmax function has property \(\sum_{c=1}^{C}\sigma(F^{c}(\mathbf{x}\,;\mathbf{\theta}))=1\), so we interpret it as the probability mass function where \(\sigma(F^{c}(\mathbf{x}\,;\mathbf{\theta}))^{c}\) represents the probability of class c. For all the analysis in this paper, we evaluate the network function and calculate kernels at the final trained weight vector.

Kernel methods.Kernel methods implicitly map the data vector \(\mathbf{x}\) to a feature vector \(\rho(\mathbf{x})\) in a higher dimensional Hilbert space \(\mathcal{V}\) for which the kernel function \(\mathbf{\kappa}(\mathbf{\cdot},\mathbf{\cdot})\) evaluates the inner product in \(\mathcal{V}\). With some abuse of notation, we will write \(\mathbf{\kappa}(\mathbf{x},\mathbf{X})\in\mathbb{R}^{N}\) for the vector whose \(j\)-th entry is \(\mathbf{\kappa}(\mathbf{x},\mathbf{x}_{j})\) and \(\mathbf{\kappa}(\mathbf{X},\mathbf{X})\in\mathbb{R}^{N\times N}\) for the matrix whose \((i,j)\)-th entry is \(\mathbf{\kappa}(\mathbf{x}_{i},\mathbf{x}_{j})\).

## 3 Kernel-based surrogate models

Given an input \(\mathbf{x}\) and a neural network \(F(\mathbf{x}\,;\mathbf{\theta})\) as defined above, our surrogate modeling goal is to learn a kernel-based approximation function denoted as \(\mathrm{kGLM}(\mathbf{x})\) such that \(\sigma(\mathrm{kGLM}(\mathbf{x}))=\sigma(F(\mathbf{x},\mathbf{\theta}))\). In the following sub-sections, we first describe the framework to learn \(\mathrm{kGLM}\left(\cdot\right)\) in Section 3.1. We introduce pNTK and other kernel functions in Sections 3.2 and 4.3. Finally, we introduce the functions for evaluating the approximation quality of \(\mathrm{kGLM}\left(\cdot\right)\) in Section 3.3.

### Kernel general linear models

Given an input sample \(\mathbf{x}\), we define the approximation provided by kGLM as \(\mathrm{kGLM}(\mathbf{x})=\mathbf{W}\mathbf{\kappa}(\mathbf{x},\mathbf{X})+\mathbf{b}\), where \(\mathbf{W}\in\mathbb{R}^{C\times N}\) is a learnable weight matrix, \(\mathbf{\kappa}\) is the kernel function, and \(\mathbf{b}\in\mathbb{R}^{C}\) is a learnable bias vector.

We make decisions from kGLM by mapping the final activations computed to softmax confidences, just as we do for a standard neural network-based classification task. The parameters \(\mathbf{W}\) and \(\mathbf{b}\) are iteratively updated using an optimizer to minimize the cross entropy loss on the same training dataset upon which the neural network is trained. This setup extends to multi-classification tasks as well; the decision function for multi-classification tasks is commonly computed from the softmax vector, so it is sufficient that the softmax vector computed from the kGLM is equal to the softmax vector computed from the neural network.
For any choice of \(\mathbf{\kappa}\) and trained kGLM, the attribution for decision class \(c\) for a test input \(\mathbf{x}\) to a training datapoint \(\mathbf{x}_{i}\) is given by \(A(\mathbf{x},\mathbf{x}_{i})^{c}=\mathbf{W}_{c,i}\ \mathbf{\kappa}(\mathbf{x},\mathbf{x}_{i})+\frac{\mathbf{b}_{c}}{N}\). The \(\frac{\mathbf{b}_{c}}{N}\) term is necessary to ensure that the sum over the attributions for the entire training dataset is equal to the kGLM's final activation for class \(c\): \(\sum_{i=1}^{N}A(\mathbf{x},\mathbf{x}_{i})^{c}=\mathrm{kGLM}(\mathbf{x},\mathbf{X})^{c}\). Note that if the kGLM is a perfect surrogate model, the softmax function applied to the vector created from each class attribution will equal the neural network confidence in each class. Consequently, we will have decomposed the reasoning for the neural network's specific confidence in each class to a linear combination of similarities between \(\mathbf{x}\) and each training datapoint in \(\mathbf{X}\). The next sub-section describes the kernel function used to compute this similarity measure. It is _critically important_ for the similarity measure to be model-centric. For example, given two images \(\mathbf{x}_{i}\) and \(\mathbf{x}_{j}\), the measure should capture the similarity of their respective classification by accounting for \(\mathbf{\theta}\) rather than simply raw data level similarity between the images.

### Psuedo neural tangent kernel

We define the kernel function \(\mathrm{pNTK}\) as follows: for any data samples \(\mathbf{x}_{i}\) and \(\mathbf{x}_{j}\), define the Jacobian of the neural network with respect to \(\mathbf{\theta}\) at datapoint \(\mathbf{x}_{i}\) by

\[\mathbf{G}(\mathbf{x}_{i};\mathbf{\theta})=\sum_{c=1}^{C}\nabla_{\mathbf{\theta}}F^{c}(\mathbf{x} _{i};\mathbf{\theta}). \tag{1}\]

Then the \(\mathrm{pNTK}\) evaluated at datapoints \(\mathbf{x}_{i}\) and \(\mathbf{x}_{j}\) is given by

\[\mathrm{pNTK}(\mathbf{x}_{i},\mathbf{x}_{j})=\frac{\langle\mathbf{G}(\mathbf{x}_{i};\mathbf{ \theta}),\mathbf{G}(\mathbf{x}_{j};\mathbf{\theta})\rangle}{\|\mathbf{G}(\mathbf{x}_{i};\mathbf{\theta })\|\|\mathbf{G}(\mathbf{x}_{j};\mathbf{\theta})\|}. \tag{2}\]

We provide more details about these definitions in Appendix B, where we also compare the \(\mathrm{pNTK}\) to the limiting NTK. We explore the intuition behind this choice of kernel in Figure 1. The normalization in the denominator of Eq. (2) makes the \(\mathrm{pNTK}\) a kernel of cosine-similarity values. It has been suggested that this normalization helps smooth out kernel mass over the entire training dataset [1].

### Evaluating the fidelity of surrogate models

Finally, we introduce the measures to evaluate approximation quality of the surrogate model. Specifically, we elaborate on our use of the Kendall-\(\tau\) rank correlation measure to model the pair-wise

Figure 1: **Geometric intuition behind the pNTK.** A neural network function is evaluated at two points creating surfaces \(F(\mathbf{x}_{i}\,;\mathbf{\theta})\) and \(F(\mathbf{x}_{j}\,;\mathbf{\theta})\). These surfaces are shown with a tangent hyper plane at the same point (\(\mathbf{\theta}\)) in parameter space coinciding with the end of training. The Jacobian vector defines the tangent hyperplane’s orientation in parameter space. The \(\mathrm{pNTK}\) is a kernel whose \((i,j)\)-th element is the cosine angle between averaged Jacobian vectors. The more similar the local geometry between \(\mathbf{x}_{i}\) and \(\mathbf{x}_{j}\), the higher the value of \(\mathrm{pNTK}_{i,j}\).
relationship between the softmax outputs returned by the neural network and it's surrogate model on a sample-by-sample basis.

Kendall-\(\tau\) rank correlation.For a sequence of observations \(S_{\tau}=\{(x_{1},y_{1}),\ldots,(x_{N},y_{N})\}\) we say that a pair \((x_{i},y_{i})\) and \((x_{j},y_{j})\) with \(i\neq j\) are concordant if either both \(x_{i}>x_{j}\) and \(y_{i}>y_{j}\) or \(x_{i}<x_{j}\) and \(y_{i}<y_{j}\). Otherwise, the pair is discordant. We count the total number of concordant and discordant pairs, respectively denoted as \(\mathrm{NC}\) and \(\mathrm{ND}\). Then, Kendall-\(\tau\) (hereafter just \(\tau\)) is defined: \(\tau(S_{\tau})=\frac{(\mathrm{NC}-\mathrm{ND})}{\mathrm{NC}+\mathrm{ND}}\).

To assess the viability of a surrogate model, we compute the \(\tau\) between the softmax probability of the neuron representing the correct class, \(\sigma(F(\mathbf{x}\,;\mathbf{\theta})^{c})\), and the kGLM softmax probability for the output representing the correct class, \(\sigma(\mathrm{kGLM}(\mathbf{x}))^{c})\).

Kendall-\(\tau\) is used for a few reasons. First, \(\tau\) has a range \([-1,1]\) with \(\pm\) 1 representing a monotonic relationship and a value of 0 representing no correlation. Furthermore, if the relationship between the kGLM and neural network is strictly monotonic, then an invertible mapping function exists between the kGLM softmax probabilities and the neural network's [6]. In Appendix G, we demonstrate how to find these mapping functions with iterative optimizers [48]. Fitting and then applying these invertible mapping functions makes the kGLM linearly correlated to the neural network. Finally, rank-correlations have been suggested by contemporaneous groups [38], making this choice consistent across our methodologies.

Test accuracy differential (TAD).Previous works evaluated whether kernel models could replicate the performance of underlying neural networks [3; 26; 27; 49]. Having the same test accuracy as the underlying neural network is a necessary but insufficient condition of being equivalent with the neural network, which gives us a preference for \(\tau\). Nonetheless, we track the test accuracy differential, or TAD, as given by \(\mathrm{TestAcc}_{\mathrm{kGLM}}-\mathrm{TestAcc}_{\mathrm{NN}}\) to demonstrate that kGLM have similar performance to the underlying neural network. A value of \(0\) is preferred.

## 4 Experiments

We pursue the following research questions:

1. Does the pNTK-based \(\mathrm{kGLM}\) provide a robust approximation of the decisions by the target neural network (Section 4.2)?
2. How does the pNTK-based \(\mathrm{kGLM}\) compare with other kernel functions? We study the correlation between surrogate models and their underlying neural networks in Section 4.3.
3. How does the pNTK-based \(\mathrm{kGLM}\) compare with other kernel functions on data-attribution tasks? We use a well-established poisoned dataset and evaluate each surrogate method based on their ability to retrieve the perturbed samples in Section 4.4.

### Experimental setup

Classification neural networks with architectures and datasets (MNIST [24], FMNIST [53], CIFAR [22], and COLA [51]) shown in Table 1 are trained using standard techniques. Details such as specifics of architecture and choice of hyperparameters are available in Appendix F. Models that have a value of more than 1 in the column '# Models' in Table 1 are trained multiple times with different seeds. The ResNet18 [17], ResNet34, and MobileNetV2 [44] models were trained by an independent research group with weights downloaded from an online repository [41]. Bertbase [11] weights were downloaded from the HuggingFace [52] repository then transferred onto the COLA dataset, as is common practice for foundation models [7]. After training, we calculate the pNTK and alternative kernels using PyTorch and automatic differentiation [39]. We train a kGLM (sklearn.SGDclassifier) [40] for each \(\mathbf{\kappa}\) using the same training dataset for training the neural network model.

### Robust surrogate modeling via pNTK

We calculate the correlation between the surrogate model and underlying neural network via \(\tau\) and report the results in Table 1. Across each experiment, we observed that the kGLM with choice
of \(\mathbf{\kappa}=\mathrm{pNTK}\) achieves similar performance as the underlying neural network. We find that the efficacy of our surrogate model as measured by the correlation to the neural network changes depending on architecture and dataset; remarkably, \(\tau\) is consistently high with a lower bound value of 0.7 across all experiments indicating relatively high fidelity. We measure a statistically significant decrease in the test accuracy when using the kGLM compared to the neural network across all but one of the experiments. The largest of these decreases in performance is only 2.2%. To visualize whether we can achieve a point-for-point linear realization of the neural network, we learn a non-linear mapping from the kGLM to the neural network (Figure 2 for Bert-base while other visualizations and details available in Appendix G).

### Comparison of surrogate model viability between kernels

For ResNet18 and Bert-base models, we evaluate our choice of \(\mathrm{pNTK}\) against alternatives described below, reporting \(\tau\) and test accuracy differential in Table 2.

The formal definition of each alternative is available in Appendix A and we briefly introduce them here. Trak [38] is computed from the gradients similar to the NTK, except that the authors introduce a random projection matrix over the parameter dimension to speed up the computation. TraceIn [1] is an approximation to the influence function [21] that has become popular for data attribution; TraceIn is calculated from the gradients of the loss function, and therefore relies on the availability of the test label. To evaluate the effect of our normalization in the \(\mathrm{pNTK}\), we compare to the unnormalized pNTK [31] (\(\mathrm{pNTK}^{0}\)), which is simply the numerator of (2). The Conjugate Kernel (CK) is the kernel

\begin{table}
\begin{tabular}{l l l l l} \hline \hline Model (Dataset) & \# Models & NN test acc (\%) & TAD (\%) & \(\tau\) \\ \hline MLP (MNIST2) & 100 & 99.64(1) & +0.03(5) & 0.708(3) \\ CNN (MNIST2) & 100 & 98.4(1) & -0.2(2) & 0.857(7) \\ CNN (CIFAR2) & 100 & 94.94(5) & -2.1(5) & 0.711(3) \\ CNN (FMNIST2) & 100 & 97.95(4) & -2.2(2) & 0.882(3) \\ ResNet18 (CIFAR10) & 1 & 93.07 & -0.28 & 0.776 \\ ResNet34 (CIFAR10) & 1 & 93.33 & -0.29 & 0.786 \\ MobileNetV2 (CIFAR10) & 1 & 93.91 & -0.4 & 0.700 \\ BERT-base (COLA) & 4 & 83.4(1) & -0.1(3) & 0.78(2) \\ \hline \hline \end{tabular}
\end{table}
Table 1: **Choice of \(\mathbf{\kappa}=\mathbf{pNTK}\) robustly forms a viable surrogate model of underlying neural network.** We report the leading digit of error as a parenthetical, when available. We perform each experiment with ‘# Models’ independent seeds. For each model and dataset we train and extract the \(\mathrm{pNTK}\), train a kGLM, then calculate and report the \(\tau\) correlation between the kGLM softmax probability and neural network softmax probability for the correct class. The neural network test accuracy column shows that training terminates with a highly performant model, and the test accuracy differential (TAD) column demonstrates that the kGLM achieve comparable performance to the underlying neural network.

Figure 2: **Linear Realization of Bert-base Model.** Each panel shows an independent linearization of a Bert-base transfer model. An invertible mapping is fit between the kGLM and neural network to transform the kGLM’s final activations to the neural network’s, described in Appendix G. The coefficient-of-determination (or \(R^{2}\)) is used to assess the goodness-of-fit; with values near 1, the linearized kGLM is a good approximation for the neural network itself.
created from the Gram matrix of the embedding computed by the last hidden layer [13]. Finally, the Embedding Kernel [1] (Em) is computed from the normalized sum of hidden layer embeddings and forms an apt comparison when considering the widespread use of frozen weight embeddings, e.g., in transfer learning paradigms [10].

We find that our choice of the normalized \(\mathrm{pNTK}\) outperforms these comparison models. Because the \(\mathrm{pNTK}\) technique achieves the highest \(\tau\) across each experiment we argue it forms the best feature space representation of the kGLM. Furthermore, because we observe high correlation across many models and datasets, we believe this methodology forms a robust approximation to the supervised classification network's decision functions. Finally, we visualize the distribution of attributions for a single test input for the ResNet18 model (Figure 3). The specific point visualized is incorrectly classified by the neural network as well as each of the \(\mathrm{pNTK}\), \(\mathrm{pNTK}^{0}\), embedding, and conjugate kernel surrogate models, with the same class prediction as the neural network.

\begin{table}
\begin{tabular}{l l l l l l l l} \hline \hline Exp Name & Metric & \multicolumn{6}{c}{\(\boldsymbol{\kappa}\)} \\ \cline{3-8}  & & \(\mathrm{pNTK}\) & \(\mathrm{pNTK}^{0}\) & \(\mathrm{Trak}\) & \(\mathrm{TraceIn}\) & CK & Em \\ \hline \multirow{2}{*}{ResNet18} & \(\tau\) & **0.776** & 0.662 & -0.143 & 0.374 & 0.632 & 0.763 \\  & TAD (\%) & -0.3 & -0.52 & -83.07 & +6.66 & -0.32 & **-0.20** \\ \hline \multirow{2}{*}{Bert-base} & \(\tau\) & **0.78(2)** & 0.6(1) & 0.368(7) & 0.74(2) & 0.32(9) & 0.74(2) \\  & TAD (\%) & **+0.1(2)** & +0.6(1) & +16.54(8) & +4.9(9) & **-0.1(1)** & **-0.3(4)** \\ \hline \hline \end{tabular}
\end{table}
Table 2: **Comparison across surrogate feature spaces.** If available, we report leading digit of error as a parenthetical. Across benchmark models in both CV and NLP, we observe that the \(\mathrm{pNTK}\) forms surrogate models with the highest correlation to the underlying neural network decision function (measured by \(\tau\)) and is furthermore consistent in replicating the performance of these networks (TAD nearly 0). In contrast, the popular influence-based alternative Traceln outperforms the neural network, despite being a linear model, and contemporary work Trak is no better than random guessing for ResNet18, or outperforms when computed with 4 model instances on Bert-base.

Figure 3: **Utilizing the structure of the linear surrogate model to visualize neural network model decision making. Using modified boxplots we compare the distribution of attributions computed from each kernel. The sum of the points in each distribution equals the kGLM activation for that class; therefore, the mean value (black bar) is a good measure of the confidence that the surrogate model will have in that class (see Appendix H for details). For the specific test image shown (left), the neural network misclassifies the image as a horse when its true class is plane. Good surrogate models should similarly misclassify this point, which we observe for the \(\mathrm{pNTK}\), \(\mathrm{pNTK}^{0}\), Embedding and CK.**
### pNTK forms robust attributions in data poisoning regime

Next, we train a 21-layer CNN (details available in Appendix F.1.5) using BadNet CIFAR10 data [15; 46]. We perturb training data by placing a yellow square in the image and modify the label of these perturbed images to a target label (see Figure 4). At test time, perturbed data tricks the model into producing labels of the targeted class. We train a model on this poisoned dataset, compute each kernel function, and produce data attributions. We compare models with precision and recall (see Appendix D) of recognizing poisoned test data by attributing affected decisions correctly to poisoned training data (voted on by a five member committee formed by the highest attributed training data exemplars), and by computing \(\tau\) between the surrogate model and the poisoned model. This experimental procedure is similar to the evaluation used in Shan et al. [46]; there, the authors specifically called for evaluating different choices of kernel for this task, which we now complete.

We report the results in Table 3, and a comparison of attributions from each kernel is given in Figure 4 with more visualizations available in Appendix I. We find that both normalized and unnormalized \(\mathrm{pNTK}\) attributions perform best on precision and recall, but the normalized \(\mathrm{pNTK}\) continues to be most correlated to the neural network. In Figure 4 we visualize the top few exemplars we used in our committees for determining perturbed test data points.

## 5 Summary and conclusions

Impact of linear surrogate modeling for explainability.We have shown evidence supporting the choice of the \(\mathrm{pNTK}\) as a robust feature space for a kernel surrogate model. Our choice of a linear surrogate model model allows us to separate the terms from each training datapoint to form an attribution (Section 3.1). Note that while the choice of the linear model implies every training datapoint will contribute some weight to the decision, we observed that the highest attributed images from the \(\mathrm{pNTK}\) have relatively small mass compared to the bulk contribution (see Figure 3), suggesting that the central tendency of the distribution, rather than a few exemplar outliers, are the main source of decision making. We contrast this to visualizing the few most similar datapoints for attribution (as in Figure 4). Presenting the top most similar training images without the context of how much those images actually influence the decision compared to the bulk of training images may be misleading to human interpreters.

Contrast with alternative kernels for surrogate models.Our results consistently showed that the TraceIn kernel generates a better-performing model than the neural network itself but was typically less correlated to the model than the \(\mathrm{pNTK}\), CK, or Embedding kernels. It is worth re-emphasizing that the TraceIn function utilizes the test labels in the computation of attribution. We believe this results in high attributions to training data of the leaked correct class. When attempting to explain the neural network behavior, we believe it's better to have a surrogate model that consistently replicates the neural network decisions rather than optimizing for accuracy. We also observed that Trak did not perform well as a feature space. This is likely because Trak was designed to be calculated with many model instances, whereas we study a single model instance paradigm. Trak's performance in this regime was observed by the authors of the original paper (cf. Figure E.1 in [38]). Because we focus on different regimes and use cases, we view our works as largely complimentary of each other.

\begin{table}
\begin{tabular}{l l l l l l l} \hline \hline Method & Precision (\%) & Recall (\%) & \(\tau\) & TAD (\%) & poi. \(\tau\) & poi. TAD(\%) \\ \hline \(\mathrm{pNTK}\) & 99.99 & 100.00 & **0.643** & **+0.45** & **0.569** & **+0.09** \\ \(\mathrm{pNTK}^{0}\) & 99.99 & 99.97 & 0.344 & +0.87 & 0.125 & +0.13 \\ TraceIn & 99.89 & 99.96 & -0.197 & +9.35 & 0.021 & -7.04 \\ Trak & 0.00 & 0.00 & 0.007 & -72.35 & 0.046 & +0.14 \\ Embedding & 99.71 & 100.00 & 0.430 & -2.73 & 0.261 & -13.98 \\ CK & 1.65 & 50.61 & 0.552 & -3.50 & 0.454 & -81.25 \\ \hline \hline \end{tabular}
\end{table}
Table 3: **Poisoned data attribution forensics. We compare the attributions from different kernel functions. Notably, we do not the attribution experiment to be very discriminating between the \(\mathrm{pNTK}\), \(\mathrm{pNTK}^{0}\), TraceIn and Embedding kernels; however, our normalized pNTK maintains the highest correlation to the underlying poisoned neural network.**
Our quantitative experiments consistently show the \(\mathrm{pNTK}\) as more correlated to the neural network model compared to the unnormalized \(\mathrm{pNTK}\), Embedding kernel, and CK. We also observe qualitative differences between each kernel's attributions (Figure 3) including what training datapoints have highest attribution (Figure 4). These differences emphasize the need for this study: given each choice of kernel provides a different explanation, we need to know which kernel is the right choice. No kernel is observed to give a perfect surrogate model, so we relax our question to investigate which kernel is the best fit. Based on the rank correlation \(\tau\), the \(\mathrm{pNTK}\) is empirically the best fit.

LimitationsOur main results studied \(\mathrm{pNTK}\) in kGLM under standard supervised classification and data poisoning techniques; however, previous works have focused on support vector machine (SVM) linear kernel surrogate models instead. We think its likely that the limitations observed for SVM surrogate models extend to the kGLM models. We know of two such limitations. We found that SVM surrogate models fail to replicate neural network behavior under gradient-based adversarial attacks Appendix E. In addition, previous researchers have found that SVM surrogate models do not have the same scaling relationships as underlying neural networks [49]. Future work can endeavor to push the scope of surrogate modeling to these new regimes.

As a limitation of our direct scope, calculating the \(\mathrm{pNTK}\) for a large model with many training datapoints remains computationally expensive. The dimension in the \(\mathrm{pNTK}\) matrix is the size of the training data squared. As the number of training data grows analyses of models using the methods in this study will become impractical. Future work could attempt to address this by finding a smaller subset of the training data that still forms a good surrogate model.

Figure 4: **Visualizing poisoned data attributions.** We compare each kernel method in a data poisoning regime. On the left side of the figure, the top (bottom) image is the unperturbed (poisoned) input. The network outputs show the decision is altered by the attack. On the right, we show the top 3 attributions of each kernel on the unperturbed (poisoned) input. An effective feature space ought to capture when the model is affected by the attack. Qualitatively, we see there are similarities and differences between the conceptual features in the attributions for each kernel. For example, many of the dogs and cats attributed to are looking out of the screen (or “at the camera”) like the test image. The embedding kernel seems overall more sensitive to the background pixel color than the other kernels. We believe this is a fault of the embedding kernel because the background pixel color should not factor heavily into the classification decision on the subject of the image. We include more examples in Appendix I.
## Acknowledgments and Disclosure of Funding

A.W.E., Z.W., S.C., N.F., and T.C. were partially supported by the Mathematics for Artificial Reasoning in Science (MARS) initiative via the Laboratory Directed Research and Development (LDRD) Program at PNNL and A.D.S and T.C. were partially supported by the Statistical Inference Generates kNowledge for Artificial Learners (SIGNAL) Program at PNNL. PNNL is a multi-program national laboratory operated for the U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract No. DE-AC05-76RL0-1830.
