[論文レビュー] Score-Based Generative Models for Molecule Generation

Dwaraknath Gnaneshwar, Bharath Ramsundar+3 more|arXiv (Cornell University)|Mar 7, 2022

Machine Learning in Materials Science被引用数 7

ひとこと要約

本論文は、150万個のZINC分子のSELFIES表現を用いて訓練されたTransformerベースのスコア関数を用いたスコアベース生成モデルを提案する。モデルは多様で妥当かつ新規の化合物を生成するが、100%の妥当性と新規性を達成している一方、学習分布から著しく逸脱した分子を生成しており、複雑な化学的空間をモデル化するためのアーキテクチャの改善の余地があることを示唆している。

ABSTRACT

Recent advances in generative models have made exploring design spaces easier for de novo molecule generation. However, popular generative models like GANs and normalizing flows face challenges such as training instabilities due to adversarial training and architectural constraints, respectively. Score-based generative models sidestep these challenges by modelling the gradient of the log probability density using a score function approximation, as opposed to modelling the density function directly, and sampling from it using annealed Langevin Dynamics. We believe that score-based generative models could open up new opportunities in molecule generation due to their architectural flexibility, such as replacing the score function with an SE(3) equivariant model. In this work, we lay the foundations by testing the efficacy of score-based models for molecule generation. We train a Transformer-based score function on Self-Referencing Embedded Strings (SELFIES) representations of 1.5 million samples from the ZINC dataset and use the Moses benchmarking framework to evaluate the generated samples on a suite of metrics.

研究の動機と目的

• スコアベース生成モデルがデ・ノボ分子生成に実用的かつ効果的であるかを調査すること。
• GAN やノーマライジングフローの限界（訓練の不安定性やアーキテクチャ的制約）を、スコアベースモデルのアーキテクチャの柔軟性を活用することで克服すること。
• SELFIES表現を用いた分子データに対するスコアベースモデルの訓練に関する基礎的ベストプラクティスを確立すること。
• MOSESベンチマークスイートを用いて、妥当性、一意性、多様性、分布の忠実度という指標でモデルの性能を評価すること。
• 等長性アーキテクチャ（例：SE(3)等長モデル）をスコア関数に組み込む可能性を検討し、より優れた3次元分子生成を実現すること。

提案手法

• 分子データの潜在空間におけるSELFIES文字列の対数確率密度の勾配（＝スコア）を推定するため、Transformerベースのスコア関数を訓練する。
• 複数のノイズレベル（σ₁ = 10, σₗ = 0.01, 350レベル）を用いたノイズ除去スコアマッチングを採用し、訓練の安定化とスコア推定の向上を図る。
• 冷却されたランジエブンダイナミクスを用いてサンプリングを行う：初期にランダムノイズを生成し、徐々にノイズレベルを低下させながらスコア関数を反復的に適用してノイズ除去を行う。
• 分子を1-of-KエンコーディングされたSELFIES文字列として表現し、それをスコアネットワークの入力としてベクトルにフラット化する。
• 目的関数：∑ᵢ λ(i)𝔼[‖sθ(x,i) − ∇ₓlog pσᵢ(x)‖²] を用いてスコア関数を最適化する。ここでλ(i)は異なるノイズレベルに重みを付ける。
• 1台のV100 GPUで約12時間訓練し、T=10の時間ステップで約4時間で3万個の分子をサンプリングする。

実験結果

リサーチクエスチョン

• RQ1 SELFIES符号化された分子データで訓練されたスコアベース生成モデルは、化学的に妥当で多様な分子を効果的に生成できるか？
• RQ2 スコアベースのスコア関数が、VAE や GAN といった既存のモデルと比較して、標準的な分子生成ベンチマークでどの程度の性能を示すか？
• RQ3 生成された分子は、学習データ（ZINC 1.6M）の統計的分布とどの程度一致するか？ また、分布の乖離を引き起こす要因は何か？
• RQ4 スコアベースモデルのアーキテクチャの柔軟性を活かして、今後の等長性層（例：SE(3)等長GNN）をスコア関数に統合できるか？
• RQ5 スコアベース分子生成に最も効果的なハイパーパrameterと訓練戦略（例：ノイズスケジュール、モデルの深さ）は何か？

主な発見

• 生成された3万個の分子すべてで100%の妥当性と100%の新規性を達成しており、構造的に妥当で完全に新しい化合物を生成していることが示された。
• 1,000個での一意性スコアは88%、10,000個でのスコアは82%であり、最先端のVAEほどではないが、依然として競争力があり、ハイパーパrameterチューニングによる改善の余地があることを示唆している。
• IntDiv₁とIntDiv₂スコアがそれぞれ0.90および0.88であり、生成セットに強い化学的多様性があり、モード崩壊が見られないことが示された。
• FCD/TestとFCD/TestSFスコアがそれぞれ39.84および40.92であり、学習データから著しく逸脱していることが示され、生成された分子は実際の分子分布に近いとは言えないことがわかった。
• モデルの性能は、分布の忠実度を犠牲にした高多様性・高新規性のトレードオフを示しており、スコア関数の能力向上やアーキテクチャのインダクティブバイアスの強化の必要性を示唆している。
• 本結果は、スコアベースモデルを用いた分子生成の可能性を裏付け、今後の等長性アーキテクチャの統合による分子の対称性のより良い捉え方と分布モデリングの向上の道を開いた。