Facebook 开源了超大规模图嵌入算法,上亿个节点也能快速完成

方栗子 发自 凹非寺

量子位 报道 | 公众号 QbitAI

图,是很有用的数据结构,用节点 (Node) 和边 (Edge) 织成一张网。比如,知识图谱就是这样的网。

null

处理这样的数据,要用到图嵌入 (Graph Embedding) :把高维的网络,转换成低维的向量。处理之后,机器学习模型才能轻易食用。

如果像上图这样,只有几个节点,十几条边,图嵌入没什么难度。

如果有几十亿个节点,几万亿条边呢?传统的图嵌入方法就捉急了。

但现在,Facebook 开源了叫做 PyTorch-BigGraph (简称 PBG) 的新工具。

有了它,再大的图 (原文是 arbitrarily large,任意大) 都能快速生成图嵌入。而且,完全不需要 GPU。

开源模型除了可以自己训练之外,还有拿 7,800 万节点的 Wikidata 数据预训练过的模型,可以直接下载来用。

PBG 发表之后,小伙伴们纷纷奔走相告:

null

LeCun 老师还亲自转发了两次。

如何养成

PBG 是一个分布式系统,用 1.2 亿个节点的完整 FreeBase 知识图谱来训练的。

null

训练过程中,PBG 会吃进图上所有边 (Edge) 的大列表,每条边都是用它两端的节点来定义,一个是源 (Source) ,一个是目标 (Target) 。定义中也有两点之间的关系 (Relation Type) 。

然后,PBG 给每一个节点,输出一个特征向量 (就是嵌入) ,让两个相邻的节点在向量空间中离得近一些,让不相邻节点的离远一些。

这样一来,那些周围环境分布相近的节点,在向量空间里的位置也会彼此靠近,图原本要表达的意思就保留下来了。

另外,针对每种不同的关系,“近似度得分 (Proximity Score) ”都可以定制不同的计算方法。这样,一个节点的嵌入,就可以在不同种类的关系里共享了。

快一点,再快一点

要快速处理大规模的图数据,PBG 用了这几个法术:

一是图分区 (Graph Partitioning) ,这样就不需要把整个模型加载到内存里了。在图嵌入质量不损失的情况下,比不分区时节省了 88% 的内存占用。二是一台机器进行多线程计算。三是在多台机器上同时跑,在图上各自跑一个不相邻的区域。四是批次负采样 (Batched Negative Sampling) ,能让一台 CPU 每秒处理 100 万条边,每条边 100 次负采样。

训练完成之后,在 FB15k、Youtube、LiveJournal 等等图谱上,都测试过。

团队说,PBG 和大前辈们的图嵌入质量相当,但需要的时间明显缩短了。

你也试一试吧

你也去训练一个 PBG 吧。

反正,也不用 GPU。

不想训练的话,还有用完整 WikiData 预训练的模型。

你需要的传送门,都在这里了。

代码传送门:

https://github.com/facebookresearch/PyTorch-BigGraph

食用索引:

https://torchbiggraph.readthedocs.io/en/latest/

论文传送门:

https://arxiv.org/abs/1903.12287

博客传送门:

https://ai.facebook.com/blog/open-sourcing-pytorch-biggraph-for-faster-embeddings-of-extremely-large-graphs/