[Paper Review] Grammars and reinforcement learning for molecule optimization

Egor Kraev|arXiv (Cornell University)|Nov 27, 2018

Machine Learning in Materials Science8 references5 citations

TL;DR

This paper proposes a novel method combining a custom context-free grammar (CFG) for SMILES with reinforcement learning using a Transformer model to generate chemically valid molecules with optimized properties. By enforcing valence rules and cycle constraints via rule masking, the approach ensures validity while enabling efficient, high-quality generation, achieving state-of-the-art performance in molecular property optimization with significantly fewer model steps per atom.

ABSTRACT

We seek to automate the design of molecules based on specific chemical properties. Our primary contributions are a simpler method for generating SMILES strings guaranteed to be chemically valid, using a combination of a new context-free grammar for SMILES and additional masking logic; and casting the molecular property optimization as a reinforcement learning problem, specifically best-of-batch policy gradient applied to a Transformer model architecture. This approach uses substantially fewer model steps per atom than earlier approaches, thus enabling generation of larger molecules, and beats previous state-of-the art baselines by a significant margin. Applying reinforcement learning to a combination of a custom context-free grammar with additional masking to enforce non-local constraints is applicable to any optimization of a graph structure under a mixture of local and nonlocal constraints.

Motivation & Objective

• To automate the design of molecules with desired chemical properties using deep learning.
• To address the challenge of generating only chemically valid SMILES strings without relying on complex architectures or large pre-defined vocabularies.
• To improve sample efficiency in molecular generation by reducing the number of model steps per atom.
• To enable optimization of complex molecular properties under both local (valence) and non-local (cycle closure, length) constraints.
• To demonstrate that a simpler, grammar-guided RL approach can surpass state-of-the-art results in molecular optimization.

Proposed method

• A new context-free grammar (CFG) for SMILES is designed to guarantee correct atom valences and support complex cycle structures by construction.
• Additional masking logic is introduced to enforce non-local constraints: terminal distance masking ensures termination within the model's step limit, and cycle ID propagation ensures cycle closure and length constraints.
• The generation process uses a Transformer-based policy network that leverages full-context attention to make informed decisions at each step.
• Molecular property optimization is framed as a reinforcement learning problem using best-of-batch policy gradient, with a reward function combining property scores, SA score, and penalties for undesirable features.
• Standard CFG parsers (e.g., NLTK) are used to decompose real molecules into production rule sequences for training and validation.
• The method allows for efficient, scalable generation of large molecules by decoupling validity enforcement from property optimization.

Experimental results

Research questions

• RQ1 Can a custom context-free grammar for SMILES be designed to guarantee chemically valid molecules while supporting complex structural features like branching and aromatic cycles?
• RQ2 How can non-local constraints—such as cycle closure and length limits—be effectively enforced during sequence generation using only local attention mechanisms?
• RQ3 Can a Transformer-based policy with best-of-batch policy gradient outperform prior RNN- or graph-based models in molecular property optimization with fewer model steps per atom?
• RQ4 To what extent can grammar-based generation with rule masking reduce the need for large pre-defined vocabularies or complex message-passing architectures?
• RQ5 How do different reward shaping strategies (e.g., anchoring, SA penalty, aromatic cycle penalties) affect the diversity and drug-likeness of generated molecules?

Key findings

• The proposed method achieved a molecular score of 8.46 on the ZINC benchmark, significantly exceeding the previous state-of-the-art score of 7.16 reported by Jin et al. (2018).
• With strong anchoring and SA penalty, the top molecules showed structural diversity including six-membered aromatic rings, five-membered heterocycles with S and N, and only one chlorine atom, resembling real drug-like molecules.
• The model required substantially fewer model steps per atom than prior approaches, enabling efficient generation of larger molecules.
• Even without explicit cycle penalties, the method produced molecules with fewer aromatic cycles and more balanced atom diversity when strong anchoring was applied.
• The approach outperformed both the baseline in Kusner et al. (2017) and the complex graph-based model in Jin et al. (2018), despite its simpler architecture and grammar-based design.
• The use of rule masking with cycle IDs enabled enforcement of complex constraints (e.g., minimum/maximum cycle length) while preserving grammatical validity and parseability.