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

arXiv: https://arxiv.org/abs/2212.03411 · HTML: https://mlsanity.com/paper/2212.03411/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 57 references that are in the map.
- First posted on arXiv 7 Dec 2022.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Matching Networks for One Shot Learning](https://mlsanity.com/paper/1606.04080/) (2016)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Fine-Grained Visual Classification of Aircraft](https://mlsanity.com/paper/1306.5151/) (2013)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Learning to Compare: Relation Network for Few-Shot Learning](https://mlsanity.com/paper/1711.06025/) (2017)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [This Looks Like That: Deep Learning for Interpretable Image Recognition](https://mlsanity.com/paper/1806.10574/) (2018)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition](https://mlsanity.com/paper/1912.02413/) (2019)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Residual Attention Network for Image Classification](https://mlsanity.com/paper/1704.06904/) (2017)
- [Influence Functions in Deep Learning Are Fragile](https://mlsanity.com/paper/2006.14651/) (2020)
- [Dataset Meta-Learning from Kernel Ridge-Regression](https://mlsanity.com/paper/2011.00050/) (2020)
- [Deep k-Nearest Neighbors: Towards Confident, Interpretable and Robust Deep Learning](https://mlsanity.com/paper/1803.04765/) (2018)
- [Regularizing Class-wise Predictions via Self-knowledge Distillation](https://mlsanity.com/paper/2003.13964/) (2020)
- [Doubly Stochastic Variational Inference for Deep Gaussian Processes](https://mlsanity.com/paper/1705.08933/) (2017)
- [End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering](https://mlsanity.com/paper/2106.05346/) (2021)
- [Interpretable Deep Learning: Interpretation, Interpretability, Trustworthiness, and Beyond](https://mlsanity.com/paper/2103.10689/) (2021)
- [Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning](https://mlsanity.com/paper/2106.02584/) (2021)
- [Analyzing and Improving Representations with the Soft Nearest Neighbor Loss](https://mlsanity.com/paper/1902.01889/) (2019)
- [Interpreting CNNs via Decision Trees](https://mlsanity.com/paper/1802.00121/) (2018)

## Built on by (most cited first, 4 of 4)

- [TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023](https://mlsanity.com/paper/2307.14338/) (2023)
- [A Framework for Interpretability in Machine Learning for Medical Imaging](https://mlsanity.com/paper/2310.01685/) (2023)
- [Learning Invariant Representations with a Nonparametric Nadaraya-Watson Head](https://mlsanity.com/paper/2309.13377/) (2023)
- [SIC: Similarity-Based Interpretable Image Classification with Neural Networks](https://mlsanity.com/paper/2501.17328/) (2025)

## Abstract (arXiv)

In this paper, we empirically analyze a simple, non-learnable, and nonparametric Nadaraya-Watson (NW) prediction head that can be used with any neural network architecture. In the NW head, the prediction is a weighted average of labels from a support set. The weights are computed from distances between the query feature and support features. This is in contrast to the dominant approach of using a learnable classification head (e.g., a fully-connected layer) on the features, which can be challenging to interpret and can yield poorly calibrated predictions. Our empirical results on an array of computer vision tasks demonstrate that the NW head can yield better calibration with comparable accuracy compared to its parametric counterpart, particularly in data-limited settings. To further increase inference-time efficiency, we propose a simple approach that involves a clustering step run on the training set to create a relatively small distilled support set. Furthermore, we explore two means of interpretability/explainability that fall naturally from the NW head. The first is the label weights, and the second is our novel concept of the ``support influence function,'' which is an easy-to-compute metric that quantifies the influence of a support element on the prediction for a given query. As we demonstrate in our experiments, the influence function can allow the user to debug a trained model. We believe that the NW head is a flexible, interpretable, and highly useful building block that can be used in a range of applications.