A Transformer-based Generative Model for De Novo Molecular Design

In the scope of drug discovery, the molecular design aims to identify novel compounds from the chemical space where the potential drug-like molecules are estimated to be in the order of 10^60 - 10^100. Since this search task is computationally intractable due to the unbounded search space, deep learning draws a lot of attention as a new way of generating unseen molecules. As we seek compounds with specific target proteins, we propose a Transformer-based deep model for de novo target-specific molecular design. The proposed method is capable of generating both drug-like compounds (without specified targets) and target-specific compounds. The latter are generated by enforcing different keys and values of the multi-head attention for each target. In this way, we allow the generation of SMILES strings to be conditional on the specified target. Experimental results demonstrate that our method is capable of generating both valid drug-like compounds and target-specific compounds. Moreover, the sampled compounds from conditional model largely occupy the real target-specific molecules' chemical space and also cover a significant fraction of novel compounds.


Introduction
Small molecule drug design aims to identify novel compounds with desired chemical properties. From the computational perspective, we consider this task as an optimization problem, where we search for the compounds that will maximize our quantitative goals in chemical space. However, this optimization task is computationally intractable because the searching space is extremely large. It has been estimated that the range of potential drug-like molecules is between 10 23 and 10 60 [1], but only about 10 8 molecules have ever been synthesized [2]. Numerous computational methods, such as virtual screening, combinatorial libraries, and evolutionary algorithms, have been developed to search the chemical space in silico and in vitro. In the past few years, recent works have demonstrated that machine learning, especially deep learning methods, could produce new small molecules [3,4,5] with biological activity.
Prior work is inspired by the recent advance in natural language processing (NLP) since the chemical structures can be formed in SMILES (Simplified Molecular Input Line Entry System) string. Thus the small molecule drug design problem can be translated to the string generation problem. A number of deep learning techniques have been successfully applied in natural language processing (NLP), including human-like text generation. For instance, the GPT series [6,7,8] uses an autoregressive language model to produce human-like text by training from the massive unlabeled human-written text. The text generated from the GPT model is of high quality and hard to be distinguished from human-written content. Therefore, we are optimistic about learning drug-like structures using deep learning from the existing large-scale chemical structures.
In the scope of small molecule drug design, it is essential to enable the generations to be guided by predefined conditions, such as the target protein. Here we formulate the molecular design problem as a conditional sequential generation given the target protein and propose a conditional Transformer architecture that auto-regressively generates target-specific compounds. We first train a conditional transformer on the MOSES dataset [9] without target information, and after training, the conditional transformer is able to generate drug-like structures. Then, the conditional transformer is fine-tuned on three target-specific datasets (EGFR, HTR1A, and S1PR1 targets). Our experiment results show that the transformer is capable of generating compounds that are similar to the ones in the training set but are still novel structures. We believe the conditional transformer to be a useful tool for de novo molecule design.

Preliminary
Transformer [10] is proposed to address sequential modeling by attention mechanism. Many neural sequence transduction models consist of an encoder and a decoder. The encoder first takes a sequence of tokens (x 1 , ..., x m ) and transforms them to a sequence of latent representations z = (z 1 , ..., z m ) (e.g., "memories"). The decoder will generate an output sequence (t 1 , ..., t n ) one by one conditioning on z. An intuitive way of sequential generation is in an auto-regressive manner [11], which means consuming all the previously generated tokens while generating the next one. While the traditional Transformer is in an encoder-decoder manner, we define our de novo SMILES generation task as a conditional generator, and we use a decoder-only design. Nevertheless, we briefly introduce the complete design of the Transformer to make this paper self-contained and refer readers to [10] for more details.
An attention mechanism mimics the process of querying a set of key-value pairs, where the output is a weighted sum over the values, and each weight is based on the matching of the key and query. The multi-head attention projects the keys, values, and queries h times and performs attention in parallel.

The formal definition of Multi-head Attention is as follows
Multi-head Attention (MHA) We first define some annotations: query matrices Q i = QW Q i , key matrices K i = KW K i , and value matrices V i = V W V i (i = 1, ..., h).
Encoder The encoder has a stack of identical layers. Each has two sub-layers: a multi-head attention component, followed by a feed-forward network; a residual connection is deployed around each of the two sub-layers, followed by layer normalization.
Decoder The encoder also has a stack of identical layers. Each has three sub-layers (two of them are the same as the encoder) including an extra sub-layer performing multi-head attention over the output (e.g., latent representations z) of the encoder.

Unsupervised Pre-training
Even though many deep learning tasks rely on supervised learning and human-labeled dataset, we are able to form an unsupervised learning task for drug discovery and overcome the challenge of expensive and hard-to-manage human labeling. As a popular supervised learning model, the sequence-to-sequence [12] (or seq2seq) model has been enjoying massive success in many natural language processing applications. The seq2seq models are usually trained end-to-end with a large number of training pairs (e.g., article-summary pairs for text summarization). To leverage the large number of unlabeled training data in the chemical space, instead of a seq2seq encoder-decoder model, we consider a decoder-only model and form our task as predicting the next token given previously generated tokens (as shown in Figure 1).
Inspired by the success of unsupervised Transformer on NLP (e.g., GPT [6] and GPT2 [7]), we propose to use the Transformer-based model for drug discovery. We form our drug discovery task by taking advantage of SMILES' sequential structure and transform it into a sequential drug-like structure generation. We formalize our task as (1) an auto-regressive token generation, (2) the generation follows SMILES grammar (e.g., atom type, bond type, and size of molecules) by observations from the training set, and (3) the generation produces more variations that are not previously observed.

Incorporating Target Information
One of the major challenges in our task is to generate target-specific SMILE sequences. The model needs to memorize both valid drug-like structures and target-specific information.
The Transformer structure is compatible with sequential dependencies, and it has been proven to be successful in memorizing drug-like structures. The challenge lies in capturing target-specific information and further target-specific generation. We propose taking advantage of the Multi-head Attention in the Transformer decoder and imposing the target-specific embeddings to the Key and Value of the attention operations. We denote our Transformer model with imposed conditional embeddings (e.g., target-specific embeddings) as cTransformer.
We draw inspiration from Transformer-based encoder-decoder structure for sequence-to-sequence translation. Taking machine translation as an example, the Transformer encoder-decoder model performs very well for this task. Specifically, the encoder memorizes the input sentence and stores them in "memories", and the decoder first attends to the previously generated tokens (first mha 2 ) and then performs multi-head attention over the output "memories" of the encoder (second mha).
When attending to the "memories" of the encoder using second multi-head attention, the queries are from the first multi-head attention and values and keys from the "memories" of the encoder. The intuition is to enable the decoder to attend over the input sequence [10]. We enable the SMILES sequence generation to condition on the specific target by feeding target-specific embeddings (denoted as e t ) to a decoder-only Transformer (shown in Figure 2). We use target-specific embeddings as keys and values of second mha, which allows every position of the decoder to attend to the target-specific embeddings and ensures the subsequent token generations are conditioned on the target-specific embeddings. It is worth noticing that our target-specific design is orthogonal to the decoder and can be easily removed by setting the condition embeddings as zero embeddings.
As shown in Figure 2, instead of fetching memory from an encoder, our decoder only design initialize the memory based on condition embeddings. For base model pre-training, e t is initialized with zero embeddings; when targets are involved, e t is initialized with target-specific embeddings.
The training process of our task can be summarized as follows 1. We first pre-train the base model of cTransformer by setting the target-specific embeddings as zero embeddings (following Figure 1 (a) without feeding targetspecific information). 2. To feed target-specific information, we fine-tune cTransformer with <compound, target> pairs and enforce conditions of the corresponding target by feeding target-specific embeddings to the attention layer of the Transformer decoder as "memories" (as shown in Figure 1 (a) and 2). 3. We can generate drug-like structure by autoregressively sampling tokens from the trained decoder (following Figure 1 (b)). Optionally, we can designate the desired target by feeding a target-specific embedding.

Experiments
We evaluate our transformer on two main tasks. The first task is molecular generative modeling (using the base model of cTransformer). This experiment validates our argument that the transformer is advantageous for small molecular generation tasks (Section 4.1). The second task is the targetspecific generation, and we show that the proposed method can generate target-biased compounds (Section 4.2).

Molecular generative model
Dataset We use the MOSES molecular dataset from Polykovskiy [9] to perform unsupervised pre-training on our cTransformer with the target-specific embedding initialized as zero. It contains 1,760,739 drug molecules extracted from ZINC clean Lead Collection [13], including 1,584,664 training molecules and 176,075 testing molecules. We follow the same train/test split as in [9].
Evaluation Metrics We evaluate the generated compounds in various aspects of molecule generation proposed in [9]. Besides basic metrics such as chemical validity and diversity, we compare the distribution of drug-likeness properties between generated and real compounds. Our metrics include: • Fraction of valid (Valid) and unique molecules report valid and unique SMILES strings. The validity is checked using a molecular structure parser (RDKit). • Unique@1K and Unique@10K for the first 1000 and 10000 valid molecules in the generated set. In general, validity measures whether the model captures enough chemical constraints (e.g., valence); Uniqueness measure whether the model overlaps with trained molecules. • Fragment similarity (Frag) compares distributions of BRICS fragments [14] in generated and reference sets. This metric measures how similar are the scaffolds present in generated and reference datasets.
• Similarity to a nearest neighbor (SNN) calculates the average Tanimoto similarity (also known as the Jaccard index) between fingerprints of a molecule from the generated set and its nearest molecule in the reference dataset based on [15] [16].
• Molecular weight (MW) is the sum of atomic weights in a molecule.
• Synthetic Accessibility Score (SA) estimates the ease of synthesis (synthetic accessibility) of drug-like molecules based on molecular complexity and fragments contributions [18].
• Quantitative Estimation of Drug-likeness (QED) measures the drug-likeness based on desirability in a value between zero (all properties unfavorable) to one (all properties favorable) [19].

Results
The performance of the different approaches is summarized in Table 1. Our method (cTransformer) achieves state-of-the-art results in the Fraction of valid (Valid), Unique@1k, Unique@10k, Fragment similarity, and Similarity to the nearest neighbor. Moreover, we compare distributions of four molecular properties in generated and test datasets in Figure 3: molecular weight (MW), octanol-water partition coefficient (logP), quantitative estimation of drug-likeness (QED), and synthetic accessibility score. Our model closely matches the real data distribution. This validates that our method (cTransformer) is capable of generating drug-like molecules.

Target-specific molecular generation
Dataset The target dataset was acquired from [23] for training our target-specific Transformer cTransformer. It contains 1381, 795, and 3,485 molecules corresponding to EGFR, S1PR1, and HTR1A, respectively.

Results
We evaluate if our approach is capable of generating target-specific compounds by visualizing the chemical space. The hypothesis is that the compounds that can potentially interact with the same protein target would populate the same sub-chemical space. To evaluate the overlapping of chemical  space, we first generate 1000 compounds for each target, then the 1024-bit FCFP6 fingerprint vectors are calculated for the generated compounds and the target-specific training datasets. We use the Uniform Manifold Approximation and Projection (UMAP) to construct 2D projections. These projections are illustrated in Figure 4. Each point corresponds to a molecule and is colored according to its target label. The dark and light colors represent the generated compounds and ground-true targetspecific molecules, respectively. The visualization of chemical space in Figure 4 demonstrates the generated target-specific molecules (dark color) and real target-specific molecules (light color) occupy the same sub-chemical space. These results show that our cTransformer can generate compounds that are similar to the ones in the training set but are still novel structures.

Conclusion
In this study, we first present a Transformer-based random molecular generator and compare it with a number of baseline models using standard metrics. We demonstrate that the Transformer-based molecular generation achieves state-of-the-art performances in generating drug-like structures. To incorporate the protein information, we present a target-specific molecular generator by feeding the target-specific embeddings to a Transformer decoder. We visualize the chemical space, and the new target-specific compounds largely populate the original sub-chemical space.