[论文解读] 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

一句话总结

本文提出 ChemRLformer，一种基于 Transformer 和基于文本的分子表示（SMILES/SELFIES）的强化学习框架，用于发现高价值分子。通过结合在对齐数据集上的预训练，以及针对性的算法组件（如爬山式经验回放缓冲区和 KL 正则化），ChemRLformer 在 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.

研究动机与目标

• 开发一种更高效且简洁的基于文本表示的强化学习方法，用于分子设计。
• 识别在基于文本的强化学习中，哪些架构和算法选择在分子生成任务中显著提升性能。
• 评估在多样化数据集上的预训练以及奖励函数设计对分子优化结果的影响。
• 分析 Transformer 在在线强化学习中用于分子生成时的稳定性和性能。
• 证明通过消融驱动设计的简化方法可在高价值分子发现中超越复杂的先前方法。

提出的方法

• 该方法采用基于 Transformer 的策略网络，通过强化学习以逐 token 的方式生成 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 神经架构选择（如 Transformer 与 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