Root-aligned SMILES: a tight representation for chemical reaction prediction

Chemical reaction prediction, involving forward synthesis and retrosynthesis prediction, is a fundamental problem in organic synthesis. A popular computational paradigm formulates synthesis prediction as a sequence-to-sequence translation problem, where the typical SMILES is adopted for molecule representations. However, the general-purpose SMILES neglects the characteristics of chemical reactions, where the molecular graph topology is largely unaltered from reactants to products, resulting in the suboptimal performance of SMILES if straightforwardly applied. In this article, we propose the root-aligned SMILES (R-SMILES), which specifies a tightly aligned one-to-one mapping between the product and the reactant SMILES for more efficient synthesis prediction. Due to the strict one-to-one mapping and reduced edit distance, the computational model is largely relieved from learning the complex syntax and dedicated to learning the chemical knowledge for reactions. We compare the proposed R-SMILES with various state-of-the-art baselines and show that it significantly outperforms them all, demonstrating the superiority of the proposed method.

both composed of multiple stacked multihead attention layers consisting of a multihead attention module and a position-wise feed forward module.
Before passing into the encoder, SMILES tokens are embedded to continuous vector representations. The multihead attention module consists of multiple scaled-dot product layers that run in parallel. A single scaled-dot product calculation works as follows: where Q, K, V represent query, key, value matrix, respectively; W Q , W K , W V are all trainable parameters; d k means the dimension of K. Depending on where Q, K, V come from, multihead attention can be the self-attention mechanism or cross-attention mechanism.
After the attention calculation of each head, they can be concatenated as follows: The position-wise feed forward module is a simple fully connected layer that utilizes the concept of residual block and works as follows: After the calculation of feed forward module, updated token vectors can be passed to another multihead attention layer. We use the vanilla transformer architecture composed of 6 layers for both encoder and decoder with 8 attention heads for all experiments. During the inference stage, the transformer takes the product SMILES and decoded reactant as the input to predict the probability of the next reactant SMILES token, which can be represented by a conditional probability distribution: where m is the maximum number of reactant tokens. The "bos" (begin of sentence) token is the beginning of reactant tokens. When the last predicted token is "eos" (end of sentence), the decoding process completes.  Figure S1: The workflow of model prediction with the beam search and data augmentation. The example applies 3× augmentation, and sets beam size = 5, topk = 5.
Data augmentation with R-SMILES We successively perform data augmentation and root alignment on the training data, and only perform data augmentation on the test data.
When inferring on the validation and test data, we input multiple SMILES of a molecule respectively and get multiple sets of outputs correspondingly. After removing invalid SMILES that cannot be recognized by Rdkit 2 and converting all outputs to canonical SMILES, we refer to Tekto et al.'s approach 3 that scores these outputs uniformly as follows: where augmentation represents the augmentation times of the test set, and beam represents the beam size. After scoring uniformly, we can select outputs with top-K scores as the final result.
Here we show an example of performing data augmentation for the test data in Fig-ure S1. When the SMILES "C(COC(C=C)=O)(Cl)(Cl)Cl" was input, first we performed SMILES enumeration to get three different SMILES representing the same molecule. Then we started model prediction using the beam search strategy with the beam size of 5, and got the top-5 prediction for each SMILES. To score uniformly, we converted them to the canonical SMILES and removed those invalid SMILES. According to the ranking of the initial results, we can give each output an initial score 1.0 k . For example, the rank-1 prediction "C=CC(=O)OCC(Cl)C(Cl)(Cl)Cl" of "C(COC(C=C)=O)(Cl)(Cl)Cl" was scored one.
If the prediction was an invalid SMILES, we would score it zero. After getting the scores for each output, we can score them uniformly by adding the scores of the same output.
For example, since "C=CC(=O)Cl.OCC(Cl)(Cl)Cl" was the rank-1 prediction for two inputs "O(C(=O)C=C)CC(Cl)(Cl)Cl" and "O=C(C=C)OCC(Cl)(Cl)Cl", its final score was two. Therefore, we can acquire the scores for all the output and got a uniform ranking, i.e, the final result. There is a scoring trade-off in the test set data augmentation, that is, how to weigh the number of prediction occurrences against the ranking of predictions. Suppose there are two predictions, one that is predicted by one input and ranks first, and the other that is predicted by two different inputs and ranks second and third respectively, which of them should get a higher final score? This problem can be expressed by the following equation: where α is the weighing parameter. The higher α, the more important the ranking and the less important the number of appearances, and vice versa. We tested different values of α and show the accuracy of the validation set in the Table S1. It can be seen that for all datasets, the best results are obtained when the α value is equal to or greater than one, which demonstrates the ranking is more important than the number of occurrences.
Moreover, when the α value is one, most of the top-K accuracies are the highest. In fact, Eq. 5 is the case where α takes the value of one in Eq. 6. Training settings We use the same data split as previous researchers 4-6 for all the datasets.
During the pretraining stage. Depending on whether it is a forward or retrosynthesis prediction, products or reacatants in the training set of USPTO-FULL are used for self-supervised training, where molecules in the test set of USPTO-50K and USPTO-MIT are removed. We apply 20× augmentation at training and test sets of USPT0-50K, and 5× augmentation at training and test sets of USPTO-MIT and USPTO-FULL. We set the embedding and hidden size as 512 except that the dimension of Q , K , V is 64. We also use the Adam optimizer and a varied learning rate with 8,000 warm up steps. The input and output share the same vocabulary, but their embedding layers are independent. We conducted all experiments with one NVIDIA GeForce RTX 3090 GPU, and the approximate training time and steps for different tasks are displayed in Fig. S2. The more detailed settings can be found at https://github.com/otori-bird/retrosynthesis.

Analysis of Pretrained Transformer
We used a similar approach to the masked language model of BERT 7 for pretraining. Specifically, 15% of tokens in SMILES are masked. Every masked token has an 80% probability of being replaced with the "unknown" token, a 10% probability of being replaced with any token in the vocabulary, and keeps unchanged for the rest of the cases. After pretraining, we can see in Figure S2 and Table S2 that the training time has been dramatically reduced, which helps a lot in the case of very limited computational resources.

Limitations for ring-opening and ring-forming reactions
The root alignment strategy does not always work well. Taking the ring-opening reaction  Table S3.
It can be seen that the accuracy of R-SMILES, as pointed out by the reviewer, is not so high as that of other reactions. To make it clearer, we also calculated the edit distance between the input and the output SMILES for these reactions. Compared with that of non-ring reaction R-SMILES, the edit distance of ring reactions is significantly larger. These results again verify our main motivation in this work that large distance between input and output strings will degrade the reaction prediction performance. We will be devoted to dealing with this problem in our future work by trying to align more than one atom.

Comparison of the attention of all layers
To give a more detailed discussion about attention mechanism, we fed the reactant or the product of the reaction "C=CC(=O)Cl.OCC(Cl)(Cl)Cl>>C=CC(=O)OCC(Cl)(Cl)Cl" to four different models and visualized attention maps of all layers in Fig. S4. For these four models, one is trained with R-SMILES for forward reaction, one trained with R-SMILES for retrosynthesis, one trained with canonical SMILES for forward reaction, and one trained with SMILES for retrosynthesis. With SMILES (Fig. S4a, c)