TransE:将关系解释为向量空间中的平移
什么是知识图谱嵌入
知识图谱嵌入是将知识图谱中的实体和关系映射到连续低维向量空间的技术。其目标是让语义相近的实体在向量空间中距离更近,同时通过向量运算捕获实体间的关联模式。嵌入后的向量可以直接输入机器学习模型,广泛应用于链接预测、实体分类和推荐系统等任务。
知识图谱通常以三元组 (头实体, 关系, 尾实体) 的形式表示事实,例如 (北京, 位于, 中国)。嵌入模型需要为每个实体和每种关系学习一个向量,使得对正确三元组的某种评分函数值最大(或最小化损失)。
TransE 的核心思想
TransE 是最经典且直观的知识图谱嵌入模型之一。它提出一个简洁的假设:关系可以解释为向量空间中头实体到尾实体的平移。如果一个三元组 (h, r, t) 为真,那么头实体向量加上关系向量应当接近尾实体向量,即:
[ \mathbf{h} + \mathbf{r} \approx \mathbf{t} ]
这里的 h、r、t 分别是头实体、关系、尾实体的嵌入向量,维度均为 (d)。这种平移不变性使得 TransE 能够轻松处理一对一关系,同时对各类对称、反对称和组合关系具有一定的建模能力。
评分函数与训练目标
TransE 使用 L1 或 L2 范数定义三元组的能量(或距离):
[ f(h, r, t) = |\mathbf{h} + \mathbf{r} - \mathbf{t}|_p ]
其中 (p) 通常取 1 或 2。正确三元组的能量应尽可能低,而错误三元组的能量尽可能高。模型通过基于边际的合页损失进行训练:
[ \mathcal{L} = \sum_{(h,r,t) \in S} \sum_{(h',r,t') \in S'_{(h,r,t)}} \max\left(0, \gamma + f(h, r, t) - f(h', r, t')\right) ]
- (S):正确三元组集合。
- (S'):负采样集合,通过替换正确三元组的头实体或尾实体(但不同时替换)生成。
- (\gamma > 0):边际超参数,控制正确与错误三元组之间的最小分数间隔。
优化目标是最小化所有正负样本对的合页损失之和,通常配合随机梯度下降和实体关系向量的归一化约束(如 (|\mathbf{h}| \leq 1),(|\mathbf{t}| \leq 1))来防止训练发散。
训练过程的完整流程
- 初始化:为每个实体和每种关系随机生成服从均匀分布或正态分布的低维向量,并对实体向量做归一化。
- 小批量采样:从训练集中抽取一个批次的正三元组。
- 负采样:对每个正三元组,随机替换头实体或尾实体(使用均匀采样或伯努利采样策略)生成一个负样本,同时确保负样本不在知识图谱内。
- 计算损失:根据当前嵌入向量计算正负三元组的能量,代入合页损失公式。
- 梯度更新:计算损失关于实体和关系向量的梯度,使用 SGD 或 Adam 等优化器更新向量,然后对实体向量重新归一化。
- 迭代:重复步骤 2-5 直到收敛或达到预设轮次。
TransE 的局限性
尽管 TransE 模型简单高效,但其平移假设在处理复杂关系时存在明显不足:
- 一对多/多对一/多对多关系:例如一个头实体 “莎士比亚” 与多个尾实体(《哈姆雷特》《麦克白》)通过 “创作” 关系连接。TransE 会迫使所有尾实体向量都相近,导致表达混淆。
- 对称关系:对于 “同事” 这样的对称关系,TransE 要求 (\mathbf{h} + \mathbf{r} \approx \mathbf{t}) 且 (\mathbf{t} + \mathbf{r} \approx \mathbf{h}),这迫使关系向量 (\mathbf{r} \approx 0),丢失了语义。
- 复杂关系模式:自反关系、传递关系或组合模式难以用单一平移精确刻画。
为了解决这些问题,后续研究提出了 TransH(关系超平面上的平移)、TransR(关系空间映射)、TransD 等改进模型,但 TransE 仍以其概念清晰和计算高效而被广泛用于基线对比和教学。
代码实战:使用 PyKeen 训练 TransE
接下来使用 PyKeen 库快速搭建并训练 TransE 模型,预测缺失链接。
环境准备
通过 pip 安装核心依赖:
pip install pykeen torch
加载数据集与模型
以下代码使用 Nations 数据集(国家间关系)训练 TransE,并给出基本评估。
from pykeen.pipeline import pipeline
from pykeen.datasets import Nations
# 加载数据集
dataset = Nations()
training_triples = dataset.training
testing_triples = dataset.testing
# 配置并运行 pipeline
result = pipeline(
dataset=dataset,
model='TransE',
model_kwargs=dict(
embedding_dim=50,
scoring_fct_norm=2, # L2 范数
margin=1.0,
),
training_kwargs=dict(
num_epochs=200,
batch_size=64,
learning_rate=0.01,
),
optimizer='Adam',
optimizer_kwargs=dict(weight_decay=0.0),
negative_sampler='basic',
negative_sampler_kwargs=dict(num_negs_per_pos=1),
evaluator='rankbased',
evaluator_kwargs=dict(filtered=True),
random_seed=42,
)
解释结果与评估指标
训练结束后,result 对象包含模型、损失曲线和评估指标。常用的排序指标有:
- MRR (Mean Reciprocal Rank):所有测试三元组排名的倒数平均值,越接近 1 越好。
- Hits@k:正确实体出现在前 k 名的比例,通常取 k=1, 3, 10。
可以快速查看指标:
metrics = result.metric_values.to_dict()
print(f"MRR: {metrics['both.realistic.mrr']:.4f}")
print(f"Hits@10: {metrics['both.realistic.hits_at_10']:.4f}")
利用嵌入进行预测
模型训练后,可直接用来预测缺失的尾实体或头实体。
from pykeen.predict import predict_tail
# 预测 (usa, exports_to, ?) 的尾实体
prediction_df = predict_tail(
model=result.model,
head_label='usa',
relation_label='exports_to',
triples_factory=dataset.training,
).df
# 显示前5个最可能的尾实体及其评分
print(prediction_df.head(5)[['tail_label', 'score']])
评分越小(距离越小)表示模型认为该实体越可能是正确答案。
可视化嵌入
可以利用主成分分析或 t-SNE 将实体向量降至二维或三维,观察聚类结构。
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
import torch
# 获取实体嵌入矩阵
entity_embeddings = result.model.entity_representations[0](indices=None).detach().cpu().numpy()
entity_labels = dataset.entity_to_id.keys()
# PCA 降维
pca = PCA(n_components=2)
reduced = pca.fit_transform(entity_embeddings)
# 绘图(仅展示部分实体避免重叠)
plt.figure(figsize=(10,8))
for i, label in enumerate(list(entity_labels)[:30]):
plt.scatter(reduced[i,0], reduced[i,1])
plt.text(reduced[i,0], reduced[i,1], label, fontsize=9)
plt.title("TransE Entity Embeddings (PCA)")
plt.show()
通过散点图可以直观看到语义相似的国家是否在向量空间位置靠近。
超参数调优建议
- 嵌入维度:通常取 50-300。简单图谱可较低,复杂图谱需更高维度。
- 边际 γ:较大边际迫使正负样本距离更远,但可能增加优化难度,常用 1.0、2.0 或 0.5。
- 学习率:Adam 常用 0.001-0.01,SGD 需要更大学习率但配合衰减。
- 范数类型:L2 更平滑,L1 产生稀疏梯度,可根据经验选择。
- 负采样策略:基本均匀负采样可能产生大量弱负样本,可使用伯努利采样(按关系头尾基数比例)或基于嵌入的对抗采样提升训练效率。
总结
TransE 通过将关系视为向量平移,用极简公式捕获了知识图谱中实体与关系的结构信息。它易于实现、解释性强,是嵌入学习的绝佳起点。虽在复杂关系类型上表现有限,但其核心思想启发了众多后续工作。初学者在上手一个知识图谱项目时,用 TransE 建立性能基线、验证数据质量与嵌入维度等设置,依然是非常明智的选择。