知识图谱通过三元组(头实体,关系,尾实体)描述客观世界的事实,例如(姚明,效力于,NBA)。这种符号化表示虽然直观,却让计算机很难计算语义相似度。Knowledge Graph Embedding 的目标是把每个实体和关系转换为一个固定维度的数值向量,使“姚明”和“篮球运动员”在向量空间中的距离更近,从而为推荐系统、智能问答和搜索引擎提供知识推理能力。

Node.js 虽然不像 Python 那样拥有丰富的机器学习生态,但通过合理的建模和优化,同样能够完成知识图谱嵌入的训练与推断。本文不依赖 TensorFlow.js 等重型框架,而是用纯 JavaScript 实现一个简化版的 TransE 模型,从数据准备到训练推理逐一拆分,帮助读者掌握嵌入的本质。
一、Knowledge Graph Embedding 的核心思想
知识图谱嵌入的第一要务是定义评分函数。对于任意三元组 (h, r, t),我们希望合法的三元组得分更高,损坏的三元组得分更低。TransE 模型采用了非常直观的平移假设:头实体向量加上关系向量应近似等于尾实体向量,用公式表示就是 h + r ≈ t。因此,评分函数可以直接定义为 h + r 与 t 之间的欧氏距离。距离越小,说明三元组越正确。
这个假设在表达对称关系时存在天然短板,但对大部分常见关系来说非常有效。例如“首都”关系,向量空间中的“北京”加上“首都”关系向量,会落在“中国”附近。借助这一特性,TransE 可以用极少的参数学习到丰富的语义信息。通常在完成向量化后,我们还能进行实体类比推理,比如“北京 - 中国 + 日本 ≈ 东京”。
训练 TransE 需要构造正负样本。正样本直接来自原始知识图谱中的三元组;负样本则通过随机替换头实体或尾实体生成。例如将(北京,首都,中国)的尾实体替换成“法国”,得到负样本(北京,首都,法国)。模型通过最大化正样本得分与负样本得分之间的间隔来学习参数,最终使正确三元组的向量关系成立。
二、Node.js 环境准备与数据预处理
在开始实现之前,需要先把原始知识图谱数据转换成程序可处理的格式。常见的数据集如 Freebase、DBpedia 通常以 CSV 或 JSON 格式提供。下面使用一个 JSON 数组存储三元组,每个元素包含 head、relation、tail 三个字段。首先读取数据,并为每个实体和关系分配唯一的数字 ID,这是建立嵌入表的前提。
在 Node.js 中,使用 Map 类型来维护实体名与 ID 的对应关系非常高效。实体名通常是一个不规则的字符串,通过 Map 可以在 O(1) 时间内完成查找。需要注意的是,数据预处理阶段要过滤掉不完整的三元组,并统计实体总数和关系总数,以便后续初始化向量。
const fs = require('fs');
// 三元组数组,每一项为 [head, relation, tail]
const triples = [];
// 读取 JSON 格式的图谱数据
const rawData = JSON.parse(fs.readFileSync('./graph.json', 'utf8'));
for (const item of rawData) {
// 简单过滤空值
if (item.head && item.relation && item.tail) {
triples.push([item.head, item.relation, item.tail]);
}
}
// 构建实体和关系的索引映射
const entityMap = new Map();
const relationMap = new Map();
for (const [h, r, t] of triples) {
if (!entityMap.has(h)) {
entityMap.set(h, entityMap.size);
}
if (!entityMap.has(t)) {
entityMap.set(t, entityMap.size);
}
if (!relationMap.has(r)) {
relationMap.set(r, relationMap.size);
}
}
const entityCount = entityMap.size;
const relationCount = relationMap.size;
console.log(`实体数量: ${entityCount}`);
console.log(`关系数量: ${relationCount}`);
有了 ID 映射之后,还需要将原始三元组转换为数字索引形式,因为嵌入表只能通过数字索引访问。转换后的三元组数组会直接用于后续训练。在实际项目中,如果图谱规模达到千万级,直接使用 JavaScript 的数组存储所有三元组可能占用太多内存,此时可以考虑使用 Int32Array 来压缩存储,但在小型实验中普通数组已经足够。
负采样是知识图谱嵌入训练中至关重要的一步。负样本的质量会直接影响模型的效果。最简单的采样方式是从实体集合中均匀随机选取一个实体,替换当前三元组的头实体或尾实体。需要注意,替换后可能仍然构成一个真实存在的三元组,这种负样本被称为“假负例”。虽然 TransE 对少量假负例不敏感,但在工程实践中可以通过维护一个哈希集合来过滤,从而提高训练质量。
function sampleNegative(head, relation, tail) {
// 随机选一个实体
const chosenEntity = Math.floor(Math.random() * entityCount);
const rand = Math.random();
// 一半概率替换头实体,一半概率替换尾实体
if (rand < 0.5) {
return [chosenEntity, relation, tail];
} else {
return [head, relation, chosenEntity];
}
}
三、实现 TransE 训练过程
有了正负样本,就可以实现训练主流程。TransE 的参数只有实体向量和关系向量。初始化时通常采用均匀分布或正态分布,并且需要将实体向量的 L2 范数归一化为 1。关系向量是否需要归一化取决于具体实现,但大多数场景下会对关系向量也做归一化,以避免训练过程中出现梯度爆炸。
下面给出一个简化但完整的 TransE JavaScript 实现。为了便于阅读,代码中省略了部分优化细节,例如小批量迭代时的随机打乱和梯度裁剪。核心训练逻辑是:对每个正样本产生负样本,计算两个样本的得分差值,如果差值小于 margin 则执行梯度更新。这里使用随机梯度下降,学习率固定为 0.01。
class TransE {
constructor(entityCount, relationCount, dim, margin, lr) {
this.entityCount = entityCount;
this.relationCount = relationCount;
this.dim = dim;
this.margin = margin;
this.lr = lr;
this.entityVec = new Map();
this.relationVec = new Map();
// 初始化实体向量并归一化
for (let i = 0; i < entityCount; i++) {
const vec = this.randomVector();
this.normalize(vec);
this.entityVec.set(i, vec);
}
// 初始化关系向量并归一化
for (let j = 0; j < relationCount; j++) {
const vec = this.randomVector();
this.normalize(vec);
this.relationVec.set(j, vec);
}
}
randomVector() {
const vec = new Array(this.dim);
for (let i = 0; i < this.dim; i++) {
// 均匀分布范围 [-0.5, 0.5]
vec[i] = Math.random() - 0.5;
}
return vec;
}
normalize(vec) {
let sum = 0;
for (const v of vec) {
sum += v * v;
}
const norm = Math.sqrt(sum);
for (let i = 0; i < vec.length; i++) {
vec[i] /= norm;
}
}
score(headId, relId, tailId) {
const head = this.entityVec.get(headId);
const rel = this.relationVec.get(relId);
const tail = this.entityVec.get(tailId);
let sum = 0;
for (let i = 0; i < this.dim; i++) {
const diff = head[i] + rel[i] - tail[i];
sum += diff * diff;
}
return Math.sqrt(sum);
}
train(triples, epochs, batchSize) {
for (let epoch = 0; epoch < epochs; epoch++) {
let totalLoss = 0;
// 将三元组转换为数字索引
const indexed = triples.map(([h, r, t]) => [
entityMap.get(h),
relationMap.get(r),
entityMap.get(t)
]);
// 随机打乱并选取一个小批量
const shuffled = indexed.sort(() => Math.random() - 0.5);
const batch = shuffled.slice(0, batchSize);
for (const [h, r, t] of batch) {
const [hn, , tn] = sampleNegative(h, r, t);
const posScore = this.score(h, r, t);
const negScore = this.score(hn, r, tn);
const loss = Math.max(0, this.margin + posScore - negScore);
if (loss > 0) {
// 梯度下降更新,这里只做示意
// 实际需要分别对 h, r, t, hn, tn 的向量分量求梯度并更新
const head = this.entityVec.get(h);
const rel = this.relationVec.get(r);
const tail = this.entityVec.get(t);
const negHead = this.entityVec.get(hn);
const negTail = this.entityVec.get(tn);
for (let i = 0; i < this.dim; i++) {
const diffPos = head[i] + rel[i] - tail[i];
const diffNeg = negHead[i] + rel[i] - negTail[i];
head[i] -= this.lr * diffPos;
rel[i] -= this.lr * diffPos;
tail[i] += this.lr * diffPos;
negHead[i] += this.lr * diffNeg;
negTail[i] -= this.lr * diffNeg;
}
totalLoss += loss;
}
}
// 每轮结束后归一化实体向量
for (const vec of this.entityVec.values()) {
this.normalize(vec);
}
console.log(`Epoch ${epoch + 1}, Loss: ${totalLoss.toFixed(4)}`);
}
}
}
上述代码中,sampleNegative 返回的负样本只替换了头实体或尾实体,没有替换关系。这是 TransE 的标准做法,因为关系集合通常远小于实体集合,采样关系容易导致负样本过于稀疏。另外,梯度更新时对所有向量分量直接使用原始差值,相当于在欧氏距离损失上执行梯度下降。实际项目中可以加入正则化项和更精细的批次归一化。
训练完模型后,所有实体和关系都被映射到低维向量空间。我们可以利用这些向量完成链接预测任务。例如已知(?,首都,中国),需要预测头实体。做法是遍历所有实体,计算每个实体的向量加上关系向量后与尾实体向量的距离,距离最小的实体就是预测结果。
function predictHead(relId, tailId) {
let bestEntity = -1;
let bestScore = Infinity;
for (let i = 0; i < entityCount; i++) {
const score = model.score(i, relId, tailId);
if (score < bestScore) {
bestScore = score;
bestEntity = i;
}
}
return bestEntity;
}
// 使用示例
// const model = new TransE(entityCount, relationCount, 64, 1.0, 0.01);
// model.train(triples, 100, 128);
// const predicted = predictHead(relationMap.get('首都'), entityMap.get('中国'));
四、性能优化与应用场景
纯 JavaScript 实现的 TransE 在小规模图谱上运行顺畅,但在大规模数据集上面临性能压力。最直接的优化手段是使用 TypedArray 存储嵌入向量。TypedArray 在内存中连续存放,访问速度远高于普通数组,并且能够有效减少垃圾回收的负担。例如用 Float64Array 存储所有实体向量,通过实体 ID 乘以维度得到向量起始位置,彻底摆脱 Map 的查找开销。
另一个优化方向是引入 Node.js 的 worker_threads 模块进行多线程训练。TransE 的训练过程本质上是频繁的矩阵运算,可以将实体向量划分到多个 worker 中,每个 worker 负责处理一部分实体,训练过程中定期同步梯度。此外,为了加快采样速度,可以预先生成一批负样本缓存到内存中,避免训练时频繁调用随机函数导致的性能抖动。
知识图谱嵌入在生产环境中有着广泛的应用。在推荐系统里,用户和商品可以被视为实体,用户的行为视为关系,嵌入向量可以直接用于计算用户与商品之间的相似度。在智能问答中,嵌入模型可以将自然语言中的实体短语映射到图谱实体,再通过向量推理找到答案。Node.js 版本虽然在大规模预训练上不如 Python 生态成熟,但非常适合部署到内存受限的微服务中,因为 Node.js 实例本身启动快,嵌入推理的接口可以用简单的 HTTP 服务暴露给上层应用。
总体而言,从符号到向量的转换让知识图谱具备了计算能力。本文用 Node.js 实现了一个可运行的 TransE 模型,代码覆盖了数据预处理、负采样、训练和推理全流程。如果你需要处理更复杂的语义关系,可以考虑升级到 TransH 或 RotatE,它们的核心思想仍然是在向量空间中定义关系运算。理解 TransE 是学习更高级知识图谱嵌入模型的最好起点。