在人工智能领域，知识图谱 (Knowledge Graphs, KGs) 充当着机器的结构化记忆。它们以三元组的形式——(头实体，关系，尾实体)——存储着海量数据，例如 (巴黎, 是…的首都, 法国)。这些图谱支撑着从搜索引擎侧边栏到推荐系统和问答机器人等各种应用。

然而，存在一个根本性的问题: 知识图谱往往是不完整的。现实世界的数据杂乱无章，常常缺失某些关系。这催生了知识图谱补全 (Knowledge Graph Completion, KGC) 这一领域，它利用算法来预测缺失的链接，例如推断 (? , 运行系统, iOS) 暗示着 Apple。

历史上，研究人员不得不面临两种选择:

  1. 快速但浅显: 结构化模型，将实体视为向量空间中的点 (效率高，但忽略了实体丰富的文本描述) 。
  2. 智能但缓慢: 基于描述的模型 (如 KG-BERT) ，它阅读实体的文本描述 (准确率高，但计算速度慢得令人痛苦) 。

在这篇文章中，我们将深入探讨一篇提出“两全其美”解决方案的研究论文。这篇论文名为 Joint Pre-Encoding Representation and Structure Embedding for Efficient and Low-Resource Knowledge Graph Completion，介绍了一种名为 PEMLM 的模型。通过巧妙地将文本编码与训练循环分离，并将其与结构化数据融合，作者在取得最先进 (SOTA) 结果的同时，将推理速度提高了 30倍 , 并将内存使用量减少了 60% 。

让我们来看看他们是如何做到这一点的。

现代 KGC 的瓶颈

要理解为什么需要 PEMLM，我们首先需要看看现有方法的局限性。

两大阵营

基于嵌入的模型 (如 TransE 或 RotatE) 通过观察图的几何结构来学习。如果 国王 - 男人 + 女人 = 女王，则模型理解了这种结构。然而，这些模型将实体视为抽象的 ID。它们无法仅通过名字知道“苹果 (水果) ”和“苹果 (公司) ”在语义上是不同的；它们只能通过它们的连接关系来区分。

基于描述的模型 (如 KG-BERT) 使用预训练语言模型 (PLMs) 来阅读文本描述。例如，阅读“史蒂夫·乔布斯联合创办了苹果公司”可以帮助模型预测链接，即使图结构很稀疏。缺点是什么？它们极其沉重。为了预测一个链接，这些模型通常必须针对每一个候选实体将长文本序列输入 BERT。面对成千上万的实体，计算成本呈爆炸式增长，使得在有限资源下对大型图谱进行训练几乎不可能。

PEMLM 的解决方案

这篇论文的作者提出了 PEMLM (预编码掩码语言模型) 。 他们的核心见解简单而有力: 不要每次都重新阅读文本。

PEMLM 没有在训练循环中将原始文本输入模型，而是预先处理所有文本描述 一次，将其转换为丰富的语义向量。在训练期间，模型使用的是这些预计算好的向量。这种转变极大地降低了计算负载，同时保留了语言模型的语义理解能力。

此外，他们注意到仅靠文本并不总是足够的。有时，图结构包含了文本所遗漏的线索。为了解决这个问题，他们引入了 PEMLM-F , 这是一个融合框架，将预编码的文本表示与结构化嵌入结合在一起。

PEMLM 的架构

该架构分为两个明显的阶段: 预编码阶段和训练阶段 。 这种分离是模型高效的关键。

Figure 1: An overview of the PEMLM for link prediction. PEMLM consists of a pre-trained Description Encoder and a Triplet Encoder. The pre-Encoding and training stages can be conducted independently.

如图 1 所示，描述编码器 (左) 处理原始文本以创建嵌入。然后将这些嵌入传递给三元组编码器 (右) ，这是实际学习预测链接的组件。

1. 描述编码器 (预编码)

假设我们有一个实体 \(e\) 及其描述 \(des_e\)。例如，实体可能是 Gary Rydstrom，描述是 “Gary Roger Rydstrom is an American sound designer…”。

作者使用标准的 BERT 模型对该句子进行分词 (Tokenize) 。在 BERT 中，会添加特殊标记 [CLS] (开始) 和 [SEP] (结束) 。序列被输入到编码器中。作者没有使用每个单词的输出，而是对最后一个隐藏层应用 平均池化 (Mean Pooling) , 以获得该实体的单个紧凑向量表示。

生成语义表示 \(u\) 的公式为:

Equation for mean pooling of description encoding.

对图中的每个实体和关系重复此过程。一旦完成，BERT 模型就可以关闭或从内存中丢弃。我们留下了代表图中每一项“语义灵魂”的静态矩阵:

Matrices showing the collection of entity and relation embeddings.

这里，\(E\) 是所有实体向量的矩阵，\(R\) 是所有关系向量的矩阵。这个预编码步骤只需几分钟，却能节省后续数天的计算时间。

2. 三元组编码器 (训练)

现在我们有了预编码向量，如何训练模型来预测链接？作者将其视为 掩码语言建模 (Masked Language Modeling, MLM) 任务，类似于 BERT 的预训练方式，但针对图三元组进行了调整。

构建输入

知识图谱三元组由头实体 (\(h\)) 、关系 (\(r\)) 和尾实体 (\(t\)) 组成。为了预测缺失的尾实体，输入序列如下所示: [CLS], 头实体表示, 关系表示, [MASK], [SEP]。

但是，这里有个问题。由于我们向第二个编码器输入的是 向量 (来自预编码步骤) 而不是原始文本标记，模型失去了顺序的概念。它无法天生知道头实体在关系之前。

为了解决这个问题，作者向向量添加了 位置嵌入 (Position Embeddings) 。 这明确地告诉模型哪个向量代表头实体，哪个是关系，哪个是掩码。

Figure 2: Composition of input embedding showing tokens + position embeddings.

在数学上，输入序列 \(u^{input}\) 构建如下:

Equation showing the input sequence construction with CLS, Head, Relation, Mask, and SEP.

预测

模型将此序列通过三元组编码器 (另一种基于 Transformer 的架构) 。它查看对应于 [MASK] 标记位置的输出向量。这个向量代表了模型 认为 缺失实体应该是什么样子的。

Equation for extracting the output at the mask position.

最后，该输出通过一个密集分类层 (标准的神经网络层) ，以预测缺失实体在图中所有可能实体上的概率。

Equation for the softmax classification probability.

这种设置有效地将链接预测问题转化为多类分类问题。复杂性从 \(O(N)\) (逐个扫描每个候选者) 降低到 \(O(1)\) (一次前向传播即可针对所有候选者进行分类) 。

整合结构: PEMLM-F

虽然基于文本的方法很强大，但作者发现它有时在处理“一对多”关系 (例如，一个父母有多个孩子) 时会遇到困难。图结构模型对此处理得很好。

为了弥补这一差距，他们引入了 PEMLM-F (融合版) 。 该变体在文本模型旁边运行一个结构化嵌入模型 (基于著名的 TransE 算法) 。

结构化组件

TransE 模型将关系视为空间中的平移: \(h + r \approx t\)。模型试图最小化头实体加关系与尾实体之间的距离。

Equation for the TransE scoring function.

评分函数使用余弦相似度来查看预测与目标的接近程度:

Equation for cosine similarity scoring.

融合模块

这里的创新点在于他们如何结合这两者。他们不仅仅是平均结果。相反，他们获取语义文本向量 (\(u\)) 和结构化向量 (\(v\)) 并将它们拼接起来。

Figure 3: The architecture of PEMLM-F showing the fusion of description and structure embeddings.

如图 3 所示，融合模块拼接了这两个表示:

Equation for concatenating u and v.

然后，一个可学习的 MLP (多层感知机) 对这些特征进行加权，以创建一个新的、融合后的表示 \(s\):

Equation for the learnable fusion layer.

这个新的融合向量 \(s\) 是被输入到三元组编码器中的内容。这使得模型能够动态地学习何时依赖文本语义，何时依赖图结构。

最终的损失函数结合了分类损失 (来自掩码语言模型) 和对比损失 (来自结构化模型) ，并由超参数 \(\alpha\) 进行平衡:

Equation for the total loss function combining cross-entropy and structural loss.

实验与结果

研究人员在三个标准数据集上测试了 PEMLM: FB15k-237 (来自 Freebase 的通用知识) 、WN18RR (来自 WordNet 的词汇关系) 和 UMLS (医学/生物医学数据) 。

准确率表现

结果令人印象深刻。如下面的表 1 所示，PEMLM-F 取得了最先进的结果，尤其是在 WN18RR 和 UMLS 上。

Table 1: Link prediction results on FB15k-237, WN18RR and UMLS. Comparing PEMLM against baselines.

在 WN18RR 数据集上，PEMLM-F 达到了 50.9% 的 Hits@1 分数 , 显着优于之前的基于描述的模型如 KG-BERT (仅得 9.5%) ，甚至优于强大的联合模型如 Pretrain-KGE。这表明融合策略在精确定位正确实体 (Hits@1) 方面非常有效，而不仅仅是让结果“接近” (Hits@10) 。

效率: 游戏规则的改变者

最引人注目的结果是效率的提升。高准确率通常伴随着高资源成本，但 PEMLM 打破了这一趋势。

Table 4: Comparisons of the time, memory, and MRR required for training and inference.

表 4 突出了资源消耗的巨大差异:

  • 推理时间: KG-BERT 在测试集上运行推理需要 4天 。 PEMLM 仅需 1分钟 。 这是一个惊人的改进。
  • 训练内存: PEMLM 需要的内存大约 少 60% (3.6GB vs 8.5GB) ，使其在消费级 GPU (如 RTX 3080 甚至更小的显卡) 上运行成为可能。
  • 训练时间: 它的训练是以分钟而不是小时计算的。

这种效率来自于预编码策略。通过不强制繁重的 BERT 模型在每个训练步骤中处理文本，计算瓶颈被移除了。

为什么融合很重要

作者还分析了融合模型在 哪里 有帮助。他们按关系类型细分了表现: 一对一、一对多、多对一和多对多。

Table 5: MRR performance of PEMLM and PEMLM-F on different types of relations.

表 5 显示，融合模型 (PEMLM-F) 在 1-N (一对多) 关系中提供了最大的提升。这是有道理的: 兄弟姐妹或整体的多个部分的语义描述在文本上可能非常相似。结构化嵌入有助于在向量空间中独特地分离这些实体，从而允许模型更准确地对它们进行排序。

Alpha (\(\alpha\)) 的作用

参数 \(\alpha\) 控制模型赋予结构化损失相对于文本分类损失多大的权重。

Figure 4: MRR on WN18RR with different alpha values.

图 4 显示，在 WN18RR 数据集上，性能在 \(\alpha = 2\) 左右达到峰值。这表明在这个特定数据集上，结构信息非常有价值，应该与文本预测一起被赋予较大的权重。

结论与启示

PEMLM 论文提出了一个令人信服的论点，即解耦文本编码与图训练。通过将实体描述视为预计算的特征而不是实时输入，研究人员在不牺牲准确性的情况下释放了巨大的效率增益。

主要收获:

  1. 预编码行之有效: 你不需要端到端地微调语言模型就能获得很好的 KGC 结果。冻结的、池化后的嵌入是足够且更快的。
  2. 融合至关重要: 文本不能捕捉一切。整合几何图结构 (TransE) 特别有助于处理复杂的一对多关系。
  3. 效率促进规模化: 将推理时间从数天减少到数分钟，为在更大、更真实的知识图谱上使用这些模型打开了大门。

虽然该模型依赖于初始文本描述的质量 (垃圾进，垃圾出) ，但它代表了“绿色 AI”迈出的重要一步——用更少的资源获得更好的结果。对于该领域的学生和从业者来说，PEMLM 提供了一个蓝图，用于构建无需超级计算机即可运行的高性能图模型。