REINVENT-Transformer: Molecular De Novo Design through Transformer-based Reinforcement Learning

Authors: pengcheng  xu
tianfan  fu
wenhao  gao
Venue: KDD_org_2024_Workshop_AIDSH
Published: 2024-06-05 12:05:40

Abstract: In this work, we introduce a method: REINVENT-Transformer to fine-tune a Transformer-based generative model for molecular de novo design. Leveraging the superior sequence learning capacity of Transformers over Recurrent Neural Networks (RNNs), our model can generate molecular structures with desired properties effectively. In contrast to the traditional RNN-based models, our proposed method exhibits superior performance in generating compounds predicted to be active against various biological targets, capturing long-term dependencies in the molecular structure sequence. The model's efficacy is demonstrated across numerous tasks, including generating analogues to a query structure and producing compounds with particular attributes, outperforming the baseline RNN-based methods. Our approach can be used for scaffold hopping, library expansion starting from a single molecule and generating compounds with high predicted activity against biological targets.

Main paper: 
 introduction: 
 Navigating the vast chemical space, which contains 10 60 -10 100 possible molecules, is a critical challenge in drug discovery [34, 45] . Early de novo design algorithms [6, 17] and RNN-based models [14, 32] have partially addressed this complexity. However, the Transformer architecture has proven superior, especially in handling the long-term dependencies necessary for modeling complex molecules, due to its:(1) Parallelization: Processes all tokens simultaneously, unlike step-by-step processing in RNNs, enhancing efficiency. (2) Long-term Dependency Handling: Employs multi-head selfattention mechanisms to capture long-range interactions. (3) Scalability: Better suited for longer sequences, a key advantage in molecular design.This work introduces a novel approach by incorporating the Decision Transformer in de novo molecular design. Leveraging "oracle feedback reinforcement learning," our model optimizes towards molecules with high predicted activity, providing precision and enhancing success rates in drug discovery [32] . This integration sets a new standard in the field, emphasizing the transformative potential of Transformer-based architectures in molecular design. 
related works: 
 Early de novo design algorithms primarily focused on structurebased methods, aiming to develop ligands that fit the binding pocket of a target, as highlighted in works by [6] and [17] . These methods often resulted in molecules with suboptimal drug metabolism and pharmacokinetic (DMPK) properties, presenting challenges in synthetic tractability. Ligand-based approaches, not relying on the 3D structure of the target, were introduced to address some of these limitations, involving a virtual library of chemical structures evaluated using a scoring function [21, 33] . However, the effectiveness of ligand-based methods compared to structure-based ones is not definitive [4] .Generative models, including RNN-based methods, have been employed in de novo design [10, 19, 35, 43] . These models learn the probability distribution over chemical structures and are further fine-tuned using reinforcement learning (RL) [24] , achieving significant improvements.The Transformer architecture, known for its self-attention mechanism, addresses the challenges of long-term dependencies in sequence data [39] . Motivated by its success, we propose using Transformer-based architectures for molecular de novo design in place of RNNs.Molecular assembly strategies include string-based approaches like SMILES and SELFIES [29, 40] , and graph-based methods representing molecular structures [26, 46] . Synthesis-based strategies focus on generating synthesizable molecules [7, 9, 10, 15] .Optimization algorithms used in molecular design include Genetic Algorithms (GAs) and Bayesian optimization (BO), which mimic natural evolutionary processes and build surrogates for the objective function, respectively [8, 28, 30, 31] . Variational autoencoders (VAEs) and reinforcement learning (RL) have also been used to map molecules from a latent space and refine models for enhanced molecule generation [19, 26, 32] . Recent advancements like Pasithea and Differentiable scaffolding tree (DST) utilize gradientbased optimization [11, 36] .The evolution from RNN-based methods to Transformer-based models reflects a desire to better handle complex chemical structures and optimize them more effectively. The Transformer architecture, particularly its self-attention mechanism, effectively handles sequence data [39] . Combined with RL for fine-tuning [12, 14, 24, 32, 42, 46] , this integration aims to improve molecular design processes, promising more sophisticated automated systems. 
methodology: reinvent-transformer: 
 Our method, named REINVENT-Transformer, first pre-trains the real 2D molecule dataset based on the transformer. Then, based on the RL paradigm, fine-tuning is performed on the molecular attributes to be optimized.This study focuses on single-objective molecular optimization for designing small organic molecules with significant scalar properties in therapeutic development. The molecular design task is an optimization problem:𝑚 * = arg max 𝑚∈ M O (𝑚),where 𝑚 is a molecular structure and M is the vast chemical space of potential candidates, approximately 10 60 [5] . We assume access to the actual value of a targeted property, O (𝑚) : M → R, evaluated by an oracle, O, which is an opaque mechanism providing a scalar value for specific chemical or biological attributes. These oracles, whether experimental or high-fidelity simulations, are costly. Hence, an efficient optimization algorithm within feasible resource constraints is vital, significantly aiding automated molecular design in advanced automated chemical design (ACD) [18] or functiondriven autonomous synthesis [16] .The transformer is used for pre-training on real 2D molecules. Specifically, it treats the prediction of a 2D molecule as a sequence prediction and lets the transformer predict the next atom based on the molecular sequence history. The pre-training of the transformer is based on maximum likelihood.Training data Overview: Segmentation and Binary Coding of SMILES A Simplified Molecular Input Line Entry System (SMILES) [41] represents a molecule as a character sequence with atoms and symbols for ring closure, opening, and branching. SMILES are typically tokenized by single characters, except for two-character atoms like "Cl" and "Br" and special cases in square brackets (e.g., [nH] ) treated as single tokens. This tokenization approach identified 86 tokens in the training data.A molecule can have multiple SMILES representations. Canonicalization algorithms [40] ensure consistent SMILES for the same molecule, though different implementations may produce varied SMILES. Transformers Overview Transformers are a neural network architecture designed to process sequential data, while also accounting for the importance of each input in relation to the others, despite their position in the sequence [39] . They manage to do this by the introduction of an attention mechanism that assesses the significance of each input in the sequence (Figure 1 ). At any given step 𝑡, the transformer state at 𝑡 is influenced by all previous inputs 𝑥 1 , . . . , 𝑥 𝑡 -1 and the current input 𝑥 𝑡 . The transformer's ability to selectively focus on the parts of the input sequence that are most relevant for each step makes them especially well suited for tasks in the field of natural language processing. Sequences of words can be encoded into one-hot vectors with a length equivalent to our vocabulary size 𝑋 . We may add two extra tokens, GO and EOS, to signify the beginning and end of a sequence, respectively.Learning to model the data Training a Transformer for sequence modeling typically involves using maximum likelihood estimation to predict the next token 𝑥 𝑡 in the target sequence, given tokens from the previous steps (Figure 1 ). The model generates a probability distribution at every step, representing the likely next character, and the objective is to maximize the likelihood assigned to the correct token:EQUATIONThe cost function 𝐽 (Θ), often applied to a subset of all training examples known as a batch, is minimized with respect to the network parameters Θ. Given a predicted log likelihood log 𝑃 of the target at step 𝑡, the gradient of the cost function with respect to Θ is used to update Θ. This method of fitting a neural network is called back-propagation. Changing the network parameters affects not only the immediate output at time 𝑡, but also influences the information flow into subsequent transformer states. Generating new samples Once a Transformer has been trained on target sequences, it can be used to generate new sequences that adhere to the conditional probability distributions learned from the training set. The first input is the GO token, and at every timestep following, we sample an output token 𝑥 𝑡 from the predicted probability distribution 𝑃 (𝑋 𝑡 ) over our vocabulary 𝑋 . The sampled 𝑥 𝑡 is then used as our next input. The sequence is considered finished once the EOS token is sampled.In this part, we load the pre-trained transformer network and finetune it based on RL. Here, our task is to generate some specific molecules with good attributes. Therefore, we use the generated molecules to measure the properties of the corresponding molecules through Oracle, and use them as rewards to finetune the neural network.Agent Decision-Making and Markov Decision Processes Consider an Agent choosing an action 𝑎 ∈ A(𝑠) in state 𝑠 ∈ S, where S is the set of states and A(𝑠) is the set of actions. The policy 𝜋 (𝑎 | 𝑠) maps states to action probabilities. Reinforcement learning often uses Markov decision processes (MDPs), where the current state is sufficient for decision-making [37] . This can extend to partially observable MDPs with partial environment views. The reward 𝑟 (𝑎 | 𝑠) measures action effectiveness, and the long-term return 𝐺 (𝑎 𝑡 , 𝑆 𝑡 ) = 𝑇 𝑡 𝑟 𝑡 is the cumulative reward from time 𝑡 to 𝑇 . For molecular desirability, we consider the return of a complete SMILES sequence.Reinforcement learning aims to improve the policy to maximize expected return E[𝐺]. Tasks ending at step 𝑇 are episodic [37] , such as SMILES generation ending with an EOS token.States and actions for training can be generated by the agent (on-policy) or others (off-policy) [37] . Two RL strategies are valuebased and policy-based [37] . Value-based RL learns a value function to derive a policy, while policy-based RL directly learns the policy. For our problem, policy-based methods are preferred because:• They can learn an optimal stochastic policy [37] , aligning with our goal. • Fine-tuning a prior sequence model with a scoring function needs minimal changes, reducing gradient estimate variance.Negative Log-Likelihood (NLL) and Loss Function To assess the likelihood of sequence generation by the agent, we use the Negative Log-Likelihood (NLL). The NLL is calculated as follows:EQUATION)This measure is crucial in understanding the generative model's performance [3] . The augmented likelihood and loss function are then computed to adjust the agent's generation process:𝑁 𝐿𝐿(𝑺) Augmented = 𝑁 𝐿𝐿(𝑺) Prior -𝝈 * 𝑀𝑃𝑂 (𝑺) score loss = 𝑁 𝐿𝐿(𝑺) Augmented -𝑁 𝐿𝐿(𝑺) Agent 2 .Scoring Functions for Molecular Sequences REINVENT-Transformer utilizes scoring functions to evaluate and guide the generation of molecular sequences. These functions are formulated as either a weighted product or a weighted sum:𝑆 (𝑥) = 𝑖 𝑝 𝑖 (𝑥) 𝑤 𝑖 1/ 𝑖 𝑤 𝑖 , 𝑆 (𝑥) = 𝑖 𝑤 𝑖 * 𝑝 𝑖 (𝑥) 𝑖 𝑤 𝑖 .This scoring approach is designed to balance various molecular properties during the generation process, facilitating the production of molecules with desired characteristics [3] . 
experiment 4.1 experimental setup: 
 Dataset. For methods requiring a database, we use the ZINC 250K dataset [22, 23] , consisting of around 250K molecules. This dataset is significant in pharmaceuticals and is used in Screening [1] , MolPAL [20] , and pretraining generative models like VAEs [27] and LSTMs [44] . Essential fragments for JT-VAE [26] , MIMOSA [13] , and DST [11] are also derived from it.Baseline. We compare eight baseline methods for performance evaluation, including REINVENT [32] , Graph-GA [25] , SELFIES-REINVENT [14] , GP BO [38] , STONED [31] , SMILES-LSTM HC [8] , SMILES-GA [8] , SynNet [15] , DoG-Gen [7] , and DST [11] . The implementations come from the PMO benchmark [14]1 .Metric. In order to evaluate both optimization capability and sample efficiency, following [14] , we use the area under the curve (AUC) of the top-𝐾 average property value in relation to the number of oracle calls.Our result is shown in Table . 1. From the table, we can observe that our method is better than the baseline method on multiple Oracles, which proves the effectiveness of the transformer in our problem. Our experiments mainly follow the benchmark paper [14] .The evaluation results depict a thorough comparison between the REINVENT-Transformer (referred to as REINVENT-Trans) and other prominent models across multiple oracles. Randomly selected SMILES generated by different models can be seen in Table 4 . And the corresponding chemical structures are shown in figure 5 .REINVENT-Transformer consistently excels in molecular generation, achieving top results in properties like 'Albuterol_Similarity', 'Mestranol_Similarity', 'QED', 'Scaffold_Hop', and 'Sitagliptin_MPO'. This indicates the model's strength in capturing intricate molecular patterns and optimizing desired properties.1. Versus REINVENT (SMILES and SELFIES): REINVENT-Transformer often outperforms REINVENT (SMILES), though REIN-VENT scores slightly better on 'Osimertinib_MPO'. SELFIES representation in REINVENT doesn't always enhance performance, highlighting the impact of model architecture and representations.2. Graph-based Models: 'Graph GA' and 'GP BO' show strong performance in 'Amlodipine_MPO' and 'Celecoxib_Rediscovery' oracles, respectively, but aren't consistently top-performing, suggesting their limited generalizability.3. Genetic Algorithms: STONED (using SELFIES) achieves the highest score in 'Fexofenadine_MPO', demonstrating the potential of genetic algorithms in specific optimization tasks despite their stochastic nature.Generation Comparison with REINVENT-SMILES In order to better investigate in the ability of our method in long sequence generation, we did the following ablation study.The box plot visualizes the distribution of evaluation scores across different molecular lengths for both the REINVENT-Transformer method and the baseline REINVENT method.Based on the figure 2 , we can derive the following observations: 1. In general, REINVENT-Transformer will generate longer average length of molecules than REINVENT.2. The REINVENT-Transformer method consistently achieves higher average scores.3. The spread (interquartile range) of scores for the REINVENT-Transformer method remains relatively consistent across molecular lengths, indicating stable performance. In conclusion, the REINVENT-Transformer method outperforms the baseline REINVENT method, particularly in the context of longer molecular sequences. We set a threshold=50 for the length of the generated molecular string. If the generated string is longer than the threshold, it will be considered as "long", other it's considered as "short". From Figure 3 , we can see the our method REINVENT-Transformer has better average score when generating long sequences.We plotted the auc_topk curve and the number of oracle calls is the x-axis. From the figure as follows, we can see that our method REINVENT-Transformer converges faster than REINVENT method. From Fig. 4 , the evolution of the average accuracy for the top 100 predictions is evident. Upon examination, across equivalent number of oracle calls, the mean accuracy of REINVENT-Transformer consistently surpasses that of REINVENT. This indicates a more expedient convergence rate for the REINVENT-Transformer compared to REINVENT.The avg_top100 curve initially displays a steep incline, eventually plateauing post approximately 6000 oracle calls. Notably, beginning from the 2500th oracle call, the performance differential between REINVENT-Transformer and REINVENT significantly widens. It is also observed that the REINVENT-Transformer possesses a higher standard deviation relative to REINVENT, suggesting potential variability in its performance. Despite this, the difference between the average top100 accuracy and the standard deviation for REINVENT-Transformer remains superior to the mean accuracy of REINVENT, reaffirming the enhanced efficacy of the REINVENT-Transformer method. 
conclusion: 
 Navigating the vast chemical space in molecular design remains challenging, but the introduction of the REINVENT-Transformer marks a significant advancement by harnessing strengths such as parallelization and long-term dependency handling in the Transformer architecture. Our experimental findings demonstrate its superior performance across multiple oracles, particularly in tasks  
 
 Appendix: The vast expanse of chemical space, encompassing an order of magnitude from 10 60 -10 100 possible synthetically feasible molecules [34, 45] , presents formidable obstacles to drug discovery endeavors. In this colossal landscape, the task of pinpointing a molecule that simultaneously meets the prerequisites for bioactivity, drug metabolism and pharmacokinetic (DMPK) profile, and synthetic accessibility becomes an undertaking similar to the proverbial search for a needle in a haystack. Pioneering de novo design algorithms [6, 17] have attempted to address this by employing virtual strategies to design and evaluate molecules, thereby condensing the vast chemical space into a more navigable realm for exploration.Traditional de novo molecule design models, based on Recurrent Neural Networks (RNNs), have proven effective in molecule generation tasks. However, RNNs possess inherent architectural limitations, notably in their capability to capture long-term dependencies in sequential data, which can be particularly detrimental when modeling complex molecular structures. Recently, the Transformer architecture has emerged as a powerful alternative to RNNs in sequence modeling tasks across various domains. Some of the key advantages of Transformers over RNNs include:(1) Parallelization: Unlike RNNs which process sequences step-by-step, Transformers process all tokens in the sequence simultaneously, allowing for better computational efficiency. (2) Long-term Dependency Handling: Transformers utilize multi-head self-attention mechanisms, which can capture long-range interactions in the data, making them particularly well-suited for modeling intricate molecular structures. (3) Scalability: Transformers are inherently more scalable, allowing for the processing of longer sequences, which is a considerable advantage in molecular design.In light of these advantages, our work introduces a novel approach by integrating the Transformer architecture, specifically the Decision Transformer, for de novo molecular design. By leveraging the inherent strengths of Transformers, our model exhibits enhanced performance in generating molecular structures with desired attributes.Furthermore, we emphasize the incorporation of the "oracle feedback reinforcement learning" method. Pretraining models on large datasets is beneficial, but downstream tasks often require fine-tuning on specific objectives. By integrating feedback from an oracle during the reinforcement learning phase, our approach can efficiently navigate the solution space, optimizing towards molecules with high predicted activity. Such oracle-guided optimization provides an added layer of precision, facilitating the generation of molecules that not only conform to structural constraints but also exhibit high bioactivity, thereby increasing the potential success rate in drug discovery endeavors.Drawing inspiration from previous work that employed RNNs and reinforcement learning for molecular optimization [32] , our approach distinguishes itself by the adoption and fine-tuning of the Transformer architecture, ensuring superior handling of long-sequence data and paving the way for innovative breakthroughs in the realm of molecular design.In summary, this work presents a fresh perspective on molecular de novo design, underscoring the potential of Transformer-based architectures, complemented by oracle feedback reinforcement learning, to revolutionize drug discovery methodologies. We envision that our approach will not only set a new benchmark in molecular generation tasks but will also inspire future research in leveraging advanced machine learning architectures for complex scientific challenges.Early de novo design algorithms primarily focused on structure-based methods, aiming to develop ligands that precisely fit the binding pocket of a target [6, 17] . While effective in certain aspects, these methods often resulted in molecules with suboptimal drug metabolism and pharmacokinetic (DMPK) properties and posed challenges in synthetic tractability. Ligand-based approaches were introduced to overcome some of these limitations, involving the creation of comprehensive virtual libraries of chemical structures evaluated using scoring functions [21, 33] . However, as noted by [4] , the effectiveness of ligand-based methods compared to structure-based ones is not definitive, with both approaches having unique advantages and limitations depending on the specific requirements of the drug design process.Recently, generative models such as RNN-based methods have been successfully applied to de novo design of molecules [10, 19, 35, 43] . These models have shown promise in learning the underlying probability distribution over large sets of chemical structures, effectively reducing the search space to reasonable molecules. Further improvements were achieved through fine-tuning using reinforcement learning (RL) techniques [24] , demonstrating considerable enhancements over initial models.Despite these advancements, challenges such as capturing long-term dependencies in sequence data persist. The Transformer architecture [39] , known for its self-attention mechanism and ability to handle long sequences, has been highly successful in various sequence prediction tasks across domains. Motivated by these successes, we propose the use of Transformer-based architectures in place of RNNs for molecular de novo design.Molecular assembly strategies play a crucial role in representing and manipulating chemical structures. String-based approaches like SMILES and SELFIES [29, 40] provide efficient representations of molecules. Graph-based methods offer intuitive two-dimensional representations of molecular structures, with nodes and edges representing atoms and bonds, respectively [26, 46] . Synthesis-based strategies aim to generate only synthesizable molecules, ensuring that the design aligns with real-world applications [7, 9, 10, 15] .Various optimization algorithms have been utilized for molecular design. Genetic Algorithms (GAs) mimic natural evolutionary processes and have been applied in molecule generation using both SMILES and SELFIES representations [8, 31] . Bayesian optimization (BO) builds a surrogate for the objective function, with applications such as BOSS and ChemBO in the molecular domain [28, 30] . Variational autoencoders (VAEs) offer a generative approach, mapping molecules to and from a latent space, with notable methods including SMILES-VAE and JT-VAE [19, 26] . Reinforcement Learning (RL) techniques, like REINVENT, have also been applied to tune models for molecule generation [32] . Recent advancements in gradient ascent methods, such as Pasithea and Differentiable scaffolding tree (DST), have leveraged gradient-based optimization for molecular design [11, 36] .The evolution of molecular design methodologies has progressively addressed various challenges and limitations. The transition from RNN-based methods to more advanced generative models underscores a quest for improved handling of complex chemical structure representations and optimization. While RNNs brought significant progress, their inherent difficulty in capturing long-term dependencies in sequential data has been a notable shortcoming. The Transformer architecture addresses this gap through its self-attention mechanism, allowing for more nuanced and effective handling of sequence data, which is critical in molecular where long-range interactions within molecules play a pivotal role.The integration of reinforcement learning (RL) for fine-tuning generative models has further enhanced the field [24] . RL's ability to iteratively improve models based on a feedback loop aligns well with the demands of molecular design, where continuous refinement based on molecular properties is essential. The combination of RL with generative models has been shown to enhance the ability to navigate the vast chemical space more effectively, achieving better results in molecule generation [12, 14, 32, 42, 46] .In light of these advancements and existing limitations, our work proposes an approach that integrates the Transformer architecture with advanced RL techniques. This proposal is underpinned by the Transformer's superior handling of sequential data and the iterative refinement capability of RL. By merging these two powerful technologies, we aim to address the existing challenges in molecular de novo design, such as the need for better sequence representation and optimization. This integration promises to enhance the effectiveness and efficiency of molecular generation processes, moving closer to achieving more sophisticated and automated molecular design systems.A transformer block is a parameterized function classEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONThe notation softmax 𝑗 indicates we take the softmax (defined in Equation 9) over the 𝑑-dimensional vector indexed by 𝑗. The LayerNorm function [2] is defined for z ∈ R 𝑘 by LayerNorm(z; 𝛾, 𝛽) = 𝛾 (z -𝜇 z )EQUATIONEQUATION)The set of parameters, denoted by 𝜃 , comprises the elements of the weight matrices 𝑊 and the LayerNorm parameters 𝛾 and 𝛽, as specified on the right-hand side. The input x ∈ R 𝑛×𝑑 represents a set of 𝑛 entities, each characterized by 𝑑 attributes (typically, though not exclusively, sequences of 𝑑-dimensional vectors of length 𝑛). It is important to note that the output z ∈ R 𝑛×𝑑 retains the same format as the input x ∈ R 𝑛×𝑑 . A transformer is an amalgamation of 𝐿 distinct transformer blocks, each equipped with unique parameters:𝑓 𝜃 𝐿 •• • •• 𝑓 𝜃 1 (x) ∈ R 𝑛×𝑑 .Key hyperparameters in a transformer include 𝑑, 𝑘, 𝑚, 𝐻 , and 𝐿, with typical configurations being 𝑑 = 512, 𝑘 = 64, 𝑚 = 2048, 𝐻 = 8. While the initial research suggested 𝐿 = 6, more recent studies tend to employ a greater number of these blocks.We frame the problem of generating a SMILES representation of a molecule with specified desirable properties via a Transformer as a partially observable Markov decision process. In this framework, an Agent must decide on an action 𝑎 ∈ A(𝑠) to take given a particular state 𝑠 ∈ S, where S denotes the set of possible states and A(𝑠) represents the set of potential actions for that state. The policy 𝜋 (𝑎 | 𝑠) of an Agent associates a state to the likelihood of each action executed within.Many reinforcement learning problems are modeled as Markov decision processes, indicating that the current state provides all essential information to inform our action choice, with no additional benefit from knowing past states' history. While this is more of an approximation than a fact for most real-life challenges, we extend this concept to a partially observable Markov decision process where the Agent interacts with a partial environment representation.Let 𝑟 (𝑎 | 𝑠) be the reward serving as an indicator of the effectiveness of an action taken at a certain state, and the long-term return 𝐺 (𝑎 𝑡 , 𝑆 𝑡 ) = 𝑇 𝑡 𝑟 𝑡 represents the cumulative rewards collected from time 𝑡 to time 𝑇 [37] . As molecular desirability is only meaningful for a completed SMILES, we will only consider a complete sequence's return.The main objective of reinforcement learning is to enhance the Agent's policy to increase the expected return E[𝐺] based on a set of actions taken from some states and the obtained rewards. A task with a definitive endpoint at step 𝑇 is known as an episodic task [37] , where 𝑇 corresponds to the episode's length. SMILES generation is an example of an episodic task, which concludes once the EOS token is sampled.The states and actions used for Agent training can be produced by the agent itself or through other means. If the agent generates them, the learning is called on-policy, and if generated by other means, it is off-policy learning [37] .Reinforcement learning commonly employs two different strategies to determine a policy: value-based RL and policy-based RL [37] . In value-based RL, the aim is to learn a value function that describes a given state's expected return. Once this function is learned, a policy can be established to maximize a certain action's expected state value. In contrast, policy-based RL aims to learn a policy directly.For the problem we are addressing, we believe policy-based methods are the most suitable for the following reasons:• Policy-based methods can explicitly learn an optimal stochastic policy [37] , which aligns with our objective.• Our method starts with a prior sequence model. The goal is to fine-tune this model based on a specific scoring function. Since the prior model already embodies a policy, fine-tuning might require only minimal changes to the prior model. • The episodes in our case are short and fast to sample, reducing the impact of the variance in the estimate of the gradients.In our approach, we use the probability distributions learned by a pre-trained Transformer model as our initial prior policy. We refer to the network using the prior policy simply as the Prior, and the network whose policy has been modified as the Agent. The task is episodic, starting with the first step of the Transformer and ending when the EOS token is sampled. The sequence of actions 𝐴 = 𝑎 1 , 𝑎 2 , . . . , 𝑎 𝑇 during this episode represents the SMILES generated, and the product of the action probabilities 𝑃 (𝐴) = 𝑇 𝑡 =1 𝜋 (𝑎 𝑡 | 𝑠 𝑡 ) represents the model likelihood of the sequence formed.We introduce a scoring function 𝑆 (𝐴) ∈ [-1, 1] that rates the desirability of the sequences formed. The goal is to update the agent policy 𝜋 from the prior policy 𝜋 Prior to increase the expected score for the generated sequences while remaining anchored to the prior policy. We define an augmented likelihood log 𝑃 (𝐴) U as a prior likelihood modulated by the desirability of a sequence: log 𝑃 (𝐴) U = log 𝑃 (𝐴) Prior + 𝜎𝑆 (𝐴) where 𝜎 is a scalar coefficient. The return 𝐺 (𝐴) of a sequence 𝐴 can be seen as the agreement between the Agent likelihood log 𝑃 (𝐴) A and the augmented likelihood:𝐺 (𝐴) = -[log 𝑃 (𝐴) U -log 𝑃 (𝐴) A ] 2The goal of the Agent is to learn a policy which maximizes the expected return, achieved by minimizing the cost function 𝐿(Θ) = -𝐺. This approach can be described as a REINFORCE algorithm with a final step reward, allowing for effective optimization of the Agent's policy towards generating desirable molecular structures.This metric, which we refer to as AUC top-K, is defined as:Given a sequence of molecules {𝑀 1 , 𝑀 2 , . . . , 𝑀 𝑁 } generated by a method, and an oracle function 𝑂 (𝑀) that returns the property value of a molecule, the top-𝐾 average property value at any point in the sequence is given by:EQUATIONMethod Overview Technical Details Advantage Disadvantage REINVENT [32] A method employing a policybased reinforcement learning approach to instruct RNNs to produce SMILES strings.Formulates molecular design as a Markov decision process with states representing partially generated molecules and actions as string manipulations. Rewards based on properties of interest.Adaptable to generate other string representations like SELFIES.Heavily reliant on the design of rewards.A genetic algorithm that manipulates molecular representations using graphs, with graph matching and atom/fragment mutations.Introduces crossover operations based on graph representations, unlike stringbased genetic algorithms.Offers a richer set of operations for exploring diverse chemical spaces.Increased complexity due to graph-based operations.An extension of REIN-VENT for generating SELFreferencing Embedded Strings (SELFIES).Uses a policy-based RL approach for SELFIES representation, ensuring syntactical validity.Produces molecules with fewer syntactical errors.Still dependent on reward system definition.GP BO [38] Combines Gaussian process Bayesian optimization with Graph-GA methods.Leverages GP acquisition function integrated with Graph-GA techniques for sampling.Balances exploration and exploitation effectively.Higher computational costs due to GP and GA interplay.A modified genetic algorithm that manipulates tokens within SELFIES strings.Interacts directly with tokens in SELFIES strings, differing from traditional stringbased GAs.Direct approach potentially reduces invalid chemical representations.Limited to SELFIES, may not generalize to other representations. SMILES-LSTM HC [8] Iterative learning method using LSTM to understand the molecular distribution in SMILES strings.Employs a variant of the cross-entropy method, fine-tuning the model with highscoring molecules.Iteratively refines the generative process.Slow convergence if initial model is suboptimal.Genetic algorithm based on SMILES context-free grammar.Implements genetic mutations and crossovers based on SMILES grammar.Exploits SMILES structure for effective exploration.Confined to SMILES grammar nuances, potentially missing novel structures. SynNet [15] Synthesis-based genetic algorithm operating on binary fingerprints and decoding to synthetic pathways.Focuses on the synthesizability of generated molecules.Prioritizes synthesizability, ensuring lab producibility of molecules.Limited diversity in molecular space exploration due to synthesis emphasis.DoG-Gen [7] Tailored to learn the distribution of synthetic pathways.Represents synthetic pathways as DAGs, using an RNN generator for modeling. Emphasizes synthesizability.Structured approach to learning synthetic pathways.Issues in capturing very long sequences with RNNs if not designed effectively. DST [11] Differentiable Scaffolding Tree method for molecular optimization using gradient ascent.Abstracts molecular graphs into scaffolding trees, using a graph neural network for gradient estimation.Direct optimization of molecular structures through gradient computation.Possible loss of information due to abstraction to scaffolding trees. https://github.com/wenhao-gao/mol_opt 



Reviews: review: title: title: Advanced Approach
main_review: main_review: The paper presents a cutting-edge application of Transformer-based models for molecular de novo design. The ability of Transformers to handle long-term dependencies in molecular sequences marks a substantial improvement over traditional RNN-based models.The research provides an in-depth comparison with existing methods, showcasing the superior performance of the REINVENT-Transformer in generating molecules with desired properties. This comprehensive benchmarking against various oracles highlights the model’s efficacy. The practical applications of the REINVENT-Transformer in drug discovery are well articulated. Its ability to generate compounds with high predicted activity against biological targets and perform tasks like scaffold hopping and library expansion is highly valuable.The experimental setup and evaluation are robust, with clear improvements demonstrated over baseline methods. The use of multiple metrics, such as AUC for top-K property values, provides a well-rounded assessment of the model’s performance. The REINVENT-Transformer sets a new standard in molecular de novo design. Its success in leveraging Transformers for this task suggests that further developments and refinements could lead to even greater advancements in drug discovery.
Limitations: The model’s predictions are based on computational evaluations. Experimental validation of the generated compounds’ properties in a laboratory setting is crucial to confirm their real-world applicability. The model’s ability to handle highly diverse molecular structures, particularly those with unconventional features, is not fully explored. Future work could focus on enhancing its generalizability across different types of molecules.
title: title: Review of "REINVENT-Transformer: Molecular De Novo Design through Transformer-based Reinforcement Learning"
main_review: main_review: Summary
The manuscript integrates transformers with reinforcement learning, presenting a novel method for molecular design that enhances the efficiency and accuracy of molecular structure generation. Additionally, the manuscript demonstrates the model's potential and superior performance across a variety of molecular design tasks through a series of experiments. The work is solid and well-executed.

Strengths
- The manuscript successfully applies a framework combining transformers with reinforcement learning fine-tuning to the field of molecular design, showcasing an exemplary AI for science application.
- The inclusion of case studies provides a clear demonstration of the model's interpretability, which significantly aids in understanding the practical implications of the research.

Weaknesses
- The manuscript could improve its formatting. For example, in the "Related Works" section, the text about transformers exceeds the text box limits, and in Section 3.2, the citation markers extend beyond the line width.

Question
- I am not very familiar with the field of chemical molecules, but I wonder why transformers might even surpass graph-based methods in tasks involving molecules, which naturally form graph structures. Could the authors provide some insight into this?