A Simple Nadaraya-Watson Head for Explainable and Calibrated Classification

Alan Q. Wang, Mert R. Sabuncu

Keywords: Image Classification, Nonparametric, Interpretability, Explainability, Calibration

TL;DR: We present a simple, nonparametric replacement to the fully-connected head in the image classification setting based on the Nadaraya-Watson (NW) estimator, which can be shown to be interpretable and well-calibrated.

Abstract:

We propose a simple, non-learnable, and nonparametric prediction head to be used with any neural network architecture. The proposed head can be viewed as a classic Nadaraya-Watson (NW) model, where 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 than its parametric counterpart, while having comparable accuracy and with minimal computational overhead. 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. In addition to using the weights as a means of interpreting model predictions, we further present an easy-to-compute ``support influence function,'' which 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.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Deep Learning and representational learning