[论文解读] SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive Knowledge Graphs

Hongyu Ren, Hanjun Dai+5 more|arXiv (Cornell University)|Oct 28, 2021

Advanced Graph Neural Networks参考文献 39被引用 7

一句话总结

SMORE 是一个可扩展的框架，适用于大规模知识图谱中的单跳知识图谱补全和多跳推理，通过双向拒绝采样和异步 CPU-GPU 计算实现最先进的训练吞吐量。它在 Freebase（8600 万个实体，3.38 亿条边）上实现了多跳推理，相比先前框架提速 2.2 倍，并且在仅使用 2GB GPU 内存的情况下，8 块 GPU 上实现近乎线性的扩展性能。

ABSTRACT

Knowledge graphs (KGs) capture knowledge in the form of head--relation--tail triples and are a crucial component in many AI systems. There are two important reasoning tasks on KGs: (1) single-hop knowledge graph completion, which involves predicting individual links in the KG; and (2), multi-hop reasoning, where the goal is to predict which KG entities satisfy a given logical query. Embedding-based methods solve both tasks by first computing an embedding for each entity and relation, then using them to form predictions. However, existing scalable KG embedding frameworks only support single-hop knowledge graph completion and cannot be applied to the more challenging multi-hop reasoning task. Here we present Scalable Multi-hOp REasoning (SMORE), the first general framework for both single-hop and multi-hop reasoning in KGs. Using a single machine SMORE can perform multi-hop reasoning in Freebase KG (86M entities, 338M edges), which is 1,500x larger than previously considered KGs. The key to SMORE's runtime performance is a novel bidirectional rejection sampling that achieves a square root reduction of the complexity of online training data generation. Furthermore, SMORE exploits asynchronous scheduling, overlapping CPU-based data sampling, GPU-based embedding computation, and frequent CPU--GPU IO. SMORE increases throughput (i.e., training speed) over prior multi-hop KG frameworks by 2.2x with minimal GPU memory requirements (2GB for training 400-dim embeddings on 86M-node Freebase) and achieves near linear speed-up with the number of GPUs. Moreover, on the simpler single-hop knowledge graph completion task SMORE achieves comparable or even better runtime performance to state-of-the-art frameworks on both single GPU and multi-GPU settings.

研究动机与目标

• 解决大规模知识图谱中多跳推理缺乏可扩展框架的问题。
• 实现在不生成所有训练实例的前提下，高效生成多跳查询的在线训练数据。
• 协同优化算法与系统级设计，以在大规模 KG 上实现高吞吐量和低内存使用。
• 在一个统一且可扩展的框架中同时支持单跳和多跳推理任务。

提出的方法

• 引入双向拒绝采样，将在线训练数据生成的复杂度降低一个平方根因子。
• 采用异步调度，重叠 CPU 上的数据采样与 GPU 上的嵌入计算，以及频繁的 CPU-GPU I/O 操作。
• 使用单机、非分区图设计，以保持跨关系的多跳路径完整性。
• 应用稀疏嵌入和优化器更新，以最小化训练期间的 GPU 内存使用。
• 支持多种查询嵌入模型（例如 GQE、Q2B、BetaE），并在单 GPU 和多 GPU 设置下提供统一的训练流水线。
• 通过双向遍历高效采样负样本实体，避免穷举搜索，提升可扩展性。

实验结果

研究问题

• RQ1 多跳推理能否在包含超过 8000 万个实体和 3 亿条边的大规模知识图谱上实现高效扩展？
• RQ2 如何使多跳查询的在线训练数据生成足够高效，以充分利用 GPU 计算能力？
• RQ3 统一框架能否在最小系统开销下同时支持单跳知识图谱补全和复杂多跳推理？
• RQ4 为实现大规模 KG 上的高吞吐量和低内存使用，需要哪些系统优化？
• RQ5 与穷举或随机采样相比，双向拒绝采样在训练效率和模型性能方面表现如何？

主要发现

• SMORE 在 Freebase 上训练 400 维嵌入时，仅使用 2GB GPU 内存，相比先前的多跳 KG 框架，吞吐量提升 2.2 倍。
• 训练速度随着 GPU 数量的增加近乎线性扩展，表明通信开销极低，且异步设计有效。
• 在 ogbl-wikikg2 数据集上，BetaE 在基线中表现最佳，而 Q2B 在 FB400k 和 Freebase 上达到最先进性能。
• 双向采样在性能上优于或匹配穷举遍历，同时将大规模图上的训练时间从数月缩短至数天。
• 即使采用非分区图设计，SMORE 在单跳链接预测任务上的性能也与最先进的框架（如 Marius 和 PBG）相当或更优。
• 该系统成功扩展至 Freebase（8600 万个节点，3.38 亿条边），其规模是此前多跳推理所支持 KG 的 1500 倍。