Neural representation and generation for RNA secondary structures

1 INTRODUCTION . 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 subfield within machine learning , with standard benchmarks ( Sterling & Irwin , 2015 ) , a set of well-known baseline approaches ( Gómez-Bombarelli et al. , 2018 ; Kusner et al. , 2017 ; Liu et al. , 2018 ; Jin et al. , 2018 ) , and high-profile 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 first step towards the development of deep generative models for complex biomolecules , focusing on the representation and generation of RNA structures . RNA plays a crucial 1e.g . LambdaZero project for exascale search of drug-like molecules . role in protein transcription and various regulatory processes within cells which can be influenced by its structure ( Crick , 1970 ; Stefl 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. , 2017 ) . However , these prior works generally focus on optimizing sequences to conform to a specific secondary structure . In contrast , our goal is to define 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 : 1 . Unsupervised generation : Generating stable , valid , and diverse RNAs that exhibit complex secondary structures . 2 . Semi-supervised learning : Learning latent representations of RNA structure that correlate with known RNA functional properties . 3 . Targeted generation : Generating RNAs that exhibit particular functional properties . These three tasks build upon each other , with the first 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 field , we introduce and evaluate three generative models for RNA . All three models build upon variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ) augmented with normalizing flows ( 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 B . 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 . 2 TASK 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 reflect the distribution in this training dataset ; ( b ) learn latent representations that reflect 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 first 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 reflective of the interaction with proteins . Therefore , key assessment for the latent space includes AUROC for the classification 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 first 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 specific modeling choices , to be introduced in section 3 , are invariant to different task formulations , and flexible enough to handle different representations of RNA secondary structures . We refer readers to appendix C for detailed explanation for the dataset and evaluation metrics on the generated molecules and latent embeddings . 3 METHODS . 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 : 1 . 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 . 2 . A probabilistic decoder network pθ ( x|z ) , which defines 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 . 3 . A parameterized prior pψ ( z ) , which defines a prior distribution over latent states and is learned based on a continuous normalizing flow ( 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 ) : L = −Eqφ ( z|x ) [ pθ ( x|z ) ] + β KL ( qφ ( z|x ) |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.1 ENCODING RNA SECONDARY STRUCTURES . The input to the encoder is a structured RNA molecule , with its sequence given by an ordered array of nucleotides x1 . . . xL , with xi ∈ { A , C , G , U } , where L is the length of the sequence , and its secondary structure , either rep- resented 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 base- pairs 2 . We use xuv to denote edge features between nucleotides u and v ; ( 3 ) or as a hypergraph T — a depth-first 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 hS , and then we obtain mean µS and log variance log σS from hS through linear transformations , and draw latent encoding zS 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 . ( 2018 ) . For much longer RNAs , it is conceptually beneficial 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 : v̂t−1uv = σ ( W g local [ xu ||xuv ] +W g msg ∑ w∈N ( u ) vt−1wu ) ( 1 ) vtuv = GRU ( v̂ t−1 uv , v t−1 uv ) ( 2 ) 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 final nucleotide level embedding is given by : hu = σ ( W g emb [ xu || ∑ v∈N ( u ) v T vu ] ) . Before pooling the nucleotide level embeddings into the graph level , we pass h1 . . . hL through a single bidirectional LSTM layer , obtaining ĥ1 . . . ĥL at each step , and hg = max ( { ĥi|i ∈ 1 ... L } ) . The latent encoding zG is similarly obtained from hG using the reparameterization trick . Hierarchical encoding of the RNA hypergraph . To encode the junction tree T of RNA , we employ a type of GRU specifically 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 : xĜi = [ L ( Ĝi ) ||maxu∈I ( Ĝi ) hu ] . Notably , maxu∈Ĝi hu is a max-pooling over all nucleotides assigned to Ĝi , and nucleotide embedding hu 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.3 vtĜi , Ĝj = T-GRU ( xĜi , { v t−1 Ĝk , Ĝi | Ĝk ∈ N ( Ĝi ) } ) ( 3 ) hĜi = σ ( W t emb [ xĜi || ∑ Ĝ∈N ( Ĝi ) hĜ ] ) ( 4 ) with details of T-GRU provided in the appendix D , and compute the embeddings for subgraphs with Eq . 4 . Further , we obtain a depth-first traversal of the subgraph embeddings hĜ1 . . . hĜD′ 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 final tree level representation hT is again given by the max-pooling over the bi-LSTM outputs . Likewise , latent encoding zT is obtained from hT . 2We do not differentiate the number of hydrogen bonds , which can be different depending on the base-pairs . For example , G-C has three hydrogen bonds whereas A-U only contains two .

Summary: This paper proposes 3 deep generative models based on VAEs (with different encoding schemes for RNA secondary structure) for the generation of RNA secondary structures. They test each model on 3 benchmark tasks: unsupervised generation, semi-supervised learning and targeted generation.  This paper has many interesting contributions — a comparison of VAE models that use different RNA secondary structure encoding schemes, including traditional dot-bracket notation and a more complex hierarchical encoding, and they also introduce various decoding schemes to encourage valid secondary structures.