[論文レビュー] 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を提案する。これはトークナイザーとデコーダーの再考に基づくもので、単純なGNNベースのトークナイザー（SGT）とリマスクデコーディングを備えた表現力のあるデコーダーを導入することで、部分グラフレベルの再構築に裏付けられた表現学習の向上により、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のパフォーマンスに与える影響を評価すること。
• 表現力のあるデコーダーとリマスクデコーディング戦略がエンコーダーの表現品質を向上させる可能性を検討すること。
• 単純なGNNベースのトークナイザーと効果的なデコーディング戦略を統合した新しい事前学習フレームワークSimSGTを開発すること。

提案手法

• 非線形活性化関数をGNN層から除去することで、効率的な部分グラフレベルのトークナイゼーションを可能にする単純なGNNベースのトークナイザー（SGT）を提案する。
• 部分グラフレベルのトークンをマスクされたグラフ表現から再構築する強力なデコーダーを備えたグラフオートエンコーダーを採用する。
• デコーダーを元のグラフから分離することでエンコーダーが再構築に集中するのではなく、意味のある表現を学習することを可能にするリマスクデコーディングを導入する。
• グラフ断片化関数に従って、GNNベースのトークナイザーを用いて分子部分グラフから固定長のトークンを生成する。
• 標準的なMGMの実践に従い、エンコーダーに入力する前にノードやエッジをランダムに破損させるグラフマスキングを適用する。
• 再構築されたトークンと元のトークン間の対照学習目的関数を用いて、モデルをエンドツーエンドで最適化し、デコーダーがマスクされた部分グラフを回復するように学習させる。

実験結果

リサーチクエスチョン

• RQ1 ノード、エッジ、モチーフ、部分グラフといったトークナイゼーションの粒度の選択が、分子表現学習におけるマスクドグラフモデリングのパフォーマンスにどのように影響するか？
• RQ2 リマスクデコーディングを組み合わせたより表現力のあるデコーダーが、学習された分子表現の質に与える影響は何か？
• RQ3 非線形性を含まない単純なGNNベースのトークナイザーは、より複雑な化学的インスピレーションを受けて設計されたトークナイザーを上回ることができるか？
• RQ4 リマスクデコーディングによってエンコーダーを直接の再構築から分離させることで、下流のパフォーマンスはどの程度向上するか？
• RQ5 SimSGTは、多様な分子性質予測ベンチマークにおいて、既存の自己教師付き事前学習手法と比較してどのように差をつけるか？

主な発見

• 部分グラフレベルのトークナイゼーションはノードレベルのトークナイゼーションよりも下流のパフォーマンスを顕著に向上させ、MoleculeNetでは平均ROC-AUCで0.8%の向上を示した。
• 十分に表現力のあるデコーダーとリマスクデコーディングを組み合わせることで、複数のデータセットで平均パフォーマンスが1.5–2.0%向上し、エンコーダー表現の質が向上することが裏付けられた。
• 非線形性を含まない単一のGNN層で構成される提案されたSGTトークナイザーは、より複雑なGNNベースや化学的インスピレーションを受けて設計されたトークナイザーと同等またはそれ以上の性能を達成した。
• SimSGTはMoleculeNet分類ベンチマークの全8つのタスクで既存手法を上回り、平均ROC-AUCが75.8%に達し、報告された最高値を記録した。
• リマスクデコーディングはエンコーダーの焦点を再構築から解放し、より強固で汎用性の高い分子表現を学習可能にした。
• アブレーションスタディの結果、部分グラフレベルのトークナイザーと表現力のあるデコーダーの両方が、MGMにおける高性能を達成する上で不可欠な要素であることが確認された。