N EURAL REPRESENTATION AND GENERATION FOR RNA SECONDARY STRUCTURES
Anonymous authors
Paper under double-blind review
A BSTRACT
Our work is concerned with the generation and targeted design of RNA, a type of genetic macromolecule that can adopt complex structures which inﬂuence their cellular activities and functions.
The design of large scale and complex biological structures spurs dedicated graph-based deep generative modeling techniques, which represents a key but underappreciated aspect of computational drug discovery.
In this work, we investigate the principles behind representing and generating different
RNA structural modalities, and propose a ﬂexible framework to jointly embed and generate these molecular structures along with their sequence in a meaningful latent space.
Equipped with a deep understanding of RNA molecular structures, our most sophisticated encoding and decoding methods operate on the molecular graph as well as the junction tree hierarchy, integrating strong inductive bias about RNA structural regularity and folding mechanism such that high structural validity, stability and diversity of generated RNAs are achieved.
Also, we seek to adequately organize the latent space of RNA molecular embeddings with regard to the interaction with proteins, and targeted optimization is used to navigate in this latent space to search for desired novel RNA molecules.
I NTRODUCTION
There is an increasing interest in developing deep generative models for biochemical data, especially in the context of generating drug-like molecules.
Learning generative models of biochemical molecules can facilitate the development and discovery of novel treatments for various diseases, reducing the lead time for discovering promising new therapies and potentially translating in reduced costs for drug development (Stokes et al., 2020).
Indeed, the study of generative models for molecules has become a rich and active subﬁeld within machine learning, with standard benchmarks (Sterling &
Irwin, 2015), a set of well-known baseline approaches (G ´ omez-Bombarelli et al., 2018; Kusner et al., 2017; Liu et al., 2018; Jin et al., 2018), and high-proﬁle cases of real-world impact 1 .
Prior work in this space has focused primarily on the generation of small molecules (with less than 100 atoms), leaving the development of generative models for larger and more complicated biologics and biosimilar drugs (e.g., RNA and protein peptides) an open area for research.
Developing generative models for larger biochemicals is critical in order to expand the frontiers of automated treatment design.
More generally, developing effective representation learning for such complex biochemicals will allow machine learning systems to integrate knowledge and interactions involving these biologically-rich structures.
In this work, we take a ﬁrst step towards the development of deep generative models for complex biomolecules, focusing on the representation and generation of RNA structures.
RNA plays a crucial role in protein transcription and various regulatory processes within cells which can be inﬂuenced by its structure (Crick, 1970; Steﬂ et al., 2005), and RNA-based therapies are an increasingly active area of research (Pardi et al., 2018; Schlake et al., 2012), making it a natural focus for the development of deep generative models.
The key challenge in generating RNA molecules—compared to the generation of small molecules—is that RNA involves a hierarchical, multi-scale structure, including a primary sequential structure based on the sequence of nucleic acids as well as more complex secondary and tertiary structures based on the way that the RNA strand folds onto itself.
An effective generative model for RNA must be able to generate sequences that give rise to these more complex emergent structures.
There have been prior works on optimizing or designing RNA sequences—using reinforcement learning or blackbox optimization—to generate particular RNA secondary structures (Runge et al., 2019; Churkin et al., 2018).
However, these prior works generally focus on optimizing sequences to conform to a speciﬁc secondary structure.
In contrast, our goal is to deﬁne a generative model, which can facilitate the sampling and generation of diverse RNA molecules with meaningful secondary structures, while also providing a novel avenue for targeted RNA design via search over a tractable latent space.
Key contributions.
We propose a series of benchmark tasks and deep generative models for the task of RNA generation, with the goal of facilitating future work on this important and challenging problem.
We propose three interrelated benchmark tasks for RNA representation and generation:
Unsupervised generation: Generating stable, valid, and diverse RNAs that exhibit complex secondary structures.
Semi-supervised learning: Learning latent representations of RNA structure that correlate with known RNA functional properties.
Targeted generation: Generating RNAs that exhibit particular functional properties.
These three tasks build upon each other, with the ﬁrst task only requiring the generation of stable and valid molecules, while the latter two tasks involve representing and generating RNAs that exhibit particular properties.
In addition to proposing these novel benchmarks for the ﬁeld, we introduce and evaluate three generative models for RNA.
All three models build upon variational autoencoders (VAEs) (Kingma & Welling, 2014) augmented with normalizing ﬂows (Rezende & Mohamed, 2015;
Kingma et al., 2016), and they differ in how they represent the RNA structure.
To help readers better understand RNA structures and properties, a self-contained explanation is provided in appendix A.
The simplest model (termed LSTMVAE) learns using a string-based representation of RNA structure.
The second model (termed GraphVAE) leverages a graph-based representation and graph neural network (GNN) encoder approach (Gilmer et al., 2017).
Finally, the most sophisticated model (termed
HierVAE) introduces and leverages a novel hierarchical decomposition of the RNA structure.
Extensive experiments on our newly proposed benchmarks highlight how the hierarchical approach allows more effective representation and generation of complex RNA structures, while also highlighting important challenges for future work in the area.
T ASK DESCRIPTION
Given a dataset of RNA molecules, i.e. sequences of nucleotides and corresponding secondary structures, our goals are to: (a) learn to generate structurally stable, diverse, and valid RNA molecules that reﬂect the distribution in this training dataset; (b) learn latent representations that reﬂect the functional properties of RNA.
A key factor in both these representation and generation processes is that we seek to jointly represent and generate both the primary sequence structure as well as the secondary structure conformation.
Together, these two goals lay the foundations for generating novel RNAs that satisfy certain functional properties.
To meet these goals, we create two types of benchmark datasets, each one focusing on one aspect of the above mentioned goals:
Unlabeled and variable-length RNA.
The ﬁrst dataset contains unlabeled RNA with moderate and highly-variable length (32-512 nts), obtained from the human transcriptome (Aken et al., 2016) and through which we focus on the generation aspect of structured RNA and evaluate the validity, stability and diversity of generated RNA molecules.
In particular, our goal with this dataset is to jointly generate RNA sequences and secondary structures that are biochemically feasible (i.e., valid), have low free energy (i.e., stable), and are distinct from the training data (i.e., diverse).
We will give an extended assessment of the generation aspect under different circumstances, e.g., when constraining the generation procedures with explicit rules.
Labeled RNA.
The second dataset is pulled and processed from a previous study on in vitro RNAprotein interaction, which features labeled RNAs with shorter and uniform length (40 nts)
(Cook et al., 2017).
With this dataset, our objective is slightly expanded (to include obj.
a), so that the latent space is adequately organized and reﬂective of the interaction with proteins.
Therefore, key assessment for the latent space includes AUROC for the classiﬁcation of protein binding, which is crucial for the design of desired novel RNA molecules.
Essentially, this creates slight variations in the task formulation, with the ﬁrst dataset suited to unsupervised learning of a generative model, while the second datasets involves additional supervision (e.g., for a semi-supervised model or targeted generation).
Our speciﬁc modeling choices, to be introduced in section 3, are invariant to different task formulations, and ﬂexible enough to handle different representations of RNA secondary structures.
We refer readers to appendix B for detailed explanation for the dataset and evaluation metrics on the generated molecules and latent embeddings.
M ETHODS
In this section, we introduce three different generative models for RNA.
All three models are based upon the variational autoencoder (VAE) framework, involving three key components:
A probabilistic encoder network q φ ( z | x ) , which generates a distribution over latent states given an input representation of an RNA
We
experiment with
three
different types of input encodings for RNA sequence and secondary structures (see Figure S1: a dot-bracket annotated string, a graph with adjacency matrix representing base-pairings, and a graph augmented with a hierarchical junction tree annotation for the secondary structure.
A probabilistic decoder network p θ ( x | z ) , which deﬁnes a joint distribution over RNA sequences and secondary structures, conditioned on a latent input.
As with the encoder network, we design architectures based on a linearized string decoding and a graph-based hierarchical junction-tree decoding approach.
A parameterized prior p ψ ( z ) , which deﬁnes a prior distribution over latent states and is learned based on a continuous normalizing ﬂow (CNF) (Chen et al., 2018).
For all the approaches we propose, the model is optimized via stochastic gradient descent to minimize the evidence lower bound (ELBO):
[Equation]
) | p ψ ( z )) where β is a term to allow KL-annealing over the strength of the prior regularization.
In the following sections, we explain our three different instantiations of the encoder (section 3.1), decoder (section 3.2), as well as our procedures to structurally constrain the decoding process using domain knowledge (section 3.3) and our procedures to avoid posterior collapse (section 3.4).
3. E NCODING RNA SECONDARY STRUCTURES
The input to the encoder is a structured RNA molecule, with its sequence given by an ordered array of nucleotides x 1 . . .
x L , with x i ∈ { A, C, G, U } , where L is the length of the sequence, and its secondary structure, either represented as (1) a dot-bracket string S = ˙ x 1 . . . ˙
x L with ˙ x i ∈ { ., ( , ) } ; (2) or as a graph G with two types of edges — covalent bonds along the RNA backbone, and hydrogen bonds between the basepairs 2 .
We use x uv to denote edge features between nucleotides u and v ; (3) or as a hypergraph
T — a depth-ﬁrst ordered array of subgraphs ˆ
G 1 . . .
ˆ G D with L ( ˆ G i ) ∈ { S, H, I, M } indicating the subgraph label, and I ( ˆ G i ) =
{ j | j ∈ { 1 . . .
L }} indicating the assignment of nucleotides to each subgraph.
Encoding RNA secondary structure as sequence.
First, we obtain a joint encoding over the nucleotide and the dot-bracket annotation, using the joint sequence-structure vocabulary { A, C, G, U } × { ., ( , ) } .
Then, these one-hot encodings are processed by a stacked bidirectional
LSTM (Hochreiter & Schmidhuber, 1997), followed by a multi-head self-attention module (Vaswani et al., 2017) to weigh different positions along the RNA backbone.
A global max-pooling is used to aggregate the information into h S , and then we obtain mean µ S and log variance log σ S from h S through linear transformations, and draw latent encoding z S from N ( µ S , σ S ) using the reparameterization trick (Kingma & Welling, 2014).
Learning graph representation of RNA secondary structure.
To encode the graph view G of an
RNA secondary structure, we pass rounds of neural messages along the RNA structure, which falls into the framework of Message Passing Neural Network (MPNN) as originally discussed in Gilmer et al.
(2017) and similarly motivated by Jin et al.
For much longer RNAs, it is conceptually beneﬁcial to pass more rounds of messages so that a nucleotide may receive information on its broader structural context.
However, this may introduce undesired effects such as training instability and over-smoothing issues.
Therefor , we combine our MPNN network with gating mechanism, which is collectively referred as the G-MPNN:
[Equation]
where [ . . .
|| . .
denotes concatenation, σ denotes the activation function and GRU indicates the gated recurrent unit (Cho et al., 2014).
Then, after T iterations of message passing, the ﬁnal nucleotide level embedding is given by: h u = σ ( W gemb [ x u ||
(cid:80) v ∈ N ( u ) v Tvu ]) .
Before pooling the nucleotide level embeddings into the graph level, we pass h 1 . . .
h L through a single bidirectional
LSTM layer, obtaining ˆ h 1 . . .
ˆ h L at each step, and h g = max( { ˆ h
i | i ∈ 1 ...L } ) .
The latent encoding z G is similarly obtained from h G using the reparameterization trick.
Hierarchical encoding of the RNA hypergraph.
To encode the junction tree T of RNA, we employ a type of GRU speciﬁcally suited to tree-like structures, which has previously been applied in works such as GGNN (Li et al., 2016) and JTVAE (Jin et al., 2018).
We refer to this tree encoding network as T-GRU, and the format of its input is shown in Figure 1.
One major distinction between our RNA junction tree and the one used for chemical compounds (Jin et al., 2018) is that an RNA subgraph assumes more variable nucleotide composition such that it is impossible to enumerate based on the observed data.
Therefore, we need to dynamically compute the features for each node in an RNA junction tree based on its contained nucleotides, in a hierarchical manner to leverage the nucleotide level embeddings learnt by G-MPNN.
Considering a subgraph ˆ G i in the junction tree T , we initialize its node feature with: x ˆ G i =
[ L ( ˆ G i )
|| max u ∈I (ˆ G i )
h u ] .
Notably, max u ∈ ˆ G i h u is a max-pooling over all nucleotides assigned to ˆ G i , and nucleotide embedding h u comes from G-MPNN.
To compute and pass neural messages between adjacent subgraphs in the RNA junction tree T , we use the T-GRU network in Eq.
v t ˆ G i , ˆ G j = T - GRU( x ˆ G i , { v t − 1 ˆ G k , ˆ G i | ˆ G k ∈ N ( ˆ G i ) } ) (3)
[Equation]
with details of T-GRU provided in the appendix C, and compute the embeddings for subgraphs with
Eq. 4.
Further, we obtain a depth-ﬁrst traversal of the subgraph embeddings
h
ˆ G 1 . . .
h ˆ G D (cid:48) which is also the order for hierarchical decoding to be discussed later.
This ordered array of embeddings is processed by another bi-directional LSTM , and the ﬁnal tree level representation h T is again given by the max-pooling over the bi-LSTM outputs.
Likewise, latent encoding z T is obtained from h T .
3. RNA MOLECULAR GENERATION
Decoding linearized sequence and structure.
In this setting, the decoder simply autoregressively decodes a token at each step, from the joint sequence-structure vocabulary mentioned before in section 3.1, plus one additional symbol to signal the end of decoding.
To simplify the design choice, we use a single-layered forward-directional LSTM, and its hidden state is initialized with the latent encoding z , which can be either z S , z G or z T .
These three types of prediction are interleaved into the procedures of decoding the junction tree structure of RNA and the nucleotide
Hierarchically decoding hypergraph and nucleotide segments.
The input to this more sophisticated hierarchical decoder are latent encodings z G which contains order and basic connectivity information of the nucleotides, and z T which contains higher order information about the arrangements of nucleotide branches and their interactions.
We give a concise description of the decoding procedures here, along with a detailed algorithm in appendix D. On a high level, we hierarchically decode the tree structure in a depth-ﬁrst manner, and autoregressively generate a nucleotide segment for each visited tree branch.
For these purposes, we interleave three types of prediction (Figure 2).
Denote the current tree node at decode step t and at the i -th visit as ˆ G t,i , whose features include (1) its node label L ( ˆ G t,i ) and, (2) a summary over the already existing i − 1 nucleotide segments max { h l,ju | u ∈ ˆ G t,i and l < t and j < i } , with l denoting the nucleotide is decoded at step l , and j indicating the nucleotide belongs to the j -th branch (this feature is simply zeros when i = 1 ).
Then, its local feature x ˆ G t,i is deﬁned as the concatenation of (1) and (2).
We make use of a notion called node state : h ˆ G t,i , which is obtained by: h ˆ G t,i =
T - GRU( x ˆ G t,i , { v ˆ G , ˆ G t,i | ˆ G ∈ N ( ˆ G t,i ) } ) .
Note its similarity to Eq. 3, and h ˆ G t,i is used to make:
• topological prediction in Figure 2 (A), to determine if the decoder should expand to a new tree node or backtrack to its parent node, based on MLP topo ( h ˆ G t,i ) ; • tree node prediction in Figure 2 (B), on condition that a new tree node is needed due to a possible topological expansion.
This procedure determines the label of the new tree node from the set of { S, H, I, M } , based on MLP node ( h ˆ G t,i ) ; • nucleotide segment decoding in Figure 2 (C), using a single-layered LSTM, whose initial hidden state is MLP dec ([ h ˆ G t,i ||
z T ||
z G ]) .
The start token is the last nucleotide from the last segment.
Our hierarchical decoder starts off by predicting the label of the root node using z T , followed by topological prediction on the root node and decoding the ﬁrst nucleotide segment.
The algorithm terminates upon revisiting the root node, topologically predicted to backtrack and ﬁnishing the last segment of the root node.
The decoded junction tree naturally represents an RNA secondary structure that can be easily transformed to the dot-bracket annotation, and the RNA sequence is simply recovered by connecting nucleotide segments along the depth-ﬁrst traversal of the tree nodes.
3. S TRUCTURALLY CONSTRAINED DECODING
To better regulate the decoding process so that generated RNAs have valid secondary structures, a set of constraints can be added to the decoding procedures at the inference stage.
Essentially, a valid
RNA secondary structure needs to observe the following rules: (1) base-pairing complementarity, which means only the canonical base-pairs and Wobble base-pairs are allowed, i.e. [A-U], [G-C] and [G-U]; (2) hairpin loop should have a minimum of three unpaired nucleotides, i.e. for any two paired bases at position i and j , | i − j | > 3 ; (3) each nucleotide can only be paired once, and overlapping pairs are disallowed.
We will translate the above rules into speciﬁc and applicable constraints
Linearized decoding constraints.
Since the linearized decoder simply proceeds in an autoregressive fashion, constraints can be easily enforced in a way that at each step, a nucleotide with an appropriate structural annotation is sampled by making use of masks and re-normalizing the probabilities.
Likewise, a stop token can only sampled when all opening nucleotides have been closed.
More details to follow in appendix E.
Hierarchical decoding constraints.
The speciﬁc set of constraints for hierarchical decoding is discussed in appendix F. Overall, considering the different natures of the three associated types of prediction, each one should require a set of different strategies, which are once again applicable by adding proper masks before sampling.
As shown in the algorithm in appendix D, the set of constraints are applied to line 13, 24 and 14 with marked asterisk.
3. A VOIDING POSTERIOR COLLAPSE
As discussed in a line of previous works, VAEs with strong autoregressive decoders are susceptible to posterior collapse, an issue where the decoder simply ignores the latent encoding of the encoder (
He et al., 2019).
Therefore, to avoid posterior collapsing, we make use of a carefully chosen KL annealing schedule during training to help the encoder adapt its information content in the latent encoding and in coordination with the decoder.
This schedule is detailed in section 4.
We also learn a parameterized prior as suggested in Chen et al.
but using a CNF instead, following a similar implementation to Yang et al.
(2019), with details given in appendix G.
Our KL annealing schedule is chosen based on empirical observations, as to our knowledge, there has yet to exist any principled methods of selecting such schedule.
We have used diagnostic metrics such as mutual information (He et al., 2019) and active units (Burda et al., 2016) along with a validation set to select a proper KL annealing schedule which is to be described later in section
R ESULTS
We consider three modes of evaluation: (1) unsupervised RNA generation; (2) generation using semi-supervised VAE models and (3) targeted RNA design from an organized latent space.
Results are presented below, and relevant hyperparameters can be found in Table S1.
Unsupervised RNA generation.
Here, we evaluate generated RNAs from models trained on the unlabeled RNA dataset for 20 epochs using a KL annealing schedule including 5 epochs of warm-up, followed by gradually increasing the KL annealing term to 3e-3 (for LSTMVAE and GraphVAE), or 2e-3 (for HierVAE).
The KL annealing schedule was chosen using a validation set of 1,280 RNAs.
Table 1 compares the generation capability of different models, from the posterior as well as the prior distribution, and in scenarios such as applying structural constraints to the decoding process or not.
It clearly shows that our most advanced model, HierVAE which employs a hierarchical view of the structure in its encoding/decoding aspects, achieves the best performance across different evaluation
[Table]
regimes, generating valid and stable RNAs even when the decoding processed is unconstrained.
It is also observed that despite having structural constraints, the validity of our generated RNAs are always slightly below 100%.
This can be explained by the threshold hyperparameter which sets the maximum number of steps for topological prediction as well as the maximal length of each nucleotide segment, as shown in Algorithm 1 in appendix D.
To further demonstrate the beneﬁts of model training from structural constraints, we sample RNAs from the prior of an untrained HierVAE model.
With structural constraints, the validity amounts to 66.34% with an extremely high free energy deviation of 22.613.
Without structural constraints, the validity translates to a mere 9.37% and the model can only decode short single stranded RNAs as it lacks the knowledge of constructing more complex structures.
This comparison illustrates that model training is essential for obtaining stable RNA folding.
The junction tree hierarchy of RNAs developed in our work shares certain modelling similarities with the probabilistic context free grammar (Dowell & Eddy, 2004) used by covariance models (CM) (Eddy & Durbin, 1994). Infernal (Nawrocki & Eddy, 2013) is one of the representative works based on CM, which is capable of sampling RNA secondary structures from a CM built around a consensus secondary structure for a conserved RNA family. However, due to the lack of homologous sequences in our dataset, Infernal is seriously limited and can only sample single stranded RNAs.
Figure 3 illustrate RNA structures generated using HierVAE from a randomly chosen short path through the latent space.
Notably, latent encoding provided by HierVAE translates smoothly in the RNA structure domain: nearby points in the latent space result in highly similar, yet different, structures.
The generated structures are particularly stable for short and medium-size RNAs, and slightly less so for longer RNAs with highly complex structures.
A side-by-side comparison between generated RNA secondary structures and MFE structures in Figure S3 shows that generated structures can evolve smoothly in the latent space along with their corresponding MFE structures.
We also visualize neighborhoods of a Cysteine-carrying transfer RNA and a 5S ribosomal RNA in ﬁgure S and S5.
Supervised RNA generation.
We then evaluate our generative approaches in a semi-supervised setting using seven RBP binding data sets from RNAcompete-S. First, we compare the efﬁcacy of different representational choices while excluding the generative components, i.e. we jointly train VAE encoders followed by simple MLP classiﬁers on top of the latent encodings for binary classiﬁcation on RBP binding.
Table S3 shows that incorporating RNA secondary structures is overall beneﬁcial for the classiﬁcation accuracy, except for RBMY where a model with access to RNA sequence alone (LSTM-SeqOnly) has the best performance.
Notably, different choices for representing RNA secondary structures do not lead to large variation in performance, with the exception of HuR and SLBP, where graph based representations have an advantage over the linearized structural representation.
On the other hand, [Table] sequence based models often have comparable performance to models incorporating RNA secondary structures, which may be explained by its capability of inferring RNA secondary structures from short
RNA sequences.
It is also worth exploring other in-vitro selection protocols such as HTR-SELEX which can select RNAs with higher binding afﬁnities than RNAcompete-S that only involves a single selection step.
Next, we train full generative models (encoder, decoder, latent CNF and MLP embedding classiﬁer), and show the results in Table 2.
Since our strategy for targeted RNA design makes use of seed molecules in the latent space, we mainly sample RNAs from the posterior distribution of these semi-supervised VAE models.
Therefore, we select a KL annealing schedule that tends to retain more information in the latent encodings, i.e. setting maximum β to 5e-4 and training 10 epochs.
Results are promising in that classiﬁcation AUROC measured by the held-out test set is comparable to the fully supervised classiﬁcation models in Table S3, and much better compared to models only using ﬁxed and pretrained VAE embeddings as shown in Table S2.
Also, RNA structures generated from the posterior distribution, even under the setting of unconstrained and deterministic decoding, have high success rates, very stable conformation and good reconstruction accuracy.
[Table]
Targeted RNA design.
We next studied the task of designing RNAs with high RBP binding afﬁnity.
Starting from the latent encodings of 10,000 randomly chosen RNA molecules that have negative labels in each RNAcompeteS test set, and use activation maximization to gradually alter the latent encodings so that the predicted binding probability from the embedding classiﬁers increases.
These embedding classiﬁers have been trained jointly with the VAE models with accuracy reported earlier (Table 2).
Then, we use separately trained full classiﬁers (also earlier shown in
Table S3)
as proxy of oracles for evaluating the “ground truth” probability of RBP binding.
Table 3, report the success rate (fraction of RNAs whose “ground truth” RBP binding probability was improved), along with the average improvement in binding probabilities.
An example of a trajectory of optimized
RNAs is shown in Fig.
S6.
R ELATED WORK
Over the years, the ﬁeld of computational drug discovery has witnessed the emergence of graphcentric approaches.
One of the earliest method, proposed in G´omez-Bombarelli et al.
(2018), is deﬁned on the linearized format of molecular structures and represents a family of methods that rely on sequential models to represent and generate SMILES strings of chemical compounds.
Later methods have sought to construct more chemical priors into the model, via (1) leveraging graph based representation and generation techniques, (2) enforcing direct chemical constraints to the decoding process, (3) considering a multi-scale view of the molecular structures, or (4) using reinforcement learning to integrate more training signal of the molecular structure and function.
As a result, greater success has been achieved by models such as Kusner et al.
; Liu et al.
; Jin et al.
You et al.
(2018) at generating and searching valid and more useful chemical compounds.
Graph representation learning is at the heart of these more recent approaches, to help understand the rules governing the formation of these molecular structures, as well as the correspondence between structures and functions.
Duvenaud et al.
(2015) were among the ﬁrst to apply GNN to learn molecular ﬁngerprints, and the general neural message passing framework for molecules is proposed in Gilmer et al.
(2017), which demonstrate the power of MPNN across various molecular benchmarking tasks.
These prior works on molecular MPNN, together with other GNN architectures developed in other areas, such as considering relational edges (Schlichtkrull et al., 2018) and attention (Velickovic et al., 2018), have laid the foundation for the success of these deep generative models.
Despite the fact that RNA molecules can adopt complex structures, dedicated graph representation learning techniques have been scarce, with some recent works beginning to leverage graph related learning techniques to predict RNA folding (Chen et al., 2020; Singh et al., 2019) and to represent
RNA molecular structures (Yan et al., 2020; Oliver et al., 2020).
Prior to our work, the design of
RNA has mostly focused on the inverse design problem, which is to conditionally generate an RNA sequence whose MFE secondary structure corresponds to an input secondary structure.
Therefore, the line of prior works have predominantly relied on sequential techniques, with some representative methods based on reinforcement learning (Runge et al., 2019), or more classically framed as a combinatorial optimization problem and solved with sampling based techniques (Churkin et al., 2018).
These prior works are mainly concerned with querying from an energy model with ﬁxed thermodynamic parameters and ﬁxed dynamics of RNA folding, which is in itself limited compared to learning based approaches (Chen et al., 2020; Singh et al., 2019), and are unable to model a joint distribution over RNA sequences and possible folds.
C ONCLUSION AND FUTURE WORKS
In this work we propose the ﬁrst graph-based deep generative approach for jointly embedding and generating RNA sequence and structure, along with a series of benchmarking tasks.
Our presented work has demonstrated impressive performance at generating diverse, valid and stable RNA secondary structures with useful properties.
For future works, there are several important directions to consider.
First, it would be beneﬁcial to obtain non-coding RNA families from the RFAM database (Kalvari et al., 2017) which would enable our models to learn more biologically-meaningful representation indicative of RNA homology and functions, as well as to learn evolutionarily conserved RNA structural motifs that would help our models to generate more stable RNA secondary structures correlated with annotated RNA functions.
In that context, a detailed comparison to Infernal and other probabilistic context-free grammar models would be meaningful.
On the methodological aspect, in light of the recent advances in protein sequences pretraining across a large evolutionary-scale (Rives et al., 2019; Elnaggar et al., 2020), our models for RNAs can certainly beneﬁt from such procedures as well with the data collected from RFAM.
After the pretraining step, reinforcement learning can be used to ﬁnetune the generative component of our model with customizable rewards deﬁned jointly on RNA structural validity, folding stability and functions such as binding to certain proteins.
Essentially, reinforcement learning would help our model overcome limitations posed by existing datasets and to actively query a larger design space.
On the evaluation side, it would be of great interest to analyze our models for any potential RNA tertiary structural motifs and to compare them with the those deposited in the CaRNAval (Reinharz et al., 2018) or RNA 3D motifs database (Parlea et al., 2016).
Our models would also need modiﬁcations to allow non-canonical interactions and pseudoknots, which are common in RNA tertiary structures.
All in all, the representation, generation and design of structured RNA molecules represent a rich, promising, and challenging area for future research in computational biology and drug discovery, and an opportunity to develop fundamentally new machine learning approaches.
A B ACKGROUND : RNA S TRUCTURE AND K EY P ROPERTIES
The representation of an RNA molecule starts from its primary sequence structure —i.e., a single chain of nucleotides (adenine (A), cytosine (C), guanine (G) and uracil (U)).
RNA sequences are ﬂexible and can fold onto themselves, enabling the formation of bonds between complementary nucleotides (Watson-Crick base-pairs [A-U,
G-C], and Wobble base-pairs [G-U]), hence stabilizing the molecule 3 .
The set of pairs of interacting nucleotides in an RNA forms its so-called RNA secondary structure .
In computational analyses of RNA, it is standard to assume that a secondary structure is nested : if [ i, j ] and [ k, l ] form base pairs with i < k , then either l < j (nesting) or k > j (non-overlapping).
This enables simple string or planar graph representations (Figure S1 a, b).
The nested structure assumption means that secondary structures can be modelled by a probabilistic context free grammar (Dowell & Eddy, 2004), or by the closely related junction tree structure (Figure S1 c) (Sarrazin-Gendron et al., 2020), where each hypernode corresponds to a particular secondary substructure element: (1) stem : consecutive stacked base-pairs locally forming a double-stranded structure; (2) hairpin loop : unpaired regions closed by a base-pair; (3) internal loop : unpaired regions located between two stems; (4) multiloop : unpaired regions at the junction of at least three stems.
Edges link elements that are adjacent in the structure.
Validity and stability of RNA folding.
The notion of free energy of RNA secondary structures can be used to characterize the stability of a particular conformation.
Given an RNA sequence, there are combinatorially many valid RNA secondary structures which all need to obey a set of constraints (summarized in section 3.3).
However, some structures are more stable than the others by having lower free energy.
Therefore, these structures are more likely to exist (hence more useful) in reality due to the stochastic nature of RNA folding.
The free energy of an RNA secondary structure can be estimated by an energy-based model with thermodynamic parameters obtained from experiments (Mathews et al., 2004), wherein the minimum free energy (MFE) structure can be predicted, up to a reasonable approximation (Lorenz et al., 2011).
B D ATASET AND M ETRICS
The unlabeled dataset is obtained from the complete human transcriptome which is downloaded from the Ensembl database (Aken et al.
; version GRCh38).
We slice the transcripts into snippets with length randomly drawn between 32 and 512 nts, and use RNAfold to obtain the MFE structures.
We randomly split the dataset into a training set that contains 1,149,859 RNAs, and 20,000 held-out
RNAs for evaluating decoding from the posterior distribution.
More information on the structural diversity and complexity of this dataset is shown in Figure S2, which should present signiﬁcant challenges for our algorithms.
The labeled dataset is pulled from a previous study on sequence and structural binding preference of RNA binding proteins (RBP), using an in vitro selection protocol called RNAcompete-S (Cook et al., 2017) which generates synthesized RNA sequences bound or unbound to a given RBP.
RNAs in this experiment are of uniform length, i.e. 40 nts, and offer a rich abundance of RNA secondary structures compared to its predecessor protocols such as RNAcompete (Ray et al., 2009; 2013).
Since no benchmark has been ever established since its publication, we randomly sample 500,000 positive sequences bound to an RBP, and the same amount of negative sequences from the pool of unbound sequences, to curate a dataset for each of the seven RBPs investigated in the paper.
Then, 80% of all
RNAs are randomly selected to the train split, and the rest goes to the test split.
Our evaluation scheme for the generated RNA secondary structures includes the following metrics:
• validity : percentage of generated RNA secondary structures that conform to the structural constraints speciﬁed in section 3.3.
• free energy deviation (FE DEV) : difference of free energy between the generated RNA secondary structure and the MFE structure of the corresponding sequence, which quantiﬁes the gap of both structures from an energy perspective.
A lower FE DEV should indicate higher stability of generated RNAs.
• free energy deviation normalized by length (Normed FE DEV) :
FE DEV divided by the length of generated RNA, which distributes the contribution of total FE DEV to each base.
• 5-mer sequence diversity : entropy of the normalized counts of 5-mer substrings, which directly measures the diversity of RNA sequences, and indirectly for RNA secondary structures when this metric is combined with FE DEV, since monolithic structures of diverse sequences would lead to high FE DEV.
C T REE ENCODING GRU
Following Eq.3, T-GRU computes a new message v t ˆ G i , ˆ G j from ˆ G i and ˆ G j , based on the features in ˆ G i denoted by x ˆ G i , as well as neural messages from neighboring subgraphs to ˆ G i , i.e. { v t − 1 ˆ G k , ˆ G i | ˆ G k ∈
N ( ˆ G i ) } .
The internal structure of T-GRU is equivalent to the tree encoder employed in Jin et al.
which is essentially a neural analogue of the
belief propagation
algorithm on junction
trees.
Nevertheless, we write down the message passing formulas of T-GRU here:
[Equation]
S
[Equation]
D
A
1 Given: z T , z G , M TI, M SI a 2 Initialize: stack ← [ ] 3 function decode( z T , z G ) root ← sample(MLP node ( z T )) ; root .add
incoming message( z T ) ; stack.
push(( root, 0 )) ; t t
M TI stack.
do
; end return
E D ETAILS FOR APPLYING RNA STRUCTURAL CONSTRAINTS TO LINEARIZED DECODING PROCEDURES
When decoding from the joint vocabulary of sequence and dot-bracket structure ( { A, C, G, U } × { ., ( , ) } ), whenever a nucleotide nuc i with a left bracket is sampled at step i , we append them to a stack, i.e. { (nuc i 0 , i 0 ) . . .
(nuc i , i ) } .
Then, at decode step j ,
• if | i − j | ≤ 3 , a proper mask will be added to the categorical logits of the vocabulary, to avoid sampling any nucleotides with right brackets, which means only an unpaired nucleotide or one that comes with a left bracket can be sampled; • if | i − j | > 3 , a mask will be applied to make sure that only a nucleotide complementary to nuc i can be sampled with the right bracket.
Sampling nucleotides with other forms of structures are allowed.
As soon as a nucleotide with a closing right bracket is sampled, we pop out (nuc i , i ) from the stack.
The special symbol for stop decoding can only be sampled when the stack has become empty.
F D ETAILS FOR APPLYING RNA STRUCTURAL CONSTRAINTS TO HIERARCHICAL DECODING PROCEDURES
Additional constraints to be enforced during the hierarchical decoding process to ensure the validity of the decoded RNA secondary structure.
Recall in section 3.2 that three types of predictions are involved with the hierarchical decoding, therefore, each type is associated with its own set of rules.
All set of rules can be observed by adding proper masks to the categorical logits before sampling, which are detailed below.
Constraints for making topological prediction, when the current node is
• stem node, then the algorithm always expands to a new node upon its ﬁrst visit, or backtracks to its parent node upon re-visit; • hairpin node, then the algorithm always backtracks; • internal loop, then the algorithm acts similarly as for stem node; • multi-loop, then the algorithm always expands upon ﬁrst visit and the next re-visit.
Further re-visits to the same multi-loop node are not regulated.
Constraints for predicting new
tree node, when the current node is • stem node, then its child node when exists can be either a hairpin loop, an internal loop, or a multi-loop; • hairpin node, internal loop or multi-loop, then its child node must be a stem node.
Constraints for decoding nucleotide segment.
Due to the property of non-empty intersection between adjacent subgraphs, the start token for decoding a segment at the current node, is always the last nucleotide decoded at the last node.
Therefore, without explicitly mentioning, the algorithm needs to decode at least one new nucleotide at each segment.
When the current node is
• stem node, and if it is upon its ﬁrst visit (i
e. decoding the ﬁrst segment of a stem), then there is no for constraints.
Otherwise, upon its re-visit, the algorithm needs to decode exactly the complementary bases and in the reverse order, according to the ﬁrst decoded segment; • hairpin node, then the decoder needs to decode at least four nucleotides before seeing the stop symbol, unless the hairpin is also the root node.
• internal loop node, and if it is upon its ﬁrst, then constraint is not necessary.
Otherwise, upon its revisit, the algorithm needs to decode at least one unpaired nucleotide on condition that the ﬁrst decoded internal loop segment does not contain any unpaired nucleotides; • multi-loop node, then there is no need for constraints.
G D ETAILS FOR PARAMETERIZING PRIOR DISTRIBUTION USING NORMALIZING FLOW
A normalizing ﬂow involves a series of bijective transformation with tractable Jacobian logdeterminant, to map an observed datapoint x ∼ p θ ( x ) from a complex distribution to a simpler one, such as the standard normal distribution.
Considering the simpliﬁed case where we have a single bijective function f θ : Z → X to map some simple latent variables z to observed datapoint x , then, using the change of variable theorem, the likelihood of the observed datapoint can be evaluated as:
[Equation]
where p z ( . ) denotes some simple base distribution, e.g. N (0; I ) .
Then, it becomes clear the efﬁciency of this scheme heavily relies on the efﬁciency of inverting the forward mapping f θ as well as computing its Jacobian log-determinant.
In this project, we use a type of continuous normalizing ﬂow (CNF) which simpliﬁes the above mentioned computation (Chen et al., 2018).
Consider a time continuous dynamics f ψ ( z ( t ) , t ) of some intermediate data representation z ( t ) , and again z ( t 0 ) ∼ p
z ( . ) , the transformation of variable, along with its inverse mapping, can be expressed as:
[Equation]
and the change of probability density can be expressed as:
[Equation]
Note that the invertibility issue is no longer a concern under some mild constraints (Chen et al., 2018).
Also, Eq. S9 only involves a more light-weight trace operation on the Jacobian rather than evaluating its log-determinant.
Therefore, we learn a parameterized prior using a CNF, and observe the decomposition of the KL term in the VAE objective:
[Equation]
Therefore, during training our CNF parameterized with ψ works on the transformation of complex latent encodings z ∼ q φ ( z | x ) to some simple z ( t 0 ) ∼ N (0; I ) , with an exact likelihood described by
Eq.
S9 and integrated into Eq.
S10 for the complete training objective.
During inference, we simply sample z t 0 ∼ N (0; I ) , and use our CNF to reversely transform it to z ∼ p ψ ( . ) which should be closer to the approximate posterior.
Our speciﬁc parameterization of the CNF follows from Yang et al.
(2019) and Grathwohl et al.
interleaving two hidden concatsquash layers of dimensionality 256 with Tanh non-linearity.
H I NFORMATION OF THE UNLABELED RNA DATASET
I H YPERPARAMETERS
[Table]
J RNA COMPETE -S CLASSIFIERS ON PRETRAINED AND FIXED VAE EMBEDDINGS
[Table]
K E ND - TO - END RNA COMPETE -S CLASSIFIERS
[Table]
L A LTERNATIVE H IER VAE TRAINING ON RNA COMPETE -S
[Table]
M C OMPARISON OF GENERATED RNA SECONDARY STRUCTURES TO MFE STRUCTURES
N N EIGHBORHOOD VISUALIZATION OF A C YSTEINE - CARRYING TRANSFER -RNA
[Figure]
[Figure]