[논문 리뷰] Wasserstein Embedding for Graph Learning

Soheil Kolouri, Navid Naderializadeh외 2명|arXiv (Cornell University)|2020. 06. 16.

Advanced Graph Neural Networks참고 문헌 64인용 수 4

한 줄 요약

WEGL은 워셔스타인 거리 기반으로 노드 임베딩 간의 유사도를 활용하여 고정 크기의 벡터 표현을 생성함으로써, 기존 방법이 요구하는 이중 비교의 제곱형 복잡도에서 벗어나 선형 복잡도로 감소시키는 선형적이고 계산적으로 효율적인 그래프 임베딩 프레임워크를 제안한다. 이는 그래프 수에 대해 선형 복잡도를 가지며, 최신 기준 성능을 달성하면서도 빠른 그래프 수준의 예측을 가능하게 한다.

ABSTRACT

We present Wasserstein Embedding for Graph Learning (WEGL), a novel and fast framework for embedding entire graphs in a vector space, in which various machine learning models are applicable for graph-level prediction tasks. We leverage new insights on defining similarity between graphs as a function of the similarity between their node embedding distributions. Specifically, we use the Wasserstein distance to measure the dissimilarity between node embeddings of different graphs. Unlike prior work, we avoid pairwise calculation of distances between graphs and reduce the computational complexity from quadratic to linear in the number of graphs. WEGL calculates Monge maps from a reference distribution to each node embedding and, based on these maps, creates a fixed-sized vector representation of the graph. We evaluate our new graph embedding approach on various benchmark graph-property prediction tasks, showing state-of-the-art classification performance while having superior computational efficiency. The code is available at https://github.com/navid-naderi/WEGL.

연구 동기 및 목표

• 모든 그래프 간의 이중 유사도 계산이 필요한 그래프 커널 방법의 높은 계산 비용을 해결하기 위해.
• 대규모 그래프 데이터셋에서 그래프 신경망(GNNs)과 커널 기반 방법의 확장성 한계를 극복하기 위해.
• 워셔스타인 거리 기반으로 메트릭의 정확성을 유지하면서도 어떤 후속 기계학습 모델에도 적용 가능한 그래프 임베딩 방법을 개발하기 위해.
• 워셔스타인 기반 그래프 유사도 추정에서 그래프 수에 대해 복잡도를 제곱형에서 선형형으로 감소시키기 위해.
• 대규모 데이터셋에서 그래프 성질 예측 작업에 대한 효율적인 학습 및 추론을 가능하게 하기 위해.

제안 방법

• WEGL은 각 그래프의 노드 임베딩에서 기준 노드 임베딩 분포로의 몽제 맵(Monge maps)을 계산하여 운반 계획을 유도한다.
• 최적의 운반 맵을 사용하여 각 그래프를 힐버트 공간에 임베딩함으로써 고정 크기의 벡터 표현을 도출한다.
• 임베딩된 그래프 간의 유클리드 거리는 노드 임베딩 분포 간의 2-워셔스타인 거리를 근사하며, 메트릭의 구조를 유지한다.
• 이 방법은 모든 그래프 간의 이중 거리 계산을 피하기 위해 선형 최적 운반을 활용하여, 복잡도를 O(M²)에서 O(M)로 감소시킨다.
• 노드 임베딩는 워셔스타인 거리 계산의 입력 분포로 사용되며, 기준 분포는 그래프 간의 풀링 또는 학습된 분포로부터 유도된다.
• 명시적인 힐버트 공간 임베딩 덕분에 SVM, 랜덤 포레스트, 기울기 부스팅 트리 등 어떤 후속 분류기와도 호환 가능하다.

실험 결과

연구 질문

• RQ1 그래프 임베딩에서 계산 복잡도를 감소시키면서도 최신 기준 성능을 달성할 수 있는가?
• RQ2 유클리드 거리가 노드 임베딩 분포 간의 2-워셔스타인 거리를 근사하는 방식으로 힐버트 공간에 그래프를 임베딩할 수 있는가?
• RQ3 워셔스타인 기반 그래프 유사도에 대해 선형 복잡도 접근이 제곱형 복잡도 방법보다 확장성과 성능 면에서 뛰어나게 되는가?
• RQ4 제안된 임베딩 프레임워크는 커널 방법에 국한되지 않고 어떤 표준 기계학습 분류기와도 사용 가능한가?
• RQ5 대규모 그래프 데이터셋에서 GNN 및 그래프 커널 기반 베이스라인과 비교해 WEGL은 효율성과 정확도 면에서 어떻게 성능을 내는가?

주요 결과

• WEGL은 ogbg-molhiv 및 TUD 그래프 분류 작업을 포함한 여러 벤치마크 데이터셋에서 최신 기준 또는 경쟁력 있는 분류 정확도를 달성한다.
• IMDB-BINARY 데이터셋에서 WEGL은 GBDT를 사용해 75.2%의 정확도를 기록했으며, GIN 및 기타 베이스라인을 모두 초월한다.
• PROTEINS 데이터셋에서 WEGL은 GBDT를 사용해 92.9%의 정확도를 기록했으며, 다음으로 좋은 방법보다 뚜렷이 뛰어나다.
• 특히 그래프 수가 많은 데이터셋에서, WEGL의 학습 시간은 WWL 및 GIN보다 수 개월 정도 빠르며, 이는 선형 복잡도 덕분이다.
• WEGL의 추론 속도는 CPU 기반 GIN보다 빠르며, GPU 가속된 GIN조차도 뛰어나며, WWL보다도 현저히 신속하다.
• 분자의 성질 예측 및 사회망 분석을 포함한 다양한 데이터셋에서 낮은 계산 오버헤드로도 높은 성능를 유지한다.