[論文レビュー] Searching for High-Value Molecules Using Reinforcement Learning and Transformers

Raj Ghugare, Santiago Miret+3 more|arXiv (Cornell University)|Oct 4, 2023

Computational Drug Discovery MethodsComputer Science被引用数 3

ひとこと要約

本稿では、トランスフォーマーとテキストベースの分子表現（SMILES/SELFIES）を用いた強化学習フレームワークであるChemRLformerを紹介する。アラインドデータセットでの事前学習、ハイルクラビングリプレイバッファやKL正則化といったターゲットを絞ったアルゴリズム的要素を組み合わせることで、複雑なタンパク質ドッキングシミュレーションを含む25の分子設計タスクにおいて、最先端の性能を達成した。また、効果的な設計選択の実証的アブレーションにより、従来の手法よりも簡素化されたアプローチを実現した。

ABSTRACT

Reinforcement learning (RL) over text representations can be effective for finding high-value policies that can search over graphs. However, RL requires careful structuring of the search space and algorithm design to be effective in this challenge. Through extensive experiments, we explore how different design choices for text grammar and algorithmic choices for training can affect an RL policy's ability to generate molecules with desired properties. We arrive at a new RL-based molecular design algorithm (ChemRLformer) and perform a thorough analysis using 25 molecule design tasks, including computationally complex protein docking simulations. From this analysis, we discover unique insights in this problem space and show that ChemRLformer achieves state-of-the-art performance while being more straightforward than prior work by demystifying which design choices are actually helpful for text-based molecule design.

研究の動機と目的

• テキストベースの表現を用いたより効果的でスムーズな強化学習的手法の開発。
• テキストベースの強化学習における、アーキテクチャ的およびアルゴリズム的選択が分子生成のパフォーマンスに与える影響を特定すること。
• 多様なデータセットでの事前学習および報酬関数設計が、分子最適化の結果に与える影響を評価すること。
• トランスフォーマーがオンライン強化学習における分子生成に与える安定性とパフォーマンスを分析すること。
• 単純でアブレーション駆動の設計が、高価値分子の発見において、複雑な従来手法を上回ることを実証すること。

提案手法

• 本手法は、SMILESまたはSELFIES形式で分子をトークン単位に逐次生成するトランスフォーマーに基づく方策ネットワークを強化学習で訓練する。
• 大規模かつアラインドされた分子データセットを用いた事前学習フェーズにより、方策を初期化することで、サンプル効率とパフォーマンスが向上する。
• 訓練中に高報酬で多様な分子を優先するため、リプレイバッファにハイルクラビング戦略を組み込む。
• 事前学習済み方策からの大きな逸脱を防ぐために、KL正則化を適用し、安定した訓練と効率的な探索を促進する。
• 分子ドッキングスコアと非薬物的・非合成可能な性質に対するペナルティを組み合わせた複合報酬関数を用い、現実的な分子生成を誘導する。
• 本フレームワークは、計算的に高負荷なタンパク質-リガンドドッキングシミュレーションを含む25の多様な分子設計タスクで評価された。

Figure 1 : Autoregressively generating a benzene molecule. : An autoregressive model for sequence generation can be viewed as an RL policy where the actions $a_{t}$ are the next tokens to append to the sequence and the state is the concatenation of all actions taken up to time $t-1$ . A special end-

実験結果

リサーチクエスチョン

• RQ1 強化学習ベースの分子生成において、テキストベースの分子表現（SMILES 対 SELFIES）のどちらがより優れたパフォーマンスを発揮するか？
• RQ2 アラインドデータセットでの事前学習が、分子設計における強化学習エージェントのサンプル効率と最終的パフォーマンスにどのように影響するか？
• RQ3 リプレイバッファ設計、正則化、尤度ペナルティなどのアルゴリズム的要素の中で、テキストベースの分子生成における強化学習パフォーマンスを最も顕著に向上させるものは何か？
• RQ4 全結合ネットワークがなぜテキストベースの強化学習で報酬ハッキング行動を示すのか、そしてその対処法は何か？
• RQ5 ニューラルアーキテクチャの選択（例：トランスフォーマー 対 RNN）が、分子設計におけるオンライン強化学習の安定性とパフォーマンスに与える影響は何か？

主な発見

• ChemRLformerは、タンパク質ドッキングシミュレーションを含む25の分子設計タスクにおいて、より単純なアーキテクチャであるにもかかわらず、従来手法を上回る最先端のパフォーマンスを達成した。
• より小さな高品質なデータセット（例：ZINC 1M）での事前学習が、より大きな低品質なデータセットでの学習をはるかに上回り、データ品質が数量よりも重要であることを示した。
• ハイルクラビングリプレイバッファ設計は顕著なパフォーマンス向上をもたらしたが、標準的なリプレイバッファはほとんど利益をもたらさなかった。
• 事前学習済み方策に対するKL正則化はパフォーマンス向上に寄与しなかったため、この設定では必要ではないことが示唆された。
• 全結合ネットワークは報酬ハッキング行動を示し、報酬関数を悪用するために、長大な炭素および窒素原子のシーケンスを生成した。
• 理論的利点があるにもかかわらず、PPOベースのアルゴリズムは、分子最適化においてヴァナイル方策勾配法を下回った。これは、この分野では単純な強化学習アルゴリズムがより安定的かつ効果的である可能性を示している。

Figure 3 : Comparison between SELFIES and SMILES: The SELFIES representation makes it relatively difficult for ChemRLformer agents to explore effectively leading to generally lower performance on pytdc and docking while scoring higher on diversity. Scores are reported for the transformer model and a