[논문 리뷰] SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive Knowledge Graphs

Hongyu Ren, Hanjun Dai외 5명|arXiv (Cornell University)|2021. 10. 28.

Advanced Graph Neural Networks참고 문헌 39인용 수 7

한 줄 요약

SMORE는 대규모 지식 그래프에서 단일 힙 지식 그래프 보완과 다중 힙 추론을 모두 위한 확장 가능한 프레임워크이다. 이는 이중 방향 거부 샘플링과 비동기 CPU-GPU 계산을 통해 최신 기준의 훈련 처리량을 달성한다. SMORE는 Freebase(86M 개체, 338M 간선)에서 다중 힙 추론을 수행하며 기존 프레임워크 대비 2.2배 빠른 성능을 기록하고, 8개 GPU에서 거의 선형적인 확장성을 보이며 GPU 메모리가 2GB 뿐이어도 가능하다.

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.

연구 동기 및 목표

• 대규모 지식 그래프에서 다중 힙 추론을 위한 확장 가능한 프레임워크의 부족을 해결하기 위해.
• 모든 훈련 인스턴스를 사전에 생성하지 않고도 다중 힙 쿼리에 대한 효율적인 온라인 훈련 데이터 생성을 가능하게 하기 위해.
• 고성능과 낮은 메모리 사용량을 달성하기 위해 알고리즘적 설계와 시스템 수준의 최적화를 동시에 고려하기 위해.
• 통합된 확장 가능한 프레임워크 내에서 단일 힙과 다중 힙 추론 작업을 모두 지원하기 위해.

제안 방법

• 온라인 훈련 데이터 생성의 복잡도를 제곱근 정도로 감소시키기 위해 이중 방향 거부 샘플링을 도입한다.
• CPU 기반의 데이터 샘플링과 GPU 기반의 임베딩 계산, 빈번한 CPU-GPU I/O를 겹치기 위해 비동기 스케줄링을 활용한다.
• 관계 간의 다중 힙 경로 무결성을 유지하기 위해 분할되지 않은 단일 머신 기반 그래프 설계를 사용한다.
• 훈련 중 GPU 메모리 사용량을 최소화하기 위해 희소 임베딩과 최적화 업데이트를 적용한다.
• GQE, Q2B, BetaE 등의 다수의 쿼리 임베딩 모델을 지원하며, 단일 GPU 및 다중 GPU 환경에서 통합된 훈련 파이프라인을 제공한다.
• 모든 엔티티를 전수 검색하지 않고도 이중 방향 탐색을 통해 부정 샘플 엔티티를 효율적으로 샘플링함으로써 확장성을 향상시킨다.

실험 결과

연구 질문

• RQ1 8000만 개 이상의 개체와 3억 개 이상의 간선을 가진 대규모 지식 그래프에서 다중 힙 추론을 효율적으로 확장할 수 있는가?
• RQ2 다중 힙 쿼리에 대한 온라인 훈련 데이터 생성이 얼마나 효율적으로 이루어져야 GPU 컴퓨팅 용량을 완전히 활용할 수 있는가?
• RQ3 통합된 프레임워크가 최소한의 시스템 오버헤드로 단일 힙 지식 그래프 보완과 복잡한 다중 힙 추론을 모두 지원할 수 있는가?
• RQ4 대규모 지식 그래프에서 고처리량과 낮은 메모리 사용량을 달성하기 위해 어떤 시스템 최적화가 필요한가?
• RQ5 이중 방향 샘플링이 훈련 효율성과 모델 성능 측면에서 전수 탐색 또는 무작위 샘플링보다 어떻게 비교되는가?

주요 결과

• SMORE는 Freebase에서 400차원 임베딩을 훈련할 때 기존의 다중 힙 지식 그래프 프레임워크보다 2.2배 더 높은 처리량을 기록하며, GPU 메모리가 오직 2GB 뿐이어도 가능하다.
• GPU 수에 따라 훈련 속도가 거의 선형적으로 증가하여, 통신 오버헤드가 최소화되고 비동기 설계가 효과적으로 기능함을 보여준다.
• ogbl-wikikg2 데이터셋에서 BetaE는 베이스라인 대비 최고의 성능을 기록했고, Q2B는 FB400k와 Freebase에서 최신 기준의 성능을 달성했다.
• 이중 방향 샘플링은 전수 탐색과 동등하거나 그 이상의 성능을 기록하면서도, 대규모 그래프에서 훈련 시간을 수개월에서 수일로 단축시켰다.
• 비분할 그래프 설계임에도 불구하고, Marius와 PBG와 같은 최신 기준의 프레임워크와 비교해도 단일 힙 링크 예측 성능은 유사하거나 더 우수했다.
• 시스템은 Freebase(86M 노드, 338M 간선)까지 성공적으로 확장되었으며, 이는 이전에 다중 힙 추론에서 지원된 지식 그래프보다 1500배 이상 큰 규모이다.