[論文レビュー] Saturn: Sample-efficient Generative Molecular Design using Memory Manipulation

Jeff Guo, Philippe Schwaller|arXiv (Cornell University)|May 27, 2024

Advanced biosensing and bioanalysis techniques被引用数 6

ひとこと要約

Saturn は Mamba SSM を Augmented Memory と SMILES 増強と組み合わせ、固定の oracle 予算下で目的指向分子設計の最新のサンプル効率を達成し、MPO ドッキング課題で多くのベースラインを上回る。

ABSTRACT

Generative molecular design for drug discovery has very recently achieved a wave of experimental validation, with language-based backbones being the most common architectures employed. The most important factor for downstream success is whether an in silico oracle is well correlated with the desired end-point. To this end, current methods use cheaper proxy oracles with higher throughput before evaluating the most promising subset with high-fidelity oracles. The ability to directly optimize high-fidelity oracles would greatly enhance generative design and be expected to improve hit rates. However, current models are not efficient enough to consider such a prospect, exemplifying the sample efficiency problem. In this work, we introduce Saturn, which leverages the Augmented Memory algorithm and demonstrates the first application of the Mamba architecture for generative molecular design. We elucidate how experience replay with data augmentation improves sample efficiency and how Mamba synergistically exploits this mechanism. Saturn outperforms 22 models on multi-parameter optimization tasks relevant to drug discovery and may possess sufficient sample efficiency to consider the prospect of directly optimizing high-fidelity oracles.

研究の動機と目的

• 厳格なサンプル予算の下で高忠実度のオラクルを直接最適化する必要性を動機付ける。
• Augmented Memory とデータ拡張が、アーキテクチャを超えてサンプル効率をどのように向上させるかを調査する。
• MPO タスクにおける高度なバックボーン（RNN、デコーダー Transformer、Mamba SSM）の利点を探る。
• ドッキングベースの MPO タスクでの Saturn の性能向上と、物理ベースのオラクルへの移植性を示す。

提案手法

• 分子を SMILES としてモデル化し、言語モデルをバックボーンとする強化学習として生成を枠組み化する。
• 高報酬系列へエージェントを誘導するために、SMILES 増強を用いた Augmented Memory を適用する。
• 上位の SMILES のリプレイバッファを用い、増強され再利用されたものとエージェントの確率の平方誤差（式（Eq. 4））でエージェントを更新する。
• バッファを更新しモード崩壊を緩和するため、任意の親集団として遺伝的アルゴリズムを組み込む。
• 繰り返しの高コスト評価を避けるためオラクル評価をキャッシュし、スキャフォルドの過剰な表現を防ぐ多様性フィルターを適用する。
• RNN からデコーダー Transformer へ、さらに Mamba SSM までのバックボーンを評価し、固定の MPO 目的と oracle_budget を前提とする。

Figure 1: Saturn generative workflow. All generated SMILES and their rewards are stored in the Oracle Cache after canonicalization. A genetic algorithm can be optionally applied using the replay buffer as the parent population. Augmented Memory is used to update the agent numerous times.

実験結果

リサーチクエスチョン

• RQ1 Memory-based augmentation と経験再生は、固定予算内で高忠実度オラクルの直接最適化を可能にするか？
• RQ2 バックボーンのアーキテクチャ（RNN、デコーダー Transformer、Mamba）が、目的指向分子設計のサンプル効率に与える影響は？
• RQ3 Augmented Memory、データ拡張、および Mamba の組み合わせは、ドッキングおよび物理ベースのタスクでベースラインよりも優れた MPO パフォーマンスをもたらすか？
• RQ4 Saturn はさまざまな生物学的ターゲットとドッキングベースの目的に対して、そのサンプル効率をどのように移転させるか？

主な発見

Model	Aug. Rounds	Yield (↑)	IntDiv1 (↑)	Scaffolds (↑)	OB 1 (↓)	OB 10 (↓)	OB 100 (↓)	Repeats
RNN	5	107±58	0.814±0.036	101±54	480±118 (10)	721±109 (10)	916±53 (4)	7±7
RNN	6	121±80	0.791±0.040	107±68	493±214 (10)	713±15 (10)6	895±107 (5)	12±11
RNN	7	144±107	0.776±0.026	117±86	467±186 (10)	684±136 (10)	871±116 (6)	38±82
RNN	8	120±95	0.734±0.128	104±85	481±288 (10)	653±145 (8)	854±54 (5)	18±28
RNN	9	141±104	0.783±0.048	112±72	453±211 (10)	654±154 (9)	871±104 (6)	59±95
RNN	10	106±76	0.76±0.056	84±63	510±201 (10)	733±122 (9)	913±64 (5)	43±47
Decoder	5	154±93	0.748±0.052	122±70	439±151 (10)	679±128 (10)	907±92 (8)	90±90
Decoder	6	116±94	0.748±0.039	86±64	517±165 (10)	728±158 (10)	904±126 (5)	73±42
Decoder	7	108±85	0.747±0.051	71±50	510±222 (10)	740±127 (9)	868±48 (4)	126±63
Decoder	8	108±94	0.708±0.109	72±57	538±164 (10)	742±116 (9)	887±87 (4)	150±72
Decoder	9	78±83	0.687±0.116	51±55	614±244 (10)	790±150 (8)	890±62 (3)	242±139
Decoder	10	120±128	0.691±0.042	74±73	663±170 (9)	768±169 (8)	805±65 (4)	344±218
Mamba	5	69±38	0.764±0.052	54±28	542±93 (10)	807±76 (10)	988±17 (3)	178±90
Mamba	6	138±46	0.759±0.039	110±42	456±89 (10)	693±75 (10)	919±36 (7)	286±137
Mamba	7	174±95	0.737±0.059	127±83	427±177 (10)	643±102 (10)	858±77 (7)	395±147
Mamba	8	209±95	0.751±0.030	137±60	461±151 (10)	617±135 (10)	817±71 (8)	482±214
Mamba	9	202±98	0.735±0.032	137±80	389±112 (10)	631±102 (10)	841±92 (8)	518±237
Mamba	10	306±57	0.714±0.035	206±34	387±148 (10)	555±66 (10)	761±58 (10)	1110±636

• Saturn は Augmented Memory と SMILES 増強を用いる Mamba により、固定オラクル予算下で MPO ドッキング課題において優れたサンプル効率を達成し、22モデルを上回った。
• Augmented Memory は増強された SMILES を高報酬領域へ絞り込み、より非現実的な系列に対してより大きな更新を行い、効率的な学習を可能にする。
• Mamba は Hop-and-locally-explore な挙動を示し、化学空間を方向付けて横断し、局所的に類似した分子を生成することで効率を高める。
• アーキテクチャ全体を通じて、Mamba は Yield と Oracle Burden の指標で RNN およびデコーダー Transformer のベースラインを一貫して上回る（1,000 oracle 予算下で）。
• Saturn は DRD2、AChE、MK2 ターゲットに対する物理ベースのドッキング MPO へサンプル効率を移転し、しばしば素朜な Augmented Memory を上回り GA が多様性を回復できることを示す。
• GEAM 対 HIT/Novel hit ベンチマークと比較して、Saturn（Saturn-GA 使用） は一部ケースで競争力のあるまたは優れた結果を低分散で達成し、厳格フィルタ付きヒットをより少ないオラクル呼び出しで見つけられる。

Figure 2: a. Average maximum token probability across agent states. Augmentation pushes the agent action distribution towards a delta distribution. b. Augmented Memory (10 augmentation rounds) makes the likelihood of generating SMILES in the buffer more likely. c. Top: On average, augmented forms of