A Flexible Nadaraya-Watson Head Can Offer
Explainable and Calibrated Classification

Nonparametric models in deep learning have received much attention in previous work. Deep Gaussian Processes (Damianou & Lawrence, 2013), Deep Kernel Learning (Wilson et al., 2015), and Neural Processes (Kossen et al., 2021) build upon Gaussian Processes and extend them to representation learning. Other works have generalized $k$-nearest neighbors (Papernot & McDaniel, 2018; Taesiri et al., 2022), decision trees (Zhang et al., 2018), density estimation (Fakoor et al., 2020), and more general kernel-based methods (Nguyen et al., 2021; Zhang et al., 2021; Ghosh & Nag, 2001) to deep networks and have explored the interpretability that these frameworks provide.

Closely-related but orthogonal to nonparametric models are attention-based models, most notably self-attention mechanisms popularized in Transformer-based architectures in natural language processing (Vaswani et al., 2017) and, more recently, computer vision (Dosovitskiy et al., 2020; Jaegle et al., 2021; Parmar et al., 2018). Retrieval models in NLP learn to look up relevant training instances for prediction (Sachan et al., 2021; Borgeaud et al., 2021) in a nonparametric or semiparametric manner. Many works have recognized the inherent interpretable nature of attention and have leveraged it in vision (Chen et al., 2018; Guan et al., 2018; Wang et al., 2017; Xu et al., 2015). More recently, nonparametric transformers (NPT) (Kossen et al., 2021) applied attention in a nonparametric setting. NPT leverages self- and cross-attention between both attributes and datapoints, which are stacked up successively in an alternating fashion to yield a deep, non-linear architecture.

Building upon and in contrast to these prior works, we present a comparatively simple nonparametric head which is computationally-efficient and performant, and which additionally lends itself naturally to being interpretable.

Metric learning seeks to quantify the similarity between datapoints for use in downstream tasks (e.g., unsupervised pretraining (Wu et al., 2018) or retrieval (Chen et al., 2021)). Of particular relevance to our method is metric learning in the context of low-shot classification (Snell et al., 2017; Vinyals et al., 2016; Sung et al., 2017; Santoro et al., 2016). In particular, Matching Networks (Vinyals et al., 2016) for low-shot and meta-learning are most similar to our model in the sense that it uses an NW-style model to compute predictions, where both the query image and support set are composed of examples from unseen classes. In contrast, we analyze the NW head from a nonparametric perspective in the typical, many-shot image classification setting.

Following classical work by Cook & Weisberg (1982), Koh & Liang (2017) propose the influence function as a means of explaining black-box neural network predictions. Let $h$ denote a function parameterized by $\theta$. $\mathcal{D}=\{z_{i}:(x_{i},y_{i})\}_{i=1}^{n}$ is a set of training examples and $L(h_{\theta}(z))$ indicates the loss for a training example $z$. $\theta^{*}$ is the set of optimal parameters which minimizes the empirical risk. Similarly, the parameter set associated with up-weighting a training example $z$ by $\epsilon$ is found by solving:

Koh & Liang (2017) define influence as the change in the loss value for a particular test point $z_{t}$ when a training point $z$ is up-weighted. Using a first-order Taylor expansion and the chain rule, influence can be approximated by a closed-form expression:

where $H_{\theta^{*}}$ represents the Hessian with respect to $\theta^{*}$. $\mathcal{I}(z_{t},z)/n$ is approximately the change in the loss for the test-sample $z_{t}$ when a training sample $z$ is removed from the training set. However, this result is based on the assumption that the underlying loss function is strictly convex in $\theta$ and that the Hessian matrix is positive-definite. Follow-up work has questioned the efficacy of influence functions in practical scenarios where these assumptions are violated (Basu et al., 2020).

In this work, we define the support influence function which follows naturally from the NW head; it computes the change in the loss directly on the prediction vector itself (and not indirectly through the model parameters $\theta$). Furthermore, it is easy-to-compute, requires no approximations, and does not require any assumptions on convexity or differentiability.

Consider a “support set” of $N_{s}$ examples and their associated labels, $\mathcal{S}=\{z_{i}:(x_{i},y_{i})\}_{i=1}^{N_{s}}$. This support set can be a subset, or the whole, of the training dataset. In this paper, we will focus on the image classification setting, where $x$’s are images and $y$’s are integer class labels from a label set $\mathcal{C}=\{1,...,C\}$. We will use $\vec{y}$ to denote the one-hot encoded version of $y$. We note that the NW head can be used with any input data type, not just images, and can be readily extended to the regression setting.

For a query image $x$, the Nadaraya-Watson prediction is computed as the weighted sum of support labels $\vec{y}_{i}$, where the weights quantify the similarity of the query image with each support image $x_{i}$:

The weights can be thought of as attention kernels (Vinyals et al., 2016) and are similar to attention mechanisms in deep learning (Vaswani et al., 2017; Dosovitskiy et al., 2020; Kossen et al., 2021). Classically, these weights are defined as a function of handcrafted kernel functions $\kappa$:

where $\kappa(\cdot,\cdot)$ quantifies the non-negative similarity between its two arguments (Bishop, 2006). Higher similarity values correspond to more similar pairs of datapoints.

In the NW head, we define the weights as:

Here, $g_{\theta}$ is a feature extractor with parameters $\theta$, which we implement as a neural network. $\tau$ is a temperature hyperparameter; unless otherwise stated, we set $\tau=1$. $\|\cdot\|_{2}$ denotes Euclidean distance. Since $\sum_{i=1}^{N_{s}}w_{\theta}(x,x_{i})=1$ and $\vec{y}_{i}$’s are one-hot encoded vectors, the output $f$ of the estimator can be interpreted as a conditional probability for the class label given the image and support set.

Let $\mathcal{D}$ be a training set of image and label pairs. Our objective to minimize is:

where $L$ is the cross-entropy loss. This indirectly encourages similar-looking images to have higher similarity score, since more similar images will tend to have the same labels.

We optimize this objective using stochastic gradient descent (SGD) by sampling a mini-batch size $N_{b}$ of query-support pairs and computing gradients with respect to the mini-batch. Note that our model assumes no dependency between support set samples; indeed, we are free to sample any arbitrary $\mathcal{S}$ and query $x$, as long as the class of $x$ is within the support classes. Therefore, the support set size $N_{s}=|\mathcal{S}|$ is a hyperparameter analogous and orthogonal to the batch size $N_{b}$.

The ability to select and manipulate the support set at inference time unlocks a degree of flexibility not possible with parametric classification models. In particular, the label set $\mathcal{C}$ of the support set directly determines the label set of the resulting prediction. For example, if the user has prior knowledge of which classes a particular query is most likely to be, the support set can be restricted to only contain those classes.

Note that we assume the model has been trained with random support sets drawn from the training data for each mini-batch. To characterize the effect of the support set on inference, in our experiments we implement the following “inference modes”:

1. 1.

   Full. Use the entire training set: $\mathcal{S}=\mathcal{D}$.

2. 2.

   Random. Sample uniformly at random over the dataset, such that each class is represented $k$ times: $\mathcal{S}\sim\mathcal{D}$ and $|\mathcal{S}|=k|\mathcal{C}|$.

3. 3.

   Cluster. Given the trained model, we first compute the features for all the training datapoints. Then, we perform $k$-means clustering on the features of the training datapoints for each class. These $k$ cluster centroids are then used as the support features for each class. Note that in cluster mode, the support set does not correspond to observed datapoints.

4. 4.

   Closest cluster (CC). Same as Cluster, except the support features correspond to real training datapoints that are closest to the cluster centroids in terms of Euclidean distance.

Each mode except Full assumes that each class is represented $k$ times in $\mathcal{S}$. One can, obviously, modify this assumption to reflect any class imbalance that might exist in the training data.

Koh & Liang (2017) define influence as the change in the loss as a result of upweighting a particular training point. With an NW head, quantifying the influence of a support image to the given query is straightforward and can be computed in closed-form. Indeed, the proposed support influence computes the change in the loss directly on the prediction vector itself (and not indirectly through the model parameters $\theta$), and furthermore does not require approximations, nor assumptions on convexity or differentiability.

Denote by $\mathcal{S}_{-z_{s}}=\mathcal{S}\setminus z_{s}$, i.e. the support set with an element removed. For a fixed query $x$:

We define the support influence of $z_{s}$ on $z$ as the change in the cross-entropy loss incurred from removing a support element $z_{s}$:

where superscript denotes indexing into a vector and $\mathbbm{1}$ is the indicator function. The derivation is provided in the Appendix.

Datasets. We experiment with an array of computer vision datasets with different class diversity and training set size. For general image classification, we experiment with Cifar-100 (Krizhevsky, 2009). For fine-grained image classification with small to medium training set size (less than 12k training samples), we experiment with CUB-200-2011 (Bird-200) (Wah et al., 2011), Stanford Dogs (Dog-120) (Khosla et al., 2011), Oxford Flowers (Flower-102) (Nilsback & Zisserman, 2008), and FGVC-Aircraft (Aircraft-100) (Maji et al., 2013). For a large-scale fine-grained image classification task (500k training samples), we experiment with iNaturalist-10k (Grant Van Horn, 2021). Additional dataset details are provided in Appendix A.1.

Network Architecture and Hyperparameters. For the feature extractor $g_{\theta}$, we use a convolutional neural network (CNN) followed by a linear projection to an embedding space of dimension $d$. For Cifar-100, Bird-102, Dog-120, Flowers-102, and Aircraft-100 dataset, we experiment with the ResNet-18 (He et al., 2015) and DenseNet-121 (Huang et al., 2016) architectures. For iNaturalist-10k, we experiment with the ResNet-50 architecture, following prior work (Zhou et al., 2019). For all feature extractors, we set $d=128$ and $\tau=1$. We train each model 3 times with different random initializations. Additional model and hyperparameter details are provided in Appendix A.2.

Baseline. We compare the NW head to a standard parametric model with a single-layer fully-connected (FC) head, and refer to it throughout this work as FC. For FC, we use the same feature extractor $g_{\theta}$ and keep all relevant hyperparameters identical to NW.

Metrics. We measure performance using error rate (the fraction of incorrectly classified test cases) and expected calibration error (ECE) (Guo et al., 2017; Naeini et al., 2015). ECE approximates the difference in expectation between confidence and accuracy, computed by partitioning the model predictions into bins and taking a weighted average of each bins’ difference between confidence and accuracy.

Tables 1, 2, and 5 show error rate and ECE performance for all models and datasets, and Fig. 4 further shows reliability diagrams. We generally observe that the NW head achieves an error rate that is comparable and at times better than its parametric counterpart. Calibration is particularly improved in small to medium-sized datasets with relatively fine-grained differences between classes.

Table 5 shows inference times on Bird-200 for both parametric and the NW-head models, where we assume that all support embeddings, cluster centroids, and closest cluster embeddings have been precomputed. We observe that inference times are nearly identical to the FC head, except Full. On a small dataset like Bird-200, Full corresponds to a two-fold increase in computation time.

In Fig. 2, we plot various metrics across support set size $k$. Generally, performance improves as support set size increases. We attribute this to the fact that more support images leads to more robust predictions, as more support examples per class can average out potential errors. At the limit, Full mode typically performs the best, at the cost of inference time.

Cluster mode often improves over Random mode and is on par with Full mode (at larger $k$) while having favorable inference time. However, this comes at the cost of loss of interpretability, since the cluster centroids do not correspond to actual images in the support set. Closest cluster addresses this deficiency, often with minimal sacrifice to performance compared to Cluster.

Fig. 3 visualizes the support images used in Closest cluster for two classes, with $k=1$ and $k=9$. These images are the prototypical examples of the class and can be interpreted as a summary or compression of the class and dataset. We observe that for $k=1$, the support image is a prototypical example of the class. For $k=9$, the images have a variety of poses and backgrounds, allowing the model to have a higher chance of computing similarities with a wide range of class prototypes. Further research might explore leveraging these prototypes.

Overall, our experiments indicate that the NW head has advantages in accuracy and calibration in the datasets we tested. The NW head pulls ahead in performance particularly for fine-grained classification tasks, and pulls ahead even further for small-medium sized datasets. Computational load is higher than FC, but is only two times higher for our most expensive inference mode on smaller datasets with less than 10k training samples.

In the following section, we restrict our attention to Full mode and explore the interpretability and explainability aspects of the NW head.