Data-Efficient Graph Grammar Learning for Molecular Generation

The problem of molecular generation has received significant attention recently . Existing methods are typically based on deep neural networks and require training on large datasets with tens of thousands of samples . In practice , however , the size of class-specific chemical datasets is usually limited ( e.g. , dozens of samples ) due to labor-intensive experimentation and data collection . This presents a considerable challenge for the deep learning generative models to comprehensively describe the molecular design space . Another major challenge is to generate only physically synthesizable molecules . This is a non-trivial task for neural networkbased generative models since the relevant chemical knowledge can only be extracted and generalized from the limited training data . In this work , we propose a data-efficient generative model that can be learned from datasets with orders of magnitude smaller sizes than common benchmarks . At the heart of this method is a learnable graph grammar that generates molecules from a sequence of production rules . Without any human assistance , these production rules are automatically constructed from training data . Furthermore , additional chemical knowledge can be incorporated in the model by further grammar optimization . Our learned graph grammar yields state-of-the-art results on generating high quality molecules for three monomer datasets that contain only ∼20 samples each . Our approach also achieves remarkable performance in a challenging polymer generation task with only 117 training samples and is competitive against existing methods using 81k data points . 1 INTRODUCTION . The rise of computational approaches has started to have a significant impact on the discovery of materials and drugs . Recent advances in machine learning , especially deep learning ( DL ) , have driven rapid development on generating novel molecular structures ( Maziarka et al. , 2020 ; Xu et al. , 2019 ) .Various forms of generative models including generative adversarial networks ( De Cao & Kipf , 2018 ; Maziarka et al. , 2020 ) , variational autoencoders ( VAEs ) ( Jin et al. , 2018 ; 2020 ; Liu et al. , 2018 ; Sattarov et al. , 2019 ) , and reinforcement learning ( You et al. , 2018 ) have been exploited to represent the complicated molecular design space and generate new molecules . They typically formulate molecular generation as a problem of distribution learning , where the generative model first learns to reproduce the distribution of a training set before generating new molecular structures . Generative models have also managed to integrate certain chemical constraints ( e.g. , valency restrictions ) ( Jin et al. , 2018 ; Liu et al. , 2018 ) and shown promising results on the common benchmarks ( Irwin et al. , 2012 ; Ramakrishnan et al. , 2014 ) . However , DL-based generative models face a serious limitation : they require large amounts of training data to achieve reasonable performance . In practice , molecule data is not always abundant ( Stanley et al. , 2021 ; Altae-Tran et al. , 2017 ; Subramanian et al. , 2016 ) ; for instance , the focus may be on a specific type of molecules fulfilling certain ingredient requirements . Particularly in the context of polymers , large amounts of training data are not available , and therefore DL models use manually constructed or generated data ( St. John et al. , 2019 ; Jin et al. , 2020 ; Ma & Luo , 2020 ) . Real data sets , as used in one of the state-of-the-art papers on polyurethane property prediction , have as little as 20 samples ( Menon et al. , 2019 ) . In such scenarios , designing a pure DL-based model is challenging . Recent renewed interest in formal grammars ( Kajino , 2019 ; Krenn et al. , 2019 ; Nigam et al. , 2021 ; Guo et al. , 2021 ) provides an alternative to pure DL methods . In formal language theory , a grammar is a set of production rules describing how to generate valid strings according to the language ’ s syntax . A chemical grammar may thus be considered as an interpretable and compact design model that simultaneously serves as a molecular representation and a generative model ; even domain-specific constraints can be explicitly incorporated into the rules . Recent examples range from string-based ( Krenn et al. , 2019 ; Nigam et al. , 2021 ) to hypergraph-based ( Kajino , 2019 ) and polymer-specific grammars ( Guo et al. , 2021 ) . Grammar-based generative models do not rely on large training datasets and easily extrapolate to generate molecules outside the distribution of training samples . Yet , they have two major drawbacks . First , current approaches require that chemical grammars are manually designed by human experts ( Krenn et al. , 2019 ; Nigam et al. , 2021 ; Dai et al. , 2018 ) . This is a tedious process that heavily relies on expertise in chemistry . Moreover , the existing grammars are also very fine-grained ( i.e. , rules are mostly attaching atoms ) in order to cover the syntax of general molecules instead of a specific dataset . Hence , it is not straightforward how to incorporate the bias of given data ( e.g. , certain molecular substructures ) into such grammars . The second drawback is that it remains challenging to integrate chemistry knowledge beyond simple constraints ( e.g. , valency restrictions ) into the grammar . Current grammar construction approaches fail to capture more abstract or complex aspects such as the diversity or synthesizabilty of the generated molecules , which constrains their practicality . Solutions such as Kajino ( 2019 ) resort to costly optimization on the level of molecules ( or their latent representation ) after the grammar is built or learned . For a more detailed delimitation of related work , see Sec . 2 . In this paper , we propose a generative model combining complex graph grammar construction with a relatively simple and effective learning technique . In particular , our grammar incorporates substructures of varying sizes ( i.e. , above atom level ) and the construction process directly optimizes various chemical metrics ( e.g. , distribution statistics and synthesizability ) while satisfying specific chemical constraints ( e.g. , valency restrictions ) . Moreover , we have the benefits of symbolic knowledge representation : explainability and data efficiency . Our evaluation focuses on polymers , particularly their monomer building blocks . Thus , we curate new , realistic polymer datasets gathered from literature that represent specific classes of monomers . Note that our model works for arbitrary molecules . Framework . Figure 1 outlines our approach . Given molecules and domain-specific metrics to be optimized , we iteratively construct and evaluate a graph grammar as our generative model . We consider the construction as a minimum spanning forest problem and combine it with optimization of the metrics , via a learnable function Fθ determining which rules to construct . Results . Our model successfully deals with extreme settings – from a DL perspective – learning meaningful production rules based on only ∼10 samples ( e.g. , of a specific class ) ; this is of significant importance in a practical setting ( Stanley et al. , 2021 ; Altae-Tran et al. , 2017 ) but has been ignored so far . In particular , our model is able to generate members of a specific monomer class with a decent success rate . No previous state-of-the-art system achieves similar performance in our experiments . Generally , our approach works on training data that is orders of magnitude smaller than the amount needed by DL-based systems to produce meaningful results . Given ∼100 samples , we achieve performance comparable to state-of-the-art systems trained on 81k samples , across a wide range of common and new evaluation metrics . Besides , our grammar optimization method can adjust to any user-defined domain-specific metrics . The learned grammars can capture domainspecific knowledge explicitly , e.g. , the characteristic functional groups of a given class of polymers can be extracted from the production rules . 2 RELATED WORK . Generative Models for Molecules . Existing models can be categorized based on their molecule representation . It is common to use SMILES strings ( Weininger , 1988 ) . Yang et al . ( 2017 ) , Olivecrona et al . ( 2017 ) , and Gómez-Bombarelli et al . ( 2018 ) use recurrent neural networks to generate the strings in a sequential manner . For instance , Gómez-Bombarelli et al . ( 2018 ) propose a VAE based on sequence-to-sequence encoders and decoders . Guimaraes et al . ( 2017 ) , Sanchez-Lengeling et al . ( 2017 ) , and Dai et al . ( 2018 ) additionally apply optimization objectives to enforce similarity or semantic correctness to generate valid SMILES . Alternatively , molecules can be treated as graphs . Simonovsky & Komodakis ( 2018 ) , Ma et al . ( 2018 ) , and De Cao & Kipf ( 2018 ) generate the graphs in a single step , directly outputting adjacency matrices and node labels . GraphNVP ( Madhawa et al. , 2019 ) uses two steps to generate the node labels based on the adjacency matrix ; it exploits a model to encode molecules and uses the same model whose layers are reversed to generate new molecules . You et al . ( 2018 ) ; Li et al . ( 2018 ) ; Samanta et al . ( 2020 ) ; Liu et al . ( 2018 ) ; Liao et al . ( 2019 ) ; Jin et al . ( 2018 ; 2020 ) build molecule graphs iteratively based on either atom nodes or substructures . For example , JT-VAE ( Jin et al. , 2018 ) first generates a tree and then expands some of its nodes by attaching substructures based on a vocabulary mined from given molecules ; Similarly , HierVAE ( Jin et al. , 2020 ) also uses a vocabulary , but improves upon the combinatorial attachment process by directly applying a multi-resolution representation , allowing for generating much larger and diverse molecules . All these methods depend on DL and require large training datasets . Moreover , the resulting models are usually not interpretable . Our approach alleviates this obstacle and can deal with practically common setting of only dozens of available data points . Generative Models using Graph Grammars . Our approach belongs to the class of models generating new graphs based on the production rules of a graph grammar . These models are non-parametric , interpretable , and can incorporate meaningful graph properties into the rules . One option is to use grammars that are manually designed by human experts ( Dai et al. , 2018 ; Nigam et al. , 2021 ) . For instance , STONED ( Nigam et al. , 2021 ) generates new molecules based on a string-based representation of given molecules by replacing tokens ( representing atoms ) in accordance with the SELFIES grammar ( Krenn et al. , 2019 ) . While the simplicity of this approach is attractive , the generated molecules are not necessary reasonable from a chemical perspective ( e.g. , they might not be synthesizable ) . This problem can be overcome by mining the grammar from real data . Recently , several automatic techniques that explicitly construct a graph grammar have been proposed in the context of large-scale graphs to facilitate understanding and analysis ( Sikdar et al. , 2019 ; Aguinaga et al. , 2018 ; Hibshman et al. , 2019 ) . These works are domain-independent and do not allow specialization of the constructed grammar to reflect domain-specific knowledge . Closest to our method is MHG ( Kajino , 2019 ) , a framework that constructs a molecular hyperedge-replacement grammar based on Aguinaga et al . ( 2018 ) . From the given data , MHG learns a fine-grained grammar where the rules iteratively attach single atoms and therefore incorporates hard chemical constraints ( e.g. , valency restrictions ) . MHG then applies a VAE conditioned on the grammar to the molecules ’ tree decompositions in order to also learn soft constraints ( e.g. , stability ) . Finally , it uses Bayesian optimization to guide molecule generation . The fine-grained rules in MHG allows for rich diversity . However , our experiments show that these rules make it hard to capture distribution-specific properties , e.g. , characteristic substructures of a specific molecular class . Such a drawback limits the practical use of MHG . In contrast , our approach focuses on subgraphs instead of individual atoms only . It directly incorporates domain-specific knowledge into grammar construction and therefore avoids complicated post-processing to achieve high generation performance .

Summary: This paper introduces a new grammar-based generative model for the molecule generation task. The graph grammar is defined as a set of production rules that operate on module graphs (i.e. the molecule graphs are not linearized as done in some previous work). The graph grammar is learned by iteratively contracting hyperedges (i.e.  edges that can connect multiple nodes, defined by simple chemistry-inspired rules) into non-terminal nodes and the submission proposes to learn how to sample the hyperedges by a REINFORCE algorithm that optimizes for several molecule generation metrics.   The evaluation is done both on small and large molecule generation datasets. Notably, the proposed method achieves strong performance while being very data-efficient.