TSSR: Two-Stage Swap-Reward-Driven Reinforcement Learning for Character-Level SMILES Generation

📄 arXiv: 2601.04521v1 📥 PDF

作者: Jacob Ede Levine, Yun Lyan Luo, Sai Chandra Kosaraju

分类: cs.LG, cs.AI

发布日期: 2026-01-08

备注: Under Review

💡 一句话要点

提出TSSR：一种双阶段交换奖励驱动的强化学习方法，用于字符级SMILES生成。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 分子生成 SMILES生成 强化学习 药物发现 化学信息学

📋 核心要点

  1. 现有化学语言模型在生成SMILES字符串时，容易累积token错误，导致生成大量无效或化学上不合理的分子。
  2. TSSR通过双阶段强化学习框架，分别奖励语法修复和化学合理性提升的token交换，从而引导模型生成更有效的分子。
  3. 实验表明，TSSR在句法有效性、化学有效性和新颖性方面均有显著提升，同时保持了药物相似性和可合成性。

📝 摘要（中文）

可靠、有效和多样化的分子设计是现代药物发现的基础。改进的分子生成技术能够更有效地探索化学空间，寻找潜在的药物候选者，并降低早期设计成本。然而，当前将分子生成为SMILES字符串的化学语言模型容易出现累积的token错误，导致许多样本无法解析或化学上不合理，而旨在防止失败的硬约束又会限制探索。为了解决这个问题，我们引入了TSSR，一种双阶段、交换奖励驱动的强化学习框架，用于字符级SMILES生成。第一阶段奖励修复语法的局部token交换，促进从无效字符串到可解析字符串的转换。第二阶段提供来自RDKit诊断的化学感知反馈，奖励减少价、芳香性和连接性问题。该奖励分解为可解释的项（交换效率、错误减少、与有效性的距离），与模型无关，并且不需要特定于任务的标签或手工制作的语法。我们在MOSES基准上评估了TSSR，使用GRU策略，通过PPO在纯强化学习（P-RL）和从预训练化学语言模型开始的微调强化学习（F-RL）中进行训练，每次运行评估10,000个生成的SMILES。在P-RL中，TSSR显著提高了句法有效性、化学有效性和新颖性。在F-RL中，TSSR在提高有效性和新颖性的同时，保留了药物相似性和可合成性。Token级别的分析表明，语法编辑和化学修复共同作用以减少RDKit检测到的错误。TSSR将稀疏的终端目标转换为更密集和更可解释的奖励，在不降低多样性的情况下提高了句法和化学质量。TSSR与数据集无关，可以适应各种强化学习方法。

🔬 方法详解

问题定义：现有化学语言模型在生成SMILES字符串时，面临着token错误累积的问题，导致生成大量无效或化学上不合理的分子。硬约束虽然可以防止无效分子的生成，但会限制化学空间的探索。因此，需要一种方法能够在保证分子有效性的同时，提高生成分子的多样性和新颖性。

核心思路：TSSR的核心思路是将分子生成过程分解为两个阶段，并分别设计奖励函数来引导模型学习。第一阶段关注语法有效性，奖励修复语法的token交换；第二阶段关注化学合理性，奖励减少化学错误的token交换。通过这种方式，模型可以逐步学习生成有效的分子，同时保持探索化学空间的能力。

技术框架：TSSR是一个双阶段的强化学习框架。第一阶段，模型通过PPO算法学习一个GRU策略，该策略的目标是最大化修复语法的奖励。奖励函数基于token交换的效率和与有效字符串的距离。第二阶段，模型继续通过PPO算法学习，但此时的奖励函数基于RDKit的诊断结果，奖励减少价、芳香性和连接性问题的token交换。整个框架是模型无关的，可以与不同的强化学习算法和分子生成模型结合使用。

关键创新：TSSR的关键创新在于其双阶段的奖励函数设计。传统的强化学习方法通常使用稀疏的终端奖励，例如只有当生成完全有效的分子时才给予奖励。TSSR将奖励分解为更密集和可解释的项，包括交换效率、错误减少和与有效性的距离。这种设计使得模型更容易学习，并且可以更好地平衡有效性和多样性。此外，TSSR不需要任务特定的标签或手工制作的语法，使其更具通用性。

关键设计：TSSR使用GRU作为策略网络，并使用PPO算法进行训练。奖励函数的设计是关键，第一阶段的奖励函数鼓励模型进行有效的token交换，以修复语法错误。第二阶段的奖励函数基于RDKit的诊断结果，鼓励模型减少化学错误。具体来说，奖励函数包括交换效率、错误减少和与有效性的距离等项。这些项的权重需要根据具体任务进行调整。

📊 实验亮点

实验结果表明，TSSR在MOSES基准测试中显著提高了句法有效性、化学有效性和新颖性。在纯强化学习（P-RL）中，TSSR的各项指标均优于基线方法。在微调强化学习（F-RL）中，TSSR在提高有效性和新颖性的同时，保留了药物相似性和可合成性。Token级别的分析表明，语法编辑和化学修复共同作用以减少RDKit检测到的错误。例如，TSSR可以将有效性从随机初始化的1%提高到超过80%。

🎯 应用场景

TSSR可应用于药物发现领域，用于生成具有特定性质的候选药物分子。通过提高生成分子的有效性和多样性，TSSR可以加速药物筛选过程，降低研发成本。此外，TSSR还可以应用于材料科学领域，用于设计具有特定功能的材料分子。该方法具有广泛的应用前景，可以促进相关领域的发展。

📄 摘要（原文）

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.