[논문 리뷰] PyTorch-BigGraph: A Large-scale Graph Embedding System

Adam Lerer, Ledell Wu외 5명|arXiv (Cornell University)|2019. 03. 28.

Advanced Graph Neural Networks참고 문헌 40인용 수 204

한 줄 요약

PyTorch-BigGraph (PBG)는 파티셔닝, 분산 실행, 배치 음수 샘플링, 그리고 하이브리드 파라미터 서버 접근법을 사용하여 그래프 임베딩을 수십억 개의 노드와 트릴리언 엣지로 확장하고, 확장성 향상과 함께 경쟁력 있는 성능을 달성합니다.

ABSTRACT

Graph embedding methods produce unsupervised node features from graphs that can then be used for a variety of machine learning tasks. Modern graphs, particularly in industrial applications, contain billions of nodes and trillions of edges, which exceeds the capability of existing embedding systems. We present PyTorch-BigGraph (PBG), an embedding system that incorporates several modifications to traditional multi-relation embedding systems that allow it to scale to graphs with billions of nodes and trillions of edges. PBG uses graph partitioning to train arbitrarily large embeddings on either a single machine or in a distributed environment. We demonstrate comparable performance with existing embedding systems on common benchmarks, while allowing for scaling to arbitrarily large graphs and parallelization on multiple machines. We train and evaluate embeddings on several large social network graphs as well as the full Freebase dataset, which contains over 100 million nodes and 2 billion edges.

연구 동기 및 목표

• 산업 규모의 그래프에서 수십억 개의 노드와 트릴리언 엣지를 가진 그래프 임베딩의 확장 필요성에 대해 동기를 부여한다.
• PyTorch-BigGraph (PBG)와 그 확장성 메커니즘을 소개한다.
• PBG가 대규모 그래프까지 확장하면서 표준 벤치마크에서 경쟁력 있는 성능을 달성함을 입증한다.
• Freebase, LiveJournal, YouTube, Twitter에서의 결과를 통해 확장성 및 효율성을 검증한다.
• 대규모 그래프 임베딩 파이프라인에 대한 오픈 소스 릴리스 및 실용적 가이드를 제공한다

제안 방법

• 관계별 변환을 갖는 인접 행렬의 요인화 스코어링 함수 f를 사용한 다중 관계 그래프 임베딩 프레임워크를 도입한다.
• 메모리 제약에 맞추고 분산 실행을 가능하게 하려면 인접 행렬을 P 버킷으로 분할하고 한 번에 하나의 버킷에서 학습한다.
• 락-서버와 파티션-서버 기반의 병렬화를 사용해 여러 기계에서의 학습을 조정하고 비동기 파라미터 서버를 통해 작은 글로벌 파라미터 집합을 공유한다.
• 메모리 대역폭을 줄이고 배치 내 음수를 재사용하여 메모리 효율적인 배치 음수 샘플링을 구현한다.
• Adagrad와 경계 기반 랭킹 목적식을 최적화하고, 샘플링 및 관계 연산자(예: TransE, DistMult, ComplEx, RESCAL)를 구성 가능하게 한다.
• 단일 머신 및 분산 학습 모두를 지원하며, 분할된 엔터티 유형과 관계별 엣지 구성의 모듈식 처리를 포함한다

실험 결과

연구 질문

• RQ1 수십억 개의 노드와 트릴리언 엣지를 가진 대규모 그래프를 파티셔닝과 분산 학습으로 임베딩하더라도 임베딩 품질이 손실되지 않는가?
• RQ2 파티션 전략, 음수 샘플링, 관계 연산자가 대규모 실제 그래프 및 지식 그래프에서 확장성과 모델 정확도에 어떤 영향을 미치는가?
• RQ3 여러 기계와 파티션으로 PBG를 확장할 때 학습 시간, 메모리 사용량, 수렴의 트레이드오프는 무엇인가?
• RQ4 다양한 그래프 유형(소셜 네트워크 vs 지식 그래프)을 서로 다른 스코어링 연산자(TransE, ComplEx 등)로 사용할 때 PBG의 성능은 어떤가?
• RQ5 전체 Freebase 그래프와 같이 대규모 임베딩의 릴리스가 가능한가? 품질 및 자원 사용 측면에서 실용적인가?

주요 결과

모델 / 구성	MRR( Raw )	MRR( Filtered )	Hits@10	비고
RESCAL	0.189	0.354	0.587	FB15k baseline
TransE	0.265	0.594	0.785	PBG with TransE (FB15k)
HolE	0.232	0.524	0.739	FB15k
ComplEx	0.242	0.790	0.872	FB15k
R-GCN+	-	-	-	Knowledge-graph related baseline
StarSpace	-	-	0.838	Baseline comparison
Reciprocal ComplEx-N3	-	-	0.910	Baseline comparison
PBG (TransE)	0.265	0.594	0.785	PBG variant on FB15k
PBG (ComplEx)	0.242	0.790	0.872	PBG variant on FB15k

• PBG는 표준 벤치마크에서 기존 임베딩 시스템의 성능과 일치하면서도 수십억 노드와 트릴리언 엣지의 그래프까지 확장한다.
• 메모리 사용을 현저히 감소시키는 파티셔닝이 임베딩 품질의 약간의 저하만으로도(예: 전체 Freebase 그래프의 88% 감소) 크게 감소를 줄인다.
• 여러 기계에 걸친 분산 실행은 대규모 그래프에서 거의 선형에 가까운 확장과 함께(예: 8대 기계에서 Freebase의 벽시계 시간 4배까지의 속도 향상) 상당한 속도 향상을 제공한다.
• 배치 내 음수를 재사용하여 학습 속도를 크게 향상시키는 배치 친화적 배치 음수 샘플링은 음수의 수가 늘어나도 효율성을 유지한다.
• 실제 대규모 그래프(LiveJournal, YouTube, Twitter)에서 경쟁력 있는 연결 예측 및 노드 분류 성능을 보여주며, 임베딩은 다운스트림 작업의 특징으로도 활용 가능하다.
• Freebase 지식 그래프에서 TransE 및 ComplEx를 사용한 PBG는 기존 기준선에 비해 경쟁력 있는 MRR 및 Hits@10을 달성하며, 분할 및 분산 학습을 통해 전체 Freebase 그래프로의 확장이 가능하다