Leveraging Partial SMILES Validation Scheme for Enhanced Drug Design in Reinforcement Learning Frameworks

TL;DR

This work tackles catastrophic forgetting during reinforcement learning fine-tuning of SMILES-based molecular generators by introducing PSV-PPO, which integrates real-time Partial SMILES Validation through a PSV truth table into the PPO framework. The approach adds PSV-specific losses (entropy, Hellinger distance, GPS, and TPC) and leverages an experience replay regime to stabilize learning while encouraging diverse, valid molecule generation. Experiments on PMO, docking, and GuacaMol benchmarks show PSV-PPO reduces invalid structures, maintains competitive optimization performance, and benefits from a prespecified pretraining regime on ChEMBL with GPT-2, achieving strong scores across tasks. The method is general, enabling broader application in drug design and potentially incorporating additional domain knowledge to further enhance sequential SMILES generation models.

Abstract

SMILES-based molecule generation has emerged as a powerful approach in drug discovery. Deep reinforcement learning (RL) using large language model (LLM) has been incorporated into the molecule generation process to achieve high matching score in term of likelihood of desired molecule candidates. However, a critical challenge in this approach is catastrophic forgetting during the RL phase, where knowledge such as molecule validity, which often exceeds 99\% during pretraining, significantly deteriorates. Current RL algorithms applied in drug discovery, such as REINVENT, use prior models as anchors to retian pretraining knowledge, but these methods lack robust exploration mechanisms. To address these issues, we propose Partial SMILES Validation-PPO (PSV-PPO), a novel RL algorithm that incorporates real-time partial SMILES validation to prevent catastrophic forgetting while encouraging exploration. Unlike traditional RL approaches that validate molecule structures only after generating entire sequences, PSV-PPO performs stepwise validation at each auto-regressive step, evaluating not only the selected token candidate but also all potential branches stemming from the prior partial sequence. This enables early detection of invalid partial SMILES across all potential paths. As a result, PSV-PPO maintains high validity rates even during aggressive exploration of the vast chemical space. Our experiments on the PMO and GuacaMol benchmark datasets demonstrate that PSV-PPO significantly reduces the number of invalid generated structures while maintaining competitive exploration and optimization performance. While our work primarily focuses on maintaining validity, the framework of PSV-PPO can be extended in future research to incorporate additional forms of valuable domain knowledge, further enhancing reinforcement learning applications in drug discovery.

Figures (7)

• Figure 1: The training flow of the PSV-PPO algorithm. The process begins with the prior model generating molecular structures and their associated probabilities (1). Both the rewards and the Partial SMILES Validation (PSV) Table are computed concurrently, which minimizes the computational overhead (2). The experience replay memory pool is updated with the scored molecules (3). Samples are then drawn from the experience replay memory pool (4) and replayed through the current model (5). The PSV-PPO algorithm calculates multiple loss functions, including HD Loss, TPC Loss, Entropy Loss, GPS Loss, Value Loss, and Policy Loss (6), which are then accumulated (7) and used to perform a backward pass to update the model parameters (8). The PSV scheme ensures that the model learns to generate valid molecular structures by providing real-time feedback on token validity during the generation process.
• Figure 2: Visualization of the Partial SMILES Validation (PSV) framework used in our study. (a) PSV Truth Table illustrating the validity assessment of each candidate token in a partially generated SMILES string. (b) Box plot showing the distribution of the number of valid candidate tokens based on the preceding token in the SMILES string. (c) Overall distribution of valid candidate tokens across all states derived from 10,000 sampled molecules.
• Figure 3: Ablation study results for the Fexofenadine MPO and Osimertinib MPO tasks, illustrating the impact of various components within the PSV-PPO framework. The plots compare PSV-PPO with two ablation models: PSV-PPO_WO_PSV, which uses the standard PPO entropy and KL divergence losses without integrating the PSV table, and PSV-PPO_WO_PL, which excludes the GPS and TPC losses. The score plots demonstrate the optimization effectiveness of each model across epochs. The valid rate plots reveal how the absence of the PSV table significantly reduces the validity of generated molecules. The duplication plots in both the experience replay memory and during generation highlight the role of GPS and TPC losses in maintaining diversity and preventing mode collapse in the generated molecules.
• Figure 4: Docking Performance: Comparison of molecular docking performance on the fa7 target across different reinforcement learning methods. The plot presents scores for Top 1, Top 10, Top 100, and Diversity (Div.), evaluating HC (green), PPO (blue), and PSV-PPO (orange). PSV-PPO achieves competitive performance, demonstrating its effectiveness in generating diverse and high-quality molecules.
• Figure 5: Illustration of the top-scoring molecules generated by PSV-PPO for tasks 11-20 on the GuacaMol benchmark.