Uncertainty Estimation via Nearest Neighbor Distance Weighted Softmax

Abstract

Trustworthy prediction in deep learning models is important for safety-critical applications in the real world. However, deep learning models often suffer from the problem of miscalibration. Approaches that require multiple stochastic inferences can particularly mitigate this problem, but the expensive cost of inference makes them impractical. In this study, we propose $k$NN-weighted Softmax which is an uncertainty estimation method that uses the distances from neighbor examples. The method scales the logit according to the distance between the input example and its neighbors in the training data, which only requires a single forward inference. Experiments on text classification and named entity recognition show that our proposed method outperforms the baselines in uncertainty estimation.