[論文レビュー] One Transformer Can Understand Both 2D & 3D Molecular Data

Shengjie Luo, Tianlang Chen+5 more|arXiv (Cornell University)|Oct 4, 2022

Machine Learning in Materials Science被引用数 4

ひとこと要約

本論文では、2次元分子グラフと3次元幾何構造を別々のが共有される注意機構を用いて統合的に符号化する、Transformerベースの統合モデル「Transformer-M」を提案する。自己注意機構におけるバイアス項として、2Dまたは3Dの構造的符号化を動的に活性化することで、2Dおよび3Dの分子性質予測タスクの両方で優れた性能を達成し、アーキテクチャの再訓練なしに異なるモodal間で一般化することを示した。

ABSTRACT

Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One can view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer-based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M develops two separated channels to encode 2D and 3D structural information and incorporate them with the atom features in the network modules. When the input data is in a particular format, the corresponding channel will be activated, and the other will be disabled. By training on 2D and 3D molecular data with properly designed supervised signals, Transformer-M automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M can simultaneously achieve strong performance on 2D and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https://github.com/lsj2408/Transformer-M.

研究の動機と目的

• 2次元分子グラフと3次元幾何構造の両方をアーキテクチャの特化なしに処理できる汎用的なディープラーニングモデルの開発。
• 既存のモデルが1つの分子データフォーマットに特化して訓練されているため、モダリティ間でのゼロショット転送が不十分であるという制限を解消すること。
• 1つのモデルにモダリティに適応した注意機構を設計することで、2次元および3次元データの共同学習を可能にすること。
• 共有パラメータとモダリティ固有の構造的符号化を用いることで、1つのニューラルネットワークが分子データのモダリティ間で効果的に学習および一般化できることを示すこと。

提案手法

• 2次元および3次元の構造的符号化にそれぞれ別々のが共有される注意機構を備えた標準的なトランスフォーマーベースを採用する。
• 2次元チャネルでは、次数符号化、最短経路距離符号化、エッジ符号化を入力バイアスとして用い、トポロジーおよび結合の特徴を捉える。
• 3次元チャネルでは、原子間の3次元ユークリッド距離をガウスカーネルベースの符号化を用いて注意のバイアス項として使用する。
• 推論時、入力モダリティに応じて有効なチャネル（2Dまたは3D）のみを用い、共有された注意およびフィードフォワード層は両方のモダリティで共通して使用する。
• ジョイントトレーニング戦略により、1バッチ内で2次元、3次元、または両方のデータタイプから学習可能であり、非アクティブなチャネルは無効化される。
• 3次元位置のノイズ除去を目的とした事前学習をPCQM4Mv2データセットで実施し、その後、下流の2次元および3次元データセットで微調整する。

実験結果

リサーチクエスチョン

• RQ1 1つのニューラルネットワークアーキテクチャが、2次元分子グラフと3次元幾何構造の両方から意味的な表現を効果的に学習できるか？
• RQ2 注意機構におけるモダリティ固有の構造的符号化が、分子データフォーマット間のモデル一般化に与える影響はいかほどか？
• RQ3 2次元および3次元データのジョイント事前学習は、多様な分子性質予測タスクにおける下流性能をどの程度向上させるか？
• RQ4 モダリティ間で共有されたパラメータを用いることで、モダリティ固有のモデルよりも優れた一般化が達成できるか？

主な発見

• PDBBindコアセットにおいて、Transformer-Mはピアソン積率相関係数（R）0.830を達成し、最良のベースライン（SIGN）および事前学習済みのGraphormerモデルを上回った。
• QM9データセットでは、2D-3Dのジョイント事前学習戦略により、スクラッチからの学習に比べてMAEが最大30％低減され、6つの分子ターゲットすべてで改善が見られた。
• DFTで生成された3次元コンフォーマーを用いることで顕著な性能向上（例：Rが0.804から0.830に上昇）が得られたが、RDKitで生成されたコンフォーマーでは僅かな改善にとどまった。
• 3次元位置のノイズ除去を目的とした事前学習が有効であることが示され、3次元分子表現学習の観点からその目的関数の有効性が裏付けられた。
• 3次元データで学習した場合でも2次元タスクで優れた性能を維持しており、強力なクロスモーダル転送性を示した。
• アブレーションスタディにより、モダリティ固有の構造的符号化が不可欠であることが確認され、いずれのチャネルも無効化すると、対応するモダリティでの性能が低下した。