[논문 리뷰] Rethinking Tokenizer and Decoder in Masked Graph Modeling for Molecules

Zhiyuan Liu, Yaorui Shi외 5명|arXiv (Cornell University)|2023. 10. 23.

Advanced Graph Neural Networks인용 수 5

한 줄 요약

이 논문은 분자의 표현 학습을 위한 새로운 마스킹 그래프 모델링 프레임워크인 SimSGT를 제안한다. 이는 토크나이저 및 디코더 구성 요소를 재고하는 것으로, 단순한 GNN 기반 토크나이저(SGT)와 재마스킹 디코딩을 갖춘 표현력 있는 디코더를 도입함으로써 MoleculeNet 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성한다. 이는 부분 그래프 수준의 재구성에 기반한 향상된 표현 학습을 통해 기존의 자기지도 학습 방법을 능가한다.

ABSTRACT

Masked graph modeling excels in the self-supervised representation learning of molecular graphs. Scrutinizing previous studies, we can reveal a common scheme consisting of three key components: (1) graph tokenizer, which breaks a molecular graph into smaller fragments (i.e., subgraphs) and converts them into tokens; (2) graph masking, which corrupts the graph with masks; (3) graph autoencoder, which first applies an encoder on the masked graph to generate the representations, and then employs a decoder on the representations to recover the tokens of the original graph. However, the previous MGM studies focus extensively on graph masking and encoder, while there is limited understanding of tokenizer and decoder. To bridge the gap, we first summarize popular molecule tokenizers at the granularity of node, edge, motif, and Graph Neural Networks (GNNs), and then examine their roles as the MGM's reconstruction targets. Further, we explore the potential of adopting an expressive decoder in MGM. Our results show that a subgraph-level tokenizer and a sufficiently expressive decoder with remask decoding have a large impact on the encoder's representation learning. Finally, we propose a novel MGM method SimSGT, featuring a Simple GNN-based Tokenizer (SGT) and an effective decoding strategy. We empirically validate that our method outperforms the existing molecule self-supervised learning methods. Our codes and checkpoints are available at https://github.com/syr-cn/SimSGT.

연구 동기 및 목표

• 마스킹 그래프 모델링(MGM)에서 그래프 토크나이저와 디코더의 미비하게 다뤄진 역할을 탐구한다.
• 노드, 엣지, 무늬, 부분 그래프와 같은 다양한 토크나이저의 정밀도가 MGM 성능에 미치는 영향을 평가한다.
• 표현력 있는 디코더와 재마스킹 디코딩 전략이 인코더 표현 품질 향상에 미치는 잠재적 영향을 탐색한다.
• 단순한 GNN 기반 토크나이저와 효과적인 디코딩 전략을 통합하여 향상된 최종 성능을 달성하는 새로운 사전학습 프레임워크인 SimSGT를 개발한다.

제안 방법

• 비선형 활성화 함수를 제거한 단순한 GNN 기반 토크나이저(SGT)를 제안하여 부분 그래프 수준의 토크나이징을 효율적으로 구현한다.
• 강력한 디코더를 갖춘 그래프 오토인코더를 사용하여 마스킹된 그래프 표현에서 부분 그래프 수준의 토큰을 재구성한다.
• 재마스킹 디코딩을 도입하여 디코더를 원본 그래프에서 분리함으로써 인코더가 재구성 작업에 집중하지 않고 의미 있는 표현을 학습하도록 한다.
• GNN 기반 토크나이저를 사용하여 분자의 부분 그래프에서 고정 길이의 토큰을 생성하며, 그래프 분할 함수에 의해 토크나이징 과정을 이끌어낸다.
• 표준 MGM 관행에 따라 인코더에 입력하기 전에 노드 또는 엣지를 무작위로 손상시켜 그래프 마스킹을 적용한다.
• 재구성된 토큰과 원본 토큰 간의 대비 학습 목표를 통해 모델을 종합적으로 최적화하며, 디코더는 마스킹된 부분 그래프를 복구하도록 학습한다.

실험 결과

연구 질문

• RQ1 노드, 엣지, 무늬, 부분 그래프와 같은 토크나이저의 정밀도 선택이 분자의 표현 학습에서 마스킹 그래프 모델링(MGM) 성능에 어떤 영향을 미치는가?
• RQ2 재마스킹 디코딩을 적용한 더 표현력 있는 디코더는 학습된 분자의 표현 품질에 어떤 영향을 미치는가?
• RQ3 비선형성을 포함하지 않는 단순한 GNN 기반 토크나이저가 MGM에서 더 복잡한 화학 기반 토크나이저를 능가할 수 있는가?
• RQ4 재마스킹 디코딩을 통해 인코더를 직접 재구성 작업에서 분리함으로써 최종 성능 향상에 얼마나 기여하는가?
• RQ5 다양한 분자의 성질 예측 벤치마크에서 SimSGT는 기존의 자기지도 사전학습 방법과 비교해 어떻게 성능을 내는가?

주요 결과

• 노드 수준의 토크나이징에 비해 부분 그래프 수준의 토크나이징이 최종 성능 향상에 크게 기여하며, MoleculeNet에서 평균 ROC-AUC가 0.8% 향상된다.
• 충분히 표현력 있는 디코더와 함께 재마스킹 디코딩을 적용하면 인코더의 표현 품질이 향상되며, 여러 데이터셋에서 평균 성능이 1.5~2.0% 향상됨을 입증한다.
• 단일 GNN 레이어로 구성된 제안된 SGT 토크나이저는 더 복잡한 GNN 기반 및 화학 기반 토크나이저보다 경쟁력 있거나 더 우수한 성능을 달성한다.
• SimSGT는 MoleculeNet 분류 벤치마크의 여덟 가지 모두에서 기존 방법을 능가하며, 평균 ROC-AUC 75.8%를 기록하여 보고된 바 중 최고 수준이다.
• 재마스킹 디코딩은 인코더의 초점을 재구성 작업에서 분리하여 더 강력하고 일반화 능력이 뛰어난 분자의 표현을 학습하도록 한다.
• 제거 분석 결과, 부분 그래프 수준의 토크나이저와 표현력 있는 디코더가 MGM에서 고성능을 내기 위해 필수적인 구성 요소임을 확인한다.