第3章：TransE模型详解——原理、损失函数与代码骨架
TransE是我刚接触知识图谱嵌入时学的第一个模型。说实话，它简单得让人有点不放心——把实体和关系都映射到同一个向量空间，关系就当平移向量用。但就是这套朴素的思想，至今还在很多工业级项目里扛大梁。下面我会把它的原理、损失函数、训练流程和实战中踩过的坑一次性说清楚。
3.1 核心思想：h + r ≈ t
TransE的核心假设就一句话：头实体向量 + 关系向量 ≈ 尾实体向量。对于每个三元组 (h, r, t)，我们希望经过关系 r 的“平移”后，头实体能尽可能靠近尾实体。
h + r ≈ t；对于错误的三元组（负样本），h + r 应该远离 t。
举个例子：三元组 (阿司匹林, 治疗, 头痛)。我们希望“阿司匹林”的向量加上“治疗”这个关系向量后，得到的向量和“头痛”的向量在空间中离得很近。如果换成负样本 (阿司匹林, 治疗, 心脏病)，这个距离就会很大。
这个模型为什么有效？因为药物知识图谱里的关系很多都是偏序或一对一的——比如“副作用”、“相互作用”、“治疗”。在这些场景下，平移假设天然成立。我在实际项目中测过，TransE在大规模稀疏关系上的表现反而比后来那些花哨的模型更稳。
3.2 损失函数：Hinge Loss + Margin
TransE用的损失函数是基于Margin的Hinge Loss。不扯公式符号，直接看逻辑：
loss = max(0, γ + d(h + r, t) - d(h' + r, t'))其中：
γ是Margin（超参数，通常0.5~2.0）d是距离函数（L1或L2范数，我个人偏好L1，梯度更稳定）(h, r, t)是正样本三元组(h', r, t')是负样本三元组（通常替换头实体或尾实体）
损失函数的含义很清楚：让正样本的距离 d(h+r, t) 尽可能小，同时让负样本的距离 d(h'+r, t') 尽可能大。当正样本距离 + Margin ≤ 负样本距离时，损失为0，模型认为已经“学到位了”。
3.3 训练流程与代码骨架
训练流程不复杂，但细节决定成败。我把它拆成5步：
- 初始化： 对每个实体和关系，随机初始化一个低维向量（维度d，通常50~200）。用均匀分布
[-6/sqrt(d), 6/sqrt(d)]初始化，记得对每个向量做L2归一化。 - 采样mini-batch： 从训练集中随机选一批正样本三元组，对每个正样本生成一个负样本（替换头或尾）。
- 前向传播： 计算正样本距离和负样本距离，得到Hinge Loss。
- 反向传播： 用到SGD或Adam更新嵌入向量。
- 归一化： 每轮迭代后将所有实体向量重新归一化到单位球面（防止向量范数无限增大）。
直接看代码更清爽。下面是一个简化但可运行的训练循环骨架：
import torch
import torch.nn as nn
class TransE(nn.Module):
def __init__(self, num_entities, num_relations, dim, margin=1.0):
super().__init__()
# 实体和关系嵌入表
self.entity_emb = nn.Embedding(num_entities, dim)
self.relation_emb = nn.Embedding(num_relations, dim)
self.margin = margin
self.dim = dim
# 初始化：均匀分布，范围 6/sqrt(dim)
nn.init.uniform_(self.entity_emb.weight, -6/(dim**0.5), 6/(dim**0.5))
nn.init.uniform_(self.relation_emb.weight, -6/(dim**0.5), 6/(dim**0.5))
def distance(self, h, r, t):
# h, r, t: [batch_size, dim]
# 使用L1距离（也可用L2）
return torch.norm(h + r - t, p=1, dim=1)
def forward(self, pos_triplets, neg_triplets):
# pos_triplets: [batch, 3] (h_idx, r_idx, t_idx)
# neg_triplets: [batch, 3]
h_pos = self.entity_emb(pos_triplets[:, 0])
r_pos = self.relation_emb(pos_triplets[:, 1])
t_pos = self.entity_emb(pos_triplets[:, 2])
h_neg = self.entity_emb(neg_triplets[:, 0])
r_neg = self.relation_emb(neg_triplets[:, 1])
t_neg = self.entity_emb(neg_triplets[:, 2])
d_pos = self.distance(h_pos, r_pos, t_pos)
d_neg = self.distance(h_neg, r_neg, t_neg)
# Hinge Loss: max(0, margin + d_pos - d_neg)
loss = torch.relu(self.margin + d_pos - d_neg).mean()
return loss
# 训练循环示例
model = TransE(num_entities=10000, num_relations=50, dim=100, margin=1.0)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(30):
total_loss = 0
for batch_pos, batch_neg in dataloader: # 假设有dataloader
optimizer.zero_grad()
loss = model(batch_pos, batch_neg)
loss.backward()
optimizer.step()
# 每步后归一化实体向量（关键！）
with torch.no_grad():
model.entity_emb.weight.data = nn.functional.normalize(
model.entity_emb.weight.data, p=2, dim=1
)
total_loss += loss.item()
print(f"Epoch {epoch}, Loss: {total_loss/len(dataloader):.4f}")
3.4 知识图谱嵌入核心概念对比
为了帮你理清TransE在整个图谱嵌入家族中的位置，我整理了一个对比表：
| 模型 | 核心思想 | 关系建模方式 | 参数复杂度 | 适用场景 |
|---|---|---|---|---|
| TransE | h + r ≈ t | 平移向量 | O(d) | 药品关系、稀疏图谱 |
| TransH | 实体在关系超平面投影后平移 | 超平面 + 平移向量 | O(2d) | 一对多/多对一关系 |
| TransR | 实体通过关系矩阵变换后平移 | 投影矩阵 + 平移向量 | O(d²+d) | 复杂多变关系 |
| DistMult | h · r · t（点积） | 对角矩阵 | O(d) | 对称关系（如“相互作用”） |
你看，TransE的参数最少，就两个查找表。这也是它在工业界受欢迎的原因——训练快，内存省。
3.5 TransE的优缺点：实战视角
先说说优点：
- 简单直接： 理解成本低，调参维度少（就Margin、学习率、嵌入维度这几个）。
- 训练速度快： 没有矩阵乘法，基本就是向量加减和范数计算。我在处理100万三元组、100维嵌入时，单机CPU一小时就能收敛。
- 适合大规模图谱： 尤其是药物知识图谱这种实体多、关系稀疏的场景，TransE不容易过拟合。
再说缺点，这部分我要多唠叨几句，因为新手很容易被它“骗”过去：
- 无法处理复杂关系： 一对多（一个药有多个适应症）、多对一（多个药治一个病）、自反关系（相似），TransE都会出问题。你想想看，如果阿司匹林和布洛芬都是治头痛的，那它们的向量差异就只能靠“治疗”这个关系来吸收，但TransE的关系向量是固定的——这是一个死结。
- 对称关系搞不定： 药物间的“相互作用”是双向的，但TransE要求 h + r ≈ t 且 t + r ≈ h，这只有在 r=0 时才成立。所以TransE学出来的对称关系向量会趋近于0，失去表达能力。
- 伪负样本问题： 随机替换头尾实体时，可能生成一个实际上在知识图谱中存在的正确三元组（数据不全导致的）。这会让模型对正负样本的区分变模糊。
3.6 一张图总结TransE核心流程
下面这张SVG图我把TransE从输入到损失计算的全链路画出来了，方便你对照代码理解：
这个流程里有个隐藏细节：右下角的实体向量归一化箭头是虚线回馈到嵌入层的。嗯，这一步千万不能省。我记得第一次跑实验时忘了加，结果loss直接崩到负值去了。
好，TransE的核心内容就这些。代码骨架我已经给了，你可以直接在自己的药物知识图谱数据集上试一试。调参时只要盯着两个指标：验证集上的Mean Rank（越低越好）和Hit@10（越高越好），基本就能快速收敛到可用状态。