Probabilistic Transformer: Modelling Ambiguities and Distributions for RNA Folding and Molecule Design
Anonymous Author(s)
AfﬁliationAddress
Abstract
Our world is ambiguous and this is reﬂected in the data we use to train ouralgorithms.
This is especially true when we try to model natural processes wherecollected data is affected by noisy measurements and differences in measurementtechniques.
Sometimes, the process itself can be ambiguous, such as in the case of
RNA folding, where a single nucleotide sequence can fold into multiple structures.
This ambiguity suggests that a predictive model should have similar probabilisticcharacteristics to match the data it models.
Therefore, we propose a hierarchicallatent distribution to enhance one of the most successful deep learning models, the
Transformer, to accommodate ambiguities and data distributions.
We show thebeneﬁts of our approach on a synthetic task, with state-of-the-art results in RNAfolding, and demonstrate its generative capabilities on property-based moleculedesign, outperforming existing work.
Introduction
Transformer models [1] are the architecture of choice for many applications.
Next to a wide rangeof NLP applications, such as language modelling [2, 3, 4] and machine translation [1, 5], they arealso very effective in other disciplines, such as computer vision [6, 7], biology [8, 9], and chemistry[10, 11].
An additional challenging application for which transformers are promising is RNA folding,where the goal is to model a secondary structure (represented in dot-bracket notation [12]) based ona given sequence of nucleotides.
RNA data is highly ambiguous since it is collected with differenttechniques, resolutions, protocols, and even contains natural ambiguities since there exist multiplestructures for the same RNA sequence in the cell [13, 14] and the same structure can be caused bymultiple sequences.
Similar to RNA structures, molecules can be represented as sequences by usingthe simple molecular line entry system (SMILES) [15] and transformers arise as the architecture ofchoice in the ﬁeld [11, 16].
Although Transformers have outstanding performance, the deterministic core of the architecturecould harm performance in real-world applications like RNA folding.
If collected data containsnoisy labels or ambiguous samples, a vanilla Transformer model can only express uncertainties in thesoftmax output but not in the latent space.
When sampling, sequential interdependencies can only bemodelled in a decoder setting but not in an encoder-only setting.
We address these limitations byproposing a probabilistic Transformer (ProbTransformer) that models a hierarchical latent distributionand performs sampling in the latent space.
The ProbTransformer can represent ambiguities in thesedistributions and reﬁne the sampled latent vectors within the computational graph.
This is in linewith recent ﬁndings in cognitive science that suggest that the human brain both represents probabilitydistributions and performs probabilistic inference [17, 18].
Our approach is based on the idea ofcombining the transformer architecture with a conditional variational auto-encoder (cVAE)
[19], but in a hierarchical fashion similar to the hierarchical probabilistic U-Net [20].
Therefore, we introduce anew probabilistic layer and incorporate it after the attention and feed-forward layer (Section 3).
In thisway, we preserve the global receptive ﬁeld through the attention mechanism and remain independentof other enhancements in the transformer ecosystem [21].
To train the latent distributions, we makeuse of the generalized ELBO with constrained optimization [22] and introduce an annealing techniqueto adapt a hyperparameter online.
We see our contributions in three aspects:
• The introduction of the ProbTransformer, a novel hierarchical probabilistic architecture enhancement to the Transformer ecosystem.
• Our training procedure using GECO, the analysis of the sensitivity of its hyperparameter κ , andthe introduction of the online adaption technique kappa annealing which could be beneﬁcial forvariational training with ELBO in general.
• A comprehensive empirical analysis that veriﬁes the ProbTransformer’s capability to learn andrecover data distributions on a novel synthetic sequential distribution task, assesses its capability
of handling data ambiguities in practice by achieving state-of-the-art performance in RNA folding, and demonstrates its generative character by outperforming existing work in Molecule Design.
We ﬁrst clarify notation and recap the cVAE [19] (Section 2), and then introduce the ProbTransformer(Section 3).
We then present our experiments on the sequential distribution task (Section 4.1), RNAfolding (Section 4.2), Molecule Design (Section 4.3), as well as our ablation study (Section 4.4),discuss related work (Section 5) and conclude (Section 6).
Background
Transformer Notation The Transformer [1] is a self-attention based sequence-to-sequence modelintroduced as an encoder-decoder architecture.
However, both encoder-only and decoder-onlyversions are very successful by themselves [2, 23], and in our work, we focus on either of these twoversions.
An encoder or decoder has N blocks and each block consists of a multi-head (masked)attention followed by a feed-forward layer.
Around both of these layers there are residual connections[24] followed by layer normalization [25].
We use the parameterization with D model dimensions inthe residuals and attention, H heads per attention, D ff latent dimensions in the feed-forward layer,and n ∈ { 1 , . . .
, N } for the block in the Transformer.
In addition, we deﬁne X (cid:98) =( x 1 , . . .
, x S ) asthe input sequence of length S , and Y (cid:98) =( y 1 , . . .
, y S )
as the target sequence, ˆ Y (cid:98) =(
ˆ y 1 , . . .
, ˆ y S ) as thepredicted sequence.
Conditional Variational Auto-Encoder (cVAE)
The cVAE [19] is a deep conditional generativemodel and an extension of the Variational Auto-Encoder [26, 27].
During inference, the cVAE aimsto generate a distribution for an output y conditional on an input x .
More speciﬁcally, given an inputx , a latent variable z is drawn from a (conditional) prior model p θ ( z | x ) and used as an additionalinput to the generation model p ρ ( y | x , z ) in order to generate the prediction ˆ y .
The cVAE is trainedby minimizing the negative evidence lower bound (ELBO), L ELBO :
[Equation]
where q ψ ( z | x , y ) describes the posterior model (called “recognition network” in [19]).
All models areneural networks, and the prior and posterior models each output a mean and variance of a Gaussiandistribution which represents the distribution of the latent z .
During training, z post ∼ q ψ ( z | x , y )is sampled from this posterior model and used as input to the generation model p ρ ( y | x , ˆ z ) , whoseoutput ˆ y is compared to the true y with a cross-entropy loss.
This objective at training time can beviewed as a reconstruction task, which is an easier task than prediction.
The KL divergence term aimsto align the (conditional) prior model p θ ( z | x ) and the posterior model q ψ ( z | x , y ) .
The two lossesare added and then used to train the prior, posterior, and generation models jointly, employing thereparameterization trick [28].
Probabilistic Transformer
Building on concepts of the cVAE, we enhance the Transformer model to a probabilistic Transformer(ProbTransformer) by adding M ∈ { 1 , . . .
, M } probabilistic feed-forward layers (prob layer) to the existing blocks.
We introduce our approach for an encoder-only model to simplify notation.
It alsoapplies to the decoder-only model but not directly to an encoder-decoder model because our trainingsetup requires the alignment of the source and target sequences in the posterior input.
Whether all oronly a selection of blocks are enhanced with prob layers is a design decision and examined empiricallyin Section 4.4.
In a block, we place the prob layer after the attention and feed-forward layer.
Theprob layer parameterizes a univariate Gaussian N ( µ, σ 2 ) used to sample a latent vector z ∈ R D z foreach position in the sequence.
We denote by X and Y the input and target sequences of length S , by Z m (cid:98) =(
z m, 1 , . . .
, z m,S )
the position-wise sampled latent sequence of block m , by Z (cid:98) =(
Z 1 , . .
, Z M )all latent sequences of all blocks, and by Z <m (cid:98) =(
Z 1 , . .
Z m − 1 ) all latent sequences of all blocksbefore m .
At inference time, we sample Z m and add it to the computation graph of the current block,similar to sampling from the (conditional) prior in the cVAE.
However, in contrast to the cVAE, oursampling is conditioned hierarchically: the z vector also depends on previously sampled z vectors ofall positions in all lower Transformer blocks:
[Equation]
A sequential relation is achieved by the attention mechanism: while the samples z m are drawnindependently for each position, the following attention operations relate them and turn the positionwise samples into a joint distribution over sequences reﬁned in higher blocks.
As a result of this hierarchical composition of the ProbTransformer and this sequential relation, priormodel p θ ( z m | Z <m , X ) and the generation model p ρ ( Y | X, Z ) effectively become a single modelp
Φ ( Y | X ) , which strongly differs from the modeling strategy of the cVAE.
This model approximatelymarginalizes over latent variables
Z : P ( Y | X ) =
(cid:82) Z p ( Z | X ) × p
( Y | X, Z )
d Z , and we can easilysample from it by hierarchically sampling p ( Z | X ) decomposed as
[Equation]
followed by sampling from p ( Y | X, Z ) .
We dub the combined model Φ the predictive model.
Atinference time, we can sample different predictions from this predictive model ˆ Y =
Φ( Y | X ) .
However, we can also use the mean of the respective (Gaussian) distributions instead of samplingfrom them.
We denote this as mean inference in contrast to sample inference .
In the following, weintroduce the prob layer architecture in detail before explaining our training setup and the learningobjective.
3.1 Position-wise Probabilistic Feed-Forward Network
[Figure]
The prob layer is applied to each position of the latent space separately and identically.
Figure 1 provides a visual description.
Itconsists of a linear layer that transforms from the model dimension
D model to the distribution dimension D z of the probabilistic latentspace, followed by an activation function.
Further, two linear layers with D z × D
z weight matrices generate the latent representation of a conditional Gaussian distribution with mean µ m ∈ R D zand log variance 1 log σ 2 m ∈ R D z for the prob layer in block m :
[Equation]
As mentioned before, during sample inference , we can sample from the distribution (Equation 6)and during mean inference we use z n (cid:98) = µ m .
An additional linear layer is used to compute the layer’soutput y m ∈ R D model (Equation 7).
Similar to the attention and feed-forward layer, we employ aresidual connection [24] followed by a layer normalization [25].
3.2 Training Setup and Learning Objective
We optimize the ELBO as the standard practice for conditional training [19].
This requires avariational posterior Q ψ ( Z | X, Y ) that depends on both the input sequence X and target sequence
Y .
We model this posterior with a separate ProbTransformer Ψ with the same architecture as thepredictive model except of an additional input embedding for the target sequence Y and no outputgeneration layers, since we only make use of the latent sampling Z post and not of its actual modeloutput.
During training, we ﬁrst run the posterior model Ψ and sample in each prob layer the latent
Z post ∼ Q ψ .
In a second step, we run the predictive model Φ but use the latent realization Z postminstead of sampling an own Z m in each prob layer.
Figure 2 illustrates this training setup.
[Figure]
The negative ELBO loss L ELBO (Equation 1) is composed of a reconstruction loss L rec and a
Kullback-Leibler divergence D KL between the prior distribution P conditioned on the latent Z postdrawn from Q and the posterior distribution Q .
The reconstruction loss is the cross-entropy betweenthe predicted sequence ˆ Y and the true sequence Y while using the latent Z post from the posteriormodel:
[Equation]
(8) and D KL is the sum of KL divergences between the hierarchical decomposition, p m and q m , of thedistributions P and Q :
[Equation]
With the default L ELBO loss, we discovered instability in the training and performance issues, evenwith a weighting factor β [29].
Therefore, we follow [20] to avoid convergence issues with the classic
ELBO objective and make use of the Generalized ELBO with Constrained Optimization (GECO)objective [22], optimizing the Lagrangian
[Equation]
The Lagrange multiplier λ balances the terms and is initialized to 1 , added to the learnable parameterspace and updated as a function of the exponential moving average of the reconstruction loss [20].
Inthe beginning of the training, there is high pressure on the reconstruction loss due to an increasing λ .
Once the desired reconstruction loss L rec ≈ κ is reached, the pressure moves to the KL-term.
Pleaseﬁnd more information about the training dynamics in Appendix A.κ was set to a constant in the original work of [22], but this is problematic since it is a sensitivehyperparameter due to the training dynamics.
Speciﬁcally, if κ is chosen too small , possible failuresmodes include that L rec never reaches it (with the pressure staying only on the reconstruction loss),that κ will be reached by over-ﬁtting, or that the posterior distribution collapses which leads to high
D KL ( Q ψ (cid:107) P θ ) and destabilizes the training.
On the other hand, if κ is chosen too large , the modelcan underﬁt or the reconstruction loss drops below κ , leading to a negative loss value and harmingthe training success signiﬁcantly.
To address this issue we introduce kappa annealing and adjustκ during training.
Let L c = 1 K (cid:80)
Kk =
L rec ( X k , Y k ) − κ be the mean constraint difference for K training samples in one epoch of training or a deﬁned number of update steps.
In our experiments,we ﬁnd that initializing κ slightly too large and updating it every epoch withκ [Equation] (cid:26) κ
[Equation]
if L [Equation] 0 and λ ≤κ otherwise (11) leads to more stable training behaviour, improved ﬁnal performance, and reduces the need forexpensive tuning of κ .
Experiments
We demonstrate the beneﬁts of the ProbTransformer in three experiments: (1) using a novel syntheticsequential distribution task, we show the advantages of distributions in the latent space over thestandard softmax output distribution of a vanilla Transformer.
(2) We show the beneﬁts of the
ProbTransformer in dealing with ambiguous data in RNA folding.
(3) We use the ProbTransformer toimprove property-based Molecule Design by sampling from the latent space instead of a softmaxoutput.
In an additional ablation study, we provide insights on the impact of kappa annealing andhierarchical prob layers.
Our source code will be publicly available upon acceptance.
4.1 Synthetic Sequential Distribution Task
In real-world generation tasks, the true distribution is not available and is only implicitly accessible bya dataset.
To provide insights on the quality of distribution learning, we created a synthetic sequentialdistribution task and compare the ProbTransformer to two sampling methods of a vanilla Transformer.
Data We design the task to map a sequence of tokens from a source vocabulary x ∈ V ∗ i to a sequenceof target tokens from a target vocabulary y ∈ V ∗ o with the same length.
The tokens in the sourcesequence are used to build ‘phrases” P .
Each phase exists of l tokens sampled with replacement(similar to the combination of words to phrases in a sentence).
We generate randomly for each sourcetoken in each phrase a unique distribution p ( y | x, P ) over the target tokens depending on the currentphrase.
Further, we design the distribution sparsely so that no more than k tokens from the targetvocabulary have a non-zero probability.
Training data are generated by sampling input sequences fromall phases (with replacement) and sampling the target sequence from its corresponding distribution.
Setup We used an encoder-only model.
For the ProbTransformer we enhanced each block with aprob layer.
We conﬁgured the task and the model to run on one GPU within a few hours.
Pleaseﬁnd more information about the task and the conﬁguration in Appendix B. For the inference of thevanilla Transformer we used MC dropout (using dropout during inference time, based on [30]) andsampling from the softmax distribution.
For the inference of the ProbTransformer we sample fromthe predictive model and use the token with the highest output value.
We generate 50 realizations persample to create a predictive distribution of the target vocabulary.
[Table]
Results We measure the generative performance of a model or sampling method with twometrics: The Validity describes the percentageof predicted tokens whose probability in the truedistribution is not zero.
Second, we measure the
KL-divergence between the sampled distributionand the true distribution since sampling shouldreproduce the true target token distribution.
Asshown in Table 1, the ProbTransformer outperforms the MC dropout and softmax sampling,demonstrating its strong performance in probabilistic sequence modelling.
Please note that this taskdoes not consider interdependence sampling, where the sampling of one token in a sequence dependson the realization of others, while the next task on RNA folding does.
4.2 RNA Folding
An RNA’s structure inﬂuences its function drastically [31] and the functional importance of RNAarguably is on par with that of proteins [32]. Since cellular RNAs typically have extensive secondary
[Table]
structures but limited tertiary structures [33], RNA folding is typically modeled as a function F : Υ ∗ → Γ ∗ , where Υ and Γ denote sequence and secondary structure alphabets, respectively. However, this is a simpliﬁed view on the RNA folding process, which ignores the fact that RNAsrespond to cellular conditions by altering their structures dynamically, resulting in an ensembleof structures occurring with different probabilities [13, 14] Further, real-world applications oftenrequire structure analysis of very similar sequences, e.g. resulting from single or multiple pointmutations [34].
This can hardly be captured with common approaches during training becauseambiguous data is typically removed to avoid overﬁtting [35, 36, 37].
We address these issues withprobabilistic modelling using our ProbTransformer and by keeping ambiguous training data, whileimplicitly measuring overﬁtting by explicitly removing similarity to the test data. Since we canexpress RNA structures in dot-bracket notation [12], the source and target have the same length,making RNA folding a suitable problem for an encoder-only Transformer model.
Data We collect a large data pool of publicly available data from recent publications [38, 39, 37, 40,41] and split the predeﬁned validation and test sets, VL0 and TS0 [37], from the pool.
We deriveda testset, TSsameStruc , from 149 samples of TS0 that share the same structure with at least oneother sequence and uniformly sampled a disjoint set of 20 sequences with more than one annotatedsecondary structure from the remaining pool to produce an ambiguous testset, TSsameSeq .
Sampleswithout pairs and with a sequence similarity greater than 80% to the test and validation sets wereremoved from the pool.
However, in contrast to previous work [37, 39, 42], we kept all remainingsamples for training to capture the inﬂuence of point mutations on the structure.
The ﬁnal dataconsists of 52007 training and 1299 validation samples, and 1304 samples in the testset TS0 and 46samples in the testset TSsameSeq.
We observe 48092, 1304, and 20 unique sequences with up to7 different structures for a single sequence, and 27179, 1204, and 46 unique structures with up to582 different sequences for a single structure in the training, TS0, and TSsameSeq sets, respectively.
This indicates that many sequences map to the same structure and different structures for the samesequence.
We refer to Appendix C.1 for more information about the data.
Setup We use a 6-block encoder-only architecture with D model of 512 for the vanilla Transformerand ProbTransformer.
We train them for 200 epochs or 1M training steps.
We compare them againstthe state-of-the-art deep learning algorithms, SPOT-RNA [-[36], MXFold
 [38],-][37], MXFold2 [39], and a commonlyused dynamic programming approach, RNAfold [43] on the testset TS0 and its subset TSsameStrucusing mean inference.
To analyse the capabilities of the ProbTransformer to reconstruct structuredistributions, we also infer the model 5, 10, 20, 50, and 100 times using sample inference on
TSsameSeq, and compare against three commonly-used algorithms specialized on the prediction ofstructure ensembles via sampling from the Boltzmann distribution based on experimentally derivedthermodynamic parameters: RNAsubopt [43], RNAshapes [44], and RNAstructure [45].
In contrastto previous work [37, 41, 39, 42], we do not use ensembling nor limitations on the accepted basepairs or secondary structures for post-processing, but independently train a CNN head on the outputof the ProbTransformer that maps the string representation to an adjacency matrix representing thestructure for evaluations on TS0 and TSsameStruc.
Although we are aware of the problems with exactevaluations via F1 Score due to the dynamic nature of RNA structures [46], we report F1 Score forreasons of comparability to previous work; however, we further report the number of solved structures(the ultimate goal of the task), and the average Hamming distance to achieve a better measure ofdistance.
For more details on the setup, CNN head, and metrics, we refer to Appendix C.2.
Results Table 2 summarizes the results of all approaches on TS0 and TSsameStruc.
Both attentionbased approaches generally achieve strong results, but we observe that the ProbTransformer outper- [Figure] forms the vanilla Transformer across all measures on both sets.
The ProbTransformer achieves thebest performance in terms of F1 Score, solves more than eight times the structures compared to thenext best approach (MXFold2) on both sets and achieves four times lower Hamming distance on
TSsameStruc compared to SPOT-RNA, indicating that it has learned to handle ambiguous sequences.
Our model is further capable of accurately reconstructing challenging structures, exemplarily shownfor two structures that contain long-range interactions in Figure 3. A functional important classof base pairs are so-called pseudoknots [47, 48], non-nested base pairs present in around 40% of
RNAs [41] that are overrepresented in functional important regions [49, 48] and known to assistfolding into 3D structures [50]. While RNAFold and MXFold2 cannot predict this kind of basepairs due to the underlying nearest neighbour model, Figure 3 as well as further results shown in Appendix C.3.2 indicate that the ProbTransformer can predict pseudoknots more accurately than SPOT-RNA if these are contained in the structures (in particular, see Figure 9) and further rarelypredicts pseudoknots if the structure does not contain non-nested base pairs.
However, a detailedanalysis on the quality of pseudoknot predictions would be out of the scope of this work.
Wheninferring the model multiple times on TSsameSeq, the ProbTransformer is the only model in ourevaluation that could reconstruct more than one structure for a given sequence (for two sequences thiswas already achieved when inferring the model only ﬁve times). The average Hamming distances ofthe best predictions on TSsameSeq for every approach are summarized in Table 15 in Appendix C.3.2,showing the performance of the ProbTransformer can improve the Hamming distance by up to 44%,indicating that the ProbTransformer reconstructs the overall structure ensemble very well.
Additionalresults and example predictions for RNA folding are reported in Appendix C.3.
4.3 Molecule Design
In the ﬁeld of generative chemistry [51], deep generative models are employed to explore the chemicalspace [52, 53, 54].
However, biological applications typically require that the designed molecules havecertain desired properties.
For example, to penetrate the blood-brain barrier in order to interact withreceptors of the nervous system, a molecule typically requires a certain permeability score [55, 56].
Conditional generation then refers to the problem of exploring the chemical space conditioned onmolecule properties; a well-suited task to evaluate the ProbTransformer in a decoder-only settingagainst the state-of-the-art transformer decoder model, MolGPT [11].
Data We use the GuacaMol [57] training data and the evaluation protocol provided by [11].
Formore information about the data we refer to Appendix D.1.
Setup We employ the same architecture as [11], enhanced with probabilistic layers, and trainour model for 20 epochs or 300k training steps.
During training, the model implicitly learns the [Table] properties of the training data by conditioning generation on molecular properties together with theinput SMILES.
At inference, novel molecules with multiple desired property values are generated byproviding the model with a start token alongside with the desired property values while predictingthe next token until either a maximum length is reached or the model produces an end-token.
Wecondition the generation of molecules on three properties: the synthetic accessibility score (SAS) , thepartition coefﬁcient (logP) , and the topological polar surface area (TPSA) , using the same domainsof values described in [11].
For each combination of property values, the model generates a totalof 10000 molecules.
Results are reported in terms of the mean average deviation (MAD) and thestandard deviation (SD) relative to the range of the desired property values, as well as the validity ofthe generated compounds, their uniqueness in terms of the internal diversity of valid predictions, andtheir novelty compared to the training data.
For more information about the measures and the generalsetup, we refer to Appendix D.2.
Results
The results for the conditional generation of molecules are summarized in Table 3.
Weobserve a high validity score and a novelty of 1.0, indicating that our ProbTransformer has learnt theunderlying SMILES grammar very well and does not suffer from overﬁtting to the train data.
For themain task of conditional generation, the ProbTransformer clearly outperforms MolGPT across allmeasures, which highlights its ability to control multiple molecular properties during generation.
Thelargest improvement can be observed for the SD and MAD scores of TPSA, an important measurefor drug delivery in the body, improving the standard deviation (SD) by 57.5% and the mean averagedeviation (MAD) score by nearly 35%.
We do not observe improvements in uniqueness comparedto MolGPT but note that nearly perfect uniqueness could, e.g., be achieved by adding carbons aposteriori [58].
Additional results and prediction examples are shown in Appendix D.3.
4.4 Ablation Studies
We perform two ablation studies to provide more evidence for the usefulness of hierarchical probabilistic construction and kappa annealing.
For this ablation, we use the RNA Folding task describedin Section 4.2 with the same architecture and the training steps.
We provide more details in E.
Hierarchical probabilistic design Due to the possibility of having a prob layer in any block, wetested different setups from one prob layer (middle) up to all blocks enhanced with a prob layer.
In Figure 4, we measure the performance improvement in Hamming distance compared to a vanilla
Transformer on the TS0 and TSsameStruct sets.
Performance saturates at 4 prob layers for the TS0set but keeps improving on the TSsameStruct set.
This is in line with our hypothesis that especiallyambiguous samples proﬁt more from the hierarchical architecture. [Table]
Kappa Annealing We assess the effect ofkappa annealing when initializing κ with different initial values.
In this ablation, we compare aconstant κ to kappa annealing on the Hammingdistance, see Table 4.
In these results, kappaannealing always improves the performance andsubstantially stabilizes ﬁnal performance acrossa range of initialization values.
Figure 4 showsthe adaption of κ over a training run, demonstrating that it smoothly converges to a smallervalue.
Initializing with such a small value would yield much worse performance.
Related Work
Related
Transformer Models A series of works enhanced the encoder-decoder Transformer architecture with a distribution in the latent space between encoder and decoder.
During inference, themodel samples from these distributions.
There is work in the ﬁeld of music representation [59] whichdoes not use a posterior model.
Work in neural machine translation [60, 61] has a separate posteriorencoder and enhances the ELBO objective.
This is similar to the work of Lin et al.
[62] whichintroduces a variational decoder layer for dialogue modelling and has a separate posterior encoder.
Other work in text generation [63, 64, 65] uses the encoder for the predictive and posterior model.
In contrast to our work, neither of the methods mentioned above model hierarchical distributions oruse the GECO objective.
The work of Pei et al.
[66] introduces a hierarchical stochastic multi-headattention mechanism aiming at uncertainty estimation.
In contrast, we preserve the original attentionmechanism by adding a new layer.
The work of Liu et al.
[64] also introduces a β scheduling inβ -VAE ELBO objective [29] similar to our kappa annealing but focused on the KL loss instead of thereconstruction objective.
Related work in RNA folding Until recently, RNA folding was dominated by dynamic programming (DP) algorithms using either thermodynamic, statistical, or probabilistic scoring functions [67].
Learning-based approaches to the problem beneﬁt from making few assumptions on the foldingprocess and allowing previously unrecognized base pairing constraints [42], which recently led tostate-of-the-art performance using deep learning [37, 42].
We discuss these state-of-the-art deeplearning approaches in detail in Appendix C.4 and refer to a recent review [68].
Related Work in Molecule design While a plethora of deep learning-based models have beenproposed for de novo generation of molecules in the past ﬁve years [69, 70, 71, 72, 73, 74, 75,76, 77, 78, 79], only some methods yet approached the challenging task of generating moleculeswith (multiple) predeﬁned property values (conditional generation).
These include conditional
RNNs [80], cVAEs [81], conditional adversarially regularized autoencoders [82], and more recentlyalso Transformers [11].
Since [81] considers inorganic molecules and conditional generation is onlyperformed in a case study in [82], and because of the usage of different evaluation protocols in [80]and [11], a correct evaluation and fair comparison of these approaches is out of the scope of thiswork.
We, therefore, compare the ProbTransformer to the most similar and recent work in the ﬁeld,
MolGPT , which, to our knowledge, is the current state-of-the-art in the ﬁeld.
For more details on thedifferent methods, we refer to multiple excellent reviews [83, 84, 51, 53].
Discussion and Conclusion
We propose a novel probabilistic layer to enhance the transformer architecture with hierarchical latentdistributions while keeping the global receptive ﬁeld via the attention mechanism.
We showed the beneﬁts of our approach in several experiments.
To our knowledge, the ProbTransformer is the ﬁrst learned RNA folding model that can provide multiple correct structure proposalsfor a given RNA sequence, which opens the doors to novel research paths in RNA structure predictionthat are in line with experimental evidence for RNA structural dynamics from, e.g., NMR studies,such as fraying [85], bulge migration [86], and ﬂuctuating base pairs [86].
On the challenging multiobjective optimization task [54] of designing molecules with desired properties, we demonstratesuperior control over molecule properties during the generation in a decoder-only setting compared to a state-of-the-art vanilla Transformer architecture.
We want to point out that molecular researchinevitably bears the risk of misuse [87], but we strongly distance ourselves from any such attempts.
The ProbTransformer samples interdependent sequences in one forward path.
This sampling happensin the latent space, and the ProbTransformer can reﬁne or interpret a sampled latent representationin a further layer.
Compared to sampling from the softmax output distribution, this approach yieldsgreater ﬂexibility.
It also is compatible with other enhancements to the Transformer model since itonly adds a new layer but keeps everything else unchanged.
However, our approach also has limitations.
The additional layer and the posterior model increase thecomputational and memory needs up to a factor of two.
Our approach is also limited to encoder-onlyor decoder-only setups since it requires a target of the same length as the input sequence.
We have notapplied our model in natural language processing yet, e.g. as a language model, which could improvepredictive performance and save compute at inference time due to the probabilistic encoder.
Anotherfuture application could be in text-to-speech or automated speech recognition since natural speechcontains many ambiguities.
In general, our approach and especially the conditional variational trainingusing the GECO objective and kappa annealing, could be used in future probabilistic models.
Kappaannealing reduces the need for exhaustive hyperparameter optimization and increases performance.
Therefore, we expect that the positive impact of our work is not limited to RNA folding and moleculedesign but also to generative models in general.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes] We describe our new approach in detail and provide a
experimental evaluation with state-of-the-art results.
(b) Did you describe the limitations of your work?
[Yes] We mention them in Section 6.(c)
Did you discuss any potential negative societal impacts of your work?
[Yes]
We arenot aware of negative impacts of RNA folding or foundation in probabilistic modelling
but in molecule design and mention this in Section 6.(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes] We read it and our paper is conform.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[Yes] We include thesource code and scripts to build the data in the supplemental material.
We also opensource our code upon acceptance as mentioned in Section 4(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they werechosen)?
[Yes] Details about hyperparameter and data can be found in Appendices B,C,and D.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] We report all experimental results including standarddeviation on multi-seed runs in the Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Details about compute resourcescan be found in Appendices B.2,C.2, and D.2.
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cite used opensource libraries in the Appendix B.2.
[Table]
(b) Did you mention the license of the assets?
[No] We use only open source software.
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]We do not include any new assets.
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [No] [Table] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[No] Our data does not contain personally identiﬁable [Table] information or offensive content. [Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A] We do not use crowdsourcing nor conducted research with humansubjects.
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
[Table]
(IRB) approvals, if applicable? [N/A] We do not use crowdsourcing nor conductedresearch with human subjects.
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation? [N/A] We do not use crowdsourcing nor conducted [Table] research with human subjects.
[Table]
In this section, we provide insights into the training dynamics of the ProbTransformer using the RNAfolding task. In Figure 5 we see the progress of the training. The ﬁrst row pictures the hammingdistance on the validation set, in the second we see the learning rate schedule, in the third theannealing of kappa, in the fourth the cross-entropy loss L rec , in the ﬁfth the KL loss D KL , in thesixth the adjustment of κ , and at the bottom the reconstruction constraint L rec − κ .
We see that the κ decreases over time due to a low λ which allows to an increase in the pressureon the reconstruction. In Figure 6, we see the same training but with a log scale on the x-axis tofocus on the early training phase. At the beginning of the training, the reconstruction constraint is notsatisﬁed and the Lagrange multiplier λ is increasing which results in pressure on the reconstructionloss. At the same time, the KL divergence increases due to reconstruction via Z post , leading to anincrease in the initial distance of P and Q .
Once the reconstruction constraint is satisﬁed, λ decreasesand the pressure moves to the KL term. Also, the performance, measured in terms of the Hammingdistance, does not improve even when the CE loss drops, since the CE loss only trains the posteriorreconstruction. It only starts to improve when the KL divergence begins to decrease and the predictivemodel learns to create a useful internal latent representation. B Synthetic Sequential Distribution Task
This section provides more details about the synthetic sequential distribution task itself, the conﬁguration of the used Transformer and ProbTransformer model, the training process, and the results.
We design the task to map a sequence of tokens from a source vocabulary x ∈ V ∗ i to a sequenceof target tokens from a target vocabulary y ∈ V ∗ o with the same length. The tokens in the sourcesequence are used to build ‘phrases” P . Each phrase consists of l tokens sampled with replacement(similar to the combination of words in a sentence).
We randomly generate a unique distributionp ( y | x, P ) over the target tokens for each source token in each phrase, depending on the current phrase.
Further, we design the distribution sparsely so that no more than k tokens from the target vocabularyhave a non-zero probability.
The training data is generated by sampling input sequences from allphrases (with replacement) and sampling the target sequence from its corresponding distribution.
Thesize of the source and target vocabulary is 500, a phrase exists of three tokens, and we create 1000different sections. Each target token is drawn from a sparse distribution with 1 to 10 non-zero tokenprobabilities. The sequence length is uniformly drawn from a length of 15 to 90. We created 100.000training samples and 10.000 validation and test samples.
Please ﬁnd the detailed conﬁguration of thetask in Table 5.
[Table]
[Figure]
[Table] [Figure]
Figure 7 shows an example of target distribution depending on the source phrase. On the x-axis, weshow the target vocabulary consisting of number-tokens. On the y-axis, there are two phrases ofsource tokens. The yellow-green-blue color scheme represents the distribution of the target tokenmapping to a source token depending on the source phrase.
Please note that the tokens in the secondand third rows are the same but have different distributions due to the position in the phrase.
A targetsequence is sampled from this distribution, and in the optimal case, the model should be able toreproduce this distribution.
B.2 Setup
In general, we implement our models and tasks in Python 3.8 using mainly PyTroch [88], Numpy [89],and Pandas [90].
We use Matplotlib [91] for the plots in the paper.
For experiments on the synthetic sequential distribution task, we use the conﬁguration listed in Table6 for the Transformer and ProbTransformer.
For MC dropout, we employed a grid search to ﬁnd theoptimal Dropout rate (0 . 1 , 0 . 2 , · · · , 0.5 ) .
The other hyperparameters were tuned manually based onpreliminary work [1] or based on preliminary experiments.
We use SiLU [92] as activation functionin both models.
Furthermore, we use automatic mixed-precision during the training, initialize thelast linear of each layer (feed-forward, attention, or prob layer) with zero, and use a learning ratewarm-up in the ﬁrst epoch of training as well as a cosine learning rate schedule.
We use the squaredsoftplus function to ensure a positive λ value during training and update λ with a negative gradientscaling of 0 .
For the moving average of the reconstruction loss, we use a decay of 0 .
B.3 Results
We provide detailed results of our models and sampling methods with mean and standard deviation forﬁve random seeds and evaluate two additional metrics: (1) We count the different output variations.
[Table]
A perfect model creates the same diversity as nonzero probabilities in the true distribution.
Wenormalize this measure to one; high values suggest more different tokens than non-zeros tokens inthe true distribution, and smaller values suggest fewer tokens.
(2) Another measure for the distancebetween two distributions is the total variation which can deal with zero probabilities.
Please ﬁnd theresults in Table 7.
[Table]
C RNA Folding
In this section, we detail our data pipeline, the general experimental setup and evaluation protocol,and show additional results, including standard deviations for multiseed runs, for our experiments onthe RNA folding problem.
We start, however, with a brief introduction to RNA functions and theimportance of their secondary structure.
RNAs are one of the major regulators in the cell and have recently been connected to diseases likecancer [93] or Parkinson’s [94, 95].
They consequently arise as a promising alternative for thedevelopment of novel drugs, including antiviral therapies against COVID-19[96] and HIV[97], orvaccines[98].
The vast majority of RNAs that are differentially transcribed from the human genome do not encodeproteins [99, 100] and reveal the functions of these so-called non-coding RNAs (ncRNAs) is one ofthe main challenges for understanding cellular regulatory processes [31].
Similar to proteins, thefunction of an RNA molecule strongly depends on its folding into complex shapes, but unlike proteinfolding, which is dominated by hydrophobic forces acting globally, RNAs exhibit a hierarchicalfolding process [101].
In a ﬁrst step, the corresponding nucleotides of the RNA sequence connect toeach other by forming hydrogen bonds, resulting in local geometries and a distinct pairing schemeof the so-called secondary structure of RNA 2 .
The secondary structure deﬁnes the accessibilityof regions for interactions with other cellular compounds [31] and dictates the formation of the3-dimensional tertiary structure [101, 102].
However, RNA structures are highly dynamic, whichdramatically inﬂuences their functions [13, 14].
A learning algorithm that tackles the problem ofpredicting these structure ensembles is currently lacking in the ﬁeld and we consider our work amajor step in the direction of accurate RNA structure prediction.
C.1 Data
In this section, we detail the datasets used during training and for our experiments.
RNA sequencesare chains of the four nucleotides (bases) adenine , cytosine , guanine , and uracil .
{+However, RNA dataoften considers an extended nucleotide alphabet using IUPAC nomenclature 
+}
and we note that thedatasets used in this work include IUPAC nucleotides.
A RNA secondary structure is typically described as a list of pairs where a pair ( i, j ) denotes twonucleotides at the positions i and j of a RNA sequence that are connected by hydrogen bonds toform a base pair.
In the simplest case, all pairs of the secondary structure are nested, i.e. if ( i, j )and ( k, l ) describe two pairs of a secondary structure with i < k , then i < k < l < j .
A functionalimportant class of base pairs [47, 48], however, is called pseudoknots, where the nested pairingscheme is disrupted by one or more pairs of type: i < k < j < l .
Canonical base pairs are formedbetween A and U , G and C (Watson-Crick pairs) or between G and U (Wobble pairs), while all otherpairings of nucleotides are called non-canonical base pairs.
We use the dot-bracket notation [12] fordescription of secondary structures where a dot corresponds to unpaired nucleotides and a pair ofmatching brackets denotes a pair of two nucleotides.
For our experiments we collect a large pool of annotated RNA secondary structures and theircorresponding sequences from recent publications [41, 39, 37, 103, 38].
In particular, we collect102098 samples from the BpRNA [103] meta database, two versions of the RNAStralign [104]dataset provided by [41] and [39] with 28168 and 20897 samples, respectively, two versions of the
ArchiveII [105] dataset provided by [41] and [39] with 2936 and 3966 samples, the TR0, VL0, and
TS0 datasets provided by [37] with 10814, 1300, and 1305 samples, respectively, the TrainSetA [106]and the TrainSetB [106] with 3164 and 1094 samples, respectively, and all available data from the
RNA-Strand [38] database (3898 samples).
For all data provided in .bpseq , .ct or similar ﬁle formatsthat only provide base pairs, we use BpRNA [103] to consistently annotate secondary structures withour major data source, the BpRNA metadatabase.
We split the testset TS0 and the validation set VL0from the pool and uniformly sampled a novel testset, TSsameSeq, from sequences of the remainingpool as described in Section 4.2.
The highly redundant raw data consists of 177035 training samples, called
1300 validation samples and 1351 test samples.
We remove duplicates from the data as well assamples that did not contain any pairs and resolved IUPAC nucleotides that occur very rarely inthe data with a random corresponding canonical nucleotide (A, C, G, U), following the notationof the International Nucleotide Sequence Database Collaboration (INSDC) because these wouldunreasonably blow up the vocabulary and can hardly be learned with a deep learning approach dueto their rare occurrence in the data.
We note, however, that the ProbTransformer generally acceptsall types of IUPAC nucleotides in the input sequence.
We applied CD-HIT-EST-2D [107] to removesequences from the training data with a sequence similarity greater than 80% to the validation and testsamples, the lowest available threshold [37].
In accordance to [37], we limit the length of sequencesto 500 nucleotides to save computational budget and since especially for longer RNAs, experimentalevidence is generally still lacking because of challenges in crystallization and spectral overlap [108].
Table 8 summarizes the ﬁnal datasets used for our experiments.
[Table] [Table]
C.2 Setup
[Table]
In this section, we provide the conﬁguration of the models and training details.
We list the hyperparameters of the Transformer and ProbTransformer training in the RNA folding experiment in Table9.
We use the squared softplus function to ensure a positive λ value during training and update λwith a negative gradient scaling of 0 . 1 .
For the moving average of the reconstruction loss we usea decay of 0 .
We performed the training on one Nvidia RTX2080 GPU and the training timefor one ProbTransformer model is ∼ 63h and for one Transformer ∼ 33h.
The training time for the
ProbTransformer nearly doubles due to the posterior model.
[Table]
C.2.1 CNN Head
Although the Transformer’s and ProbTransformer’s output prediction has a high quality, its stillsometimes ﬂawed.
This hinders the evaluation of the F1 score and therefore the comparison on thismetric to related work.
Instead of manually designing an error correction heuristic we decided tolearn a simple model which takes the Transformer’s last latent and predicts an adjacency matrixwhich is use to evaluate the F1 score of our prediction.
We use a ﬁxed-size CNN without up or down scaling.
The detailed hyperparameters and trainingconﬁguration of our CNN is listed in Table 10.
The input is a concatenation of the vertical andhorizontal broadcast of the last latent from the Transformer as well as the embedded nucleotidesequence.
The output are two classes, one for a connection between nucleotides and one for noconnection.
We train our CNN head on the same training data as the Transformer and pre-computethe Transformer output to save computational resources during the CNN training, i.e. we do not trainthem jointly.
We use early stopping based on the Hamming distance of the validation set.
We performthe training on one Nvidia RTX2080 GPU and the training time is ∼ 3h.
[Table]
C.3 Results
In this section, we describe our evaluation protocol in detail and show additional results for thepredictions on the three test sets, TS0, TSsameStruc, and TSsameSeq, including standard deviations.
For drawing of RNA secondary structures, we use VARNA [109] provided under GNU GPL License.
C.3.1 Evaluation
We evaluate all approaches concerning Hamming distance, the number of solved tasks, and the F1score.
The Hamming distance is the raw count of mismatching characters in two strings of the samelength.
A dot-bracket structure with a Hamming distance of zero counts as solved.
F1 score describesthe harmonic mean of precision (PR) and sensitivity (SN) and is computed as follows:
[Equation]
where TP, FP and FN denote true positives, false positives and false negatives, respectively.
For
TSsameSeq, we only evaluate the best predictions concerning Hamming distance.
SPOT-RNA
The output of SPOT-RNA is in .ct tabular format with columns for indication ofpairs.
Deriving pseudoknots from base pairs is not trivial [110, 103] and we, therefore, convert theoutput to .bpseq format and apply BpRNA [103], the same annotation tool as we used during datageneration, to yield annotated secondary structures for all predictions of SPOT-RNA.
MXFold2 MXFold2 directly outputs secondary structures in dot-bracket format, which we evaluatedirectly.
RNAfold As for MXFold2, RNAfold’s predictions can be evaluated directly from the output indot-bracket format.
We note that RNAfold and MXFold2 are not capable of predicting pseudoknots due to their underlyingdynamic programming approach and are limited to predictions of canonical base pairs.
UFold In contrast to all other approaches, UFold cannot handle IUPAC nucleotides in the inputsequences.
When evaluating the exact same test data used for all other approaches, the recommendedwebserver of UFold at https://ufold.ics.uci.edu/ as well as the standalone version generatespredictions with different lengths compared to the inputs which cannot be evaluated.
A fair comparison with UFold thus was not possible and we decided to exclude UFold from the evaluations inthe main paper.
However, we resolved IUPAC nucleotides by uniformly sampling correspondingcanonical nucleotides for IUPAC nucleotides to create a dataset accepted by UFold.
We use thedot-bracket output of UFold for evaluations since provided .ct ﬁles resulted in errors when trying toobtain secondary structures using BpRNA similar to described for SPOT-RNA due to predictionswith nucleotides pairing with themselves.
The results of UFold on TS0 and TSsameStruc are shownin Table 11.
[Table]
UFold 58.
[Table]
ProbTransformer Regarding Hamming distance and the number of solved structures, we directlyevaluate the predictions of the ProbTransformer from the raw model outputs.
For the F1 Score,however, we use a post-processing step using the CNN head to obtain an adjacency matrix asdescribed above.
Structure Ensemble Predictions
We use the dot-bracket output of all approaches for evaluationson TSsameSeq.
The predictions with the lowest Hamming distance to the respective ground truthstructure were used for the evaluation of performance.
C.3.2 Detailed results
In this section, we provide further results for our experiments on the RNA folding problem, includingstandard derivations for multiseed runs.
TS0 We provide additional results for predictions on TS0.
In Table 12 we provide results withstandard deviations.
Figure 8, 9 and 10 show example predictions of different approaches. Figure 13shows the F1 score for different base-pairs in comparison with related work. [Figure]
[Figure]
[Figure]
[Table]
[Table]
TSsameStruc We provide additional results for predictions on TSsameStruc.
Table 14 showsresults with standard deviation.
[Table] TSsameSeq We provide results for the predictions on TSsameSeq to analyze the ability of the ProbTransformer to capture the structure distribution of RNA sequences that map to differentstructures. For all experiments, we inferred the model 5, 10, 20, 50, and 100 times using sampleinference and analyzed the raw predictions without further post-processing.
We observe that the
ProbTransformer has learned the structure distributions from the data, producing predictions closerto the desired structures as indicated by a low Hamming distance shown in Table 15.
The meanand standard deviations of the predictions are shown in Table 17, results for individual samples aresummarized in Table 16.
Remarkably, the ProbTransformer is the only model that reproduces bothtrue structures for two of the 20 RNA sequences from the raw model predictions directly using only 5inferences (with two out of three random seeds, results not shown).
[Table] [Table] [Table]
C.4 Related work RNA folding
In this section, we discuss state-of-the-art deep learning approaches for the RNA folding problem indetail.
SPOT-RNA [37] was the ﬁrst algorithm using deep neural networks for end-to-end prediction of RNAsecondary structures.
In this work, an ensemble of residual networks (ResNets) [111] and bidirectional
LSTMs [112] (BiLSTMs) [113] was pre-trained on a large set of RNA secondary structure data andthen ﬁne-tuned on a small set of experimentally-derived RNA data, including tertiary interactions.
Although the authors claimed the possibility of predicting RNA tertiary interactions, the performancefor these types of base pairs was poor and the currently available version of the algorithm excludestertiary interactions from its outputs.
We thus consider this work as RNA secondary structureprediction.
E2efold [41] uses a Transformer encoder architecture to learn the prediction of RNA secondarystructures.
The algorithm was trained on a very homologous set of RNA data and showed stronglyreduced performance when evaluated on data of other publications [39, 42], indicating strongoverﬁtting.
Since we use the same data set as the respective work, we exclude E2efold from ourevaluations.
MXFold2 [39] combines deep learning with a DP approach by using a CNN/BiLSTM architecture tolearn the scoring function for the DP algorithm.
The network is trained to predict scores close to a set of thermodynamic parameters to increase robustness.
MXFold2 is restricted to predict a reduced setof base pairs due to limitations in the DP algorithm.
UFold [42] employs a UNet [114] architecture for solving the RNA folding problem.
Similar to SPOTRNA , the authors additionally report results for predictions on data that contains tertiary interactionsafter ﬁne-tuning the model on experimental data with slightly worse overall performance comparedto SPOT-RNA .
In contrast to the previously described works, however, UFold treats an RNA sequenceas an image of all possible base-pairing maps (16 maps corresponding to 16 possible pairs) and anadditional map for pair probabilities, represented as square matrices of the provided sequence.
D Molecular Design
In this section, we provide further information on our experiments for the conditional generation ofmolecules based on multiple desired properties.
Estimations of the size of the chemical space [52] vary widely [115] (typically between 10 20 and10 100 ) with a common consensus that it contains too many molecules to be explicitly enumerated [53].
Deep generative models recently attracted huge interest in exploring this practically inﬁnite spacefor the use in drug discovery and deep learning-based molecular de novo generation has emergedas the most interesting and fast-moving ﬁeld in cheminformatics [53] during the last ﬁve years.
Inthis so-called generative chemistry [51], deep generative models are typically trained on a largepart of enumerated chemical space to learn a biased distribution of molecular representations andevaluated for their ability to generate novel compounds and explore the unseen chemical space.
Common evaluation protocols include metrics to measure e.g. the novelty of the designed compoundsconcerning the examples visited during training, their uniqueness to measure the internal diversityof predictions, and validity of the generated compounds regarding e.g. the underlying SMILESgrammar [57].
However, besides general exploration which could be achieved using uniform samplingapproaches [53], biological applications typically require that the designed molecules have certaindesired properties.
For generative models, the task is then to explore the chemical space conditionedon molecule properties (conditional generation).
D.1 Data
For sequence-based approaches, a common way of representing molecules is the simple molecularline-entry system (SMILES) [15].
This notation was originally proposed to represent moleculesas strings and uses a sequence of elements combined with special characters to enable branching,ring-closure, and different bond orders as well as indications for properties like charges [53].
Totrain the ProbTransformer, we use the training data of the GuacaMol [57] benchmark suite providedby [11].
Overall, the training data consists of 1259543 SMILES with a vocabulary of 94 uniquecharacters.
D.2 Setup
In this section, we provide the conﬁguration of the models and training details.
We list the hyperparameters of the Transformer and ProbTransformer training for the molecule design experiment in
Table 18.
Furthermore, we use automatic mixed-precision during the training, initialize the last linearof each layer (feed-forward, attention, or prob layer) with zero, and use a learning rate warm-upin the ﬁrst epoch of training.
We use the squared softplus function to ensure a positive λ valueduring training and update λ with a negative gradient scaling of 0 . 01 .
For the moving average of thereconstruction loss, we use a decay of 0 .
We performed the training on one Nvidia RTX2080 GPUand the training time for one ProbTransformer model is ∼ 25h and for one Transformer ∼ 13h.
We condition the generation of molecules on three properties:
The synthetic accessibility score (SAS) is a measure of how difﬁcult it is to synthesize a compound.
The values for SAS could generally range between 1 (easy to synthesize) and 10 (very difﬁcult tomake).The partition coefﬁcient (logP) describes the logarithm of the partition coefﬁcient of a compound.
This measure compares the solubilities of a solute in two immiscible solvents at equilibrium.
If oneof the solvents is water and the other one is non-polar, logP is a measure of hydrophobicity.
The topological polar surface area (TPSA) measures the ability of a drug to permeate cell membranesand describes the contributions of all polar atoms, such as oxygen and nitrogen and their attachedhydrogens, to the molecular surface area.
The polar surface area is a good estimator of the absorption,distribution, metabolism, excretion
and toxicology (ADMET)-relations of a compound and providesa rule-of-thumb for chemists to avoid dead-ends during the development process in drug discoverypipelines [116].
For our experiments, we follow the protocol described for MolGPT by choosing a value for eachproperty from the following domains of values.
SAS : 2.0, 4.0; logP 2.0, 6.0; TPSA : 40, 80.
The modelthen conditionally generates molecules with the task to match all chosen values.
Following MolGPT ,
[Table]
results are reported in terms of the mean average deviation (MAD) and the standard deviation (SD)relative to the range of the desired property values.
While our experiment focuses on conditionalgeneration to match the desired property values, we also report the following scores.
Validity describes the fraction of generated molecules that are expressed as valid SMILES.
Highvalidity indicates strong learning of the underlying SMILES grammar.
Uniqueness is a measure of the prediction diversity.
Uniqueness is the fraction of unique predictionsfrom all generated valid SMILES.
Although this metric is known to be not well-deﬁned [16] and canbe tricked by very simple means [58], we report uniqueness scores for reasons of comparability toprevious work in the ﬁeld.
Novelty is the fraction of valid molecules that are different from the training samples.
A high noveltyindicates strong exploration while a low novelty indicates overﬁtting.
We use rdkit [117] for computations of TPSA and logP and use the provided script by [11] forcomputations of SAS.
D.3 Results
We provide detailed results with standard derivation in Table 19.
Figure 11 shows the distributions ofproperties for all valid predictions of the ProbTransformer and MolGPT.
We observe less deviationfrom the desired property values for the predictions of the ProbTransformer compared to MolGPT.
In line with these results, we observe that the ProbTransformer generates more unique moleculeswith properties close to the desired property values compared to MolGPT as indicated in Table 20.
Example predictions of these molecules are shown in Figures 12, 13, 14, 15, 16, 17, 18, and 19.
Weuse rdkit [117] for drawing of the molecules.
[Table]
TPSA
{+LogP 
+}
S A S
ProbTransformer (TPSA, LogP, SAS) (40.0, 2.0, 2.0) (40.0, 2.0, 4.0) (40.0, 6.0, 4.0) (40.0, 6.0, 2.0) (80.0, 6.0, 4.0) (80.0, 2.0, 4.0) (80.0, 2.0, 2.0) (
[Figure]
TPSA
{+LogP 
+}
S A S
MolGPT (
[Figure]
[Figure]
[Figure]
(TPSA, LogP, SAS)
MolGPT
ProbTransformer
Mean
Std (40.0, 2.0, 2.0) 58.
D.4 Related Work
[Figure]
Inspired by progress in the ﬁeld of natural language processing (NLP), early work employed recurrentneural networks (RNNs) to produce focus libraries based on the SMILES notation [73].
Later on,these approaches were coupled with reinforcement learning (RL) to focus the generation on moleculeswith desirable properties [69, 74].
Additional methods were proposed to tackle the problem, includinggenerative adversarial networks (GANs) [70, 71], variational autoencoders (VAEs) [75, 76], andadversarial autoencoders (AAEs) [72, 77, 78, 79].
For more details on the different methods, we referthe interested reader to multiple reviews of the ﬁeld [83, 84, 51, 53].
More recently, the success of self-attention mechanisms entered the ﬁeld and novel methods weredeveloped, adding attention either to RNNs [16] or VAEs [16, 118].
Remarkably, Transformerbased VAEs showed more complex latent representations of molecules and outperformed previousstate-of-the-art VAEs [79] in the ﬁeld [16].
However, as discussed before in Section 5, only some methods yet approached the challenging taskof generating molecules with (multiple) predeﬁned property values (conditional generation) [73, 80,81, 82, 11].
E Ablation Study
We list the hyperparameters of the Transformer and ProbTransformer as well as the training conﬁguration for the ablation study in Table 21.
We performed the training on one Nvidia RTX2080 GPUand the training time for one ProbTransformer model is ∼ 56h (1 prob layer) to ∼ 80h (all prob layer)and for one Transformer ∼ 33h.
Also, we reduced the learning rate for the architecture ablation studydue to unstable training when using all prob layers.
[Table]