Rethinking Tokenizer and Decoder in Masked Graph Modeling for Molecules

TL;DR

This paper critically examines and enhances the components of masked graph modeling for molecules, introducing a novel method that improves self-supervised molecular graph representations by focusing on tokenizer and decoder design.

Contribution

The study provides a comprehensive analysis of tokenizer and decoder roles in MGM, proposing a new method with a simple GNN tokenizer and effective decoding that outperforms existing approaches.

Findings

Subgraph-level tokenizer significantly improves encoder representations.

Expressive decoders with remask decoding enhance learning.

SimSGT outperforms previous self-supervised molecular learning methods.

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.