[논문 리뷰] Memory-Based Graph Networks

Amir Hosein Khasahmadi, Kaveh Hassani외 3명|arXiv (Cornell University)|2020. 02. 21.

Advanced Graph Neural Networks참고 문헌 60인용 수 11

한 줄 요약

이 논문은 메모리 기반 그래프 네트워크(MemGNN)와 그래프 메모리 네트워크(GMN)를 제안하며, 다중 헤드 메모리 레이어를 통해 계층적 노드 표현 학습과 그래프 군집화를 동시에 수행하여, 9개의 그래프 분류 및 회귀 벤치마크 중 8개에서 최신 기술 수준(SOTA) 성능을 달성하며 분자의 화학적으로 의미 있는 부분 구조를 정확하게 식별한다.

ABSTRACT

Graph neural networks (GNNs) are a class of deep models that operate on data with arbitrary topology represented as graphs. We introduce an efficient memory layer for GNNs that can jointly learn node representations and coarsen the graph. We also introduce two new networks based on this layer: memory-based GNN (MemGNN) and graph memory network (GMN) that can learn hierarchical graph representations. The experimental results shows that the proposed models achieve state-of-the-art results in eight out of nine graph classification and regression benchmarks. We also show that the learned representations could correspond to chemical features in the molecule data. Code and reference implementations are released at: https://github.com/amirkhas/GraphMemoryNet

연구 동기 및 목표

• 그래프 신경망에서 노드 표현 학습과 그래프 군집화를 동시에 수행하는 효율적인 메모리 레이어를 개발한다.
• 이 메모리 레이어를 활용하여 계층적 그래프 표현 학습을 위한 두 가지 새로운 아키텍처—MemGNN와 GMN—을 설계한다.
• 분자 데이터에서 기능 그룹과 같은 의미 있고 해석 가능한 그래프 구조를 엔드 투 엔드로 학습할 수 있도록 한다.
• 각 풀링 단계 이후 메시지 전파를 반복하지 않음으로써 기존 풀링 기법 대비 계산 오버헤드를 줄인다.

제안 방법

• 메모리 레이어는 노드 표현을 전역적이고 추상화된 클러스터로 군집화하기 위해 학습 가능한 키의 다중 헤드 어레이를 사용하며, 연결 정보에 의존하지 않는다.
• 이전 레이어의 노드 표현을 쿼리로 사용하여 메모리 키에 어텐션을 적용함으로써, 소프트 클러스터 할당과 집계를 통해 군집화된 노드 표현을 생성한다.
• 메모리 레이어는 다수의 헤드에서 온 어텐션 스코어를 집계하기 위해 컨볼루션 연산자를 활용하여 강건하고 미분 가능한 클러스터링을 가능하게 한다.
• MemGNN는 초기 표현 학습을 위한 GNN과 스택된 메모리 레이어를 조합하여 전역 그래프 수준까지 계층적 표현을 구축한다.
• GMN은 메시지 전파를 완전히 대체하여 다중 메모리 레이어를 통해 직접 계층적 표현을 학습함으로써 더 빠른 추론을 가능하게 한다.
• 구조적 인식을 향상시키기 위해 그래프 확산(예: 랜덤 워크 또는 열 커널)을 사용해 토폴로지 임베딩을 초기화하며, 명시적 엣지 의존성 없이도 효과를 발휘한다.

실험 결과

연구 질문

• RQ1 기존 풀링 기법에 비해 메모리 기반 레이어가 노드 표현 학습과 그래프 군집화를 동시에 더 효율적으로 수행할 수 있는가?
• RQ2 학습된 메모리 키가 분자 그래프에서 기능 그룹과 같은 의미 있는 하위 구조를 대표하는가?
• RQ3 키와 헤드의 수가 다양한 그래프 크기를 가진 데이터셋 간의 모델 성능 및 일반화에 어떤 영향을 미치는가?
• RQ4 순수하게 메모리 레이어로 구성된 모델(GMN)이 그래프 수준 예측 작업에서 메시지 전파 기반 GNN보다 우월한가?
• RQ5 메모리 레이어에서 명시적 연결 정보가 없더라도, 구조적으로 복잡한 그래프에서 성능 저하가 발생하는가?

주요 결과

• MemGNN는 9개의 그래프 분류 및 회귀 벤치마크 중 8개에서 최신 기술 수준(SOTA) 성능을 달성하였으며, ESOL(RMSE = 0.52)과 리포필리시티에서 SOTA 결과를 기록하였다.
• 랜덤 이웃 샘플링을 사용하여 Collab 데이터셋에서 10겹 교차 검증 정확도 73.9%를 달성하였으며, RWR 기반 샘플링(73.1%)보다 뛰어났다.
• 학습된 클러스터의 시각화 결과, 메모리 키가 분자 내 알려진 화학적 부분 구조—예를 들어 하이드록실(OH), 카복실(COOH), 벤젠 고리—에 대응하는 것으로 확인되었다.
• 고정된 파라미터 예산 하에서 헤드 수를 늘리면 성능 향상이 관찰되었으며, 예를 들어 ESOL에서 32개의 키와 5개의 헤드 조합이 RMSE = 0.53을 기록한 반면, 160개의 키와 1개의 헤드 조합은 RMSE = 0.54를 기록하였다.
• K-평균 초기화와 비교해도 무작위 초기화된 메모리 키가 유사한 성능을 보였으며, 이는 엔드 투 엔드 학습이 온전히 의미 있는 클러스터 중심점을 학습할 수 있음을 시사한다.
• 메모리 레이어는 국소적 토폴로지에 의존하지 않고도 효과적으로 그래프를 군집화하고 표현을 집계하며, 과도한 스무딩을 방지하고 효율적인 추론을 가능하게 한다.