[論文レビュー] 3D-Transformer: Molecular Representation with Transformer in 3D Space

Fang Wu, Qiang Zhang+6 more|arXiv (Cornell University)|Oct 4, 2021

Machine Learning in Materials Science参考文献 79被引用数 10

ひとこと要約

3D-Transformerは、多スケール自己注意メカニズムと適応的ポジショナルエンコーディングを用いて空間幾何を組み込むことで、分子表現学習のための3D対応型トランスフォーマー・アーキテクチャを導入した。アテンションベースのダウンサンプリングと、さまざまな分子サイズにおける原子間距離の文脈に配慮したモデリングを活用することで、結晶性質予測およびタンパク質リガンド結合親和性タスクで最先端の性能を達成した。

ABSTRACT

Spatial structures in the 3D space are important to determine molecular properties. Recent papers use geometric deep learning to represent molecules and predict properties. These papers, however, are computationally expensive in capturing long-range dependencies of input atoms; and have not considered the non-uniformity of interatomic distances, thus failing to learn context-dependent representations at different scales. To deal with such issues, we introduce 3D-Transformer, a variant of the Transformer for molecular representations that incorporates 3D spatial information. 3D-Transformer operates on a fully-connected graph with direct connections between atoms. To cope with the non-uniformity of interatomic distances, we develop a multi-scale self-attention module that exploits local fine-grained patterns with increasing contextual scales. As molecules of different sizes rely on different kinds of spatial features, we design an adaptive position encoding module that adopts different position encoding methods for small and large molecules. Finally, to attain the molecular representation from atom embeddings, we propose an attentive farthest point sampling algorithm that selects a portion of atoms with the assistance of attention scores, overcoming handicaps of the virtual node and previous distance-dominant downsampling methods. We validate 3D-Transformer across three important scientific domains: quantum chemistry, material science, and proteomics. Our experiments show significant improvements over state-of-the-art models on the crystal property prediction task and the protein-ligand binding affinity prediction task, and show better or competitive performance in quantum chemistry molecular datasets. This work provides clear evidence that biochemical tasks can gain consistent benefits from 3D molecular representations and different tasks require different position encoding methods.

研究の動機と目的

• 既存の幾何的ディープラーニングモデルが長距離依存性を捉えることや、分子内の非一様な原子間距離を処理することに課題を抱えているのを是正すること。
• 異なる空間スケールおよび分子サイズに適応可能な3D対応型アテンションメカニズムを開発すること。
• 分子表現の集約において、固定または距離に依存するダウンサンプリング手法の欠陥を克服すること。
• 分子サイズに動的に適応するポジショナルエンコーディング戦略を設計し、表現品質を向上させること。
• 量子化学、材料科学、プロテオミクスを含む多様な科学分野において、3D-Transformerの有効性を検証すること。

提案手法

• すべての原子間相互作用を完全連結グラフとしてモデル化し、3次元空間における直接的な原子間接続を保持する。
• 受容 field を段階的に拡大することで、微細な局所的パターンと広範な文脈的依存性を捉える多スケール自己注意モジュールを導入する。
• 分子サイズに応じて異なるエンコーディング戦略を選択する適応的ポジショナルエンコーディングモジュールを採用する—小分子には相対的ポジショナルエンコーディング、大分子には学習済み絶対的エンコーディングを用いる。
• アテンションスコアを用いて代表的な原子を選択する注目型最遠点サンプリング法を提案し、バーチャルノードや距離ベースのプーリング戦略を凌駕する。
• トランスフォーマー・エンコーダーの入力特徴として3次元座標を活用し、アテンション計算に空間的インダクティブバイアスを組み込む。
• 標準的な損失関数を用いて、分子性質予測タスク上でモデルをエンドツーエンドに訓練する。

実験結果

リサーチクエスチョン

• RQ1 空間幾何と長距離依存性を明示的にモデル化することで、3D対応型トランスフォーマー・アーキテクチャは分子表現学習を向上させ得るか？
• RQ2 多スケールアテンションは、異なる原子間距離分布に応じた文脈依存パターンの捉え方をどのように向上させるか？
• RQ3 適応的ポジショナルエンコーディングは、固定エンコーディング方式と比較して、さまざまなサイズの分子でより良い性能を発揮するか？
• RQ4 アテンションベースのダウンサンプリングは、従来の距離ベースまたはバーチャルノードプーリングを上回る性能を発揮するか？
• RQ5 3D空間的情報を組み込むことで、量子化学、材料科学、プロテオミクスのような多様な科学分野での性能向上はどの程度達成されるか？

主な発見

• 3D-Transformerは結晶性質予測ベンチマークで最先端の性能を達成し、先行するSOTAモデルを上回った。
• タンパク質リガンド結合親和性予測タスクにおいて顕著な改善を示し、生物学的・化学的分野にわたる強力な一般化能力を示した。
• 量子化学データセットでは、既存のモデルと同等またはそれ以上の性能を達成し、電子構造予測における有効性を裏付けた。
• 適応的ポジショナルエンコーディングモジュールは、分子サイズに応じたエンコーディング戦略の最適化により、小分子および大分子の両方で性能向上を実現した。
• 多スケール自己注意メカニズムは、非一様な原子間距離を有する分子でさえも、長距離依存性と微細な空間的パターンを効果的に捉えることができた。
• 注目型最遠点サンプリング法は、バーチャルノードや距離ベースのプーリング手法と比較して、より情報量が多く代表的な分子埋め込みを生成した。