Description Generation Using Variational Auto-Encoders for Precursor microRNA

Abstract

Micro RNAs (miRNA) are a type of non-coding RNA involved in gene regulation and can be associated with diseases such as cancer, cardiovascular, and neurological diseases. As such, identifying the entire genome of miRNA can be of great relevance. Since experimental methods for novel precursor miRNA (pre-miRNA) detection are complex and expensive, computational detection using Machine Learning (ML) could be useful. Existing ML methods are often complex black boxes that do not create an interpretable structural description of pre-miRNA. In this paper, we propose a novel framework that makes use of generative modeling through Variational Auto-Encoders to uncover the generative factors of pre-miRNA. After training the VAE, the pre-miRNA description is developed using a decision tree on the lower dimensional latent space. Applying the framework to miRNA classification, we obtain a high reconstruction and classification performance while also developing an accurate miRNA description.

1. Introduction

In living organisms, DNA encodes all the information used by the organism to develop and survive. To make use of the information encoded in the DNA, it needs to be translated to RNA. One type of RNA is messenger RNA (mRNA), which is used to create proteins. Sometimes, depending on the cell type, it is desirable that a certain mRNA is not expressed, which is where micro RNA (miRNA) comes into play by silencing and inhibiting the expression of certain mRNAs. MiRNAs can silence mRNAs by binding to them, following which the mRNA is degraded. As such, the definition of what constitutes an miRNA is functional since it needs to carry out this task. Depending on the shape of the miRNA, it can be used to silence (multiple) different mRNAs. Since over- and under-expression of miRNA can lead to various diseases, it is important to understand the genome of miRNA better.

It is estimated that there are around 2300 types of human miRNA [1] while only around 700 are documented [2]. Finding new miRNAs often requires complex lab conditions [3,4] where potential new miRNA needs to be differentiated from other RNA. To detect novel miRNAs, the focus is typically on identifying precursor miRNAs (pre-miRNAs), which are intermediate forms processed before becoming mature miRNAs. Initially, primary miRNA is transcribed from DNA and processed by enzymes into pre-miRNA [5]. Pre-miRNA is then further processed into mature miRNA within the cell. Folded in a hairpin structure, pre-miRNA contains more nucleotides and therefore possesses more distinguishing features (see Figure 1) compared to other types of RNA. As a result, detecting novel pre-miRNAs is generally easier than detecting mature miRNAs. While some features, like the presence of a terminal loop and a large fraction of base pairs (C-G/A-U) in the stem, are prevalent in pre-miRNA, they are not enough to differentiate it from other RNA molecules and hence are not sufficient to form a structural description of pre-miRNA.

Figure 1. Artificial pre-miRNA strand [6] labeled with different properties that can be present in (non) pre-miRNA.

A number of data-driven methods have been developed to identify novel pre-miRNA. Some of the pre-miRNA detection methods are based on traditional ML methods, which use engineered features [7] such as Random Forests (RFs) [8] and Support Vector Machines (SVMs) [3,8,9,10,11], where only limited interpretability is possible through feature importance. In addition, several Deep Learning methods have been proposed, which rely on Convolutional Neural Networks (CNNs) [12,13,14,15,16,17], Recurrent Neural Networks (RNNs) [14,18,19], and transformers [20,21]. While these methods are able to achieve state-of-the-art accuracy on the miRNA classification task, interpretation is only possible through methods that use activation maps or techniques such as concept whitening [22]. As such, these black-box models do not directly allow us to form a description of pre-miRNA, which would, in turn, allow us to better understand these molecules and their roles in the cell.

In recent years, generative models such as Variational Auto-Encoders (VAEs) and Generative Adversarial Networks (GANs) have been successfully applied to a wide variety of tasks such as protein [23,24], drug [25,26,27], and DNA design [28,29]. Typically, VAEs and GANs generate data by mapping a low-dimensional latent space to high-dimensional data. The low-dimensional latent space in these models is usually entangled, meaning that one latent variable can encode multiple aspects of the generative process. To tackle this problem, different disentanglement techniques have been proposed [30,31,32], where the loss function or model architecture forces the model to learn a latent representation in which each variable only encodes one aspect of the generative process. However, it is not inherently clear which latent variable corresponds to which property of the data and how these properties could be used to describe the class of a datapoint.

In this work, we develop a method for developing pre-miRNA descriptions, using the lower dimensional latent space of a latent variable model (Figure 2). Our contributions are two-fold: (i) We propose a novel framework based on VAEs and Decision Trees (DTs), which can be used to develop interpretable descriptions of the different object classes y being modeled in terms of features f; (ii) We apply this framework to the domain of precursor microRNA where we demonstrate that it obtains biologically sound structural descriptions.

Figure 2. Overview of the proposed description generation framework. The Variational Autoencoder (VAE) is trained to encode the input into a latent space, which is then decoded. A Decision Tree (DT) is trained on the latent space to split based on features, using an SVM for each feature to assess which split maximizes information gain for the classification of pre-miRNA versus non-pre-miRNA. The resulting tree provides interpretable descriptions of pre-miRNA.

2. Methods

2.1. Data

We use the modmiRBase dataset [12], which contains 49602 RNAs from various organisms and is balanced in terms of class. The pre-miRNA sequences in the dataset were sourced from mirbase.org [33] and mirgenedb.org [34]. The non-pre-miRNA sequences were obtained from various datasets, which included existing non-pre-miRNAs with characteristics similar to pre-miRNAs [3,35,36], as well as shuffled or suboptimally folded pre-miRNAs [7]. The RNA sequences are folded using the RNAFold algorithm [37], which predicts the secondary structure by identifying stable configurations where complementary nucleotides bond. Following this, the RNA sequences are encoded as images with dimensions of 100 by 25 pixels [12]. In this encoding, each nucleotide is represented as a colored bar—adenine (A) in blue, cytosine (C) in yellow, guanine (G) in green, and uracil (U) in red. Gaps (represented in black) are added for nucleotides lacking complementary pairs. The length of each bar is inversely related to the bond strength between nucleotides, with strong bonds occurring between base pairs and weaker bonds between other pairs or gaps. Additionally, the length of each bar increases with multiple consecutive weak bonds, while gaps are consistently represented with a length of 2 pixels. This encoding captures the physical structure of the RNA, with the hairpin formations resembling actual hairpins in the images. Additionally, any bulges present in the RNA structure are represented as protruding bars, accurately reflecting the spatial arrangement of the nucleotides. An example of an image encoding can be found in Figure 3.

Figure 3. Example of RNA image encoding.

Since the shape plays an important role in the process of silencing mRNAs, we also encode the bond strength between nucleotides. This is performed in vector m of size 100, where a strong bond is assigned a value of 1, while a weak bond (or no bond) is assigned a value of 0.

For each datapoint, we also calculated multiple features, using the algorithms from [22]. Several of the features can be found in Figure 1, such as the presence of a terminal loop, the size of the terminal loop, and the length of the stem. These features are not used directly in our VAE but are used in the decision tree algorithm instead.

2.2. Model

To model the (non) pre-miRNA, we made use of the Domain Invariant Variational Autoencoder (DIVA) framework [32], with a latent space for bond strength (${\mathbf{z}}_{\mathbf{m}}$), class (${\mathbf{z}}_y$) and remaining variance (${\mathbf{z}}_{\mathbf{x}}$). Each latent space has a size of 64. For each latent space, a separate encoder based on ResNet [38] was used, with an overview of the architecture in Figure 4a. Furthermore, ${\mathbf{z}}_{\mathbf{m}}$ and ${\mathbf{z}}_y$ have conditional priors, whereas ${\mathbf{z}}_x$ has a standard normal prior. The prior for ${\mathbf{z}}_{\mathbf{m}}$ consists of a 1D convolutional network (Figure 4b), followed by a fully connected layer, while the prior of ${\mathbf{z}}_y$ consists of two fully connected layers. In addition, auxiliary classifiers were used to further enhance the disentanglement. The auxiliary classifier for $\mathbf{m}$ consists of a fully connected layer followed by 1D deconvolutional layers (Figure 4c), while the auxiliary classifier for y consists of two fully connected layers. Since the shape of the miRNA plays an important role in the mRNA silencing process, we made use of an additional auxiliary classifier for y on ${\mathbf{z}}_{\mathbf{m}}$. Finally, to organize the latent space better, Inverse Autoregressive Flow (IAF) [39] was applied to all 3 latent spaces. Each IAF consists of 8 blocks of 2 Masked Autoencoder for Distribution Estimation (MADE) [40] layers, which have a context size of 32 (obtained from the encoder), as well as a hidden size of 1080.

Figure 4. Architecture of different model components. The first number following the layer name indicates the kernel size, and the second number indicates the amount of filters. Unless mentioned otherwise, each convolutional layer uses a stride (s) of 1. Convolutional layers in the encoder use same padding (p), while other layers use valid padding unless mentioned otherwise. (a) Encoder architecture. Sampling from the final layer is followed by IAF to obtain the latent space. (b) Prior (${\mathbf{z}}_m$) architecture. (c) Classifier for m architecture. (d) Decoder for bar height/color. Both decoders are preceded by the same fully connected layer and reshape operation. The height decoder uses 2D transposed convolutions, while the decoder for color first sums over the height dimension and then uses 1D transposed convolutions. For the final layer, the color decoder gives the probability for each color (top and bottom row, 5 colors in each bar). Outputs of the two decoders are combined to obtain $\widehat{\mathbf{x}}$.

Since our data are highly structured, both the encoders and decoders were designed with the data complexity in mind. In the input representation, we one-hot encoded the 5 possible different colors to ensure that distances between different colors are always the same, which results in an image with 5 channels. Furthermore, using a regular deconvolutional image decoder could yield invalid reconstructions, as respecting the structure of the data would not be possible. To limit the amount of invalid images, we separately decode the height and color of each bar. After both height and color distributions are decoded, we combine the two distributions to obtain a distribution over both height and color, which is used to calculate the reconstruction loss. The decoder of the bar height is a 2D deconvolutional decoder. It outputs the probability for the different possible bar heights after softmax activation. To obtain the actual bars, the probability distribution is multiplied by the different bar heights, similar to the approach in StampNet [41]. The decoder for the bar color is a 1D deconvolutional decoder and outputs the color probabilities of the bars in the bottom and top rows. Finally, the heights and colors are multiplied to obtain the reconstruction. When sampling from the decoder, a discrete value should be sampled from both the height and color distributions. The architecture of both decoders can be found in Figure 4d. In both the encoders and decoders, the number of filters was chosen based on Cordero et al. [12].

An overview of the full model architecture can be found in Figure 5. In the model, each layer is followed by batch normalization and uses the ELU activation function (unless mentioned otherwise). For training the model, the loss functions in Equation (2) were used. In Equation (1), the first term represents the reconstruction loss, while the remaining terms represent the KL divergence between the encoded prior ($q_{\varphi }$) and the (learned) prior ($p/p_{\theta }$). The penalty terms for each KL divergence are represented by $\beta$. In Equation (2), each additional term represents the loss of the different auxiliary classifiers ($q_{\omega }$), with $\alpha$ being the penalty term for each classifier. The dataset was split into a training set containing 34,721 samples and a test set containing 14,881 samples, with both datasets being balanced. For training, we set each $\beta$ parameter to 0.5. Both ${\alpha }_{y1}$ and ${\alpha }_{y2}$ were set to 12, while ${\alpha }_m$ was set to 1. The hyperparameters for $\alpha$ and $\beta$ were selected to ensure the model learns a well-organized latent space while still generating high-quality samples. The model was trained using the Adam optimizer [42] with a batch size of 64 and a learning rate of 0.0005 until convergence.

$$\begin{array}{cc}\hfill {\mathcal{L}}_s(m,\mathbf{x},y)=& E_{q_{{\varphi }_m}\left({\mathbf{z}}_m\right|\mathbf{x}),q_{{\varphi }_x}\left({\mathbf{z}}_{\mathbf{x}}\right|\mathbf{x}),q_{{\varphi }_y}\left({\mathbf{z}}_y\right|\mathbf{x})}[log\left(p_{\theta }\left(\mathbf{x}\right|{\mathbf{z}}_y,{\mathbf{z}}_{\mathbf{x}},{\mathbf{z}}_m)\right)]\hfill \\ \hfill & -{\beta }_m\ast D_{KL}(q_{{\varphi }_m}\left({\mathbf{z}}_m\right|\mathbf{x})\left|\right|p_{{\theta }_m}\left({\mathbf{z}}_m\right|m))\hfill \\ \hfill & -{\beta }_x\ast D_{KL}(q_{{\varphi }_x}\left({\mathbf{z}}_{\mathbf{x}}\right|\mathbf{x})\left|\right|p\left({\mathbf{z}}_{\mathbf{x}}\right))\hfill \\ \hfill & -{\beta }_y\ast D_{KL}(q_{{\varphi }_y}\left({\mathbf{z}}_y\right|\mathbf{x})\left|\right|p_{{\theta }_y}\left({\mathbf{z}}_y\right|y))\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill {\mathcal{F}}_{\mathrm{DIVA}}(m,\mathbf{x},y)& ={\mathcal{L}}_s(m,\mathbf{x},y)\hfill \\ \hfill & +{\alpha }_{y_1}{E}_{q_{{\varphi }_y\left({\mathbf{z}}_y\right|\mathbf{x})}}[log{q}_{{\omega }_y}\left(y\right|{\mathbf{z}}_y)]\hfill \\ \hfill & +{\alpha }_{y_2}{E}_{q_{{\varphi }_y\left({\mathbf{z}}_m\right|\mathbf{x})}}[log{q}_{{\omega }_y}\left(y\right|{\mathbf{z}}_m)]\hfill \\ \hfill & +{\alpha }_m{E}_{q_{{\varphi }_m\left({\mathbf{z}}_m\right|\mathbf{x})}}[log{q}_{{\omega }_m}\left(m\right|{\mathbf{z}}_m)]\hfill \end{array}$$

(2)

Figure 5. Full DIVA model architecture. $\parallel$ represents the concatenation of the different latent spaces.

To verify that the model architecture is optimal, we performed several ablation studies where we compared the aforementioned model with a standard(/$\beta$) VAE using a decoder consisting of fully connected layers (VAE), a $\beta$-VAE with IAF applied on the latent space ($\beta$-IAF-VAE), as well as with a deconvolutional decoder and IAF (DC-$\beta$-IAF-VAE). For all models, $\beta$ was set to 0.5. Code for all models and experiments can be found on https://www.github.com/marko-petkovic/mirna (accessed on 22 September 2024).

2.3. Decision Tree

To develop a description of pre-miRNA, we create a decision tree algorithm on top of the ${\mathbf{z}}_m$ latent space, as this latent space contains both information regarding the shape and class of RNA. Since this latent space is organized, various features of the RNA should be linearly separable within this space. As such, a linear SVM can be trained to partition the space based on a feature, such as the length of the RNA.

Compared to a normal decision tree algorithm, such as CART [43] or C4.5 [44], our algorithm makes splits based on the learned latent representation rather than the features (f) themselves. To create a split using our decision tree algorithm, for each feature, we train an SVM to partition the latent space (${\mathbf{z}}_m$), with the feature being the target variable. In case a feature is continuous, we make multiple binary features out of it, where the feature is divided according to different thresholds. Alternatively, one could train a regression model rather than a classifier and decide the threshold for a binary feature afterward. For each split, we assess whether the performance of the SVM is above a threshold, and only features classified with an accuracy above the threshold are considered. Then, to decide which split to make, we asses which split yields the highest information gain based on the class of the RNA. The SVM used to make the split is stored as an attribute of the decision tree node. After training a full decision tree, we can assess which descriptions of pre-miRNA the model developed by following the paths down the tree that lead to a pre-miRNA classification. The full algorithm for making splits for the DT can be found in Algorithm 1.

Algorithm 1 MakeSplit($\mathbf{z},\mathit{y},\mathit{f}$)

Require:  $max\_depth,min\_samples,min\_acc$  1: $\mathrm{Node}\leftarrow Node(\mathbf{z},y,concepts)$  2: if current depth $\ge max\_depth$ or $len\left(y\right)<min\_samples$ or Node is pure then  3:    return Node  4: end if  5: gain ← 0, cls ← None  6: for $feature$ in f do  7:     train SVM $(X=\mathbf{z},y=feature)$  8:     split $\mathbf{z}$ based on classifier predictions in ${\mathbf{z}}_{\mathbf{0}}$, ${\mathbf{z}}_{\mathbf{1}}$  9:     calculate information gain based on split  10:   if information gain > gain and SVM accuracy $>min\_acc$ then  11:       $\mathrm{gain}\leftarrow \mathrm{information}\mathrm{gain}$  12:       cls ← SVM  13:   end if  14: end for  15: Split data into ${\mathbf{z}}_{\mathbf{0}}$, ${\mathbf{z}}_{\mathbf{1}}$, $f_0$, $f_1$, $y_0$, $y_1$ according to cls  16: Node.left_child ←MakeSplit(${\mathbf{z}}_{\mathbf{0}}$,$y_0$,$f_0$,$depth$)  17: Node.right_child ←MakeSplit(${\mathbf{z}}_{\mathbf{1}}$,$y_1$,$f_1$,$depth$)  18: return Node

To predict whether a new datapoint is a pre-miRNA, it should first be encoded by the encoder for ${\mathbf{z}}_m$. Then, the decision tree should be followed, where it is checked on which side of the split by the SVM the point lies.

After training the model, we trained a decision tree with $max\_depth=5$, $min\_samples=10$, and $min\_acc=0.8$ using the training set and features f from [22]. Each of the continuous features was partitioned into multiple binary features.

3. Results

3.1. Performance

In Table 1, the reconstruction performance of our model and the different ablations can be found. The reconstruction performance was calculated based on three different metrics: the Mean Absolute Error (MAE) of the entire image, the nucleotide accuracy, and the MAE of the length of the nucleotides. In Equations (3)–(5), the evaluation metrics are defined. In Equation (3), $y_i$ represents the true pixel value, while ${\widehat{y}}_i$ is the predicted pixel value for each of the P pixels. In Equation (4), $c_b$ is the true color label, and ${\widehat{c}}_b$ is the predicted nucleotide for the b-th bar, with B being the total number of bars. The indicator function $\mathbb{1}(c_b={\widehat{c}}_b)$ evaluates to 1 if the colors match and 0 otherwise. Finally, in Equation (5), $L_b$ and ${\widehat{L}}_b$ represent the true and predicted lengths of the b-th bar, respectively, for all B bars. As can be seen, the DC-IAF-DIVA model achieved the highest performance in terms of all three metrics.

$$\begin{array}{cc}\hfill \mathrm{MAE}& ={\displaystyle \frac{1}{P}}\sum _{i=1}^P\left|y_i-{\widehat{y}}_i\right|\hfill \end{array}$$

(3)

$$\begin{array}{cc}\hfill \mathrm{Nucleotide}\mathrm{Accuracy}& ={\displaystyle \frac{1}{B}}\sum _{b=1}^B\mathbb{1}(c_b={\widehat{c}}_b)\hfill \end{array}$$

(4)

$$\begin{array}{cc}\hfill \mathrm{MAE}\mathrm{length}& ={\displaystyle \frac{1}{B}}\sum _{b=1}^B\left|L_b-{\widehat{L}}_b\right|\hfill \end{array}$$

(5)

Table 1. Reconstruction statistics for different model types. Bold statistics represent the lowest obtained errors.

Model Name	MAE	Nucleotide Accuracy	MAE Length
VAE	0.136	0.695	0.784
$\beta$-VAE	0.055	0.908	0.478
$\beta$-IAF-VAE	0.053	0.918	0.444
DC-$\beta$-IAF-VAE	0.009	0.985	0.067
DC-IAF-DIVA	0.007	0.988	0.057

To verify that the model has learned an organized latent space, we have visualized the latent spaces using t-SNE [45] dimensionality reduction. As can be seen in Figure 6, the disentanglement of the DIVA model indeed appears to be successful, as there seems to be a clear separation between classes in ${\mathbf{z}}_m$ and ${\mathbf{z}}_y$, while this is not present in ${\mathbf{z}}_{\mathbf{x}}$. Taking a further look at the distribution of the fraction of base pairs in the stem, terminal loop length, and the stem length over the ${\mathbf{z}}_m$ latent space (Figure 7), we see that the latent space is well organized. Therefore, it can be used by our decision tree algorithm to develop a pre-miRNA description.

Figure 6. Latent spaces of DIVA model after dimensionality reduction using t-SNE. Yellow dots represent pre-miRNA.

Figure 7. ${\mathbf{z}}_m$ latent space colored according to the fraction base pairs in stem, terminal loop length, and stem length. Lighter colors indicate higher values for each variable.

3.2. Conditional Generation

In addition to inspecting the latent space embeddings, we also qualitatively verify the learned representation of ${\mathbf{z}}_m$. Here, we first encode a selected RNA image to obtain ${\mathbf{z}}_m$, ${\mathbf{z}}_y$, and ${\mathbf{z}}_{\mathbf{x}}$. Following this, the bond strength m of the sample is modified and passed through the conditional prior to obtain ${\mathbf{z}}_m^{\prime }$. Afterward, we reconstruct the RNA based on ${\mathbf{z}}_m^{\prime }$, ${\mathbf{z}}_y$, and ${\mathbf{z}}_x$. Finally, we verify whether the newly generated RNA image indeed follows the provided bond strength.

In Figure 8, we can find the results of the conditional generation experiment. In the top left, the original image can be found, with the bond strength m displayed below the picture. In each column, a different sample from the same conditional prior is reconstructed, while for each row a unique m was used. Overall, we see that there are a few (minor) errors in reconstruction, but in general, the generated images follow the provided bond strength m. As such, we can conclude that the learned latent space is well organized both in terms of features and bond strength.

Figure 8. Generated RNA images, with ${\mathbf{z}}_m$ sampled from the prior of various m. In each row, a different m was used. In the second row, the same m as in the original image was used.

3.3. Decision Tree

The obtained description from the decision tree can be found in Figure 9. Overall, the model achieved an accuracy of 0.912, a sensitivity of 0.923, and a specificity of 0.955. The concept whitening method achieved an accuracy of 0.928 on the same dataset while offering less interpretation compared to the decision tree. An overview of the performance comparison between the two methods can be found in Table 2.

Figure 9. Learned pre-miRNA description.

Table 2. Model Performance on pre-miRNA Classification.

Model	Accuracy	Sensitivity	Specificity
Decision Tree VAE	0.912	0.923	0.955
Concept Whitening [22]	0.928	-	-

In addition to the high classification performance, the decision tree learned biologically sound classification rules. The first split in the decision tree is made based on the amount of base pairs in the stem, where samples with a high amount of base pairs are considered to be pre-miRNA, which is also supported by literature [6]. Similarly, short RNA sequences are also considered unlikely to be pre-miRNA [6], which can be seen in the second split of the tree. The final feature used by the tree to discriminate pre-miRNA from other RNA is the absence of large asymmetric bulges, which is also a known property of pre-miRNA [46]. In general, given the supplied features, we see that the decision tree is able to generate a biologically valid description of pre-miRNA.

4. Discussion

In this work, we have proposed a new method for latent space interpretability in VAEs and applied it to pre-miRNA generation. The trained VAE for pre-miRNA generation has shown excellent performance, validating its suitability for the application of our interpretability method. Our analysis demonstrated that the latent space descriptions developed through the method align with existing biological knowledge of pre-miRNAs, suggesting both the VAE model and our interpretability approach perform robustly in capturing meaningful structural features.

Despite these promising results, there are limitations to our approach. The model is focused solely on pre-miRNA structures and does not account for interactions with target genes, which could provide further insights into miRNA functionality. Additionally, the interpretability method relies on existing biological knowledge for validation, meaning its utility may be limited in less well-characterized regions of the latent space or for novel pre-miRNA variants. Expanding the model to account for more diverse and less-studied pre-miRNA types could pose challenges in interpretability and generalization.

A potential direction for future work could be to explore the use of geometric GNNs instead of CNN-based models for learning pre-miRNA structures. Since miRNA is primarily distinguished from other RNA by its function, which is closely linked to its shape, GNNs may provide a more effective representation of the spatial and structural properties of miRNAs. By using a geometric representation of the RNA molecule, GNNs can more accurately capture and process its shape, potentially improving interpretability and generalization across more complex datasets compared to image-based representations. Additionally, target interaction data, such as miRNA-mRNA interactions, could be integrated to account for functional relationships that influence miRNA activity. Recent machine learning methods that predict miRNA-target interactions [47,48,49] could be leveraged to enhance the model’s understanding of these regulatory networks. Incorporating such interaction data, either as additional input features or for refining the latent space, could significantly improve the model’s ability to generalize across diverse datasets and capture the biological relevance of pre-miRNAs more effectively.