TSSR: Two-Stage Swap-Reward-Driven Reinforcement Learning for Character-Level SMILES Generation
作者: Jacob Ede Levine, Yun Lyan Luo, Sai Chandra Kosaraju
分类: cs.LG, cs.AI
发布日期: 2026-01-08
备注: Under Review
💡 一句话要点
提出TSSR：一种双阶段交换奖励驱动的强化学习方法，用于字符级SMILES生成。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: SMILES生成 强化学习 药物发现 分子设计 化学信息学
📋 核心要点
- 现有化学语言模型在生成SMILES字符串时易累积token错误，导致生成无效或化学上不合理的分子。
- TSSR通过双阶段强化学习框架，利用交换奖励驱动，分别修复语法错误和化学合理性问题。
- 实验表明，TSSR在句法有效性、化学有效性和新颖性方面均有显著提升，且能保留药物相似性和可合成性。
📝 摘要（中文）
可靠、有效和多样化的分子设计是现代药物发现的基础。改进的分子生成技术能够更有效地探索化学空间，寻找潜在的药物候选者，并降低早期设计成本。然而，当前将分子生成为SMILES字符串的化学语言模型容易出现累积的token错误，导致许多样本无法解析或在化学上不合理，而旨在防止失败的硬约束又会限制探索。为了解决这个问题，我们引入了TSSR，一种双阶段、交换奖励驱动的强化学习（RL）框架，用于字符级SMILES生成。第一阶段奖励修复语法的局部token交换，促进从无效字符串到可解析字符串的转换。第二阶段提供来自RDKit诊断的化学感知反馈，奖励价、芳香性和连通性问题的减少。该奖励分解为可解释的项（交换效率、误差减少、与有效性的距离），具有模型无关性，并且不需要特定于任务的标签或手工制作的语法。我们在MOSES基准上评估了TSSR，使用GRU策略，通过PPO在纯RL（P-RL）中从随机初始化开始训练，以及在微调RL（F-RL）中从预训练的化学语言模型开始训练，每次运行评估10,000个生成的SMILES。在P-RL中，TSSR显著提高了句法有效性、化学有效性和新颖性。在F-RL中，TSSR在提高有效性和新颖性的同时，保留了药物相似性和可合成性。Token级别的分析表明，语法编辑和化学修复共同作用以减少RDKit检测到的错误。TSSR将稀疏的终端目标转换为更密集和可解释的奖励，在不降低多样性的情况下提高了句法和化学质量。TSSR是数据集无关的，可以适应各种强化学习方法。
🔬 方法详解
问题定义：论文旨在解决化学语言模型在生成SMILES字符串时出现的token错误累积问题，这些错误导致生成的分子无效或不合理。现有方法要么依赖硬约束限制探索，要么难以有效纠正语法和化学错误，导致生成分子的质量和多样性受限。
核心思路：论文的核心思路是将SMILES生成过程分解为两个阶段，并设计相应的奖励函数来指导强化学习。第一阶段关注语法修复，奖励将无效字符串转换为可解析字符串的token交换。第二阶段关注化学合理性，奖励减少价、芳香性和连通性问题的token交换。这种分解使得奖励信号更密集、可解释，并能有效引导模型生成高质量的SMILES字符串。
技术框架：TSSR框架包含两个主要阶段：语法修复阶段和化学合理性阶段。在每个阶段，使用强化学习算法（如PPO）训练一个策略网络（如GRU）来生成SMILES字符串。在语法修复阶段，奖励函数基于token交换的效率和减少语法错误的程度。在化学合理性阶段，奖励函数基于RDKit诊断结果，奖励减少化学错误的程度。整个框架是模型无关的，可以与不同的策略网络和强化学习算法结合使用。
关键创新：TSSR的关键创新在于其双阶段交换奖励驱动的强化学习框架。与传统的单阶段强化学习方法相比，TSSR能够更有效地解决SMILES生成中的语法和化学错误。此外，TSSR的奖励函数是可解释的，并且不需要特定于任务的标签或手工制作的语法，使其更具通用性和易用性。
关键设计：TSSR的关键设计包括：1) 使用GRU作为策略网络，用于生成SMILES字符串；2) 使用PPO算法进行策略优化；3) 设计了两个阶段的奖励函数，分别用于语法修复和化学合理性；4) 奖励函数分解为可解释的项，如交换效率、误差减少和与有效性的距离；5) 使用RDKit进行化学合理性评估。
📊 实验亮点
TSSR在MOSES基准测试中表现出色。在纯RL（P-RL）中，TSSR显著提高了句法有效性、化学有效性和新颖性。在微调RL（F-RL）中，TSSR在提高有效性和新颖性的同时，保留了药物相似性和可合成性。Token级别的分析表明，语法编辑和化学修复共同作用以减少RDKit检测到的错误。
🎯 应用场景
TSSR框架可应用于药物发现领域，用于生成具有特定性质的候选药物分子。通过优化SMILES生成过程，可以更有效地探索化学空间，发现新的潜在药物，并降低药物研发的成本。此外，该方法还可以扩展到其他分子生成任务，例如材料科学中的新材料设计。
📄 摘要（原文）
The design of reliable, valid, and diverse molecules is fundamental to modern drug discovery, as improved molecular generation supports efficient exploration of the chemical space for potential drug candidates and reduces the cost of early design efforts. Despite these needs, current chemical language models that generate molecules as SMILES strings are vulnerable to compounding token errors: many samples are unparseable or chemically implausible, and hard constraints meant to prevent failure can restrict exploration. To address this gap, we introduce TSSR, a Two-Stage, Swap-Reward-driven reinforcement learning (RL) framework for character-level SMILES generation. Stage one rewards local token swaps that repair syntax, promoting transitions from invalid to parseable strings. Stage two provides chemistry-aware feedback from RDKit diagnostics, rewarding reductions in valence, aromaticity, and connectivity issues. The reward decomposes into interpretable terms (swap efficiency, error reduction, distance to validity), is model agnostic, and requires no task-specific labels or hand-crafted grammars. We evaluated TSSR on the MOSES benchmark using a GRU policy trained with PPO in both pure RL (P-RL) from random initialization and fine-tuning RL (F-RL) starting from a pretrained chemical language model, assessing 10,000 generated SMILES per run. In P-RL, TSSR significantly improves syntactic validity, chemical validity, and novelty. In F-RL, TSSR preserves drug-likeness and synthesizability while increasing validity and novelty. Token-level analysis shows that syntax edits and chemistry fixes act jointly to reduce RDKit detected errors. TSSR converts a sparse terminal objective into a denser and more interpretable reward, improving both syntactic and chemical quality without reducing diversity. TSSR is dataset-agnostic and can be adapted to various reinforcement learning approaches.