Faithful and Efficient Explanations for Neural Networks via Neural Tangent Kernel Surrogate Models

1. We thank the reviewer for their time, and echo their excitement about the results. We believe the core strengths of the paper are: a) identifying a metric Kendall-$\tau$ that is a better measure of faithfulness over the previously used test accuracy in comparison of similarity between models, and b) presentation of the evidence we collect that challenge the assumption of sparsity in data attribution/explain by example explainability work. To build upon its significance and relate to the question posed below: The choice of the structure of the surrogate model can be an implicit assumption on the kind of data attribution that NN model performs. We chose to use a kGLM because it does not assume sparsity-- we could have in principle observed that the resulting kGLM was effectively sparse, but do not. In contrast, a choice like a KNN would be assuming that the data attribution is local in the feature space and sparse. We continue this discussion below.

2. Re: In section 5, You mentioned: "...an interesting follow-on work would investigate using kernel functions in K-Nearest Neighbors surrogate models." How much argument of this paper can transfer to KNN or generally any other surrogate models on explaining NN decision?

   • An excellent question that we invite the reviewer to discuss with us. We have conducted some preliminary work investigating the roles other surrogate models (kernel-SVM and kernel-KNN) can play in investigating the NN model. So far we understand that the choice of surrogate model can have profound impact on what we might infer. For example, with the kernel-KNN we would specifically be modeling that only the local structure of the feature space matters to the classification; in effect, assuming sparsity a priori. However, it is not yet clear how we should assign a class probability/logit value to utilize the Kendall-$\tau$ measure. We do include some analysis with kernel-SVMs as well looking into adversarial examples. We infer that kSVMs are not a good surrogate for modeling the effect of adversarial examples since the robustness scales differently between the NN and the kSVM. This may be in part due to the structure of the kSVM itself: in the kSVM, only a sparse number of data points near the decision boundary matter for classification (i.e., the support vectors). Perhaps it is easier for an adversarial example to fool the kSVM given fewer samples the kSVM is utilizing than the NN (under our hypothesis that the NN is not sparse).

3. Re: also, there are some of the minor typos in the paper

   • We thank the reviewer for pointing out the mistakes and will fix them in the revision

Good afternoon, we are still preparing our response, but today we realized we were not sure of the meaning of the following question. Could the reviewer please clarify? "Given the data attribution with kernels theory in page 3, wouldn’t you be able to test directly whether a kGLM is an “ideal surrogate” (according to eq 2) by comparing across all data points the NN confidence in each class with the data attribution for each class?"

1. We thank the reviewer for their detailed comments. Before we clarify some of the specific issues raised, we want to echo one of the main strengths the reviewer identifies in the paper. Part of our contributions are the implementation of projection variants of approximate eNTK. We show in Table 4 how these projection variants achieve a massive reduction in total computation time with exponentially decaying residuals to the non-projected variant (Appendix E). We are excited to see how researchers will build off this work to conduct eNTK research enabled by our software and TRAK.
2. Re: Contributions are unclear. For point (1) however the authors also state that the trNTK was introduced in Chen et al. 2021 (end of page 2)...
   • The full story is somewhat complex and in editing to meet page limits it became less clear. Chen et al. (2021) never formally defined the trNTK and never mentions it in their main body. They only describe trNTK in a single sentence in the appendix: (Appendix G, first sentence: "We approximate the neural tangent kernel on the i1k (1/10) subset by averaging over block diagonal entries ... in the full NTK."). Therefore, we feel our wording is accurate albeit slightly confusing: we are the first to formally define the trNTK. In Chen's work, the trNTK was a (hidden) computational convenience to enable their experiment for a single data set; in contrast, our work treats the trNTK to be the central object of study as an approximation of the NTK. We believe that this is a novel contribution but at the same time we should give proper credit to Chen et al. We will look to revise the contribution section and introduction to make this point clearer.
3. Re: ... and that the random projections approach is based on Park et al. 2023 (end of page 1).
   • We copy from our comment to all on this topic for the reviewer's convenience: we would like to take this opportunity to delineate our contributions from that of Park et al 2023 (TRAK). To begin, we want to acknowledge that TRAK is software at the foundation of our projection kernel functions, so it rightly takes a prominent role in our paper. The TRAK paper references the eNTK and devotes a section to show how it is an explicit part of the TracIn kernel, but its actual calculation is not implemented in their code. You can view the lines of code here
   • Notice that the model output functions implemented are only the loss (as this serves their use case for counterfactual analysis), rather than the neural network function itself. TRAK does not compute the eNTK or any approximation of the eNTK (p-, tr-, etc.). We have been in direct communication with the maintainer of TRAK to bring the capabilities we describe in this work to the TRAK software package.
4. Is the main contribution of the paper proposing a new NTK method or evaluating empirical exercises?
   • We view the paper as having three major contributions: a) We describe methods, building off recent successes from TRAK, that use projection matrices to compute efficient approximations to the eNTK. b) We establish the faithfulness metric Kendall-$\tau$ and detail how it is superior to previous methods relying on test accuracy. c) We identify the trNTK kernel function as the best performing surrogate model to the NN across a variety of experiments. Analyzing this choice of kernel in a data attribution context, we find that many training data points influence the output on any new data input.
5. Re: The authors consider different alternative NTKs to compare to the trNTK, but never compare in the main text the methods to the eNTK or the pNTK...
   • Our goal was to develop surrogate model methods that scale to large computer vision experiments and transformer models with potentially large number of classes. This implies we will need to find an approximation to the eNTK as the eNTK scales quadratically with the number of classes. A direct comparison to the eNTK is an important topic, but we did not consider it part of this paper's scope.
   • We understand the reviewer's concern that we did not compare with the pNTK. Our rational is that we "brute-forced" the computation of several trNTK and compared them with the proj-trNTK. After observing that the proj-trNTK had exponentially decaying residuals (Appendix E) with the trNTK, we felt secure in evaluating the proj-pNTK alone.

(continued below)

1. We thank the reviewer for their insightful comments: we also believe that one of the main strengths of our paper is our use of the Kendall-$\tau$ metric because it allows us to go beyond test accuracy measures of similarity. This allows us to identify the trNTK as the most faithful kernel function. This in turn allows us to show that data attribution is actually not sparse. That conclusion could help move the field, as previous works have assumed sparsity in various data attribution/explain-by-example techniques.
2. Re: Only the rank correlation of the softmax probabilities for the correct class is considered...
   • We agree with the reviewer that a current limitation of Kendall-$\tau$ is that it does not consider the logit value of incorrect classes. We think that an ideal faithfulness measure would include this information, we will add a discussion of Kendall-$\tau$ the limitation of Kendall-$\tau$ to our limitation section; however, we feel despite our use of Kendall-$\tau$ is a significant step away from previous works that rely on test accuracy to measure the similarity of surrogate functions to the neural network.
3. Re: ... However, to be faithful enough, the surrogate model should also behave similarly to the NN for the incorrect classes.
   • We agree with the reviewer's intuition. While this set of experiments have not appeared in the paper, we conducted an investigation into a "misclassification coincidence rate" as an additional line of evidence suggesting that the kGLM can mimic the behavior of the NN on misclassifications. The misclassification coincidence rate, $R_{\textrm{miss}}$, is the ratio of the number of times the kGLM and NN predict the same class out of all times either are incorrect over the number of times either the NN or the kGLM are incorrect

$

R_{\textrm{miss}} = \frac{ |\{ i : \textrm{kGLM}(\mathbf{x}_i) = \textrm{NN}(\mathbf{x}_i) \ne y_i \}| }{ | \{i : \textrm{kGLM}(\mathbf{x}_i) \ne y_i\} \cup  \{i : \textrm{NN}(\mathbf{x}_i) \ne y_i\}| }

$

* This rate is actually very high (about 75\% coincidence of misclassifications on ResNet18 for trNTK), which is powerful evidence suggesting that the kGLM can mimic the NN overall for both correct classification and misclassifications. We will include this analysis in the revision.

4. Re: An important application of data attribution is to explain why a NN makes a wrong prediction. This is not considered in the paper. * While our main body is primarily concerned on establishing our framework for data attribution and then an initial analysis re. the assumption of sparsity of the attribution scores, our framework can be extended to study misclassification. Our current status is to try and understand holistically from all the visualization generated why the NN made a misclassification by investigating the kGLM as proxy. Take as example figures 20, 31, and 42: In figure 20, we see that the mean plane and horse attribution score are highest driving the classification in their respective logits, but with the mean attribution from the remaining classes having higher negative impact on the plane logit. Figure 31 zooms into the horse logit for the trNTK, showing that the logit horse value is dependent upon the horse class. Even though our analysis shows that a sparse number of representations do not best describe the NN predictions, we can still leverage the most similar data points to help us form hypothesis of NN behavior. In figure 42, we show the most similar images are of people riding horses. This was quite the realization, since we had not noticed that the initial test image is actually of a person siting on top of an airplane! While we still must find a way to test this hypothesis, we believe that helping researchers form these hypothesis is a valuable tool in understanding misclassification. 5. Re: Eq. (4) is confusing. In the denominator, the sqrt is applied to the inner product. However, according to Appendix C and the definition of cosine similarity, the sqrt should be applied to the sum, not the inner product. * We thank the reviewer for pointing out mistake; we will add parenthesis to Eq (4) and to ensure the equation matches the definition of cosine similarity. 6. Re: The quality of Figure 2 could be improved. * We agree; there are rendering errors that we will fix in the revision and double check that boxes do not move in front of the axis labels.

(continued below)

1. We thank the reviewer for their thoughtful review. Part of the importance of our work is to explore the space of scalable approximations to the eNTK that retain intrinsic properties of the eNTK. Our hope is that our research will enable eNTK research on larger datasets and models. As an example of the analysis this software enables, we perform investigations into data attribution problems using kernel surrogate models on large computer vision and language models. We think the software will enable many new analyses and are excited to see how other researchers build off our work.
2. Re: Surfacing similar images is a not a meaningful evaluation of attribution.
   • We wholeheartedly agree with the limitations of surfacing similar images as the reviewer states --- in fact, this is a major conclusion we draw in the paper. See the end of section 5 paragraph 1: "Therefore, presenting the top highest attribution training images without the context of the entire distribution of attribution is probably misleading." Our kernel machine surrogate framework goes beyond simply surfacing similar images to show how the surrogate model, which is deeply correlated with the underlying neural network, leverages all of the training data to make classification on new test data.
3. Re: I'd be more careful about making any claims of data attribution ... as the paper does carry out any counterfactual evaluations.
   • We acknowledge the historical context of data attribution in classical statistics, and recent research in data attribution that have utilized counterfactual experiments; however, it is not the only sense in which "data attribution" is used. We subscribe to the framework given in Park et al 2023 (TRAK): "a data attribution method computes a score for each training input indicating its importance to the output of interest." Specifically, they cite the approach of Yeh et al 2018 (i.e. Representer Points) as a form of data attribution, and we mirror the same sense of data attribution in our work. We describe how we compute our data attribution score in great detail in the comment to all reviewers.
4. Re: Overall, the contributions seems somewhat marginal. Also, the fast approximate versions implemented primarily rely on prior work (Park et al.)'s implementation, so not sure there is much to claim as contribution there (since Park et al. also used it for faster approximations to eNTK).
   • We politely but strongly disagree: Park et al. (TRAK) neither compute the eNTK nor any approximation of the eNTK. While Park et al. do reference the eNTK and devote a section to show how it is a part the TracIn kernel, the actual calculation of the eNTK itself is not implemented in their code. You can view the relevant lines of code from TRAK here. Notice that the model output functions implemented are only the loss, rather than the neural network function itself. We have been in direct communication with the maintainer of TRAK to bring the capabilities we describe in this work to the TRAK software package.
   • It is true that our fast approximation would not be possible without TRAK (and so TRAK rightfully takes a prominent position in our paper); however, building off each other's work is a success story of the open-source community and is at the foundation of AI research. We see this as a major strength rather than a weakness.
5. Re. Writing ... the paper goes into more detail than necessary in defining rank correlation / R2, etc from scratch...
   • We thank the reviewer for pointing out the confusion. We agree that the space devoted to the Kendall-$\tau$ definition could be saved by moving it to the appendix. We would instead use this space to discuss the motivation for the Kendall-$\tau$, which we believe is more important than discussing why Pearson is flawed. If it interests the reviewer, we provide additional information in the comment to all reviewers about our argument for why we use Kendall-$\tau$.
6. Re: Confused by what the rank correlation is measured over exactly. I understand it's measured between the truth model outputs and surrogate model outputs, but what is it varied over? Are you measuring across different inputs x?
   • Yes: we are varying it over the different model inputs x from the test dataset. For cases beyond binary classification (single logit), we also must make a choice for which output of the neural network/ surrogate model to use. We chose to use the output corresponding to the correct class (as given by the ground-truth labels).

(continued in next comment)

I would like to know the answer to my Question 2. It seems like a surrogate model should be fit with the outputs of the model it is emulating, not with ground-truth labels. I doubt this change would alter the results in the paper. But I think the authors should consider making this chance, if feasible, for the final paper draft.

Good Morning,

We aim to post a thread of answers to each of your points today. To answer this specific question now: it is a very reasonable suggestion. The reason we did not attempt this is we were worried about overfitting the kGLM onto the NN's outputs, so we investigated using the ground truth labels first. After finding that the fits were very reasonable, we moved ahead. We were surprised that even without using the NN's outputs there would be such high alignment between the NN and the kGLM.

We now think that even if the model were to overfit onto the NN's outputs, this would become apparent when evaluating Kendall-$\tau$ on the test data, so overfitting would be observable and therefore not represent a great risk. In addition, we recently became aware of some contemporaneous work that was accepted at Neurips23 [Tsai et al 2023] that utilize the NN's outputs to derive ideal weights.

Due to all of this, we resolve to include this experiment in our final paper and compare to the existing results. We think it would add an interesting comparison to Tsai et al.'s methodology.

Tsai, C.-P., et al., “Sample based Explanations via Generalized Representors” 2023

I don't understand the concern with "overfitting". Wouldn't you expect a reasonable surrogate to at least match the NN exactly on training points?

I appreciate the authors' willingness to add experiments training the surrogate to directly emulate the NN. I look forward to reading the full rebuttal.

1. We thank the reviewer for their time and appreciate the breadth of topics covered. They have identified what we consider to be the major strength of this paper: bringing in new ideas to modeling surrogacy and data attribution: i.e., challenging both the use of test accuracy to measure the similarity of models as well as the a priori assumption that neural network classifications can be explained using a sparse number of training data exemplars.
2. Re The paper relies on previous work to establish credibility of attribution-based scores for neural network explanation.
   • We have expanded upon our argument to establish the credibility of our attribution scores in the comment to all (point 2). We resolve to present this in the revision so our work is more self-contained.
3. Re It doesn’t seem obvious that attribution is the same as similarity for learned kernel functions.
   • In Eq (3) we describe how attribution is the product of the kernel similarity with the learned weights, so is not just the similarity. We expand upon this in the comment to all reviewers as well (point 3), since it was a shared confusion. We will be more explicit in our revision.
4. Re I find the second sentence in the abstract confusing. I expected this trend to have to do with using kernel-based models for data attribution rather than to “investigate a diverse set of neural network behavior”. Isn’t the goal of your paper exactly to apply kernel models to investigate network behavior?
   • Our explicit goal is to use kernel methods as a surrogate for neural networks, and the application we focus on in this paper is data attribution. We agree that our attempt to remain general in the second sentence of the abstract reads overly pedantic for the sake of the narrative device used in the third sentence. We will address this sentence in the revision.
5. Re The 3rd experiment on qualitative evaluation of attribution is weak. A user study is probably beyond the scope of this paper, and I believe the work is strong enough to stand without such a study. However, the paper would significantly benefit from some discussions about how these attributions could be better qualitatively evaluated in the future.
   • We share a preference for stronger quantitative evaluation as well; however, we believe the qualitative evidence we have collected is still important to share because researchers often generate new hypotheses from looking at the data in a new qualitative way.
   • We seek to publish this result without human-subjects testing. We will add a short discussion of the limitations of explainability and caution the reader that a user-study must be performed in the revision.

(continued below)

3. How is "data attribution" different from "similarity"?
   • Multiple reviewers asked about the connection between our data attribution and "similarity" between data input $\boldsymbol{x}$ and training input $\boldsymbol{x}_i$. We point out that our definition of data attribution is not equivalent to similarity. We interpret $\boldsymbol{\kappa}(\mathbf x,\mathbf x_i)$ as measuring how similar $\mathbf{x}$ and $\mathbf{x}_i$ are. Our definition of the data attribution score is given in the draft as Eq 3 and above. Measured between a data input $\boldsymbol{x}$ and training datapoint $\boldsymbol{x}_i$, the attribution is:

$

     A(\mathbf x, \mathbf x_i)^c =  \mathbf W^c_i \boldsymbol{\kappa}(\mathbf x,\mathbf x_i)+\frac{\mathbf b^c}N.

$

* A kernel function $\boldsymbol{\kappa}(\mathbf x,\mathbf x_i)$ computes the inner product between $\mathbf{x}$ and $\mathbf{x}_i$ in an appropriately-defined vector space. While the similarity is encoded in the kernel function's value, $\kappa(\boldsymbol{x},\boldsymbol{x}_i)$. In short, the learned weights can and do moderate the similarity for the attribution score.

4. Why is a rank based correlation measure a good way to measure "faithfulness" and why do we specifically use Kendall-$\tau$? * We begin by considering three properties a faithfulness measure should exhibit and will explain how these are satisfied by non-parametric rank correlation like Kendall-$\tau$. * Recall our ideal objective is to evaluate to what degree a kernel machine surrogate model is equivalent to the NN at points of interest $\boldsymbol{x} \in \mathcal{X}$:

$

\text{NN}(\boldsymbol{x},\; \boldsymbol{\theta}) = \sum_{i=1}^N{W_i K(x,x_i)} + \boldsymbol{b}

$

* We only focus on whether the kernel machine is a good surrogate when $\boldsymbol{x}$ is drawn from the population $\mathcal{X}_{\text{test}}$. Because we do not know the distribution of the test population so we are forced to sample from it, leading to the first property. **The first property:** our metric should be a single number we can calculate from the samples that synthesizes the information across each of the samples.
* **The second property:** the faithfulness standardized so as to be comparable across different surrogate models. For example, it should have a bounded maximum. 
* **The third property:** the measure of faithfulness should be invariant and maximized for the following relationship over any invertible mapping function $\Phi$: 

$

\text{NN}(\boldsymbol{x}_i) = \Phi(kGLM(\boldsymbol{x}_i))

$

* Note that the assumption that $\Phi$ is an invertible map means that it preserves any monotonic relationship between $\text{NN}(\mathbf{x}_i)$ and $\text{kGLM}(\mathbf{x}_i)$. Invertible maps are injective maps. It is this fundamental injective property that makes the kGLM have ``one-to-one'' correspondence with the NN that we care about. 
*  These three properties suggest we use a rank based correlation. First, the correlation will reduce to a single number over the set of samples across which we evaluate. Second, the correlation will be bounded in [-1,1] so that values of correlation are easily comparable across kernels. Third, rank-based correlation are invariant under invertible transformations. 
* Both Kendall-$\tau$ and Spearman-$\rho$ satisfy these properties.  We chose to use Kendall-$\tau$ over Spearman-$\rho$ because Spearman-$\rho$ is generally larger than Kendall-$\tau$ and is observed to converge faster to 1 as the error between NN and KGLM approach 0 [Fredericks and Nelson, 2007]. We provide a notebook in the supplementary materials to verify this property. It is desirable to have our faithfulness measure go to 1 slowly so that our measurement distinguishes between surrogate functions over a larger range of errors.
* **PROPOSED CHANGE:** Because these motivations are not clearly elucidated in the draft, we would include a discussion about this in our revision. We will also add a discussion of the limitations of Kendall-$\tau$

(continued below)