[论文解读] $k$NN-NER: Named Entity Recognition with Nearest Neighbor Search

Shuhe Wang, Xiaoya Li+5 more|arXiv (Cornell University)|Mar 31, 2022

Topic Modeling被引用 11

一句话总结

本文提出 $k$NN-NER，一种用于命名实体识别的检索增强框架，通过在推理过程中从缓存的训练集中检索 $k$ 个最近邻并插值其预测结果，来增强原始 NER 模型。该方法在不微调的情况下利用基于相似度的邻居搜索，实现了最先进性能，包括在中文微博命名实体识别数据集上达到 72.03 的 F1 分数，且在仅使用 40% 训练数据时仍能取得与标准模型相当的结果。

ABSTRACT

Inspired by recent advances in retrieval augmented methods in NLP~\citep{khandelwal2019generalization,khandelwal2020nearest,meng2021gnn}, in this paper, we introduce a $k$ nearest neighbor NER ($k$NN-NER) framework, which augments the distribution of entity labels by assigning $k$ nearest neighbors retrieved from the training set. This strategy makes the model more capable of handling long-tail cases, along with better few-shot learning abilities. $k$NN-NER requires no additional operation during the training phase, and by interpolating $k$ nearest neighbors search into the vanilla NER model, $k$NN-NER consistently outperforms its vanilla counterparts: we achieve a new state-of-the-art F1-score of 72.03 (+1.25) on the Chinese Weibo dataset and improved results on a variety of widely used NER benchmarks. Additionally, we show that $k$NN-NER can achieve comparable results to the vanilla NER model with 40\% less amount of training data. Code available at \url{https://github.com/ShannonAI/KNN-NER}.

研究动机与目标

• 解决命名实体识别中长尾实体识别难题，即标准模型对罕见或未见实体的处理能力较差。
• 通过检索相关训练样本，提升少样本和低资源学习能力。
• 在无需额外训练或架构修改的前提下，增强模型泛化能力。
• 证明推理阶段检索相似训练样本可优于基于记忆的训练方法。
• 表明 $k$NN-NER 在训练数据减少 40% 的情况下仍保持强劲性能，体现数据效率优势。

提出的方法

• 在训练阶段，缓存预训练原始 NER 模型（如 BERT）隐藏状态中的表示-标签对。
• 在推理阶段，对输入 token（如 'Washington'）进行编码，并使用余弦相似度在知识库中执行 $k$ 个最近邻（kNN）搜索。
• 检索 $k$ 个最相似的训练样本，并提取其标签分布。
• 使用可学习或固定加权方案，将原始 NER 模型的输出分布与 $k$ 个最近邻的标签分布进行插值。
• 将插值后的分布作为每个 token 的最终预测，从而提升对罕见或 OOV 实体的鲁棒性。
• 该方法可直接应用于多种预训练模型（BERT、RoBERTa、ChineseBERT），在英文和中文命名实体识别基准上无需微调。

实验结果

研究问题

• RQ1 从缓存的训练集中进行 $k$ 个最近邻搜索，是否能提升在长尾和罕见实体上的命名实体识别性能？
• RQ2 所提出的 $k$NN-NER 框架是否能在不增加训练的前提下提升少样本和低资源学习能力？
• RQ3 性能对 kNN 搜索中 $k$ 值的选择有多敏感？
• RQ4 $k$NN-NER 是否能在显著减少训练数据的情况下，取得与标准 NER 模型相当的结果？
• RQ5 检索增强方法是否在多样化的命名实体识别基准（包括多语言和低资源场景）中持续提升性能？

主要发现

• 在中文微博命名实体识别数据集上，$k$NN-NER 达到 72.03 的新 SOTA F1 分数，较原始 ChineseBERT-Large 模型提升 +1.23。
• 在中文 OntoNotes 4.0 数据集上，当 $k$=256 时，该方法相较原始 BERT-Base 模型实现 +1.75 的 F1 分数提升。
• 仅使用 60% 的训练数据，$k$NN-NER 在中文 OntoNotes 4.0 上即可达到全量数据下原始模型的性能，展现出强大的数据效率。
• 当 $k$=256 时性能趋于稳定，$k$=512 时无进一步增益，表明对超参数选择具有鲁棒性。
• 该方法在所有评估基准（包括 CoNLL03、OntoNotes5.0、MSRA 和 Weibo NER）上均一致提升 F1 分数，增益范围为 +0.11 至 +1.63。
• 插值机制有效利用了语义相似的训练样本，提升了对罕见或未见实体的预测性能，且无需微调。