[論文レビュー] 3D molecule generation by denoising voxel grids

Pedro O. Pinheiro, Joshua A. Rackers+7 more|arXiv (Cornell University)|Jun 13, 2023

Machine Learning in Materials Science被引用数 6

ひとこと要約

VoxMol は、ボクセル化された原子密度グリッドのノイズ除去を用いて、3次元分子の新しいスコアベースの生成モデルを提案する。Langevin MCMC サンプリングと1回のノイズ除去ステップを活用し、ポイントクラウド拡散モデルと比較して、より高速かつ単純に分子を生成でき、非条件的3次元分子生成ベンチマークで最先端の性能を達成した。

ABSTRACT

We propose a new score-based approach to generate 3D molecules represented as atomic densities on regular grids. First, we train a denoising neural network that learns to map from a smooth distribution of noisy molecules to the distribution of real molecules. Then, we follow the neural empirical Bayes framework (Saremi and Hyvarinen, 19) and generate molecules in two steps: (i) sample noisy density grids from a smooth distribution via underdamped Langevin Markov chain Monte Carlo, and (ii) recover the "clean" molecule by denoising the noisy grid with a single step. Our method, VoxMol, generates molecules in a fundamentally different way than the current state of the art (ie, diffusion models applied to atom point clouds). It differs in terms of the data representation, the noise model, the network architecture and the generative modeling algorithm. Our experiments show that VoxMol captures the distribution of drug-like molecules better than state of the art, while being faster to generate samples.

研究の動機と目的

• 既存のポイントクラウド拡散モデルの制限を克服する非条件的3次元分子生成の新規手法の開発。
• 分子を規則的なボクセルグリッド上の連続的原子密度グリッドとして表現し、より効果的な生成モデリングを可能にする。
• 原子数の予測や混合連続・カテゴリカル特徴の別々の処理を回避することで、学習と推論を単純化する。
• ポイントクラウドにおける最先端の拡散モデルと比較して、学習の複雑さと推論時間を低減しながら、競争力のある生成品質を達成する。
• ウォーク・ジャンプ・ランゲビン MCMC とピークベースの座標リファインメントにより、下流の分子解析に適した効率的なサンプリングを可能にする。

提案手法

• 分子は、各ボクセルが原子位置を中心とするガウス関数に類似した関数で原子密度を符号化する3次元ボクセルグリッドとして表現される。
• ノイズの大きい固定分散のガウスノイズを含むノイズの入ったボクセルグリッドから、元の実際の分子グリッドにマッピングするためのノイズ除去ニューラルネットワークを学習する。
• ノイズ除去ネットワークは、滑らかにしたデータ分布のスコア関数を近似し、スコアベースの生成サンプリングを可能にする。
• 分子は2段階のプロセスで生成される：(i) 滑らかなノイズ分布からのランゲビン MCMC サンプリング、(ii) クリーンな分子への1ステップのノイズ除去。
• 原子座標は、生成後、ボクセルのピーク検出と最適化ベースのリファインメントにより回復され、幾何的正確性が向上する。
• ボクセル化の際、原子半径を固定の 0.5 Å として使用しており、これが複数のデータセットで最良のパフォーマンスを発揮することが経験的に示された。

Figure 1: Voxelized molecules generated by our model and their corresponding molecular graphs. Left, samples from a model trained on QM9 dataset ( $32^{3}$ voxels). Right, samples from a model trained on GEOM-drugs ( $64^{3}$ voxels). In both cases, each voxel is a cubic grid with side length of $.2

実験結果

リサーチクエスチョン

• RQ1 ボクセル化された原子密度グリッドを用いたスコアベースの生成モデルは、非条件的3次元分子生成においてポイントクラウド拡散モデルを上回ることができるか？
• RQ2 連続的でグリッドベースの表現を用いることで、離散的原子ポイントクラウドと比較して、学習と推論が単純化されるか？
• RQ3 ボクセル化における原子半径の選択が、モデルのパフォーマンスとサンプル品質に与える影響はいかほどか？
• RQ4 固定ノイズレベルを用いたウォーク・ジャンプ・ランゲビン MCMC サンプリングは、分子の多様性と構造的正確性をどの程度保持できるか？
• RQ5 ボクセルピークからの座標リファインメントは、生成された分子の妥当性と幾何的正確性を顕著に向上させることができるか？

主な発見

• VoxMol は QM9 および GEOM-Drugs データセットで競争力ある生成品質を達成し、非条件的3次元分子生成分野で最先端の結果を再現した。
• ノイズレベル σ = 1.0 の場合、座標リファインメントを適用すると QM9 で 68.6% の妥当性と 96.3% の一意性を達成したのに対し、リファインメントなしでは 55.0% の妥当性と 100% の一意性であった。
• GEOM-Drugs データセットでは、σ = 1.0 でリファインメントを適用した場合、45.1% の妥当性と 99.7% の一意性を達成し、リファインメントなしでは 32.5% の妥当性にとどまった。
• モデルは、ノイズレベルにかかわらず、真値と同等の形状多様性を示し、NPMIR 分布がロッド型、ディスク型、および球型の形状をカバーしていた。
• ボクセル化の際、原子半径を 0.5 Å に固定すると最良のパフォーマンスが得られ、0.25 Å では学習収束が困難となり、より大きな半径ではサンプル品質が劣化した。
• ポイントクラウド拡散モデルと比較して、原子数の予測や混合特徴タイプの別々の処理を必要としないため、学習がより簡単で、分子生成がより高速である。

Figure 2: (a) A representation of our denoising training procedure. Each training sample is corrupted with isotropic Gaussian noise with a fixed noise level $\sigma$ . The model is trained to recover clean voxels from the noisy version. (b) Graphical model representation of the walk-jump sampling sc