GROVER는 분자의 그래프 구조 정보와 contextual 정보를 반영하는 그래프 pre-training 모델로, 그래프 내의 local subgraph를 마스킹하고 이를 예측하는 것을 통해 적절한 노드나 엣지의 representation을 학습한다. GROVER는 GNN과 Transformer 구조를 동시에 사용함으로써 노드 주변의 local subgraph의 정보와 동시에 각 노드 간의 global relation까지 반영할 수 있게 한다. 이번 글에서는 GROVER가 무엇인지에 대해 알아보겠다.
Motivation
분자 구조의 graph representation을 학습하는 것은 어려운 일이다. 이를 어렵게 만드는 가장 큰 이유는 분자의 특성을 모두 label로 매핑하기 위해서는 시간이 오래 걸리고 비용이 많이 드는 실험실에서의 실험이 필수적이기 때문이다. 그렇기 때문에 분자 데이터들은 label이 부족하여 일부 분자들에 대해서만 주어져 있는 경우가 많고, 이것은 모델의 generalization capability를 떨어뜨려 모델이 쉽게 overfitting되거나 학습 데이터셋에서 벗어나는 데이터를 만나면 예측 성능이 크게 떨어지는 문제를 일으켰다.
그렇기 때문에 분자 데이터에 대해서는 BERT나 GPT처럼 self-supervised learning을 하는 것이 좋은 선택지가 될 수 있다. GROVER는 단순히 분자의 노드, 엣지 타입만을 고려하는 것이 아니라 분자 그래프의 구조적인 정보(structural information)까지 모두 고려하여 self-supervised learning을 통한 pre-training model을 만드는 것을 목적으로 한다.
Method
우선 GROVER의 1) 모델 아키텍처에 대해 알아보고, 이 모델을 2) 어떤 self-supervised task에 대해 학습했는지에 대해 알아보겠다.
Model architecture (GNN Transformer)
GROVER의 모델 아키텍처는 GNN Transfomer (GTransformer)라고 부르고, 이는 이름 그대로 GNN과 Transformer 부분으로 나뉜다. GNN을 통해서 각 노드와 엣지에 해당하는 vector를 얻고, 이 vector를 Transformer의 query, key, value로 사용함으로써 각 노드와 엣지 사이의 관계를 attention 기반으로 학습하게 된다.
GNN (dyMPN)
GROVER의 GNN은 각 노드와 엣지에 해당하는 vector representation을 얻기 위해서 사용된다. 이 때, GNN을 통해서 얻어진 vector representation은 이웃하는 노드 혹은 엣지로부터 정보를 얻기 때문에 local subgraph structure 정보를 포함하고 있다. GROVER는 다양한 GNN 구조 중 dynamic message passing neural network (dyMPN)을 채택한다.
DyMPN의 가장 큰 특징은 기존의 GNN과는 달리, 얼마나 멀리 떨어진 이웃 노드까지의 정보를 받을 것인지(hop)를 고정하는 것이 아닌, 랜덤하게 결정한다는 점이다. 이를 통해 DyMPN은 고정된 receptive field (정보를 얻을 수 있는 노드의 범위)를 가지는 것이 아닌 다양한 receptive field를 가지게 된다. 이는 앞서 motivation에서 설명했던 것과 마찬가지로, 모델이 더 나은 generalization 성능을 가질 수 있게 도와준다.
GTransformer
GROVER의 Transformer는 각 노드와 엣지 representation 간의 관계를 모델링하기 위해서 사용된다. 이는 GNN이 얻던 local 정보를 넘어 그래프 내에서 모든 노드와 엣지 간의 attention을 계산함으로써 global한 정보를 얻을 수 있도록 도와준다. 각 노드와 엣지 representation이 Transformer의 key, query, value 역할을 하며 이들 간의 self-attention을 계산하게 된다.
GTransformer의 특징 중 하나는 long-range residual connection을 사용한다는 점이다. 이는 initial 노드와 엣지 feature를 마지막 layer에 직접 넣는 것을 의미하고, vanishing gradient problem을 해결하고 over-smoothing도 어느 정도 해결해 준다.
Self-supervised task
GROVER의 self-supervised task는 1) 노드와 엣지 단계에서의 contextual property prediction과 2) 그래프 전체 단계에서의 graph-level motif prediction으로 나누어진다.
Contextual property prediction
첫 번째 self-supervised task인 contextual property prediction은 타겟 노드나 엣지의 주변에 어떤 노드나 엣지가 있는지(contextual property)를 예측하는 task이다. 이 글에서는 노드 기준의 contextual property를 설명하겠다. 어떤 타겟 노드의 contextual property는 hop 내의 node type - edge type - 해당 개수로 구성된다. 예시로 다음 그림을 살펴보자.
우리의 타겟 노드가 빨간색 C 노드라고 하자. 이 노드의 hop=1 (직접적으로 이웃한 노드) contextual property는 이중결합으로 연결된 파란색 N, 단일결합으로 연결된 주황색 O이기 때문에 이는 C_N-DOUBLE1_O-SINGLE1으로 표현된다. 이렇게 모든 노드들에 대해서 contextual property를 계산한 후, 각 노드를 masking하고 이에 해당하는 contextual property를 예측하는 것이 contextual property prediction이다.
Graph-level motif prediction
두 번째 self-supervised task인 graph-level motif prediction은 어떤 그래프를 input으로 받았을 때 미리 정해진 motif 셋에서 어떤 motif를 포함하고 있는지를 예측하는 task이다. 이를 통해 우리가 학습하는 node / edge representation embedding은 각 노드와 엣지의 local 정보 뿐만 아니라 그래프 전체의 정보까지 어느 정도 포함할 수 있게 된다.
Experiments
실험 파트에서는 GROVER를 활용해 얻은 pre-trained representation을 활용하여 다양한 down-stream task들(분자의 classification과 regression)을 진행한다. 그 결과, 다음과 같이 GROVER가 다른 GNN 방법들이나 pre-training 방법들보다 더 좋은 성능을 보임을 알 수 있다.
이외에도 실험을 통해 pre-training을 하면 AUC가 지속적으로 향상되는 것을 보이기도 하고, GTransformer가 단순한 GIN이나 MPNN보다 expressive power가 얼마나 좋은지를 실험하기도 했다.
이번 글에서는 Transformer를 활용한 molceular pre-training network인 GROVER에 대해서 알아보았다. 다음 글에서는 다른 graph transformer 구조들에 대해 알아보겠다.
References
Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., & Huang, J. (2020). Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33, 12559-12571.
최근댓글