11institutetext: Eindhoven University of Technology, The Netherlands 11email:
<EMAIL_ADDRESS>
# Description Generation using Variational Auto-Encoders for precursor
microRNA

Marko Petković 11 0009-0000-4918-6027 Vlado Menkovski 11 0000-0001-5262-0605

###### Abstract

Micro RNAs (miRNA) are a type of non-coding RNA, which are 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
ML could be useful. Existing ML methods are often complex black boxes, which
do not create an interpretable structural description of pre-miRNA. In this
paper, we propose a novel framework, which 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.

###### Keywords:

Generative Models Interpretability Description Generation.

##  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
desired that 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 a 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.

Figure 1: Artificial pre-miRNA strand [2], labeled with different properties
which can be present in (non) pre-miRNA.

It is estimated that there are around 2300 types of human miRNA [1], while
only around 700 are documented [25]. Finding new miRNAs often requires complex
lab conditions [27, 30], where potential new miRNA needs to be differentiated
from other RNA. When identifying novel miRNAs, precursor miRNA (pre-miRNA) is
typically used. This is an RNA molecule that is transformed into miRNA in the
cell. Pre-miRNA is folded in a hairpin and contains more nucleotides and thus
has more features (see Figure 1). As a result, it is easier to differentiate
from other RNA. 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.

A number of data-driven methods have been developed for identifying novel pre-
miRNA. Some of the pre-miRNA detection methods are based on traditional ML
methods using engineered features [31], such as Random Forests (RFs) [20] and
Support Vector Machines (SVMs) [3, 9, 20, 27, 35], 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) [8, 10, 32, 36] and Recurrent Neural Networks (RNNs) [28, 32].
While these methods are able to achieve state-of-the-art accuracy on the miRNA
classification task, interpretation is only possible through methods using
activation maps, or techniques such as concept whitening [4]. As such, these
black-box models do not directly allow us to form a description of pre-miRNA
that would in turn allow us to understand better 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 [19], drug [7] and DNA design [22].
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 [6, 16, 18], 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. 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 different object classes $y$ being modeled, in terms of features $f$ and
(ii) apply this framework to the domain of precursor microRNA, where we
demonstrate that it obtains biologically sound structural descriptions.

##  2 Methods

###  2.1 Data

Figure 2: Example of RNA image encoding.

We use the modmiRBase dataset [8], which contains 49602 RNAs from various
organisms, and is balanced in terms of class. The pre-miRNA sequences in the
dataset were taken from mirbase.org [13] and mirgenedb.org [11], while the non
pre-miRNA sequences were taken from [14, 27, 31, 34]. The RNA sequences are
folded using the RNAFold algorithm [17] , following which they are encoded as
images of 100 by 25 pixels [8]. In the image encoding, each nucleotide is
represented as a colored bar (A: blue, C: yellow, G: green and U: red), while
gaps (black) are introduced for nucleotides without an opposing nucleotide.
The length of each bar inversely depends on the bond strength between
nucleotides, where bonds between base pairs are considered strong, and bonds
between other pairs or gaps are considered weak. The length of each bar
additionally increases with multiple subsequent weak bonds (gaps always have a
length of 2 pixels). As such, the image encoding resembles the physical shape
of the RNA. An example of an image encoding can be found in Figure 2.

Since the shape plays an important role in the process of silencing mRNAs, we
also encode the bond strength between nucleotides. This is done 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 [4]. Several of the features can be found in Figure 1, such as the
presence of a terminal loop, the size of the terminal loop, the length of the
stem etc. These features are not used directly in our VAE, but are used in the
decision tree algorithm instead.

###  2.2 Model

Figure 3: Full model architecture. $\|$ represents concatenation.

To model the (non) pre-miRNA, we made use of the DIVA model [18], with a
latent space for bond strength ($\mathbf{z_{m}}$), class ($\mathbf{z}_{y}$)
and remaining variance ($\mathbf{z}_{x}$). Each latent space has a size of 64.
For each latent space, a separate encoder based on ResNet [15] was used, with
an overview of the architecture in Figure 4(a). Furthermore, $\mathbf{z_{m}}$
and $\mathbf{z}_{y}$ have conditional priors, whereas $\mathbf{z}_{x}$ has a
standard normal prior. The prior for $\mathbf{z_{m}}$ consists of a 1D
convolutional network (Figure 4(b)), 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 follow by 1D deconvolutional layers (Figure 4(c)), while the
auxiliary classifier for $y$ consists of two fully connected layers. Since the
shape of the miRNA place an important role in the mRNA silencing process, we
made use of an additional auxiliary classifier for $y$ on $\mathbf{z_{m}}$.
Finally, to organize the latent space better, Inverse Autoregressive Flow
(IAF) [24] was applied on all 3 latent spaces. Each IAF consists of 8 blocks
of 2 Masked Autoencoder for Distribution Estimation (MADE) [12] layers, which
have a context size of 32 (obtained from the encoder), as well as a hidden
size of 1080.

(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 processes 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
$\mathbf{\hat{x}}$

Figure 4: Architecture of different model components. First number following
the layer name indicates the kernel size, the second number 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.

Since our data is highly structured, both the encoders and decoders 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 with the
different bar heights, similar to the approach in StampNet [33]. The decoder
for the bar color is a 1D deconvolutional decoder, and outputs the color
probabilities of the bars in the bottom and top row. 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 4(d).

An overview of the full model architecture can be found in Figure 3. 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 was used. In Equation LABEL:eq:diva1, the first
term represents the reconstruction loss, while the remaining terms represent
the KL divergence between the encoded prior ($q_{\phi}$) and the (learned)
prior ($p/p_{\theta}$). The penalty terms to 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 34721 samples and a test set containing 14881 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 model was trained using the Adam optimizer
[23], a batch size of 64 and a learning rate of 0.0005 until convergence.

$\begin{split}\mathcal{L}_{s}(m,\mathbf{x},y)=\text{
}&E_{q_{\phi_{m}}(\mathbf{z}_{m}|\mathbf{x}),q_{\phi_{x}}(\mathbf{z_{x}}|\mathbf{x}),q_{\phi_{y}}(\mathbf{z}_{y}|\mathbf{x})}[\log(p_{\theta}(\mathbf{x}\,|\,\mathbf{z}_{y},\mathbf{z_{x}},\mathbf{z}_{d}))]\\\
&-\beta_{m}*D_{KL}(q_{\phi_{m}}(\mathbf{z}_{m}\,|\,\mathbf{x})\;||\;p_{\theta_{m}}(\mathbf{z}_{m}\,|\,m))\\\
&-\beta_{x}*D_{KL}(q_{\phi_{x}}(\mathbf{z_{x}}\,|\,\mathbf{x})\;||\;p(\mathbf{z}_{x}))\\\
&-\beta_{y}*D_{KL}(q_{\phi_{y}}(\mathbf{z}_{y}\,|\,\mathbf{x})\;||\;p_{\theta_{y}}(\mathbf{z}_{y}\,|\,y))\\\
\end{split}$ (1)
$\begin{split}\mathcal{F}_{\text{DIVA}}(m,\mathbf{x},y)&=\mathcal{L}_{s}(m,\mathbf{x},y)\\\
&+\alpha_{y_{1}}E_{q_{\phi_{y}(\mathbf{z}_{y}|\mathbf{x})}}[\log
q_{\omega_{y}}(y\,|\,\mathbf{z}_{y})]\\\
&+\alpha_{y_{2}}E_{q_{\phi_{y}(\mathbf{z}_{m}|\mathbf{x})}}[\log
q_{\omega_{y}}(y\,|\,\mathbf{z}_{m})]\\\
&+\alpha_{m}E_{q_{\phi_{m}(\mathbf{z}_{m}|\mathbf{x})}}[\log
q_{\omega_{m}}(m\,|\,\mathbf{z}_{m})]\end{split}$ (2)

To verify that the model architecture is optimal, we performed an 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
github.com/marko-petkovic/mirna.

###  2.3 Decision Tree

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

0: $max\\_depth,min\\_samples,min\\_acc$

1: $\text{Node}\leftarrow Node(\mathbf{z},y,concepts)$

2: if current depth $\geq max\\_depth$ or $len(y)<min\\_samples$ or Node is
pure then

3: return Node

4: end if

5: gain $\leftarrow$ 0, cls $\leftarrow$ 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_{0}}$,
$\mathbf{z_{1}}$

9: calculate information gain based on split

10: if information gain $>$ gain and SVM accuracy $>min\\_acc$ then

11: $\text{gain}\leftarrow\text{information gain}$

12: cls $\leftarrow$ SVM

13: end if

14: end for

15: Split data into $\mathbf{z_{0}}$, $\mathbf{z_{1}}$, $concepts_{0}$,
$concepts_{1}$, $y_{0}$, $y_{1}$ according to cls

16: Node.left_child $\leftarrow$
MakeSplit($\mathbf{z_{0}}$,$y_{0}$,$f_{0}$,$depth$)

17: Node.right_child $\leftarrow$
MakeSplit($\mathbf{z_{1}}$,$y_{1}$,$f_{1}$,$depth$)

18: return Node

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.

Compared to a normal decision tree algorithm, such as CART [5] or C4.5 [29],
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 afterwards. For each split, we assess whether the
performance of the SVM is above a threshold, and only keep the SVMs with a
high enough performance. Then, to decide which split to make, we asses which
split yields the highest information gain based on the class of the RNA. After
training a full decision tree, we can assess which descriptions of pre-miRNA
the model developed, by following paths down the tree which lead to a pre-
miRNA classification. The full algorithm for making splits for the DT can be
found in Algorithm 1.

To perform 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 [4]. Each of the continuous features were 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 MAE on the entire image, the MAE based on
nucleotide color and the MAE in terms of length of the nucleotides. As can be
seen, our model achieved the highest performance both in terms of all three
metrics.

Table 1: Reconstruction statistics for different model types. Bold statistics represent lowest obtained errors. Model Name | MAE | MAE nucleotide | MAE length
---|---|---|---
VAE | 0.136 | 0.305 | 0.784
$\beta$-VAE | 0.055 | 0.092 | 0.478
$\beta$-IAF-VAE | 0.053 | 0.082 | 0.444
DC-$\beta$-IAF-VAE | 0.009 | 0.015 | 0.067
DC-IAF-DIVA | 0.007 | 0.012 | 0.057
Figure 5: Latent spaces of DIVA model after dimensionality reduction using
t-SNE. Yellow dots represent pre-miRNA. Figure 6: $\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.

To verify that the model has learned an organized latent space, we have
visualized the latent spaces using t-SNE [26] dimensionality reduction. As can
be seen in Figure 5, 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_{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 6), 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.

###  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_{x}}$. Following this, the bond strength $m$ of the sample is
modified, and passed through the conditional prior, to obtain
$\mathbf{z}_{m}^{\prime}$. Afterwards, 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.

Figure 7: 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.

In Figure 7, 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.

###  3.3 Decision Tree

Figure 8: Learned pre-miRNA description.

The obtained description from the decision tree can be found in Figure 8.
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.924 on the same dataset, while offering less interpretation compared to the
decision tree.

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 is considered to be an pre-miRNA, which is also
supported by literature [2]. Similarly, short RNA sequences are also
considered unlikely to be pre-miRNA [2], 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 [21]. 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. The trained VAE for pre-miRNA generation
has shown excellent performance, and was therefore suitable for application of
our method. Finally, we have shown that the developed description matches
existing knowledge of pre-miRNAs, suggesting that both the VAE and our
interpretability method perform well.

## References

  * [1] Alles, J., Fehlmann, T., Fischer, U., Backes, C., Galata, V., Minet, M., Hart, M., Abu-Halima, M., Grässer, F.A., Lenhof, H.P., et al.: An estimate of the total number of true human mirnas. Nucleic acids research 47(7), 3353–3364 (2019)
  * [2] Allmer, J.: Computational and bioinformatics methods for microrna gene prediction. In: MiRNomics: MicroRNA biology and computational analysis, pp. 157–175. Springer (2014)
  * [3] Batuwita, R., Palade, V.: micropred: effective classification of pre-mirnas for human mirna gene prediction. Bioinformatics 25(8), 989–995 (2009)
  * [4] Van den Brandt, I.: Towards Concept-based Interpretability of Pre-miRNA Detection using Convolutional Neural Networks. Master’s thesis, Eindhoven University of Technology (2021)
  * [5] Breiman, L., Friedman, J.H., Olshen, R.A., Stone, C.J.: Classification and regression trees. Routledge (2017)
  * [6] Chen, R.T., Li, X., Grosse, R.B., Duvenaud, D.K.: Isolating sources of disentanglement in variational autoencoders. Advances in neural information processing systems 31 (2018)
  * [7] Cheng, Y., Gong, Y., Liu, Y., Song, B., Zou, Q.: Molecular design in drug discovery: a comprehensive review of deep generative models. Briefings in bioinformatics 22(6), bbab344 (2021)
  * [8] Cordero, J., Menkovski, V., Allmer, J.: Detection of pre-microrna with convolutional neural networks. bioRxiv (2019). https://doi.org/10.1101/840579, https://www.biorxiv.org/content/early/2019/11/13/840579
  * [9] Ding, J., Zhou, S., Guan, J.: Mirensvm: towards better prediction of microrna precursors using an ensemble svm classifier with multi-loop features. BMC bioinformatics 11(11), 1–10 (2010)
  * [10] Do, B.T., Golkov, V., Gürel, G.E., Cremers, D.: Precursor microrna identification using deep convolutional neural networks. BioRxiv p. 414656 (2018)
  * [11] Fromm, B., Domanska, D., Høye, E., Ovchinnikov, V., Kang, W., Aparicio-Puerta, E., Johansen, M., Flatmark, K., Mathelier, A., Hovig, E., Hackenberg, M., Friedländer, M.R., Peterson, K.J.: MirGeneDB 2.0: the metazoan microRNA complement. Nucleic Acids Research 48(D1), D132–D141 (10 2019). https://doi.org/10.1093/nar/gkz885, https://doi.org/10.1093/nar/gkz885
  * [12] Germain, M., Gregor, K., Murray, I., Larochelle, H.: Made: Masked autoencoder for distribution estimation. In: International conference on machine learning. pp. 881–889. PMLR (2015)
  * [13] Griffiths-Jones, S.: The microrna registry. Nucleic acids research 32(suppl_1), D109–D111 (2004)
  * [14] Gudyś, A., Szcześniak, M.W., Sikora, M., Makałowska, I.: Huntmi: an efficient and taxon-specific approach in pre-mirna identification. BMC bioinformatics 14(1), 1–10 (2013)
  * [15] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770–778 (2016)
  * [16] Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., Lerchner, A.: beta-vae: Learning basic visual concepts with a constrained variational framework. In: International conference on learning representations (2016)
  * [17] Hofacker, I.L., Fontana, W., Stadler, P.F., Bonhoeffer, L.S., Tacker, M., Schuster, P.: Fast folding and comparison of rna secondary structures. Monatshefte für Chemie/Chemical Monthly 125(2), 167–188 (1994)
  * [18] Ilse, M., Tomczak, J.M., Louizos, C., Welling, M.: Diva: Domain invariant variational autoencoders. In: Medical Imaging with Deep Learning. pp. 322–348. PMLR (2020)
  * [19] Ingraham, J., Garg, V., Barzilay, R., Jaakkola, T.: Generative models for graph-based protein design. Advances in neural information processing systems 32 (2019)
  * [20] Jiang, P., Wu, H., Wang, W., Ma, W., Sun, X., Lu, Z.: Mipred: classification of real and pseudo microrna precursors using random forest prediction model with combined features. Nucleic acids research 35(suppl_2), W339–W344 (2007)
  * [21] Kang, W., Friedländer, M.R.: Computational prediction of mirna genes from small rna sequencing data. Frontiers in bioengineering and biotechnology 3, 7 (2015)
  * [22] Killoran, N., Lee, L.J., Delong, A., Duvenaud, D., Frey, B.J.: Generating and designing dna with deep generative models. arXiv preprint arXiv:1712.06148 (2017)
  * [23] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 (2014)
  * [24] Kingma, D.P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., Welling, M.: Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems 29 (2016)
  * [25] Kozomara, A., Birgaoanu, M., Griffiths-Jones, S.: mirbase: from microrna sequences to function. Nucleic acids research 47(D1), D155–D162 (2019)
  * [26] Van der Maaten, L., Hinton, G.: Visualizing data using t-sne. Journal of machine learning research 9(11) (2008)
  * [27] Ng, K.L.S., Mishra, S.K.: De novo svm classification of precursor micrornas from genomic pseudo hairpins using global and intrinsic folding measures. Bioinformatics 23(11), 1321–1330 (2007)
  * [28] Park, S., Min, S., Choi, H.S., Yoon, S.: Deep recurrent neural network-based identification of precursor micrornas. Advances in Neural Information Processing Systems 30 (2017)
  * [29] Quinlan, J.R.: C4. 5: programs for machine learning. Elsevier (2014)
  * [30] Saçar, M.D., Allmer, J.: Machine learning methods for microrna gene prediction. In: miRNomics: MicroRNA Biology and Computational Analysis, pp. 177–187. Springer (2014)
  * [31] Saçar Demirci, M.D., Baumbach, J., Allmer, J.: On the performance of pre-microrna detection algorithms. Nature communications 8(1), 1–9 (2017)
  * [32] Tasdelen, A., Sen, B.: A hybrid cnn-lstm model for pre-mirna classification. Scientific reports 11(1), 1–9 (2021)
  * [33] Visser, J., Corbetta, A., Menkovski, V., Toschi, F.: Stampnet: unsupervised multi-class object discovery. In: 2019 IEEE International Conference on Image Processing (ICIP). pp. 2951–2955. IEEE (2019)
  * [34] Wei, L., Liao, M., Gao, Y., Ji, R., He, Z., Zou, Q.: Improved and promising identification of human micrornas by incorporating a high-quality negative set. IEEE/ACM transactions on computational biology and bioinformatics 11(1), 192–201 (2013)
  * [35] Xue, C., Li, F., He, T., Liu, G.P., Li, Y., Zhang, X.: Classification of real and pseudo microrna precursors using local structure-sequence features and support vector machine. BMC bioinformatics 6(1), 1–7 (2005)
  * [36] Zheng, X., Xu, S., Zhang, Y., Huang, X.: Nucleotide-level convolutional neural networks for pre-mirna classification. Scientific reports 9(1), 1–6 (2019)


# Room Classification on Floor Plan Graphs using Graph Neural Networks

Abhishek Paudel  George Mason UniversityFairfaxVirginiaUSA ,  Roshan Dhakal
George Mason UniversityFairfaxVirginiaUSA and  Sakshat Bhattarai  George Mason
UniversityFairfaxVirginiaUSA

###### Abstract.

We present our approach to improve room classification task on floor plan maps
of buildings by representing floor plans as undirected graphs and leveraging
graph neural networks to predict the room categories. Rooms in the floor plans
are represented as nodes in the graph with edges representing their adjacency
in the map. We experiment with House-GAN dataset that consists of floor plan
maps in vector format and train multilayer perceptron and graph neural
networks. Our results show that graph neural networks, specifically GraphSAGE
and Topology Adaptive GCN were able to achieve accuracy of 80% and 81%
respectively outperforming baseline multilayer perceptron by more than 15%
margin.

floor plan map, graph convolutional network, graph neural network, node
classification

##  1\. Introduction

A floor plan is a map representing the layout of a floor in a building. Floor
plans generally illustrate the locations of walls, doors windows and other
structural properties in the floor. Although it has its canonical usage in
architectural design, various types of maps resembling a floor plan is widely
used in the context of indoor navigation by mobile robots. Robotic mapping
techniques like SLAM are commonly leveraged to construct a map of floor of a
building which the robot can then use for localization and navigation. Such
generated maps contain no explicit information about room usage. However,
information about room usage can be useful for robot navigation. In this
paper, we present our approach to improve room classification task on floor
plan maps of buildings by representing floor plans as undirected graphs and
leveraging graph neural networks to predict the room categories. The key
motivation towards this approach lies in the intuition that category of a room
has some predictable relationship with the neighboring rooms. This
relationship can be captured using graph neural networks that take into
account the influence of neighboring nodes in the graph through message
passing mechanism between nodes.

We generate room connectivity graph from floor plans in House-GAN dataset by
representing each room as nodes in the graph with edges between them
representing their adjacency. We experiment with the dataset using four types
of graph neural networks: graph convolutional networks (GCN)(Kipf and Welling,
2016), graph attention network (GAT)(Veličković et al., 2017),
GraphSAGE(Hamilton et al., 2017) and topology adaptive graph convolutional
networks (TAGCN)(Du et al., 2017), along with a baseline multi-layer
perceptron (MLP). Our experiments demonstrate that GraphSAGE and TAGCN
significantly outperform all other models in terms of classification accuracy.
Our experiments with deepening the layers in the network also show that
GraphSAGE and TAGCN do not underperform upon deepening the layers in the
network in contrast to what one would generally expect due to oversmoothing(Li
et al., 2019).

Figure 1. Floor plan map can be represented as an undirected graph with rooms
as nodes, and edges representing their adjacency in the map. A graph neural
network then takes this graph as input to make prediction of room categories
as a node classification task.

##  2\. Related Work

Detecting and classifying elements of floor plans have been studied as a part
of architectural floor plan analysis for many years using various rule-based
heuristics approaches and image processing approaches(Ahmed et al.,
2012)(Dominguez et al., 2012)(Dosch et al., 2000)(Macé et al., 2010)(Barducci
and Marinai, 2012). Later, machine learning approaches that use fully
connected neural networks and convolutional neural networks were also used for
tasks like vectorization, segmentation and detection (Liu et al., 2017)(de las
Heras et al., 2014)(Dodge et al., 2017)(Zeng et al., 2019). In (Barducci and
Marinai, 2012) and (Dominguez et al., 2012), elements in floor plans are
represented as a graph structure for analysis and recognition of floor plan.
Most of these aforementioned tasks have focused on semantic understanding of
floor plan in the context of architectural design.

With the introduction of graph neural networks (GNNs) in the recent years, a
new frontier of learning over graph structures has been opened. In (Hu et al.,
2021), the authors use relational graph convolutional network (RGCN) in
addition to traditional machine learning to infer the usage of rooms in public
buildings using floor plan represented as graph. However, their results show
that RGCN does not outperform traditional machine learning approaches like
random forest.

In the field of robotics, a robotic agent uses a map of the environment for
localization and navigation. These maps are broadly classified into two
categories: metric maps and topological maps(Thrun, 1998). Metric maps like
occupancy grid maps(Moravec and Elfes, 1985) represent the environment in 2D
grid distinguishing obstacles like walls and objects from free space where
robot can navigate. Topological maps represent the environment as graph where
nodes represent notable places and edges represent navigability between them
generally without metric information. In this context, making inferences about
unknown spaces is an active field of research in robotic navigation. In (Liu
and von Wichert, 2014), a method is proposed to infer semantic information
from occupancy grid map represented as a scene graph. In (Aydemir et al.,
2012), the authors propose to predict the topology and the categories of rooms
given a partial map. The map of the environment is represented as a graph with
rooms as nodes and their connectivity as edges. Our work aims to use graph
neural networks to demonstrate that inference on such graph representations
can be improved with the recent advancements in learning over graph
structures.

##  3\. Dataset

We experiment with House-GAN dataset(Nauata et al., 2020) that consists of
143,184 floor plans extracted in vector format from real floor plan images in
Lifull Home’s dataset(Repository, 2020). There are 8 categories of rooms111The
original size of the dataset and number of classes are slightly different.
This is because we perform some data cleaning and also ignore irrelevant
samples like those consisting of only one room in floor plan which would
result in empty edges when represented as a graph..

The dataset consists of the following.

  * •

Room bounding boxes

  * •

Room categories

  * •

Wall coordinates

  * •

Wall to room mapping

  * •

Presence of door in walls

  * •

Name of the original RGB image from Lifull Home’s dataset

Figure 2. Class distribution in House-GAN dataset Figure 3. Distribution of
floor plans by number of rooms

##  4\. Approach

Our approach consists of three stages. First, we generate undirected graphs
from the vectorized floor plan maps in the House-GAN dataset. Second, we
extract relevant node features from the geometric and structural properties of
rooms that are available in the floor plans. Finally, we train multiple graph
neural networks on the floor plan graphs in inductive settings as node
classification task to predict the categories of the rooms.

###  4.1. Graph Generation

We generate the undirected graphs from vectorized floor plan maps with nodes
in the graph representing the rooms, and edges representing the adjacency
between two rooms. We consider two rooms to be adjacent if the distance
between them is smaller than a certain threshold. This is done on the
normalized bounding boxes of rooms with the threshold value set to 0.03
meaning that two rooms are adjacent if they are within 3% of the length of
floor plan map.

###  4.2. Node Feature Extraction

We extract the following node features corresponding to each room in the floor
plan.

  * •

Area of the room

  * •

Length of the room

  * •

Width of the room

  * •

Number of doors in the room

  * •

Whether a room is a parent room

  * •

Whether a room is a child room

The first four features: area, length, width, and number of doors in the list
above are self-explanatory. The last two features in the list above are useful
for nested rooms i.e. room within a room. This is because some rooms in the
floor plan are inside or intersect with another room and could be useful as a
distinguishing feature, for example the closet could be inside a bedroom, the
dining room could be inside the kitchen, and so on. If the intersection of a
room with another room is greater than 70% of its area, we assume that room to
be a child room and the another room that it intersects with to be a parent
room. These are both binary features with two possible values: 0 or 1.

###  4.3. Models

We experiment with multilayer perceptron (MLP) as a baseline model, and graph
neural network (GNN) models. The multilayer perceptron model only takes node
features as input without any graph structures that define connectivity
between nodes, while the graph neural network models take both node features
and the graph structures represented as a list of edges between nodes as
input. For graph neural network models, we experiment with the following four
variants.

  * •

Graph Convolutional Network

Proposed by Kipf et. al.(Kipf and Welling, 2016), Graph Convolutional Network
(GCN) is one of the first models that proposed a general approach for
convolution on graphs in transductive settings. Intuitively, this method
aggregates features from neighboring nodes after a linear transformation to
obtain new features for the node.

  * •

Graph Attention Network

Proposed by Veličković et. al.(Veličković et al., 2017), Graph Attention
Network (GAT) improves graph convolutions by leveraging masked self-
attentional layers to address the shortcomings of prior methods and thus
enabling the model to learn different weights for different nodes in a
neighborhood.

  * •

GraphSAGE

Proposed by Hamilton et. al.(Hamilton et al., 2017), GraphSAGE extends
previous approaches for graph convolution to inductive settings by learning to
generate node embeddings by sampling and aggregating features from the node’s
local neighborhood.

  * •

Topology Adaptive Graph Convolutional Network

Proposed by Du et. al.(Du et al., 2017), Topology Adaptive Graph Convolutional
Network (TAGCN) extends previous approaches by learning a fixed-size filter to
perform convolution on graphs that resembles convolutions on images. It does
not approximate convolution, unlike previous methods.

We implement a general model architecture as shown in Figure 4 for MLP and all
GNN models for a fair comparison between them. The input layer takes in a
vector of size 6 corresponding to 6 node features. We experiment with varying
number of hidden or message passing layers each with a size of 16 and followed
by ReLU activation function. The output layer has a size of 8 corresponding to
8 categories of rooms.

Figure 4. General Model Architecture

##  5\. Experimental Setup

We split the House-GAN dataset into two parts for training and testing. We
take 120,000 floor plans for training and the remaining 23,184 floor plans for
testing. Preliminary small-scale experiments were run on a small fraction of
the dataset for hyperparameter tuning before training on the full dataset. For
large-scale training on full dataset, we use the following setup and
hyperparameters for training.

  * •

GPU: NVIDIA Tesla K80 with 8GB memory

  * •

Number of training epochs: 100

  * •

Loss function: Cross-entropy loss

  * •

Gradient descent optimizer: Adam

  * •

Batch size: 128

  * •

Learning rate: 0.004

  * •

Learning rate scheduling: Exponential with step size of 10 and gamma of 0.8

We use PyTorch and PyTorch Geometric for our experiments. Large-scale
experiments were run on ARGO, a research computing cluster provided by the
Office of Research Computing at George Mason University. The cumulative
training time to run all experiments with all models was approximately 75
hours.

##  6\. Results

The best results obtained in terms of classification accuracy for each model
are shown in Table 1. The best performance is achieved by TAGCN with 81.07%
test accuracy followed by GraphSAGE with test accuracy of 80.26%. We observe
that GCN and GAT had very low test accuracy of 54.85% and 52.50% respectively,
and performed even worse than MLP which had the test accuracy of 65.74%. It
should also be noted that although we have not used any explicit
regularization techniques in our models, the results demonstrate that
overfitting is not an issue since all the models seem to have generalized well
to the test set.

Table 1. Classification Accuracy Model | Train Accuracy | Test Accuracy
---|---|---
MLP | 0.6585 | 0.6574
GCN | 0.5469 | 0.5485
GAT | 0.5269 | 0.5250
GraphSAGE | 0.8068 | 0.8026
TAGCN | 0.8140 | 0.8107
Figure 5. Visualization of test accuracy for each model with respect to the
number of hidden/message passing layers

To study the effects of oversmoothing which is common in graph neural
networks, we train our models with a varying number of hidden/message passing
layers ranging from 2 to 12 layers. The results are shown in Figure 5. We
observe that MLP has no improvement in accuracy after 3 hidden layers.
GraphSAGE and TAGCN have considerable improvements early on with only slight
improvements as the number of layers increase. However, we do not see GCN and
GAT perform better with more message passing layers. In fact, the performance
of GAT generally keeps decreasing upon increasing the number of message
passing layers which suggests that GAT might be more susceptible to
oversmoothing. Surprisingly, both GraphSAGE and TAGCN have no decrease in
performance upon increasing the number of message passing layers. This
suggests that GraphSAGE and TAGCN are more robust to oversmoothing than other
models.

To understand why GCN and GAT were not able achieve a good performance unlike
TAGCN and GraphSAGE both of which had superior performance, we study the
characteristics of these models in terms of their abilities to capture
relational inductive biases that are inherent to the graph structures due to
the edges in the graph affecting how node features interact during message
passing in the graph neural network(Battaglia et al., 2018). To study this, we
instantiate the GNN models with random weights and obtain the node embeddings
for 10,000 samples of floor plan graphs from test split by propagating them
through the untrained models. Then, we perform non-linear dimensionality
reduction on the obtained node embeddings using t-Distributed Stochastic
Neighbor Embedding (t-SNE). We observe that GraphSAGE and TAGCN were able to
partially capture relational inductive biases even before training the network
as demonstrated in Figure 6 and 7. In both of the figures, we can see partial
clusters for some of the classes in the dataset. For example, bedroom,
bathroom and balcony have partially distinguishable clusters. We did not
observe any meaningful clusters with GCN and GAT suggesting that these models
were not quite capable of capturing relational inductive biases in the
dataset. Trivially, no meaningful clusters were observed with MLP or with
original features as one would generally expect strengthening our hypothesis
that GraphSAGE and TAGCN are better at capturing relational inductive biases
in the dataset resulting in better performance after training.

Figure 6. t-SNE on node embeddings obtained from GraphSAGE model initialized
with random weights. Partial clusters for some classes are distinguishable.
Figure 7. t-SNE on node embeddings obtained from TAGCN model initialized with
random weights. Partial clusters for some classes are distinguishable.

##  7\. Conclusion and Future Work

Overall, our experiments demonstrate that graph neural networks can be
leveraged to improve performance on room classification using floor plans
represented as undirected graphs. Specifically, we observe that TAGCN and
GraphSAGE show superior performance compared to other models like GCN, GAT and
MLP we experimented with for this task. It was a surprising finding that GCN
and GAT, which leveraged graph structures could not even perform on par with
MLP which did not leverage graph structures. Additionally, our experiments to
study relational inductive bias on untrained graph neural networks suggest
that GCN and GAT are likely not very capable for capturing relations between
the nodes in the graph resulting in subpar performance. However, further
experiments would be needed to confirm this hypothesis. We also observe
decreased performance of GAT as number of message passing layers were
increased while GCN only has decreased performance when number of layers is 11
or 12. These results are likely due to oversmoothing. However, it was
intriguing that GraphSAGE and TAGCN had either improved or similar performance
as more layers were stacked. This suggests that GraphSAGE and TAGCN are more
robust to oversmoothing than the other models. And again, more experiments
would be needed to conclusively assert this statement.

In terms of improving the performance over our current approach, it would be
interesting to explore whether adding more node features could be useful. One
could also eliminate the need for manual feature extraction by using CNN. One
possible direction for this could be to use the image of the floor plan with
specific room segmented as node property and leveraging CNN to extract
features from these images to be used as node features. This would however be
computationally more expensive than our current approach.

###### Acknowledgements.

All experiments were run on ARGO, a research computing cluster provided by the
Office of Research Computing at George Mason University. We would like to
thank the Office of Research Computing for providing the computing
infrastructure to run these experiments.

## References

  * (1)
  * Ahmed et al. (2012) Sheraz Ahmed, Marcus Liwicki, Markus Weber, and Andreas Dengel. 2012.  Automatic room detection and room labeling from architectural floor plans. In _2012 10th IAPR international workshop on document analysis systems_. IEEE, 339–343.
  * Aydemir et al. (2012) Alper Aydemir, Patric Jensfelt, and John Folkesson. 2012\.  What can we learn from 38,000 rooms? Reasoning about unexplored space in indoor environments. In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_. IEEE, 4675–4682.
  * Barducci and Marinai (2012) Alessio Barducci and Simone Marinai. 2012.  Object recognition in floor plans by graphs of white connected components. In _Proceedings of the 21st International Conference on Pattern Recognition (ICPR2012)_. IEEE, 298–301.
  * Battaglia et al. (2018) Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. 2018\.  Relational inductive biases, deep learning, and graph networks.  arXiv:1806.01261 [cs.LG]
  * de las Heras et al. (2014) Lluís-Pere de las Heras, Sheraz Ahmed, Marcus Liwicki, Ernest Valveny, and Gemma Sánchez. 2014\.  Statistical segmentation and structural recognition for floor plan interpretation.  _International Journal on Document Analysis and Recognition (IJDAR)_ 17, 3 (2014), 221–237.
  * Dodge et al. (2017) Samuel Dodge, Jiu Xu, and Björn Stenger. 2017.  Parsing floor plan images. In _2017 Fifteenth IAPR international conference on machine vision applications (MVA)_. IEEE, 358–361.
  * Dominguez et al. (2012) Bernardino Dominguez, ÁL García, and Francisco R Feito. 2012.  Semiautomatic detection of floor topology from CAD architectural drawings.  _Computer-Aided Design_ 44, 5 (2012), 367–378.
  * Dosch et al. (2000) Philippe Dosch, Karl Tombre, Christian Ah-Soon, and Gérald Masini. 2000.  A complete system for the analysis of architectural drawings.  _International Journal on Document Analysis and Recognition_ 3, 2 (2000), 102–116.
  * Du et al. (2017) Jian Du, Shanghang Zhang, Guanhang Wu, Jose M. F. Moura, and Soummya Kar. 2017.  Topology Adaptive Graph Convolutional Networks.  arXiv:1710.10370 [cs.LG]
  * Hamilton et al. (2017) William L. Hamilton, Rex Ying, and Jure Leskovec. 2017\.  Inductive Representation Learning on Large Graphs.  arXiv:1706.02216 [cs.SI]
  * Hu et al. (2021) Xuke Hu, Hongchao Fan, Alexey Noskov, Zhiyong Wang, Alexander Zipf, Fuqiang Gu, and Jianga Shang. 2021.  Room semantics inference using random forest and relational graph convolutional networks: A case study of research building.  _Transactions in GIS_ 25, 1 (2021), 71–111.  https://doi.org/10.1111/tgis.12664
  * Kipf and Welling (2016) Thomas N. Kipf and Max Welling. 2016.  Semi-Supervised Classification with Graph Convolutional Networks.  arXiv:1609.02907 [cs.LG]
  * Li et al. (2019) Guohao Li, Matthias Müller, Ali Thabet, and Bernard Ghanem. 2019.  DeepGCNs: Can GCNs Go as Deep as CNNs?  arXiv:1904.03751 [cs.CV]
  * Liu et al. (2017) Chen Liu, Jiajun Wu, Pushmeet Kohli, and Yasutaka Furukawa. 2017\.  Raster-to-vector: Revisiting floorplan transformation. In _Proceedings of the IEEE International Conference on Computer Vision_. 2195–2203.
  * Liu and von Wichert (2014) Ziyuan Liu and Georg von Wichert. 2014.  Extracting semantic indoor maps from occupancy grids.  _Robotics and Autonomous Systems_ 62, 5 (2014), 663–674.
  * Macé et al. (2010) Sébastien Macé, Hervé Locteau, Ernest Valveny, and Salvatore Tabbone. 2010.  A system to detect rooms in architectural floor plan images. In _Proceedings of the 9th IAPR International Workshop on Document Analysis Systems_. 167–174.
  * Moravec and Elfes (1985) Hans Moravec and Alberto Elfes. 1985.  High resolution maps from wide angle sonar. In _Proceedings. 1985 IEEE international conference on robotics and automation_ , Vol. 2. IEEE, 116–121.
  * Nauata et al. (2020) Nelson Nauata, Kai-Hung Chang, Chin-Yi Cheng, Greg Mori, and Yasutaka Furukawa. 2020.  House-GAN: Relational Generative Adversarial Networks for Graph-constrained House Layout Generation.  arXiv:2003.06988 [cs.CV]
  * Repository (2020) Informatics Research Data Repository. 2020\.  LIFULL HOME’S Dataset.  https://www.nii.ac.jp/dsc/idr/en/lifull/
  * Thrun (1998) Sebastian Thrun. 1998\.  Learning metric-topological maps for indoor mobile robot navigation.  _Artificial Intelligence_ 99, 1 (1998), 21–71.
  * Veličković et al. (2017) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. 2017.  Graph Attention Networks.  arXiv:1710.10903 [stat.ML]
  * Zeng et al. (2019) Zhiliang Zeng, Xianzhi Li, Ying Kin Yu, and Chi-Wing Fu. 2019\.  Deep floor plan recognition using a multi-task network with room-boundary-guided attention. In _Proceedings of the IEEE/CVF International Conference on Computer Vision_. 9096–9104.