NTK Surrogate Models for Explaining Neural Networks

Paper context

Paper title: Faithful and Efficient Explanations for Neural Networks via Neural Tangent Kernel Surrogate Models Abstract: A conference poster detailing Neural Tangent Kernel (NTK) surrogate models for explaining neural networks, featuring empirical intuition plots, attribution formulas, kernel function comparisons, and visualizations of similar training data for specific classes. Paper body (method & results): Accepted as a conference paper at ICLR 2024 FAITHFUL AND EFFICIENT EXPLANATIONS FOR NEU- RAL NETWORKS VIA NEURAL TANGENT KERNEL SUR- ROGATE MODELS Andrew Engel1 Zhichao Wang2 Natalie S. Frank3 Ioana Dumitriu2 Sutanay Choudhury1 Anand Sarwate4 Tony Chiang1,5,6 1Pacific Northwest National Laboratory 2University of California, San Diego 3Courant Institute, NYU 4Rutgers University 5University of Washington 6University of Texas, El Paso {andrew.engel,sutanay.choudhury,tony.chiang}@pnnl.gov; {zhw036,idumitriu}@ucsd.edu; nf1066@nyu.edu; ads221@soe.rutgers.edu ABSTRACT 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 intro- duced trace NTK the most consistent performer. Open source software allowing users to efficiently calculate kernel functions in the PyTorch framework is available here*. 1 INTRODUCTION Explainability remains a critical open problem for applications of deep neural networks (NNs) (Leavitt & Morcos, 2020). Explain-by-example techniques (Lai et al., 2021; Yang et al., 2020) have emerged as a major category of algorithms for explainability, including prototype examples (Chen et al., 2019), Deep K-Nearest Neighbors (Papernot & McDaniel, 2018; Wang et al., 2021; Dziedzic et al., 2022), and Representer Points (Yeh et al., 2018; Tsai et al., 2023). These techniques explain models by providing example(s) that capture model behavior on new data. Kernel functions (Alvarez et al., 2011) are a natural choice for building explain-by-example algorithms (Yeh et al., 2018); a kernel measures the similarity between individual data points via an inner product in a reproducing kernel Hilbert space (RKHS) (Hilbert, 1912; Ghojogh et al., 2021). A RKHS that faithfully represents a linearized NN feature space can be used in a kernel machine to explain (model) the NN decision as a weighted sum of similarities to training data. In this work, we investigate computationally efficient approximations to the empirical neural tangent kernel (eNTK), which is a kernel function motivated by advances in the theory of deep learning (Jacot et al., 2018). It is well established that NNs trained using gradient descent are equivalent to kernel machines (Schölkopf & Smola, 2002) with a kernel constructed from a sum over eNTK (Lee et al., 2020) computed at each gradient step (Domingos, 2020; Bell et al., 2023). Given this equivalence, we would like to evaluate the eNTK as the kernel function for an explain-by-example algorithm; however, computing eNTK is computationally expensive (Novak et al., 2022; Chen et al., 2022), *https://github.com/pnnl/projection_ntk 1 arXiv:2305.14585v5 [cs.LG] 11 Mar 2024 Accepted as a conference paper at ICLR 2024 so low computational cost approximations have been developed instead (Mohamadi & Sutherland, 2022). We are the first to define and evaluate one such approximate kernel, the trace neural tangent kernel (trNTK). Additionally, we build from the work of Park et al. (2023) to provide software to compute random-projection variants that can be computed and stored with lower time and memory cost over traditional eNTK. Using these approximations, we build low-cost and faithful surrogate models for neural network classifiers. Our methodology improves over the past evaluation of kernel surrogate models. We measure the faith- fulness of a kernel function by assessing how well a kernel generalized linear model (kGLM) (Hof- mann et al., 2007) correlates with the softmax probabilities of the original NN using a rank correlation. Previous evaluations relied on test accuracy (Mohamadi & Sutherland, 2022; Long, 2021), or having high similarity to the correct class (Hanawa et al., 2021), which are both flawed. Our approach and accompanying code-repository will allow users to evaluate how close their own NNs are to kernel machines in the PyTorch framework with limited overhead (Paszke et al., 2019). CONTRIBUTIONS We make three major contributions in this work: 1. We define and evaluate new kernel functions for faithful approximation of an underlying neural network; we are the first to analyze random projection variants that permit tuning the computational and memory expense of approximate eNTK. 2. We are the first to show that approximate eNTK kernel surrogate models are consistently correlated to the underlying neural network across experiments including ResNet18 on CIFAR10 and Bert-base on COLA. 3. We compare explanations of NN decisions generated from each kernel function through a data attribution strategy and through an explain-by-example strategy; this is the first such qualitative evaluation between approximate eNTK. RELATED WORK Surrogate Models for Explaining Neural Network Behavior. Recent work in explainable AI has focused on determining when NNs are exactly equivalent to other common ML algorithms (Lee et al., 2018; Balestriero & Baraniuk, 2018; Schmitz et al., 1999), including kernel machines. It has been shown that infinitely wide NNs are equivalent to a kernel machine with kernel function chosen as the neural tangent kernel (Jacot et al., 2018). These infinitely wide models, however, do not replicate the feature learning behavior seen in finite-width networks (Chizat et al., 2018; Yang & Hu, 2021; Wang et al., 2022). Subsequently, researchers turned to investigate properties of finite-width models with NTK computed at various checkpoints (Domingos, 2020; Bell et al., 2023) and/or after training (Long, 2021). This framework was used to explore inductive biases (Ortiz-Jiménez et al., 2021), feature learning (Radhakrishnan et al., 2022), learning dynamics (Fort et al., 2020; Atanasov et al., 2022), and adversarial faithfulness (Tsilivis & Kempe, 2023; Loo et al., 2022). Support vector machines (Vapnik, 1999) using eNTK or approximate eNTK kernels computed after training were shown to achieve the same test accuracy as the underlying NN (Atanasov et al., 2022; Long, 2021; Vyas et al., 2022; Mohamadi & Sutherland, 2022). Our work builds upon this by evaluating whether kernel machines can approximate the underlying neural network function itself, rather than simply reproduce the same test accuracy. Kernels for Explainability. Kernel functions defined from various RKHS have been proposed to explain the behavior of NN in different contexts, (Park et al., 2023; Koh & Liang, 2017; Pruthi et al., 2020; Akyürek et al., 2023), but in each of these works the kernel studied is loss-based and relies upon the availability of labels at inference time. We differ in that our goal is to model/explain the classification behavior on any new data, including unlabeled data where the loss is incalculable. Most relevant to our work, Yeh et al. (2018) (hereafter Representer Points) used a kernel formed from the NN final embedding in what we call the data attribution task (see section 2). We build from Representer Points by evaluating their assumptions under new approximate eNTK kernels. Computationally Feasible Approximations of the eNTK The computational cost of the eNTK is prohibitively high for large models and datasets. Advances on this issue have been two-pronged: 2 Accepted as a conference paper at ICLR 2024 Some groups focus on algorithmic improvements to calculate the eNTK directly (Novak et al., 2022). An alternative strategy has been to avoid eNTK calculation and instead compute kernel functions that share a similar structure to the eNTK (Mohamadi & Sutherland, 2022). One such approximate kernel was introduced quietly in Chen et al. (2022) which we refer to as the trace-NTK (trNTK). We are the first to explicitly investigate the trNTK’s properties. Finally, Park et al. (2023), hereafter TRAK, utilized random projection matrices to scale the computation of a loss-based kernel function. We modify TRAK to compute projected variants of approximate eNTK. Evaluating Kernel Attribution. In this paper, we use three evaluation strategies. The first focuses on evaluating the faithfulness of the surrogate model through rank correlation. The second evaluates surrogate model performance on a data-attribution task. We follow the methodology in Shan et al. (2022) to evaluate the model via precision and recall in tracing decisions on poisoned test data back to poisoned training data. Finally, we compare kernels qualitatively via explain-by-example. Previous work evaluated kernels through whether the attributions trace to training data of the correct class (Hanawa et al., 2021), whether surrogate models replicate NN test accuracy (Mohamadi & Sutherland, 2022; Long, 2021). These are insufficient: our goal is that kernel functions reflect the neural network behavior, but test accuracy is invariant to the specific classification on individual datapoints. Representer Points used Pearson correlation as a faithfulness measure, but Pearson correlation can conflate covariance with faithfulness (see Appendix H). We will demonstrate that our methodology is more secure measurement of faithfulness. 2 PRELIMINARIES Neural Networks for Classification. We consider the supervised classification problem with C classes. Consider a data input x ∈X ⊆Rn with n the dimensionality of inputs, and a one-hot encoded data label vector z ∈Z ⊆RC. We define a neural network F(x ; θ) : X →Y where the output space Y ⊆RC is an intermediary step in our classification called a “logit.” The NN F(x ; θ) is parameterized by the vector θ and was learned via back-propagation to minimize the cross entropy loss between the target label vector z and softmax probability vector σ(F(x ; θ)), with σ : Y →Z the softmax function. We denote the c-th scalar output of the network as F c. We interpret the predicted confidence for the c-th class for input x as σ(F(x ; θ))c. Kernel Functions. Kernel functions implicitly map the data vector x to a feature vector ρ(x) in a higher dimensional RKHS V for which the kernel function κ(·, ·) evaluates the inner product of two feature vectors in V. We will notate the data matrix X = [x1, . . . , xN] ∈RN×n with N the number of training samples. With some abuse of notation, we will write κ(x, X) ∈RN for the vector whose j-th component is κ(x, xj) and