[论文解读] Rethinking Tokenizer and Decoder in Masked Graph Modeling for Molecules

Zhiyuan Liu, Yaorui Shi+5 more|arXiv (Cornell University)|Oct 23, 2023

Advanced Graph Neural Networks被引用 5

一句话总结

本文提出 SimSGT，一种用于分子表征学习的新型掩码图建模框架，重新思考了分词器与解码器组件。通过引入一种简单的基于图神经网络的分词器（SGT）以及具备重掩码解码功能的表达能力强的解码器，SimSGT 在 MoleculeNet 基准测试中实现了最先进性能，通过基于子图级重建的改进表征学习，超越了现有自监督方法。

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 性能的影响。
• 探索表达能力强的解码器与重掩码解码策略在提升编码器表征质量方面的潜力。
• 开发一种新型预训练框架 SimSGT，集成一种简单的基于图神经网络的分词器与高效的解码策略，以提升下游性能。

提出的方法

• 提出一种简单的基于图神经网络的分词器（SGT），其图神经网络层中去除非线性激活函数，实现高效的子图级分词。
• 采用具有强大解码器的图自编码器，从掩码图表征中重建子图级分词。
• 引入重掩码解码，将解码器与原始图解耦，使编码器能够专注于学习有意义的表征，而非直接重建。
• 使用基于图神经网络的分词器从分子子图生成固定长度的分词，分词过程由图碎片化函数引导。
• 通过在输入编码器前随机破坏节点或边来实施图掩码，遵循标准的 MGM 实践。
• 通过重建分词与原始分词之间的对比学习目标端到端优化模型，使解码器学习恢复被掩码的子图。

实验结果

研究问题

• RQ1 在分子表征学习的掩码图建模中，分词粒度（节点、边、基元或子图）的选择如何影响性能？
• RQ2 使用更具表达能力的解码器并结合重掩码解码策略，对所学分子表征的质量有何影响？
• RQ3 无非线性激活的简单基于图神经网络的分词器是否能超越更复杂的化学启发式分词器？
• RQ4 通过重掩码解码将编码器与直接重建解耦，在多大程度上提升了下游性能？
• RQ5 SimSGT 在多样化的分子性质预测基准测试中，与现有自监督预训练方法相比表现如何？

主要发现

• 与节点级分词相比，子图级分词显著提升了下游性能，在 MoleculeNet 上平均 AUC 提升 0.8%。
• 表达能力足够的解码器结合重掩码解码，可提升编码器表征质量，多个数据集上平均性能提升 1.5–2.0%。
• 所提出的 SGT 分词器（单层 GNN）在性能上达到或超越更复杂的基于图神经网络和化学启发的分词器。
• SimSGT 在 MoleculeNet 的全部八个分类基准上均优于现有方法，平均 ROC-AUC 达 75.8%，为目前报告的最高值。
• 重掩码解码有效解耦了编码器对重建的关注，使其能够学习到更鲁棒、更具泛化能力的分子表征。
• 消融实验确认，子图级分词器与表达能力强的解码器均为 MGM 高性能的关键组件。