Efficient Nearest Neighbor based Uncertainty Estimation for Natural Language Processing Tasks

TL;DR

The paper introduces k-Nearest Neighbor Uncertainty Estimation (kNN-UE), a density-based UE method that also incorporates neighbor label information to calibrate NLP predictions with a single forward pass. By weighting base logits with a learned combination of neighbor distances and label agreement, kNN-UE improves confidence calibration and selective prediction across SA, NLI, and NER, in both in-domain and out-of-domain settings. The authors evaluate against strong baselines, analyze top-k and the label-term contribution, and explore efficient nearest-neighbor techniques (PQ, IVF, PCA) to address latency, offering practical recommendations and revealing trade-offs. The work demonstrates significant gains in UE metrics while providing insights into when and how to use approximate NN methods without severely compromising calibration, making it attractive for deployment in real-world NLP pipelines.

Abstract

Trustworthiness in model predictions is crucial for safety-critical applications in the real world. However, deep neural networks often suffer from the issues of uncertainty estimation, such as miscalibration. In this study, we propose $k$-Nearest Neighbor Uncertainty Estimation ($k$NN-UE), which is a new uncertainty estimation method that uses not only the distances from the neighbors, but also the ratio of labels in the neighbors. Experiments on sentiment analysis, natural language inference, and named entity recognition show that our proposed method outperforms the baselines and recent density-based methods in several calibration and uncertainty metrics. Moreover, our analyses indicate that approximate nearest neighbor search techniques reduce the inference overhead without significantly degrading the uncertainty estimation performance when they are appropriately combined.

Figures (4)

• Figure 1: Illustrations of $k$NN-UE behavior. The orange circle indicates predicted data instances and other circles indicate training data instances. $k$NN-UE gives high uncertainty when the predicted query representation is far from examples obtained from the $k$NN search (left) and the predicted label is different from the labels of neighbors (center). $k$NN-UE outputs low uncertainty only when the query representation is close to neighbors and the labels of neighbors contain many of the model's predicted label (right).
• Figure 2: A diagram of $k$NN-UE when $K = 3$ and the estimated hyperparameters are $\alpha = 0.5$, $\tau = 1.0$, $\lambda = 0.5$ and $b = 0.1$. A datastore is constructed with the representations of the training data as keys and their labels as values. The distances of the nearest examples from the test representation, and the neighbor labels are aggregated into $W_{k {\rm{NN}}} (\hat{y})$. Finally we obtain calibrated confidence by correcting the raw logits with $W_{k {\rm{NN}}} (\hat{y})$ as in Eq. \ref{['eq:knn_ue_main']}.
• Figure 3: Changes in ECE and E-AURC in SA when changing the number of neighbors of $k$NN-UE. On the x-axis, the parameter $K \in \{8, 16, 32, 64, 128\}$ is represented on a log scale.
• Figure 4: Changes in ECE and E-AURC in NER when changing the number of neighbors of $k$NN-UE. On the x-axis, the parameters $K \in \{8, 16, 32, 64, 128\}$ are represented on a log scale.