[論文レビュー] TSSR: Two-Stage Swap-Reward-Driven Reinforcement Learning for Character-Level SMILES Generation

Jacob Ede Levine, Yun Lyan Luo+1 more|arXiv (Cornell University)|Jan 8, 2026

Machine Learning in Materials Science被引用数 0

ひとこと要約

この論文は、まず局所トークン置換で構文エラーを修復し、次にRDKit診断を用いた化学整合性の評価で文字レベルのSMILES生成を改善する2段階強化学習フレームワーク「TSSR」を紹介する。学習-from-scratchとファインチューニングの双方でMOSESにおける有効性と新規性の向上を実証する。

ABSTRACT

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.

研究の動機と目的

• SMILESによる新規分子生成を信頼性・有効性・多様性の観点から動機付ける
• トークンレベルのSMILES生成を導く密度の高い、解釈可能なフィードバックを提供する
• モデルおよびデータセットに依存しないRLフレームワークを、スクラッチまたはファインチューニングで適用可能とする
• MOSESベンチマークで構文的/化学的有効性と新規性の改善を示す
• 手作り文法に頼らず、標準的なRL手法と互換性があることを示す

提案手法

• 2段階報酬を提案: 第1段階は構文を修正してパース可能なSMILESを生み出す局所トークン交換を報酬化
• 第2段階は構文修正後のRDKit検出化学問題の減少を報酬化
• スワップ効率、エラー低減、有効性までの距離などを含むモデル非依存の報酬分解を使用
• PPOで訓練されたGRUベースの化学言語モデルを2つの regimeで使用: P-RL（ランダム初期化）と F-RL（事前学習モデル）
• MOSESデータ上で、トークン事前分布をグローバルトークン頻度と標準SMILES語彙から導出
• 学習ダイナミクスを解釈するため、トークンレベルの分析とスワップ数、修正率、化学エラー低減を公開報告

Figure 1: Example a Two-Stage, Swap-Reward-driven (TSSR) reinforcement learning (RL) framework for character-level SMILES generation.

実験結果

リサーチクエスチョン

• RQ1 2段階のスワップ報酬RLフレームワークは文字レベルSMILES生成の構文的有効性を改善できるか？
• RQ2 第2段階の化学認識型フィードバックは第1段階の修正後のRDKit検出エラーを減らすか？
• RQ3 TSSRで最適化された場合、スクラッチ学習と事前学習モデルの両方で有効性と新規性は向上するか？
• RQ4 生成分子の薬物様性、合成可能性、多様性、骨格多様性にはどのような影響があるか？
• RQ5 このアプローチはデータセットやモデルに依存せず、PPOのような標準的なRLパイプラインと互換性があるか？

主な発見

• TSSRはP-RLにおける構文的有効性を大幅に改善し、未学習ベースラインと比較して化学的有効性と新規性を高める
• P-RLでは構文的有効性が6.14%から35.03%へ、化学的有効性が4.77%から9.61%へ上昇し、顕著な新規性の向上を示す
• F-RLでは有効性は平均0.83%の穏やかな増加にとどまる一方で新規性はおおよそ99.6%を維持し、全体の化学的有効性は19.20%へ上昇
• 第1段階のスワップと第2段階の修正が協働して機構を作り、構文修正により化学補正が可能となりRDKit検出エラーを減らす
• TSSRはより密で解釈可能な報酬信号を返し、多様性を損なうことなく構文的・化学的品質を向上させる
• P-RLはピーク報酬と学習効率が高い一方、F-RLは事前学習された事前知識を活用しスループットは高いが有効性の獲得はやや小さい

Figure 2: Examples of TSSR Stage Two fixes: Invalid SMILES to Chemically valid with upto 3 Fixes Each