Generative AI: Transformers For Molecular Design

Building a transformer model for generating molecules with desired physical properties. A pytorch implementation

The complete project code is accessible on github including model, data, training and testing.

The first time I was really impressed with machine learning was when Alphafold came out. These breakthroughs can only be appreciated when we consider the complexity and significant challenges of determining protein structures. One protein structure determination could be an entire PhD in molecular chemistry/biology. Breakthroughs like alphafold and the large language models like chatgpt demonstrated the immense potential of deep learning. Whilst the value of these deep learning models is unquestionable however, some of the most exciting innovations in the world have come from unlikely sources, not from statistically most likely outcomes so I don’t think AI will replace scientists just yet even if large scale automation in labs were possible.

With that being said, it is very exciting to see the developments of these models applied in research fields like molecular chemistry. One example of significant impact is drug discovery, where these models, in conjunction with first-principle models could be very promising. For example, deep learning models can be used for generating potential drug candidates for specific targets associated with a particular disease, such as proteins, genes, or RNA molecules. Once these potential candidates are generated, they can be further tested and refined using computational methods such as free energy calculations with molecular dynamics that provide key insights into properties like the stability of the molecule, ligand binding, and protein conformational changes etc. These calculations allow researchers to predict how a drug candidate will really behave in a real biological environment, which further refines the candidate selection. It is never this simple, there is of course ALOT more to drug discovery than this but I am not qualified enough to discuss them much further.

In this article, I want to demonstrate the transformer network used as a molecular generator based on desired physical properties. In a way, this is like the original transformer model intended for text translation purposes, except our input will be the set of desired physical properties and the ‘translated’ output will be the molecular structure. A set of physical properties used as inputs are not a sequence so we’ll have a significantly different network architecture for the encoder. Whilst molecules are inherently a graph and not sequential, we can treat the SMILES notation of the molecule as a sequence. In this article, I will (briefly) go through the transformer network and then implement it from scratch in pytorch. I appreciate you can now just import transformer models and take a blackbox approach but implementing those networks from scratch give a deeper understanding of the underlying principles, which is particularly important in problems that require novel challenges, customisation, and optimisation; a black box approach will not be enough and possibly soon automated.

The Transformer Network; Multihead-Attention

At the heart of the transformer network is the multi-head attention mechanism that enable the model to focus on all parts of the sequence at the same time. The recurrent nature of RNN, LSTM, GRU (and other variations of them) is great because its intake of data is inherently sequential without needing further treatment but generating one output at a time means training larger models will computationally very long and performance can also be a problem due vanishing gradient. The multi-head attention is able to overcome these limitations. To get the intuition behind attention mechanisms, consider a time series, X

Where xi is a scalar value (for language and other applications xi may be an embedding vector rather than a scalar). lets say you wanted to predict the next value in the series; which value in the current series do you pay most attention to? the most recent one? An average of all values? This is where the attention layer becomes very helpful, it learns which values and how much to focus on in order to predict future values. The attention mechanism is built upon three main components derived from this input sequence X, these include query (q), key (k) and value (v) which are calculated as follows:

Where Wq, Wv, Wk are all learnable matrices that are determined during training. Subsequently the dot-product attention (q.k) can be determined, which computes the attention scores between each query and all keys, and then uses these scores to weight the corresponding values. The attention scores are passed through a softmax function to obtain the attention weights, which are probability distributions over the keys. These attention weights along with the values matrix are used to determine an attention layer (also referred to as a head)

In case you are not comfortable with these set of equations and linear algebra, we can visualise each step of the process:

This covers the foundation of the self-attention but in practice, multi-head attention is used which is essentially running the attention mechanism described above multiple times, each with different learnable matrices (i.e., different Wq, Wk, Wv). Let’s write the code for this in pytorch, it’s worth going over the code carefully to ensure you fully understand the equations above. I have provided some comments to guide the reader

# MULTIHEAD ATTENTION 
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        try:
            assert d_model % num_heads == 0
        except Exception as e:
            logger.error("dimension of the embedding model is not divisable by number of heads")
        self.d_models = d_model
        self.num_heads = num_heads
        self.depth = d_model // num_heads

        # The query, key, value learnable matrices
        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)

        self.FCLayer = nn.Linear(d_model, d_model)
    def split_embedding_perHead(self,x):
        # x shape is (batch_size, seq_len, d_model)
        (batch_size, seq_len, d_model) = x.shape
        # let's reshape to (batch_size, seq_len, num_heads, depth)
        x = x.view(batch_size, -1, self.num_heads, self.depth)
        # changing the dimensions order to:(batch_size, num_heads, seq_len, depth)
        x = x.permute(0,2,1,3)
        return x

    def cal_attention(self,q,k,v,mask):
        qk = torch.matmul(q, k.permute(0,1,3,2))
        dk=torch.tensor(k.shape[-1], dtype=torch.float32)
        #dk is a tensor scalar!
        attention = qk/torch.sqrt(dk)

        if mask is not None:
            attention += (mask*-1e9)
        attention_weights = F.softmax(attention, dim=-1) # should be applied along the sequence which is the 3rd dimension
        output = torch.matmul(attention_weights, v)
        return output, attention_weights

    def forward(self, v,k,q,mask):
        batch_size = q.shape[0]

        q = self.split_embedding_perHead(self.Wq(q))
        k = self.split_embedding_perHead(self.Wk(k))
        v = self.split_embedding_perHead(self.Wv(v))

        attention,atten_weights = self.cal_attention(q,k,v,mask)
        attention = attention.permute(0,2,1,3).contiguous()
        attention = attention.reshape(batch_size, -1, self.d_models)

        output = self.FCLayer(attention)
        return output

Masks: Managing Varying Sequence Lengths

The code above refers to a mask which allows the attention network to ignore (or pay no attention) to some values in the sequence. Whilst generally, in time series data, the number of input (historical data window) and output is always fixed by the network, in other applications such as language, the number of inputs and outputs is not fixed. In this mini experiment of generating molecular structures based on desired properties, the output sequence will vary depending on the size and complexity of the molecule; the padding mask helps to deal with this varying sequence length by taking the sequence and adding padding to the sequence so all sequences are of the same length as required. For example consider the SMILES notation for methane (“C”) and ethanol (“CCO”), if we define a sequence length of 6 then the padding mask for these two molecules could be:

Methane: [1, 0, 0, 0, 0, 0] Ethanol: [1, 1, 1, 0, 0, 0]

Where 0 suggests padding and 1 suggest molecule defined character. The mask is multiplied with a large negative number (check the code) which results in a large negative attention score. The corresponding attention weight is essentially 0, hence, the padding mask helps the network pay attention only to the molecule and not the padding.

I mentioned earlier the computational challenge of training RNN (and any variants) and that transformers do not have the same restrictions because during training we are able to predict the entire output sequence simultaneously with the help of the look-ahead mask. This mask is used to prevent the model from accessing future tokens in a sequence, ensuring that predictions of each token in the sequence are made based only on past and present tokens. The mask is a lower triangle matrix:

[[1, 0, 0, 0, 0],
 [1, 1, 0, 0, 0],
 [1, 1, 1, 0, 0],
 [1, 1, 1, 1, 0],
 [1, 1, 1, 1, 1]]

Each element/token in the output sequence is predicted with the future tokens masked (hence the lower triangular matrix). This is only relevant during training where we have access to the target output sequence and not applicable during the inference stage where each token is generated one at a time.

The Whole Architecture

In this mini-experiment I want to use molecular properties to then predict potential molecules. The properties include polar surface area (polararea), molecular complexity (complexity), heavy atom count (heavycnt), hydrogen bond donors (hbonddonor), and hydrogen bond acceptors (hbondacc). I actually chose these properties based on availability of data but the model should work with any set of properties providing the properties chosen are good indicators of the corresponding molecular structure.

Left) A Schematic of the transformer network including an encoder layer for inputs and decoder layer for predicted outputs. Image obtained from original paper; All you need is Attention. Right) The modified transformer network where the encoder is a standard feed forward ANN layer.

Since these molecular properties input is not a sequence, the traditional encoder of the transformer network is no good and can be replaced with a simple ANN feed forward network. The decoder works with the SMILEs notation of the molecule which can be treated as a sequence so the decoder is largely unchanged.

# THE ENCODER
class EncoderLayer(nn.Module):
    def __init__(self,d_model,dff):
        super(EncoderLayer,self).__init__()
        self.FeedForwardNN = nn.Sequential(
            nn.Linear(d_model,dff),
            nn.ReLU(),
            nn.Linear(dff,dff)
        )

    def forward(self,x):
        output = self.FeedForwardNN(x)
        return output

# THE DECODER LAYER
class DecoderLayer(nn.Module):
    def __init__(self,d_model, num_heads, dff):
        super(DecoderLayer,self).__init__()
        self.MultiHAttention1 = MultiHeadAttention(d_model, num_heads)
        self.MultiHAttention2 = MultiHeadAttention(d_model, num_heads)
        self.FeedForwardNN = nn.Sequential(
            nn.Linear(d_model,dff),
            nn.ReLU(),
            nn.Linear(dff,d_model)

        )
        self.layerNorm1 = nn.LayerNorm(d_model, eps=1e-6)
        self.layerNorm2 = nn.LayerNorm(d_model, eps=1e-6)
        self.layerNorm3 = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, x, enc_output, look_ahead_mask, padding_mask):
        attn_output1 = self.MultiHAttention1(x,x,x,look_ahead_mask)
        attn_output1 = self.layerNorm1(x+attn_output1)
        attn_output2 = self.MultiHAttention2(enc_output, enc_output,attn_output1, padding_mask)
        attn_output2 = self.layerNorm2(attn_output2+attn_output1)

        Feedforward_output = self.FeedForwardNN(attn_output2)
        final_output = self.layerNorm3(attn_output2+Feedforward_output)
        return final_output

The decoder layer follows the same architecture as in the illustration with two multihead attention layers, followed by a feed forward and layer normalisation in-between. This represents a single pass through the decoder but actually the sequence is passed through this network iteratively multiple times for refinement. The complete decoder is combined below; in addition to the iterative refinement through the decoder layer, there are two other key elements worth discussing.

Firstly the target sequence is used as input to the decoder is converted to an embedding, here we are using random embedding vectors which is not ideal for accuracy and better care is necessary but for an experiment and demonstrations purposes, this is fine.

class Decoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, maximum_position_encoding):
        super(Decoder, self).__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        self.embedding = nn.Embedding(target_vocab_size, d_model) # d_model is the size of embedding vector
        self.pos_encoding = self.positional_encoding(maximum_position_encoding, d_model)
        self.dec_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, dff) for _ in range(num_layers)])

    def positional_encoding(self, position, d_model):
        angle_rads = self.get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model)
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
        pos_encoding = angle_rads[np.newaxis, ...]
        return torch.tensor(pos_encoding, dtype=torch.float32)

    def get_angles(self, pos, i, d_model):
        angle_rates = 1 / np.power(1000, (2 * (i // 2)) / np.float32(d_model))
        return pos * angle_rates

    def forward(self, x, enc_output, look_ahead_mask, padding_mask):
        seq_len = x.size(1)
        x = self.embedding(x)
        x *= torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
        x += self.pos_encoding[:, :seq_len, :]

        for i in range(self.num_layers):
            x = self.dec_layers[i](x, enc_output, look_ahead_mask, padding_mask)
        return x

Secondly, a positional encoding is used to include the sequential information of the sequence. As previously mention, the multihead attention network consider relationships between all tokens in the sequence simultaneously, regardless of their position. This is importantt because it enables the model to capture long-range dependencies but it treats all tokens equally and doesn’t inherently know the relative or absolute positions of tokens in the sequence e.g. consider a sentence, any permutable shuffling of the sentence would make no difference because the self-attention network pays no attention to positions in the sequence. But the order of tokens in a sequence is important and if you change the order, the meaning of a sentence or in this application, the molecular structure would drastically change. The positional encoding is calculated as follows:

where pos is the position of the token in the sequence and i is the index within the embedding vector. The positional encoding should have the same dimensions as the embedding matrix.

Now that we have completed all the key parts whole architecture, we can all be combined to define a transformer model:

# TRANSFORMER
class Transformer(nn.Module):
    def __init__(self,num_layers, enc_d_model, dec_d_model,
                enc_num_heads, dec_num_heads, enc_dff, 
                dec_dff, target_vocab_size, pe_target):
        super(Transformer, self).__init__()
        self.encoder = EncoderLayer(enc_d_model, enc_dff)
        self.decoder = Decoder(num_layers, dec_d_model, dec_num_heads, dec_dff, target_vocab_size, pe_target)
        self.final_layer = nn.Linear(dec_d_model, target_vocab_size)

    def forward(self, properties, target, look_ahead_mask, dec_padding_mask, training):
        enc_output = self.encoder(properties)
        enc_output_reshaped = enc_output.unsqueeze(1).repeat(1, target.shape[1],1)
        dec_output = self.decoder(target, enc_output_reshaped, look_ahead_mask, dec_padding_mask)
        ffl_output = self.final_layer(dec_output)
        return ffl_output

The next step is to define the data processing, train the model and test it. The training and processing of the data is actually fairly standard and similar to most model training in pytorch so we shant focus too much on them. For these steps, including the training data, you can find the complete code in github repository.

Evaluating The Model

Let’s discuss the model performance. I trained the model on around 5000 molecules and for about 8 epochs. The model itself is barely optimised but suprisingly, it does a decent job! Let’s look at some examples of generated molecules during testing:

Generated SMILES: <start>[O-].[O-].[O-].[O-].[Al+2]<end>
actual smiles: <start>[N-]=[N+]=O<end>
Generated SMILES: <start>C1=C(C(=C(C(=C1Cl)Cl)Cl)C(=O)Cl)Cl)C(=O)Cl<end>
actual smiles: <start>C1[C@@H]2[C@H]3[C@@H]([C@H]1[C@H]4[C@@H]2O4)[C@]5(C(=C([C@@]3(C5(Cl)Cl)Cl)Cl)Cl)Cl<end>
Generated SMILES: <start>CC1=CC(=C(C=C1)C(=O)OC2=CC=CC=C2)C<end>
actual smiles: <start>CC(C)CC(=O)O[C@@H]1CC2CC[C@]1(C2(C)C)C<end>
Generated SMILES: <start>[O-]S(=O)(=O)[O-].[Na+]<end>
actual smiles: <start>C(=O)([O-])[O-].[Mg+2]<end>
Generated SMILES: <start>CCCCCCCCCCCCCCO<end>
actual smiles: <start>CN1CCC[C@H]1C2=CN=CC=C2.Cl<end>
Generated SMILES: <start>CCCCCCCCCCCCO<end>
actual smiles: <start>CCCCCCCC1OCC(O1)C<end>
Generated SMILES: <start>CCCCCCCCCCCCCCCC(=O)OC(=O)CCCC(=O)C<end>
actual smiles: <start>CC(C)(C)C1=CC=CC=C1OP(=O)([O-])OC2=CC=CC=C2<end>
Generated SMILES: <start>CC(=O)OC(=O)C1=CC=CC=C1<end>
actual smiles: <start>C1CN1P(=O)(N2CC2)N3CC3<end>
Generated SMILES: <start>CC(=O)OC(=O)OCC<end>
actual smiles: <start>C=O.C=O.C=O.C=O.C=O.[Fe]<end>
Generated SMILES: <start>CCCCCCCCCC<end>
actual smiles: <start>C=CC1=CC=CC=C1Cl<end>
Generated SMILES: <start>CCCCCCCCCCCCCCCCO<end>
actual smiles: <start>CCCCCCCOC(=O)CCCCCC<end>
Generated SMILES: <start>C(C(C(C(C(C(CO)O)O)O)O)O)C(C(C)O)O<end>
actual smiles: <start>C(CNCCNCCNCCNCCN)N<end>
Generated SMILES: <start>CC(C)C(=O)O<end>
actual smiles: <start>CC(C)(C#C)O<end>
Generated SMILES: <start>[NH4+]<end>
actual smiles: <start>[W]<end>
Generated SMILES: <start>[NH4+]<end>
actual smiles: <start>[He]<end>
Generated SMILES: <start>C(C(=O)O)OC(=O)O.C(C(=O)O)O.C(=O)O.[Na+]<end>
actual smiles: <start>C(C(=O)[O-])C(CC(=O)[O-])(C(=O)[O-])O.N.O.[Fe+3]<end>
Generated SMILES: <start>CCCCCCCCCCCCO<end>
actual smiles: <start>CCCCCCCCCC(=O)OC<end>
Generated SMILES: <start>CCCCCCCCCC(=O)OCC<end>
actual smiles: <start>CC/C=C\C/C=C/CCOC(=O)C<end>
Generated SMILES: <start>CCCCC=O<end>
actual smiles: <start>C1=CC=NC=C1<end>
Generated SMILES: <start>CCCCCCCC(=O)O<end>
actual smiles: <start>CC1=CC(=CC=C1)C(=O)O<end>
Generated SMILES: <start>CCCCCCCCCCC<end>
actual smiles: <start>CC(C)(C)C1=CC=CC=C1<end>
Generated SMILES: <start>CC1=CC(=C(C=C1)[N+](=O)[O-])[N+](=O)[O-](Cl)Cl<end>
actual smiles: <start>C[C@]12CC[C@@H](C1(C)C)C[C@@H]2OC(=O)CSC#N<end>
Generated SMILES: <start>CCCCCCCCCC(=O)OCC(=O)O<end>
actual smiles: <start>C=CCOC(=O)C1=CC=CC=C1N<end>
Generated SMILES: <start>CC1=CC(=C(C=C1)C2=CC=CC=C2C(=C2)C(=O)OC(=O)CC(C)C)C(C)C(C)C(C)CC(C)C<end>
actual smiles: <start>C1C(CC2=CC=CC=C2C1C3=C(C4=CC=CC=C4OC3=O)O)C5=CC=C(C=C5)C6=CC=C(C=C6)Br<end>
Generated SMILES: <start>C(C(F)(F)(F)F)(F)F<end>
actual smiles: <start>C(C(F)(Br)Br)(F)(F)F<end>
Generated SMILES: <start>CC(=O)OC(=O)CC(=O)C<end>
actual smiles: <start>COC(=O)/C=C/C(=O)OC<end>
Generated SMILES: <start>CCCCCCCCCC=O<end>
actual smiles: <start>CC1=CC2=CC=CC=C2O1<end>
Generated SMILES: <start>CCCCCCCCCCCCC(=O)OC(=O)CCC<end>
actual smiles: <start>C1=CC(=CC=C1[N+](=O)[O-])OC2=C(C=C(C=C2)Cl)Cl<end>
Generated SMILES: <start>CCCCCCCCCC=O<end>
actual smiles: <start>CCCCC=C(CC)C=O<end>
Generated SMILES: <start>CC(=O)OC(=O)C<end>
actual smiles: <start>COP(=O)(C)OC<end>
Generated SMILES: <start>CC(=O)[O-].CC(=O)[O-].[Na+]<end>
actual smiles: <start>CC(C)SSSC(C)C<end>
Generated SMILES: <start>CCCCC=O<end>
actual smiles: <start>CCSCCCl<end>

Identifying a sensible evaluation metric for these models is challenging because while structural similarity could be used as a metric, evaluating how well the generated molecules meet the specified properties is key. But this is not so easy to do.

Firstly I was impressed that the model is even able to generate chemically valid SMILES strings. Strings like “CCCCCCCCCCCCCCO” which is a long carbon chain with a hydroxyl group and “CC(=O)OC(=O)C” whcih is a simple ester represent valid molecular structures. it seems the transformer model has learned some underlying principles of molecular structure, such as the connectivity of atoms and basic functional group formation.

However, there are challenges with stereochemistry and Complex Structures, e.g. the generated SMILES “C1=CC=C(C=C1)C(=O)O”, which corresponds to benzoic acid, is a much simpler structure compared to the actual target molecule; “CC(=O)NC1=CC=C(C=C1)Cl”. In general the model seems to miss specific stereochemical configurations suggesting either that we need more training data or to adjust the model so that it can capture these specific molecular structures.

Another issue seems to be that the model tends to generate simpler molecules e.g. when the target molecule has multiple chiral centers, the model might just generate a simpler, non-chiral molecule. This over simplification of molecules is most likely due to the model’s inherent bias towards generating more common, simpler structures that it has seen more frequently during training.

One potential issue that is difficult to confirm without further data is how well the desired input molecular properties map to the molecular structure. Although the generated molecules may not match the target molecules but it might have the desired input properties which would indicate the model is learning to generate molecules that adhere to the specified properties constraints. In this case, it is not the model that is limiting the performance but rather a need for data with more features that map to unique molecular structures.

That brings us to the end of this article; thank you for taking the time to read it, I hope you found it insightful! If you are interested, definitely explore these models further and apply it to different problems. It should be very exciting to see these developments in science.

Unless otherwise noted, all images are by the author