Reasoning Through Memorization:Nearest Neighbor Knowledge Graph Embeddings

这项研究提出了一种基于预训练语言模型的kNN-KGE方法，用于扩展知识图嵌入。该方法通过k近邻插值实体分布，利用预训练语言模型的掩模输出作为锚实体嵌入，结合知识存储显式记忆稀有实体或新出现的三元组。实验结果表明，该方法在低资源设置下性能显著提升，特别适用于知识图推理。

摘要

传统上，知识图嵌入方法将实体映射到特定表示空间，并通过分数函数来预测目标实体。然而，这些方法在处理罕见或未出现的实体时表现不足。我们提出了一种基于预训练语言模型的kNN-KGE知识图嵌入方法，该方法通过k近邻实体分布的线性插值，实现了对实体分布的更精细建模。通过计算实体嵌入空间与知识存储之间的距离，我们确定了最近的邻居实体。我们的方法使得显式记忆稀有实体的能力得以实现，而无需依赖隐式的模型参数。实验结果表明，我们的方法显著提升了推理性能，在仅包含少数三元组的资源有限场景中表现优异，这得益于显式地将稀有实体编码到模型参数中。

1.介绍

知识图谱（Knowledge Graphs, KGs）通过结构化的方式将事实以三元组（s, p, o）的形式组织起来，其中s和o代表实体，p则定义了实体之间的关系。由于实际应用中新增实体及其相互关联的缺失，现有的知识图谱大多尚未完善；为此，知识图谱补全问题——在缺少三元组的情况下扩展知识图谱——引起了研究者的广泛关注。目前，基于知识嵌入的方法如TransE[2]、ComplEx[30]和RotatE[28]等，通过将实体和关系映射到向量空间，并利用预定义的评分函数来预测潜在的三元组。这些方法通常采用端到端的训练方式，将所有关系知识隐含地编码到参数化神经网络的权重中。然而，这些方法的主要缺陷在于，它们在推理过程中难以有效处理稀有实体或在训练阶段未被观察到的新实体。

请注意，复杂的系统相互作用推动了人类的推理过程，例如，融合当前知识和从记忆中检索信息。增强记忆的神经网络的最新进展促使了模块化架构的设计，将计算处理与记忆存储区分开来。这些基于记忆的方法（或非/半参数方法）已在语言建模[15]和问题回答[14]等任务中得到应用，它们在表现力和适应性方面表现突出。

在此启发下，我们提出了𝑘NN-KGE方法，这是一种通过线性插值技术在实体嵌入空间中扩展知识图谱嵌入的创新方法。如图1所示，我们构建了基于预训练语言模型(plm)的实体知识存储系统，并通过计算实体嵌入空间中实体间的距离来实现近邻检索。对于缺少头实体或尾实体的三元组，我们利用[MASK]标记的输出表示作为锚点实体嵌入，在知识图谱存储中检索出最邻近的实体，并通过线性插值的方式，将锚点实体嵌入与最近邻实体分布进行融合，从而实现对稀有实体或未见三元组的显式记忆。这种方法通过显式记忆而非隐式推理，显著提升了模型的推理能力。

在两个数据集（FB15k-237[29]和WN18RR[8]）上，基于转换和归纳推理的实验结果验证了该方法的有效性。从定性的角度来看，我们观察到该方法显著地优于低资源知识图嵌入，这可能更易通过显式记忆进行访问。

我们的贡献可以总结如下:

1. 经研究可知，这是知识图嵌入领域的先驱性研究。我们的工作可能为通过外在记忆改进知识图推理提供了新的研究方向。
2. 我们提出了一种新的方法，称为𝑘NN-KGE，它能够明确地存储稀有或新兴的实体。在实际应用中，这一方法不可或缺，因为知识图谱不断扩展。
3. 通过两个基准数据集的实验分析，我们模型在性能上优于基准模型，特别适用于低资源推理场景。

2.方法

2.1 Preliminary

masking语言模型(MLM)。随机屏蔽输入中的某些标记，目的是基于上下文预测这些标记的原始词汇id[9]。大多数预训练语言模型采用MLM作为预训练任务。这些特殊标记是在训练过程中预先定义的，通过plm的编码器，可以轻松地获得输入序列中特殊符号的输出嵌入。

MEM Head类似于预训练语言模型中的词嵌入机制(MLM Head)，它将上下文信息的表征映射到词汇表中各词汇的概率分布。MEM Head由实体嵌入构成，其功能是将上下文信息的实体表征映射到知识图谱中各实体的概率分布。

2.2 Framework

在本节，我们将介绍该方法的整体架构。在2.3节中，我们提出了一种方法，通过屏蔽实体建模和扩展实体词汇，将链接预测任务转化为实体预测问题。在2.4节中，我们构建了一个知识存储系统，该系统通过锚嵌入（即[MASK]输出的表示）来检索实体信息。在2.5节中，我们详细描述了推理过程，该过程通过记忆机制而非直接在模型参数中进行推断。

2.3 Contextualized KG Representation

在本小节中，我们将BERT模型视为实体识别器，其原因在于我们将链接预测任务转化为不涉及实体建模的任务。该任务通过整合结构信息和文本描述来识别缺失的实体。

Masked Entity Modeling.

针对一个不完整的三元组（e_i, r_j, ?），已有研究通过KG嵌入或文本编码技术对三元组进行表示，并基于预定义的评分函数对这些向量进行评估。在处理预训练编码器时，现有方法如KG-BERT[34]展现了在通用性和对数据完整性处理方面的卓越表现；然而，在推理过程中，该方法需要遍历所有可能的三元组。在本文中，我们采用了一种更为简洁的实体建模方法进行链接预测，类似于屏蔽语言模型（MLM）任务。具体而言，通过实体建模的屏蔽机制，该模型能够减轻由于不正确实体存在于三元组中而产生的复杂评分函数设计、评分计算过程中的负抽样需求，从而在预测尾部或头部缺失实体时表现出色。换句话说，该方法仅需要关注于预测缺失的实体部分，而无需考虑所有可能的三元组。

在建模过程中需要注意的是，掩模实体建模这一过程具有较高的效率，相较于基于BERT的模型StAR[31]而言，其推理速度更为出色。表1详细列出了推断时间的对比结果，为评估模型性能提供了重要依据。与掩码语言建模任务具有相似性，我们需要预测每个[MASK]令牌的原始值。为了实现这一目标，我们采用相同的损失函数来优化我们的屏蔽实体模型。

由于PLM通常通过字节对编码将序列信息转化为子词单元，导致实体表示与原始的令牌信息脱离，从而使得推理过程难以有效进行。针对这一问题，我们提出了一种实体词汇扩展方法，通过将每个实体的嵌入信息融入模型，从而提升整体的表示能力。

Entity Vocabulary Expansion

在进行实体推理的过程中，子词的利用具有重要意义。具体而言，我们将实体𝑒∈E表示为语言模型词汇表中的特殊标记；进而，知识图推理可以重新表述为一个掩模实体预测任务，如图1和图2所示。为了获得这些实体嵌入（词汇表中的特殊标记），我们通过实体描述来设计提示句式，如“[MASK]的描述是那样”来获得实体嵌入。给定一个实体描述𝑑=(𝑤1,…,𝑤𝑛)的实体𝑒𝑖，在输入序列中前后添加标记[CLS]和[[SEP]的开始和结束描述，构建输入序列𝑥𝑑如下：

我们对这些实体嵌入进行随机初始化后进行优化。通过预测实体𝑒的位置，同时固定其他参数。正式来说，我们有:

2.4 Knowledge Store

基于记忆增强神经网络的最新进展[14,15]，我们开发了知识存储机制，以显式地存储实体信息。具体而言，我们首先从实体描述和实体三元组两个方面构建知识存储机制。

值得注意的是，我们采用上下文化KG构建了模型的知识存储体系，该体系从两个维度对知识进行了整合，如图2所示。此外，该方法能够明确地存储罕见实体或新兴实体，这种设计既具有灵活性又具备适应性。在介绍施工细节之前，我们首先说明了以下几点：首先，我们通过上下文化KG进行关联推理，这为知识的整合提供了有效支持；其次，我们采用动态更新机制，确保知识库的实时性；最后，我们引入了多模态数据融合技术，以提升信息处理的全面性。

From Descriptions D.

我们定义函数𝑓(·)，该函数将输入x中的实体e映射为由PLM计算得到的定长向量表示。通过预先设计的提示符，我们从实体描述中提取出实体嵌入。与2.3节相同，我们通过提示输入Xd提取实体嵌入并将其添加到知识存储中。因此，我们可以由此构造...

知识存储由描述和三元组构建的实体嵌入，支持存储多达数百万个实体嵌入。具体而言，我们定义了一个基于键-值对的推理机制(𝑘,𝑣)，其中键𝑘是基于描述和实体三元组生成的向量表示嵌入𝑓(𝑥)，而值𝑣对应的目标实体𝑒。在这里，𝑥表示描述𝑑与目标实体𝑒组成的三元组𝑡所构成的序列。我们借助开源库FAISS[12]在高维空间中实现快速的最近邻检索。

2.5 Memorized Inference

对于一个缺少头或尾实体的三元组，我们通过[MASK]位置的输出表示来推导锚实体的嵌入，以便在知识存储中找到最接近的邻居。具体而言，模型通过查询[MASK]位置的最终隐藏状态来确定最接近的知识实体，通过距离函数d(·，·)来确定最接近的知识实体，从而实现对缺失实体的补全。

的分布。

在知识存储中检索到的每个实体，我们仅选择知识存储中最近的嵌入来表示该实体，由于多个实体嵌入对应于同一个实体：