SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Here, we propose a more scalable variant we term tracking representers that accumulates changes in the global importance during kernel gradient descent updates when solving Eqn.4). Let  $\varPhi:\mathbb{R}^d\mapsto\mathcal{H}$  be a feature map corresponding to the kernel  $K$ , so that  $K(\mathbf{x},\mathbf{x}^{\prime}) = \langle \varPhi (\mathbf{x}),\varPhi (\mathbf{x}^{\prime})\rangle$ . We can then cast any function in the RKHS as  $f_{K}(\mathbf{x}) = \langle \theta ,\varPhi (\mathbf{x})\rangle$  for some parameter  $\theta \in \mathcal{H}$ . Suppose we solve the unregularized projection problem in Eqn. via stochastic gradient descent updates on the parameter  $\theta$ :  $\theta^{(t)} = \theta^{(t - 1)} - \frac{\eta^{(t)}}{|B^{(t)}|}\sum_{i\in B^{(t)}}\nabla_{\theta}\mathcal{L}(f_{\theta}(\mathbf{x}_{i}),f(\mathbf{x}_{i}))\varPhi (\mathbf{x}_{i})|_{\theta = \theta^{(t - 1)}}$ , where we use  $B^{(t)}$  and  $\eta^{(t)}$  to denote the minibatch and the learning rate. The corresponding updates to the function is then given by "kernel gradient descent" updates:  $f_{K}^{(t)}(\mathbf{x}) = f_{K}^{(t - 1)}(\mathbf{x}) - \alpha_{it}K(\mathbf{x}_{i},\mathbf{x})$ , where  $\alpha_{it} = \frac{\eta^{(t)}}{|B^{(t)}|}\sum_{i\in B^{(t)}}\frac{\partial\mathcal{L}(f_{K}^{(t - 1)}(\mathbf{x}_{i}),f(\mathbf{x}_{i}))}{\partial f_{K}^{(t - 1)}(\mathbf{x}_{i})}$ . The function at step  $T$  can then be represented as:
$$
f _ {K} ^ {(T)} (\mathbf {x}) = \sum_ {i = 1} ^ {n} \alpha_ {i} ^ {(T)} K \left(\mathbf {x} _ {i}, \mathbf {x}\right) + f _ {K} ^ {(0)} (\mathbf {x}) \text {w i t h} \alpha_ {i} ^ {(T)} = - \sum_ {t: i \in B ^ {(t)}} \frac {\eta^ {(t)}}{| B ^ {(t)} |} \frac {\partial \mathcal {L} \left(f _ {K} ^ {(t - 1)} \left(\mathbf {x} _ {i}\right) , f \left(\mathbf {x} _ {i}\right)\right)}{\partial f _ {K} ^ {(t - 1)} \left(\mathbf {x} _ {i}\right)}. \tag {8}
$$
Definition 11 (Tracking representatives). Given a finite set of steps  $T$ , we term the global importance weights obtained via tracking kernel gradient descent as tracking representatives:
$$
\alpha_ {i} ^ {*} = - \sum_ {t \in [ T ]: i \in B ^ {(t)}} \frac {\eta^ {(t)}}{| B ^ {(t)} |} \frac {\partial \mathcal {L} \left(f _ {K} ^ {(t - 1)} \left(\mathbf {x} _ {i}\right) , f \left(\mathbf {x} _ {i}\right)\right)}{\partial f _ {K} ^ {(t - 1)} \left(\mathbf {x} _ {i}\right)}. \tag {9}
$$
We note that one can draw from standard correspondences between gradient descent with finite stopping and ridge regularization (e.g. [50]), to in turn relate the iterates of the kernel gradient descent updates for any finite stopping at  $T$  iterations to regularized RKHS regression solutions for some penalty  $\lambda$ . The above procedure thus provides a potentially scalable approach to compute the corresponding global importances: in order to calculate the global importance  $\alpha_{i}^{(T)}$ , we need to simply monitor the evolution of  $\alpha_{i}^{(t)}$  when the sample  $\mathbf{x}_i$  is utilized at iteration  $t$ . In our experiment, we use the following relaxation for further speed up:
$$
\alpha_ {i} ^ {*} = - \sum_ {t \in [ T ]: i \in B ^ {(t)}} \frac {\eta^ {(t)}}{| B ^ {(t)} |} \frac {\partial \mathcal {L} \left(f ^ {(t - 1)} \left(\mathbf {x} _ {i}\right) , y _ {i}\right)}{\partial f ^ {(t - 1)} \left(\mathbf {x} _ {i}\right)}, \tag {10}
$$
where we assume the target model is trained with (stochastic) gradient descent,  $f^{(t)}(\mathbf{x}_i)$  denotes the target model at  $t^{\mathrm{th}}$  iteration during training, and  $B^{(t)}$  and  $\eta^{(t)}$  are the corresponding mini-batch and learning rate. Similar to the intuition of replacing the surrogate derivative with to target derivative, we track the target model's training trajectory directly instead of solving Eqn.(4) with kernel gradient descent.
5 Choice of Kernels for Generalized Representatives
Previously, we discussed approaches for deriving global importance given user-specified kernels, which can in general be specified by domain knowledge relating to the model and the application domain. In this section, we discuss natural choices of kernels for modern non-linear models. Moreover, we show that existing sample based explanation methods such as representer points [7] and influence functions [5] could be viewed as making particular choices of kernels when computing general representatives. We also discuss TracIn [6] as a natural extension of our framework to multiple rather than a single kernel.
5.1 Kernel 1: Penultimate-layer Embeddings
A common method for extracting a random feature map from a neural network is to use the embeddings of its penultimate layer [7, 51, 52]. Let  $\varPhi_{\Theta_1}:\mathbb{R}^d\mapsto \mathbb{R}^\ell$  denote the mapping from the input to the second last layer. The target model  $f$  can be represented as
$$
f (\mathbf {x}) = \varPhi_ {\Theta_ {1}} (\mathbf {x}) ^ {\top} \Theta_ {2}, \tag {11}
$$
where  $\Theta_2\in \mathbb{R}^\ell$  is the weight matrix of the last layer. That is, we treat the deep neural network as a linear machine on top of a learned feature map. The kernel function is then defined as  $K_{\mathrm{LL}}(\mathbf{x},\mathbf{z}) = \langle \varPhi_{\Theta_1}(\mathbf{x}),\varPhi_{\Theta_1}(\mathbf{z})\rangle ,\forall \mathbf{x},\mathbf{z}\in \mathbb{R}^d$  . This is the case with most deep neural network architectures, where the feature map  $\varPhi_{\Theta_1}$  is specified via deep compositions of parameterized layers that take the form of fully connected layers, convolutional layers, or attention layers among others. While the last-layer weight matrix  $\Theta_{2}$  may not lie in the span of  $\{\varPhi_{\theta_1}(\mathbf{x}_i)\}_{i = 1}^n$  , we may solve the its explanatory surrogate function using Eqn.(4).
Corollary 12. (Representer point selection [7]) The minimizer of Eqn.(4), instantiated with  $K_{LL}(\mathbf{x},\mathbf{z}) = \langle \varPhi_{\Theta_1}(\mathbf{x}),\varPhi_{\Theta_1}(\mathbf{z})\rangle ,\forall \mathbf{x},\mathbf{z}\in \mathbb{R}^d$  can be represented as
$$
\hat {f} _ {K} (\cdot) = \sum_ {i = 1} ^ {n} \alpha_ {i} K _ {L L} \left(\mathbf {x} _ {i}, \cdot\right), \text {w h e r e} \alpha_ {i} = - \frac {1}{n \lambda} \frac {\partial \mathcal {L} \left(\hat {f} _ {K} \left(\mathbf {x} _ {i}\right) , f \left(\mathbf {x} _ {i}\right)\right)}{\partial \hat {f} _ {K} \left(\mathbf {x} _ {i}\right)}, \forall i \in [ n ]. \tag {12}
$$
The above corollary implies that  $\hat{\Theta}_2 = \sum_{i=1}^n \alpha_i \Phi_{\theta_1}(\mathbf{x}_i)$ . In other words, the RKHS regularization in Eqn. (4) can be expressed as  $\|f_K\|_{\mathcal{H}_K}^2 = \|\Theta_2\|^2$ , which is equivalent to L2 regularization. Consequently, the representative point selection method proposed in Yeh et al. [7] can be generalized to our framework when we use last-layer embeddings as feature maps.
5.2 Kernel 2: Neural Tangent Kernels
Although freezing all layers except for the last layer is a straightforward way to simplify neural networks to linear machines, last-layer representatives may overlook influential behavior that is present in other layers. For example, Yeh et al. [53] shows that representation in the last layer leads to inferior results for language models. On the other hand, neural tangent kernels (NTK) [22] have been demonstrated as a more accurate approximation of neural networks [54-56]. By using NTKs, we use gradients with respect to model parameters as feature maps and approximate neural networks with the corresponding kernel machines. This formulation enables us to derive a generalized representative that captures gradient information of all layers.
For a neural network with scalar output  $f_{\theta}:\mathbb{R}^{d}\mapsto \mathbb{R}$  parameterized by a vector of parameters  $\theta \in \mathbb{R}^p$ , the NTK is a kernel  $K:\mathbb{R}^d\times \mathbb{R}^d\mapsto \mathbb{R}$  defined by the feature maps  $\varPhi_{\theta}(\mathbf{x})=\frac{\partial f_{\theta}(\mathbf{x})}{\partial\theta}$ :
$$
K _ {\mathrm {N T K}, \theta} (\mathbf {x}, \mathbf {z}) = \left\langle \frac {\partial f _ {\theta} (\mathbf {x})}{\partial \theta}, \frac {\partial f _ {\theta} (\mathbf {z})}{\partial \theta} \right\rangle . \tag {13}
$$
Connection to TracIn [6]: TracIn measures the change in model parameters from the start to the end of training. While it is intractable due to the need to store model parameters of all iterations, Pruthi et al. [6] used checkpoints(CP) as a practical relaxation: given model parameters  $\theta^{(t)}$  and learning rates  $\eta^{(t)}$  at all model checkpoints  $t = 0,\dots ,T$ , the formulation of TracInCP is given
$$
\begin{array}{l} \phi_ {\operatorname {T r a c I n C P}} \left(f _ {\theta}, \left(\mathbf {x} _ {i}, y _ {i}\right)\rightarrow \mathbf {x}\right) = - \sum_ {t = 0} ^ {T} \eta^ {(t)} \frac {\partial \mathcal {L} \left(f _ {\theta} (\mathbf {x} _ {i}) , y _ {i}\right) ^ {\top}}{\partial \theta} \frac {\partial f _ {\theta} (\mathbf {x})}{\partial \theta} \Big | _ {\theta = \theta^ {(t)}} \\ = - \sum_ {t = 0} ^ {T} \underbrace {\eta^ {(t)} \frac {\partial \mathcal {L} \left(f _ {\theta} \left(\mathbf {x} _ {i}\right) , y _ {i}\right)}{\partial f _ {\theta} \left(\mathbf {x} _ {i}\right)} \Big | _ {\theta = \theta^ {(t)}}} _ {\text {g l o b a l i m p o r t a n c e}} \underbrace {K _ {\mathrm {N T K} , \theta^ {(t)}} \left(\mathbf {x} _ {i} , \mathbf {x}\right)} _ {\text {l o c a l i m p o r t a n c e}}. \tag {14} \\ \end{array}
$$
When the learning rate is constant throughout the training process, TracInCP can be viewed as a generalized representative instantiated with target derivative as global importances and NTK (Eqn.(13)) as the kernel function, but uses different kernels on different checkpoints.
5.3 Kernel 3: Influence Function Kernel
The influence functions [5] can also be represented as a generalized representative with the following kernel:
$$
K _ {\operatorname {I n f}, \theta} (\mathbf {x}, \mathbf {z}) = \left\langle \frac {\partial f _ {\theta} (\mathbf {x})}{\partial \theta}, \frac {\partial f _ {\theta} (\mathbf {z})}{\partial \theta} \right\rangle_ {H _ {\theta} ^ {- 1}} = \frac {\partial f _ {\theta} (\mathbf {x}) ^ {\top}}{\partial \theta} H _ {\theta} ^ {- 1} \frac {\partial f _ {\theta} (\mathbf {z})}{\partial \theta}, \tag {15}
$$
where  $H_{\theta} = \frac{1}{n}\sum_{i = 1}^{n}\frac{\partial^{2}\mathcal{L}(f_{\theta}(x_{i}),y_{i})}{\partial\theta^{2}}$  is the Hessian matrix with respect to target model parameters. The influence function then can be written as:
$$
\phi_ {\operatorname {I n f}} \left(f _ {\theta}, \left(\mathbf {x} _ {i}, y _ {i}\right)\rightarrow \mathbf {x}\right) = - \frac {\partial \mathcal {L} \left(f _ {\theta} \left(\mathbf {x} _ {i}\right) , y _ {i}\right)}{\partial \theta} H _ {\theta} ^ {- 1} \frac {\partial f _ {\theta} (\mathbf {x})}{\partial \theta} = - \underbrace {- \frac {\partial \mathcal {L} \left(f _ {\theta} \left(\mathbf {x} _ {i}\right) , y _ {i}\right)}{\partial f _ {\theta} \left(\mathbf {x} _ {i}\right)}} _ {\text {g l o b a l i m p o r t a n c e}} \underbrace {K _ {\operatorname {I n f} , \theta} \left(\mathbf {x} _ {i} , \mathbf {x}\right)} _ {\text {l o c a l i m p o r t a n c e}}. \tag {16}
$$
Therefore, the influence function can be seen as a member of generalized representatives with target derivative global importance (Definition [10]) and the influence function kernel. Influence functions [57] were designed to measure how would the model's predictions change if a training input were perturbed for convex models trained with empirical risk minimization. Consequently, the inversed Hessian matrix describes the sensitivity of the model parameters in each direction.
In the experiment, we compare different representatives within our proposed framework on both vision and language classification tasks. We use convolutional neural networks (CNN) since they are widely recognized deep neural network architectur. We compare performance of different choices of kernels and different ways to compute global importance. Existing generalized representatives, such as influence functions [5], representative point selections [53], and TracIn [6], are included in our experiment.
Evaluation metrics: We use case deletion diagnostics [53, 57, 58],  $\mathrm{DEL}_{-}(\mathbf{x}, k, \phi)$ , as our primary evaluation metric. The metric measures the difference between models' prediction scores on  $\mathbf{x}$  when we remove top- $k$  negative impact samples given by method  $\phi$  and the prediction scores of the original models. We expect  $\mathrm{DEL}_{-}$  to be positive since models' prediction scores should increase when we remove negative impact samples. To evaluate deletion metric at different  $k$ , we follow Yeh et al. [53] and report area under the curve (AUC): AUC-  $\mathrm{DEL}_{-} = \sum_{i=1}^{m} \mathrm{DEL}_{-}(\mathbf{x}, k_i, \phi) / m$ , where  $k_1 < k_2 < \dots < k_m$  is a predefined sequence of  $k$ .
We choose  $k_{i} = 0.02in$  for  $i = 0,1,\dots ,5$  with  $n$  as the size of the training set. The average of each metric is calculated across 50/200 randomly initialized neural networks for vision/language data. For every neural network, sample-based explanation methods are computed for 10 randomly selected testing samples.
3We replace the loss function on the test point  $\mathcal{L}(f(\mathbf{x}),y)$  with the target function prediction  $f(\mathbf{x})$  to measure training point influence on the predictions.
Datasets and models being explained: For image classification, we follow Pruthi et al. [6] and use MNIST [59] and CIFAR-10 [60] datasets. For text classification, we follow Yeh et al. [53] and use Toxicity and AGnews datasets, which contain toxicity comments and news of different categories respectively. Due to computational challenges in computing deletion diagnostics, we subsample the datasets by transforming them into binary classification problems with each class containing around 6,000 training samples. The CNNs we use for the four datasets comprise 3 layers. For vision datasets, the models contain around 95,000 parameters. For text datasets, the total number of parameters in the model is 1,602,257 with over  $90\%$  of the parameters residing in the word embedding layer that contains 30,522 trainable word embeddings of dimensions 48. Please refer to the Appendix C for more details on the implementation of generalized representatives and dataset constructions.
Table 1: Case deletion diagnostics, AUC-DEL-, for removing negative impact training samples on four different datasets.  $95\%$  confidence interval of averaged deletion diagnostics on  $50 \times 10 = 500$  (or  $200 \times 10 = 2000$ ) samples is reported for vision (or language) data. Larger AUC-DEL- is better. Init, middle, and final denote initial parameters  $\theta^{(0)}$ , parameters of a middle checkpoint  $\theta^{(T/2)}$ , and final parameters  $\theta^{(T)}$  for neural networks trained with  $T$  epochs. We only use the last-layer parameters to compute influence functions as in [5, 53] since the total number of parameters are too large for text models.
The results are shown in Table[1]. We also provide deletion curves we compute AUC-DEL_ for in the Appendix.
I. Comparison of different global importance: In the first experiment, we fix the kernel to be the NTK computed on final model parameters, and we compare different methods for computing global importance in Section 4. We do not compute the surrogate derivative on the text datasets since the total numbers of parameters are too large, making the computation infeasible.
4https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge
http://groups.di.unipi.it/gulli/AG Corpus_of_news_articles.html
We observe that tracking has the best performance, followed by target derivative and then surrogate derivative. This could be due to the loss flattening when converged and the loss gradients becoming less informative. Consequently, accumulating loss gradients during training is the most effective approach. Moreover, if tracking is not feasible when training trajectories are not accessible, we may use target derivative instead of surrogate derivative as an alternative to explain neural networks since they have similar performance.
II. Comparison of different kernels: Next, we fix the global importance to tracking and compare different kernels in Section 5. We employ the tracking representatives to compute global importance since it showed the best performance in the previous experiment. We can see that the influence function kernel performs the best in the vision data sets, and the NTK-final kernel has the best performance in language data sets. Note that influence functions exhibit distinctly contrasting performances on image and text data, which could be attributed to our reliance solely on last-layer parameters for influence function computation on language datasets. This finding aligns with the conclusions of Yeh et al. [53], who suggest that the last layer gradients provide less informative insights for text classifiers.
In summary, these findings indicate that NTK-final is a dependable kernel selection due to its consistent high performance across all four datasets, while also offering a computational efficiency advantage over the influence function kernel. These results also demonstrate that accessing target model checkpoints for computing kernels is unnecessary since NTK and influence function on the final model already provide informative feature maps.
III. Comparison of different generalized representatives: Finally, we compare the new generalized representative with other existing generalized representatives. We categorize TracInCP, the influence function, and the representative point as existing generalized representatives: TracInCP can be viewed as an ensemble of generalized representatives with target derivatives using the Neural Tangent Kernel. The influence function can be expressed as the influence function kernel with the target derivative. Lastly, the representative point can be seen as a form of generalized representative that utilizes the last-layer kernel and the surrogate derivative.
We find that the Inf-final has comparable performance to TracInCP and they outperform other approaches. Although TracInCP has the best performance on MNIST, it requires accessing different checkpoints, which requires a significant amount of memory and time complexity. In contrast, the NTK and Inf tracking representatives are more efficient since they only require tracking gradient descent trajectories during training without the need for storing checkpoints.
In this work, we present generalized representatives that are the only class of sample based explanations that satisfy a set of desirable axiomatic properties. We explore various techniques for computing generalized representatives in the context of modern non-linear machine learning models and show that many popular existing methods fall into this category. Additionally, we propose tracking representatives that track sample importance along the gradient descent trajectory. In future work, it would be of interest to derive different generalized representatives by altering different global importance and choices of kernels, as well as investigating their applicability to diverse machine learning models and modalities.
We acknowledge the support of DARPA via FA8750-23-2-1015, ONR via N00014-23-1-2368, and NSF via IIS-1909816.
[1] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. Why should i trust you?: Explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pages 1135-1144. ACM, 2016.
[56] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. arXiv preprint arXiv:2210.05643, 2022.
[57] R Dennis Cook and Sanford Weisberg. *Residuals and influence in regression*. New York: Chapman and Hall, 1982.
[58] Xiaochuang Han, Byron C Wallace, and Yulia Tsvetkov. Explaining black box predictions and unveiling data artifacts through influence functions. arXiv preprint arXiv:2005.06676, 2020.
[59] Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL http://yann.lecun.com/exdb/mnist/.
[60] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[61] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[62] Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization in linear time. stat, 1050:15, 2016.
[63] James Mercer. Xvi. functions of positive and negative type, and their connection the theory of integral equations. Philosophical transactions of the royal society of London. Series A, containing papers of a mathematical or physical character, 209(441-458):415-446, 1909.