[논문 리뷰] DGL-KE: Training Knowledge Graph Embeddings at Scale

Da Zheng, Xiang Song외 7명|arXiv (Cornell University)|2020. 04. 18.

Advanced Graph Neural Networks참고 문헌 17인용 수 36

한 줄 요약

DGL-KE는 대규모 그래프에서 지식 그래프 임베딩의 확장 가능한 학습을 위한 새로운 최적화를 도입하여 다중 GPU, 다수 코어 CPU, 분산 환경에서 GraphVite 및 Pytorch-BigGraph에 비해 상당한 속도 향상을 달성합니다.

ABSTRACT

Knowledge graphs have emerged as a key abstraction for organizing information in diverse domains and their embeddings are increasingly used to harness their information in various information retrieval and machine learning tasks. However, the ever growing size of knowledge graphs requires computationally efficient algorithms capable of scaling to graphs with millions of nodes and billions of edges. This paper presents DGL-KE, an open-source package to efficiently compute knowledge graph embeddings. DGL-KE introduces various novel optimizations that accelerate training on knowledge graphs with millions of nodes and billions of edges using multi-processing, multi-GPU, and distributed parallelism. These optimizations are designed to increase data locality, reduce communication overhead, overlap computations with memory accesses, and achieve high operation efficiency. Experiments on knowledge graphs consisting of over 86M nodes and 338M edges show that DGL-KE can compute embeddings in 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 cluster with 4 machines with 48 cores/machine. These results represent a 2x~5x speedup over the best competing approaches. DGL-KE is available on https://github.com/awslabs/dgl-ke.

연구 동기 및 목표

• 수백만 개의 노드와 수십억 개의 에지를 가진 거대 지식 그래프의 확장 가능한 임베딩 학습을 촉진합니다.
• 데이터 지역성 향상, 통신 감소, KGE 학습의 계산 overlapped를 개선하기 위한 최적화를 개발하고 구현합니다.
• 대규모 학습을 가능하게 하는 DGL, PyTorch/MXNet 및 분산 KVStore 위에 구축된 오픈 소스 패키지(DGL-KE)를 제공합니다.

제안 방법

• 지식 그래프를 컴퓨팅 유닛 간에 파티션하여 파티션 간 데이터 이동을 최소화합니다.
• METIS 기반 그래프 파티셔닝을 사용해 임베딩과 트리플을 근접 배치하고 기계 간 통신을 줄입니다.
• 메모리 이동을 줄이고 더 어려운 음수를 생성하기 위해 결합 음수 샘플링과 차수 기반 샘플링을 구현합니다.
• 여러 관계의 데이터 전송을 최소화하기 위해 관계 임베딩을 분할하고 희소한 읽기/업데이트를 전파합니다.
• CPU와 GPU 간의 그래디언트 업데이트를 배치 처리와 중첩시켜 CPU-GPU 대기 시간을 숨기고 처리량을 향상합니다.
• 주기적 동기화와 분산 KVStore를 도입해 머신 간 학습을 조정합니다.

실험 결과

연구 질문

• RQ1 수십만 개의 노드와 수십억 개의 간선을 가진 그래프에서 지식 그래프 임베딩을 대규모로 효율적으로 학습하려면 어떻게 해야 합니까?
• RQ2 그래프 파티셔닝, 음수 샘플링, 데이터 배치 최적화의 어떤 조합이 속도와 임베딩 품질 사이에서 최상의 트레이드오프를 제공합니까?
• RQ3 비동기식 멀티프로세싱 및 분산 학습이 데이터 전송 및 메모리 병목을 줄이면서도 경쟁력 있는 임베딩 품질을 달성할 수 있습니까?
• RQ4 DGL-KE의 성능은 대규모 지식 그래프에서 기존 도구(GraphVite, PBG)와 비교하여 어떠합니까?
• RQ5 관계 분할 및 그래디언트 업데이트의 중첩이 학습 처리량과 모델 정확도에 미치는 실용적 영향은 무엇입니까?

주요 결과

• DGL-KE는 상당한 속도 향상을 달성합니다. 예를 들어 86M 노드 및 338M 에지의 그래프에 대해 8-GPU EC2 인스턴스에서 100분, 48코어를 가진 4대 머신 클러스터에서 30분이 소요되는 반면 경쟁 접근 방식에 비해 우수한 성능을 보입니다.
• Joint 음수 샘플링은 다중-GPU 설정에서 데이터 이동을 최대 약 40배까지 감소시키고 텐서 연산 효율을 향상시킵니다.
• 차수 기반 음수 샘플링은 Freebase와 같은 대규모 그래프에서 모델 정확도를 향상시키지만 학습 속도를 증가시키지 못할 수 있습니다.
• CPU와 GPU 간에 그래디언트 업데이트를 중첩하면 대규모 그래프의 많은 모델에서 약 40%의 속도향상을 얻으면서도 정확도를 유지합니다.
• 관계 분할은 데이터 이동을 크게 감소시키고 특히 관계 관련 매개변수가 큰 모델(예: TransR)에 대해 상당한 속도 향상을 제공합니다.
• 전반적으로 DGL-KE는 GPU 수와 거의 선형으로 확장되며 대규모 그래프에서 복잡한 KGE 모델을 수 시간 내에 학습할 수 있어 GraphVite 및 PBG보다 우수한 것으로 보고된 실험에서 나타납니다.