[Paper Review] 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 Science3 citations

TL;DR

This paper proposes Augmented Memory, a reinforcement learning algorithm that enhances sample efficiency in de novo molecular design by combining SMILES data augmentation with experience replay. By reusing high-scoring oracle results across multiple training updates and selectively purging low-diversity molecules, the method achieves state-of-the-art performance on the PMO benchmark, outperforming prior methods on 19 out of 23 tasks under a 10,000-oracle call budget.

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.

Motivation & Objective

• To address the critical challenge of sample efficiency in de novo molecular design, where expensive oracles limit the number of evaluations.
• To improve the performance of policy-based reinforcement learning models like REINVENT by leveraging experience replay and SMILES augmentation.
• To maintain molecular diversity while enhancing sample efficiency through a novel selective memory purge mechanism.
• To extend the Practical Molecular Optimization (PMO) benchmark with new implementations of Augmented Memory and BAR, enabling rigorous comparison.
• To demonstrate superior optimization efficiency in both exploitation tasks and real-world drug discovery scenarios.

Proposed method

• The method uses a replay buffer to store high-scoring molecules from previous generations, enabling multiple policy updates per oracle call.
• It applies SMILES augmentation to generate alternative string representations of the same molecule, reusing the same oracle score for multiple training updates.
• An augmented likelihood is computed by combining the prior model’s log-probability with a scaled oracle score, forming a surrogate reward.
• The algorithm performs multiple training updates per epoch: one on original samples and another on augmented SMILES, both using the same reward signal.
• A replay buffer is updated iteratively and optionally purged using a diversity filter that removes molecules with undesired scaffolds.
• The final loss combines updates from both original and augmented SMILES sequences, improving policy learning efficiency.

Experimental results

Research questions

• RQ1 Can experience replay significantly improve sample efficiency in policy-based molecular generation under a strict computational budget?
• RQ2 To what extent does SMILES data augmentation enhance learning efficiency when combined with experience replay?
• RQ3 Can a selective memory purge mechanism maintain molecular diversity while preserving sample efficiency gains?
• RQ4 How does Augmented Memory compare to state-of-the-art methods like REINVENT and BAR in terms of optimization performance across diverse molecular design tasks?
• RQ5 Does the proposed method generalize across both exploitation tasks and complex, sparse-reward drug discovery scenarios?

Key findings

• Augmented Memory achieves a new state-of-the-art performance on the PMO benchmark, outperforming the previous best method, REINVENT, on 19 out of 23 tasks.
• The method significantly improves sample efficiency by reusing oracle scores across multiple training updates via SMILES augmentation and experience replay.
• The inclusion of selective memory purge enhances diversity without sacrificing optimization performance, preventing premature convergence to low-diversity solutions.
• The algorithm demonstrates strong generalization, excelling in both controlled exploitation tasks and complex, real-world drug discovery case studies.
• Empirical results show that combining data augmentation with experience replay leads to faster convergence and higher reward accumulation within the 10,000-oracle budget.
• The method’s performance is validated through extensive ablation and comparison, confirming the critical role of experience replay and augmentation in sample-efficient molecular design.