摘要

本文介绍了一种基于最近邻（k-Nearest Neighbor, kNN）的不确定性估计方法（kNN-UE），用于自然语言处理任务中的深度神经网络（DNNs），特别是预训练语言模型（PLMs）。在安全关键应用中，可信预测至关重要，但DNNs常常在不确定性估计方面存在问题，如校准不当。传统的多随机推理方法虽能缓解此问题，但其高昂的推理成本使其不切实际。kNN-UE方法通过利用邻居的距离和标签存在比率来估计不确定性，实验表明，该方法在情感分析、自然语言推理和命名实体识别任务中，无论是在域内还是域外设置下，都优于基线方法或最近的密度基方法。此外，研究还表明，通过适当的组合，引入维度缩减或近似最近邻搜索可以减少推理开销，而不会显著降低估计性能。

原理

kNN-UE方法的核心在于利用k-最近邻搜索得到的邻居距离和标签信息来校正模型的置信度。具体来说，该方法通过计算输入示例与其邻居之间的距离，并根据模型预测标签与邻居标签的匹配比例来加权logits。这种方法仅需要模型的单次前向推理，通过这种方式，模型在预测实例存在于高密度训练数据区域且邻居中存在高比例的预测标签时，给出高置信度；反之，则给出低置信度。这种基于密度的不确定性估计方法不仅考虑了模型知识带来的认知不确定性，还考虑了数据变异带来的偶然不确定性。

流程

kNN-UE的工作流程包括以下步骤：

1. 对输入示例进行k-最近邻搜索，找到其在训练数据中的最近邻。
2. 计算输入示例与其最近邻之间的距离。
3. 统计最近邻中与模型预测标签相同的实例数量。
4. 根据距离和标签匹配比例加权logits，计算修正后的置信度。
5. 输出最终的不确定性估计结果。 例如，在情感分析任务中，kNN-UE会根据输入文本的表示在训练数据中找到相似的文本，并根据这些相似文本的情感标签来调整模型对当前文本情感的置信度。

应用

kNN-UE方法在自然语言处理领域的多种任务中展现出优越的不确定性估计性能，尤其在文本分类任务中。未来，该方法可能被应用于更广泛的领域，如文本生成任务，以及需要高可靠性预测的系统，如医疗诊断和自动驾驶等安全关键系统。此外，该方法的快速推理特性也使其适合于实时应用场景。