[论文解读] Augmented Memory: Capitalizing on Experience Replay to Accelerate De Novo Molecular Design

Jeff J. Guo, Philippe Schwaller|arXiv (Cornell University)|May 10, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文提出了一种名为增强记忆（Augmented Memory）的强化学习算法，通过结合SMILES数据增强与经验回放，提升了从零开始的分子设计中的样本效率。通过在多个训练更新中重用高分评分的模型输出结果，并选择性地清除低多样性分子，该方法在PMO基准测试中实现了最先进性能，在10,000次模型调用预算下，23项任务中有19项优于先前方法。

ABSTRACT

Sample efficiency is a fundamental challenge in de novo molecular design. Ideally, molecular generative models should learn to satisfy a desired objective under minimal oracle evaluations (computational prediction or wet-lab experiment). This problem becomes more apparent when using oracles that can provide increased predictive accuracy but impose a significant cost. Consequently, these oracles cannot be directly optimized under a practical budget. Molecular generative models have shown remarkable sample efficiency when coupled with reinforcement learning, as demonstrated in the Practical Molecular Optimization (PMO) benchmark. Here, we propose a novel algorithm called Augmented Memory that combines data augmentation with experience replay. We show that scores obtained from oracle calls can be reused to update the model multiple times. We compare Augmented Memory to previously proposed algorithms and show significantly enhanced sample efficiency in an exploitation task and a drug discovery case study requiring both exploration and exploitation. Our method achieves a new state-of-the-art in the PMO benchmark which enforces a computational budget, outperforming the previous best performing method on 19/23 tasks.

研究动机与目标

• 为解决从零开始分子设计中的关键挑战，即昂贵的模型输出限制了评估次数。
• 通过利用经验回放和SMILES增强，提升基于策略的强化学习模型（如REINVENT）的性能。
• 通过一种新颖的选择性记忆清除机制，在保持分子多样性的同时提升样本效率。
• 通过新增增强记忆和BAR的实现方式，扩展实用分子优化（Practical Molecular Optimization, PMO）基准，以支持严格比较。
• 在利用型任务和真实世界药物发现场景中，展示更优的优化效率。

提出的方法

• 该方法使用回放缓冲区存储先前生成中高分分子，实现每次模型输出调用后进行多次策略更新。
• 通过SMILES增强生成同一分子的替代字符串表示，使相同模型评分在多次训练更新中重复使用。
• 通过将先验模型的对数概率与缩放后的模型输出评分相结合，计算增强似然，形成代理奖励。
• 算法在每个训练周期内执行多次更新：一次基于原始样本，一次基于增强SMILES，两者使用相同的奖励信号。
• 回放缓冲区通过迭代方式更新，并可选择性地使用多样性过滤器清除具有不良骨架结构的分子。
• 最终损失函数结合了原始和增强SMILES序列的更新结果，从而提升策略学习效率。

实验结果

研究问题

• RQ1 在严格的计算预算下，经验回放是否能显著提升基于策略的分子生成中的样本效率？
• RQ2 当与经验回放结合时，SMILES数据增强在多大程度上能提升学习效率？
• RQ3 选择性记忆清除机制是否能在保持样本效率增益的同时维持分子多样性？
• RQ4 在多种分子设计任务中，增强记忆与REINVENT和BAR等最先进方法相比，其优化性能如何？
• RQ5 所提出方法是否在利用型任务和复杂、稀疏奖励的药物发现场景中均具有泛化能力？

主要发现

• 增强记忆在PMO基准测试中实现了新的最先进性能，在23项任务中有19项优于此前最佳方法REINVENT。
• 该方法通过SMILES增强与经验回放结合，显著提升了样本效率，实现多次训练更新中重用模型输出评分。
• 引入选择性记忆清除机制在不牺牲优化性能的前提下增强了多样性，防止过早收敛至低多样性解。
• 该算法展现出强大的泛化能力，在控制型利用任务和复杂的真实世界药物发现案例研究中均表现优异。
• 实证结果表明，将数据增强与经验回放结合可实现更快收敛和在10,000次模型调用预算内更高的奖励累积。
• 通过广泛的消融实验与对比验证，确认经验回放与数据增强在样本高效分子设计中具有关键作用。