Explaining and Improving Model Behavior with k Nearest Neighbor Representations

TL;DR

This paper introduces a kNN-based interpretability method for NLP models that identifies training examples influencing predictions, uncovers spurious correlations, detects mislabeled data, and enhances model robustness without retraining.

Contribution

It presents a novel use of kNN representations for interpretability and model improvement in NLP, especially for Natural Language Inference tasks.

Findings

kNN representations help identify training data responsible for predictions

The method uncovers learned spurious associations and mislabeled examples

Using kNN improves model robustness to adversarial inputs

Abstract

Interpretability techniques in NLP have mainly focused on understanding individual predictions using attention visualization or gradient-based saliency maps over tokens. We propose using k nearest neighbor (kNN) representations to identify training examples responsible for a model's predictions and obtain a corpus-level understanding of the model's behavior. Apart from interpretability, we show that kNN representations are effective at uncovering learned spurious associations, identifying mislabeled examples, and improving the fine-tuned model's performance. We focus on Natural Language Inference (NLI) as a case study and experiment with multiple datasets. Our method deploys backoff to kNN for BERT and RoBERTa on examples with low model confidence without any update to the model parameters. Our results indicate that the kNN approach makes the finetuned model more robust to adversarial inputs.