[論文レビュー] Learning Invariant Molecular Representation in Latent Discrete Space

Xiang Zhuang, Qiang Zhang+6 more|arXiv (Cornell University)|Oct 22, 2023

Computational Drug Discovery Methods被引用数 5

ひとこと要約

本論文は、最初にGNNで分子を符号化し、その後で残留ベクタ量子化（RVQ）モジュールを用いて離散的潜在空間で不変特徴を分離することで、不変分子表現を学習する新しいフレームワークiMoLDを提案する。この手法は、因果的で分布シフトに強い特徴を特定するためのタスクに依存しない自己教師あり目的関数を採用し、18の分子データセットにおいて分布シフト下でも最先端のOOD一般化性能を達成した。特に共変量シフトおよび概念シフトの状況下で顕著な性能を示した。

ABSTRACT

Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called ``first-encoding-then-separation'' to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at https://github.com/HICAI-ZJU/iMoLD.

研究の動機と目的

• 分子表現学習における分布外（OOD）一般化性能の低さ、特に異なるデータソースからの分布シフト下での性能向上を図ること。
• 分子構造のエンタングルメントに対処できず、環境に関する仮定を必要とする従来の「まず分離してから符号化する」手法の限界を克服すること。
• 回帰やマルチラベル分類などの多様な分子タスクに適用可能な、タスクに依存しない強力なフレームワークの開発。
• 分布シフトにかかわらず安定した不変で因果的な分子特徴を特定することで、一般化性能の向上を図ること。
• タスク固有の教師信号を必要とせず、正確な不変性の同定を可能にする自己教師あり目的関数の設計。

提案手法

• 「まず符号化してから分離する」戦略を提案：GNNが分子グラフ全体を潜在表現に符号化した後、特徴の分離を実行する。
• 潜在空間の離散化を実現するため、残留ベクタ量子化（RVQ）モジュールを導入し、過学習を低減しながらエンコーダーの表現能力を維持する。
• スコアリング用GNNを用いて、離散的潜在表現から不変部分構造を同定・抽出する。
• 環境間で不変である特徴を促進するため、タスクに依存しない自己教師あり学習目的関数を設計する。
• エンコーダー出力を学習済みコードブックに一致させるために、コミットメント損失を用いることで、安定した学習を実現する。
• アーキテクチャの変更なしに、回帰やマルチラベル分類を含む多様な下流タスクにフレームワークを適用する。

Figure 1: (a) First-Separation-Then-Encoding: the input is separated using a subgraph generator, then each subgraph is encoded respectively. (b): First-Encoding-Then-Separation: the input is encoded, then the representation is separated by a scorer.

実験結果

リサーチクエスチョン

• RQ1 「まず符号化してから分離する」戦略は、従来の「まず分離してから符号化する」アプローチに比べ、不変分子特徴をより効果的に捉えることができるか？
• RQ2 残留ベクタ量子化（RVQ）は、表現能力を損なわず、OOD一般化性能を向上させるのにどの程度有効か？
• RQ3 提案されたタスクに依存しない自己教師あり目的関数は、多様な分子タスクにおいて不変性の同定をどの程度向上させるか？
• RQ4 共変量シフトや概念シフトを含むさまざまな分布シフト下で、モデルの性能はどの程度か？
• RQ5 各構成要素（例：RVQ、不変性目的関数）が最終的なOOD一般化性能に果たす寄与度はどの程度か？

主な発見

• 提案されたiMoLDフレームワークは、共変量シフトおよび概念シフトを含む複数の分布シフト状況下で、18の実世界の分子データセットにおいて最先端の性能を達成した。
• GOOD-HIV-Scaffoldデータセットでは、VQモジュールを削除すると顕著な性能低下が生じ、OOD一般化性能においてその重要性が示された。
• より複雑なGOOD-PCBAデータセットでは、残留接続を削除すると顕著な性能低下が観察され、モデルの表現能力を維持する上でその重要性が裏付けられた。
• アブレーションスタディの結果、不変性学習目的関数 $\mathcal{L}_{\textrm{inv}}$ とコミットメント損失 $\mathcal{L}_{\textrm{cmt}}$ の両方が、強固な性能を発揮するために不可欠であることが確認された。
• t-SNE可視化では、iMoLDがより均一な特徴分布と、環境間の距離（ウォッシャーライン距離）が小さい特徴を生成しており、より優れた不変性が裏付けられた。
• iMoLDは、訓練スコアが高く、検証スコアも高い結果を示し、ERMベースラインと比較して過学習が低減され、一般化性能が優れていることが示された。

Figure 2: An overview of distribution shifts in molecules. Distribution shifts occur when molecules originate from different scaffold, size or assay environments.