[논문 리뷰] Learning Attributed Graph Representations with Communicative Message Passing Transformer

Jianwen Chen, Shuangjia Zheng외 3명|arXiv (Cornell University)|2021. 07. 19.

Computational Drug Discovery Methods참고 문헌 17인용 수 10

한 줄 요약

이 논문은 분자 그래프 표현을 향상시키기 위해 노드와 에지 간의 이중 방향 메시지 전달을 트랜스포머 기반 아키텍처를 사용해 명시적으로 모델링하는 CoMPT(의사소통 메시지 전달 트랜스포머)를 제안한다. 메시지 폭발과 과도한 스무딩을 완화하기 위해 위상적 메시지 확산 메커니즘을 통합함으로써 CoMPT는 최신 기술 수준의 성능을 달성하였으며, 그래프 수준 작업에서 평균 성능을 약 4% 향상시키고, 1H-NMR 노드 수준 예측에서 MAE 0.214를 기록하였다.

ABSTRACT

Constructing appropriate representations of molecules lies at the core of numerous tasks such as material science, chemistry and drug designs. Recent researches abstract molecules as attributed graphs and employ graph neural networks (GNN) for molecular representation learning, which have made remarkable achievements in molecular graph modeling. Albeit powerful, current models either are based on local aggregation operations and thus miss higher-order graph properties or focus on only node information without fully using the edge information. For this sake, we propose a Communicative Message Passing Transformer (CoMPT) neural network to improve the molecular graph representation by reinforcing message interactions between nodes and edges based on the Transformer architecture. Unlike the previous transformer-style GNNs that treat molecules as fully connected graphs, we introduce a message diffusion mechanism to leverage the graph connectivity inductive bias and reduce the message enrichment explosion. Extensive experiments demonstrated that the proposed model obtained superior performances (around 4$\%$ on average) against state-of-the-art baselines on seven chemical property datasets (graph-level tasks) and two chemical shift datasets (node-level tasks). Further visualization studies also indicated a better representation capacity achieved by our model.

연구 동기 및 목표

• 기존 GNN이 장거리 의존성을 포착하는 데에 한계가 있고 분자 그래프에서 에지 정보를 효율적으로 활용하지 못하는 문제를 해결하기 위해.
• 깊은 GNN과 트랜스포머 스타일 GNN에서 발생하는 과도한 스무딩과 메시지 강화 폭발 문제를 극복하기 위해.
• 노드와 에지 특징 간의 상호작용을 명시적으로 모델링함으로써 분자 표현 학습을 향상시키기 위해.
• 분자 위상 구조를 존중하는 확산 메커니즘을 통해 구조적 인덕티브 바이어스를 통합함으로써 일반화 능력과 표현력 향상시키기 위해.

제안 방법

• 노드와 에지 임베딩이 트랜스포머 기반 아키텍처를 통해 노드와 에지 간의 교차 어텐션 메커니즘을 통해 반복적으로 업데이트된다.
• 분자 그래프의 위상적 인접 행렬을 사용하여 메시지 전파를 제약하고 폭발적인 메시지 누적을 줄이기 위해 메시지 확산 메커니즘이 도입된다.
• 구조적 맥락을 유지하고 표현 품질을 향상시키기 위해 학습 가능한 노드 위치 및 에지 방향 임베딩을 사용한다.
• 메시지 전달이 의사소통 방식으로 수행되며, 노드와 에지 간에 이중 방향으로 정보를 교환함으로써 특징의 표현력을 향상시킨다.
• 학습 가능한 리더아웃 메커니즘과 노드별 회귀 헤드를 통해 그래프 수준 및 노드 수준 예측 작업을 모두 지원한다.
• 과제에 따라 교차 엔트로피 또는 평균 절대 오차(MAE) 손실을 사용하여 엔드 투 엔드로 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 노드-에지 상호작용을 명시적으로 모델링하는 트랜스포머 기반 GNN이 표준 GNN과 베이직 트랜스포머보다 분자 표현 학습에서 더 우수한 성능을 낼 수 있는가?
• RQ2 위상적 메시지 확산 메커니즘을 통합할 경우 분자 그래프 과제에서 모델 성능과 안정성에 어떤 영향을 미치는가?
• RQ3 CoMPT는 깊은 아키텍처에서 과도한 스무딩을 어느 정도 완화하고, 구분 가능한 원자 표현을 유지하는가?
• RQ4 노드 위치, 에지 방향, 메시지 확산이 각각 최종 성능에 기여하는 정도는 어떠한가?

주요 결과

• CoMPT는 일곱 개인 그래프 수준 분자 성질 예측 데이터셋에서 최신 기술 수준의 베이스라인보다 평균적으로 약 4% 향상된 성능을 기록하였다.
• 노드 수준 과제에서 CoMPT는 1H-NMR에 대해 평균 절대 오차(MAE) 0.214 ppm, 13C-NMR에 대해 1.321 ppm을 기록하여 모든 베이스라인을 초월하였다.
• 제거 실험 결과 메시지 확산 메커니즘을 제거할 경우 성능 저하가 가장 크게 발생하여, 이 메커니즘이 메시지 폭발 제어에 핵심적인 역할을 한다는 점을 입증하였다.
• t-SNE 시각화 결과 CoMPT가 과도한 스무딩을 완화하고 스케일드 원자와 명확히 구분되는 측쇄 기능기의 분리가 가능하다는 점을 확인하였다.
• 노드 위치 및 에지 방향 임베딩의 포함이 성능 향상에 기여하여, 이들이 구조적 인식에 중요하다는 점을 시사하였다.
• 다양한 랜덤 시드에서 모델이 높은 안정성을 보였으며, 반복 실험에서 표준편차가 낮았다(예: 1H-NMR MAE의 경우 0.003).