[논문 리뷰] Optimal Transport Graph Neural Networks

Benson Chen, Gary Bécigneul외 3명|arXiv (Cornell University)|2020. 06. 08.

Advanced Graph Neural Networks참고 문헌 61인용 수 22

한 줄 요약

이 논문은 워셔슈타인 거리와 학습 가능한 프로토타입 클러스터를 사용하여 그래프 임베딩을 계산하는 그래프 신경망 OT-GNN을 소개한다. 이는 보편적 근사와 분자 성질 예측 성능 향상에 기여한다. 최적 운반 이론과 매개변수화된 프로토타입, 노이즈 대비 정규화를 조합함으로써 모델은 집합 붕괴를 피하고 표준 GNN보다 더 매끄럽고 정보가 풍부한 표현을 얻는다.

ABSTRACT

Current graph neural network (GNN) architectures naively average or sum node embeddings into an aggregated graph representation -- potentially losing structural or semantic information. We here introduce OT-GNN, a model that computes graph embeddings using parametric prototypes that highlight key facets of different graph aspects. Towards this goal, we successfully combine optimal transport (OT) with parametric graph models. Graph representations are obtained from Wasserstein distances between the set of GNN node embeddings and ``prototype'' point clouds as free parameters. We theoretically prove that, unlike traditional sum aggregation, our function class on point clouds satisfies a fundamental universal approximation theorem. Empirically, we address an inherent collapse optimization issue by proposing a noise contrastive regularizer to steer the model towards truly exploiting the OT geometry. Finally, we outperform popular methods on several molecular property prediction tasks, while exhibiting smoother graph representations.

연구 동기 및 목표

• 표준 GNN이 노드 임베딩을 단순히 합하거나 평균하는 방식으로 구조적 및 의미적 정보를 손실하는 한계를 해결하기 위해.
• 입력 그래프를 추상적이고 학습 가능한 프로토타입과 비교하기 위해 최적 운반 이론(Wasserstein 거리)을 활용하는 매개변수화된 그래프 모델을 개발하기 위해.
• 특히 프로토타입이 단일 점으로 분해되는 프로토타입 붕괴 현상으로 인해 발생하는 최적 운반 기반 모델의 최적화 과제를 극복하기 위해.
• 제안된 함수 클래스가 표준 GNN 집합 기반 집합과 비교해 더 높은 표현 능력을 가지는 데 이론적 근거를 제공하기 위해.
• 표준 GNN보다 더 우수한 성능을 분자 성질 예측 벤치마크에서 입증하면서도 더 매끄럽고 해석 가능한 그래프 표현을 생성하기 위해.

제안 방법

• 모델은 입력 그래프를 노드 임베딩으로 인코딩하기 위해 표준 GNN(GIN 또는 GAT 등)을 사용하여 임베딩 공간 내의 점 클러스터를 형성한다.
• 학습 가능한 프로토타입 점 클러스터를 자유 매개변수로 도입하여, 하류 작업과 관련된 추상적 구조적 또는 기능적 특징(예: 기능기, 기능기 등)을 나타낸다.
• 그래프 수준의 표현은 입력 그래프의 노드 임베딩 점 클러스터와 각 프로토타입 점 클러스터 사이의 워셔슈타인 거리로 계산된다.
• 프로토타입이 단일 점으로 붕괴되지 않도록 보장하기 위해, 최적 운반의 기하학적 구조를 유지하면서도 엔드 투 엔드로 훈련하기 위해 노이즈 대비 정규화를 사용한다.
• 최종 그래프 임베딩은 모든 프로토타입과의 워셔슈타인 거리를 연결하여 형성되며, 이는 풍부하고 기하학적 구조를 고려한 표현을 가능하게 한다.
• 이론적으로는 OT-GNN이 정의하는 함수 클래스가 표준 GNN 집합 기반 집합과 달리 점 클러스터에서 보편 근사자임을 증명함으로써 방법의 타당성을 뒷받침한다.

실험 결과

연구 질문

• RQ1 학습 가능한 프로토타입을 사용한 최적 운반 기반 그래프 표현 학습이 분자 성질 예측에서 표준 GNN 집합 기반 집합보다 우월한가?
• RQ2 최적 운반과 결합된 매개변수화된 프로토타입을 사용할 경우, 합 또는 평균 집합과 비교해 GNN의 표현 능력이 향상되는가?
• RQ3 특히 엔드 투 엔드 훈련에서 프로토타입 붕괴 문제를 어떻게 완화할 수 있는가?
• RQ4 학습된 프로토타입은 해석 가능한 분자 부분 구조 또는 기능기와 대응하는가?
• RQ5 결과로 도출된 임베딩 공간은 표준 GNN보다 더 매끄럽고 더 강건한가?

주요 결과

• OT-GNN은 GIN 및 GAT을 포함한 최신 GNN 모델보다 여러 분자 성질 예측 벤치마크에서 뛰어난 성능을 보이며 일관된 실험적 성과를 확보했다.
• 노이즈 대비 정규화는 프로토타입 붕괴를 효과적으로 방지하여 최적 운반의 전체 기하학적 구조를 활용할 수 있도록 했다.
• 이론적 분석을 통해 OT-GNN이 정의하는 함수 클래스는 표준 GNN 집합 기반 집합과 달리 점 클러스터에서 보편 근사자임을 확인했다.
• 기준 모델 대비 더 매끄러운 그래프 임베딩 공간을 생성하여 일반화 능력과 강건성 향상을 시사했다.
• 해석 가능성 분석 결과, 학습된 프로토타입은 의미 있는 분자 부분 구조와 대응한다: 예를 들어, 한 프로토타입은 고리 구조를 가진 분자(불용성)와 가장 가까운 반면, 다른 프로토타입은 더 용해성 있는 분자들을 선호한다.
• 하이퍼파rameter 탐색 결과, 프로토타입 수(M=10,20)와 프로토타입 크기(N=10)에 관계없이 성능이 안정적이며, 최적의 성능은 특정 은닉 차원(d=5,10)과 정규화 계수(nc_coef=0.1)에서 달성된다.