Ancestral protein sequence reconstruction using a tree-structured Ornstein-Uhlenbeck variational autoencoder
Authors: Lys Sanz Moreta
Ola Rønning
Ahmad Salim Al-Sibahi
Jotun Hein
Douglas Theobald
Thomas Hamelryck
Venue: ICLR 2022 Poster (2022)
Keywords: biological sequences
variational autoencoders
latent representations
ornstein-uhlenbeck process
evolution

Abstract: We introduce a deep generative model for representation learning of biological sequences that, unlike existing models, explicitly represents the evolutionary process. The model makes use of a tree-structured Ornstein-Uhlenbeck process, obtained from a given phylogenetic tree, as an informative prior for a variational autoencoder. We show the model performs well on the task of ancestral sequence reconstruction of single protein families. Our results and ablation studies indicate that the explicit representation of evolution using a suitable tree-structured prior has the potential to improve representation learning of biological sequences considerably. Finally, we briefly discuss extensions of the model to genomic-scale data sets and the case of a latent phylogenetic tree.

Published as a conference paper at ICLR 2022
ANCESTRAL PROTEIN SEQUENCE RECONSTRUCTION USING A TREE-STRUCTURED ORNSTEIN- UHLENBECK VARIATIONAL AUTOENCODER
Lys Sanz Moreta
Ola Rønning & Ahmad Salim Al-Sibahi
Probabilistic programming group PLTC Section
Probabilistic programming group
PITTC Section
University of Copenhagen
University of Copenhagen
Copenhagen, Denmark
Copenhagen, Denmark
1ys. swn..an...... www.meeeettttloo..cooo com
{ola, ahmad @di ku. dk
Jotun Hein
Douglas Theobald
Department of Statistics University of Oxford Oxford, United Kingdom hein@stats. ox. ac. uk
Brandeis University
Biochemistry Department
MA, USA
dtthobbald@branddieeeco
Thomas Hamelryck
Probabilistic programming group SCARB PLTC Section
Departments of Biology Computer Science University of Copenhagen Copenhagen, Denmark
thamelry@bio. ku.dk
ABSTRACT
We introduce deep generative model for representation learning of biological sequences that, unlike existing models, explicitly represents the evolutionary process. The model makes use of a tree-structured Ornstein-Uhlenbeek process, obtained from a given phylogenetic tree, as informative prior for variational autoencoder. We show the model performs well on the task of ancestral sequence reconstruction of single protein families. Our results and ablation studies indicate that the explicit representation of evolution using a suitable tree-structured prior has the potential to improve representation learning of biological sequences considerably. Finally, we briefly discuss extensions of the model to genomic- scale data sets and the case of latent phylogenetic tree.
INTRODUCTION
Representation learning of biological sequences is important for data exploration and downstream tasks such as protein design (Detlefsen et al. 2020 Alley et al. 2019). Deep
generative models such as variational autoencoders (VAFs) (Kingma & Welling 2013
2019 have been espocially useful for this purpose (Riesselman et al., 2018 Greener et al. 2018. However, current models do not take evolutionary information fully into account, Le., by relating the sequences belonging to protein family in a phylogenetic tree and incorporating parameterized evolutionary models (Durbin et al.) 1998). To address this problem, replace the standard multivariate Gaussian prior of conventional VAE with a tree-structured prior that takes into account a given evolutionary tree. We propose prior based on the Ornstein-Uhlenber Gaussian process tree (Hansen 1997 Jones &

---

Published as a conference paper at ICLR 2022
Moriarty 2013
We apply the model to classic problem in phylogenetics, namely the inference of ancestral sequences.
Ancestral sequence reconstruction (ASR), iee. the inference of ancestral sequences given their descendants or leaf sequences (Pauling et al. 1963 Yang et al., 1995 Koshi &
Goldstein 1996 Joy et al.) 2016 Hochberg & Thornton 2017 Selberg et al. 2021). has
important applications including protein engineering [Cole & Gaucher 2011 Spence ct al.
2021 modeling tumour evolution (El-Kebir et al. 2015), evaluating virus diversily
vaccine design (Gaschen et al.,
2002). ; understanding drug mechanisms (Wilson et al.
2015) and reconstructing ancient proteins in vitro (Chang al., 2002 Wilson et al.) 2013
Hochberg & Thornton 2017].
As input, assume set of ng known, aligned leaf sequences and their phylogenetic tree. The task we want to address is the inference of the nA S nns- unknown, ancestral sequences (Joy et al., 2016). We show that our probabilistic model, called Draupnir, about on par with or better than the accuracy of established ASR methods for standard experimentally- derived data set (Alieva et al. 2008 Randall al. 20166) and several simulated data sets. In addition,
show that Draupnir is capable of capturing coevolution among sequence positions, unlike conventional ASR methods.
The paper is organised as follows. In Background, we briefly discuss evolution of biological sequences, ancestral sequence reconstruction and the tree-structured Orustein-Uhlenber process. In Related Work, discuss deep generative models of biological sequences. In Methods, we describe the Draupnir model, the inference of ancestral sequences and the setup of the benchmarking experiments. In Results, we discuss the quality of the latent representations, compare the accuracy of Draupnir with state-of-the-ard phylogenetic methods for ASR, and present the results of ablation experiments. We end with a brief discussion of future work, including extending the method to genomic-scale data sets and the case of latent phylogenetic tree.
BACKGROUND
2.1 PROTEIN SEQUENCES AND EVOLUTION
Biological molecules such as proteins and nuclcic acids (DNA, RNA) can be characterised by sequences of characters from alphabet of size nc, where typically ne = 21 for proteins and ne = 5 for nucleic acids (Durbin et al.) 1998). These alphabets include one character that represents gap, which is useful in aligning related sequences in multiple sequence alignment (MSA). In the course of evolution, mutations arise that cause changes in these sequences, including character substitutions, deletions and insertions. A set of ns known, homologous, extant sequences and their evolutionary relationships are naturally represented ny lcaf nodes in binary tree phylogeny, where the n. internal or ancestral nodes represent unknown, ancestral sequences (Joy et al. 2016). Among the internal nodes, the root node is the most ancient node.
Edges between two nodes in the tree are labelled by positive real numbers that represent the time difference or the amount of change between them. Such a labelled binary tree naturally defines (ns + AA)) x (ns nA)) matrix containing the pairwise distances between all nodes, called the patristic distance matria, T. In the context of biological sequences, the field of phylogenctics is concerned with the inference of the tree topology, the labels of the tree's edges and the composition of the ancestral sequences, making use methods based on heuristics (such as maximum parsimony) or probabilistie, evolutionary models (Joy et al. 2016
2.2 ANCESTRAL PROTEIN RECONSTRUCTION
The ASR problem amounts to inferring the composition of the n ancestral sequences from the ng extant sequences, making use of tractable model of evolution (Joy et all 2016). Typically, the phylogenetic tree that relates the sequences is assumed known. Standard methods to do this typically assume independent (factorized) evolution of the characters
2

---

Published as a conference paper at ICLR 2022
in the sequence, which is computationally convenient but unrealistic assumption. For example, in proteins, amino acids are involved in an intricate 3-dimensional network of interactions that lead to strong dependencies between amino acids far part in the sequence. This phenomenon is called epistasis (Hochberg & Thornton 2017], which requires coevolutionary models that go beyond the factorized assumption. Nenetheless, it has been possible to infer ancestral sequences and subsequently resurrect functional ancient, ancestral proteins in vitro [Hochberg & Thornton, 2017). The aim of this work is go beyond the assumption of independent, factorized cvolution by using model of evolution that features continuous, latent vector representations of the protein sequences. This allows us to formulate the ASR problem in the context of a deep generative model.
2.3 THE ORNSTEIN- UHLENBECK PROCESS ON A PHYLOGENETIC TREE
Typically, ASR of biological sequences is done using factorised evolutionary models that represent substitutions, insertions and deletions of the discrete characters in the sequences (Joy et al. 2016). In contrast, Draupnir aims to model the evolution of latent, continuous representations or underlying traits of the sequences. A simple diffusive process allowing for
an equilibrium distribution is the Ornstein-ul Masstein-Uhlenbeck (OU) process (Hansen 1997 Jones
& Moriarty 2013). As the OU process is Gaussian process, it has Gaussian equilibriun distribution, as well as Gaussian marginal distributions.
phylogenetic tree (TOU process) Hansen 1997 Jones & Moriarty 2013) to put the latent representations under the control of a parameterized evolutionary model. Apart from the mean, which for our purposes can be assumed be zero, the TOU process has three parameters: the variation unattributable to the phylogeny or the intensity of specific variation On. the characteristic length scale of the evolutionary dynamics X, and the intensity of inherited variation oy. The covariance function for the corresponding multivariate Gaussian distribution then given by (Hadjipantelis et al. 2012 Jones & Moriarty 2013).
We use OU process
(Thhi) A (-Th:)+
Σx.1 = o
expxp
exp
o o
where Ta. is the patristic distance between nodes k and I in the tree, and the Kronecker delta 0x,1 = 1 if == and 0 otherwise.
The TOU process and related diffusive processes on trees well-establiseed evolutionary models that have been used to model the evolution of continuous traits, such as body mass or length (Joy et al. 2016). For example, Lartillot (2014) proposes a phylogenetic Kalman filter for ancestral trait reconstruction of low-dimensional, continuous traits; Tolkoff et al. (2018) propose phylogenctic factor analysis, in which latent variable under the control of small number independent univariate Brownian diffusion processes is related to observed trails through a loading matrix; Horta al. (2021) use a multivariate TOU process and Markov chain Monte Carlo to model both continuous traits and sequences of discrete characters. To represent the latter, they make use of a pairwise Potts model.
RELATED WORK
3.1 REPRESENTATION LEARNING OF BIOLOGICAL SEQUENCES
A VAE Kingma& & Welling 2013 2019) is a probabilistic, generative model featuring latent
N n= that are independently sampled from a prior distribution, 2n ~ (2..). The latent vectors are passed to neural network (the decoder) with parameters
vectors or representations, 2
0, leading to a likelihood, Xn ~ Px(nn (NN for the data, {x} The prior is typically standard multivariate Gaussian distribution, but other priors have been used, such as distributions on the Poincaré ball to recover hierarchical structures (Mathieu et al. 2019). The posterior distribution p(z,, x,)) is intractable, but can be approximated with a variational distribution or guide, (z, (x,)),)), involving a second neural network (the encoder). Point estimates of the parameters 0 and & obtained by maximizing the
3

---

Published as a conference paper at ICLR 2022
evidence lower bound (ELBO),
Po (x, 2)
(x) Eng
log
x)
using stochastic gradient ascent (Hoffman et al. 2013
VA are increasingly used for representation learning of biological sequences Detleesen ct al. 2020). Riesselman et al. (2018) use a VAE with biologically motivated priors to evaluate the stability of mutants and to explore new regions of sequence space. Greener et al. (2018) use autoencoders to design metal-binding proteins and novel protein folds. Ding et al. (2019) show that the latent representations obtained with VAE can capture evolutionary relationships between sequences. The above models do not represent the phylogenctic tree explicitly, but typically aim to condition on some evolutionary information by training on pre-computed MSAs - an approach that has been called evo-tuning (Rao et al., 2019 Detlefsen al. 2020). Hawkins- Hooker al. (2021) use a VAE with convolutional encoder and decoder, combining upsampling and autoregression, without relying on MSA.
The above models assume that the latent vectors factor independently, which is computationally convenient but unrcalistic if the sequences are related to each other in phylogeny. A more realistic approach thus uses prior I T(22) | (K)((() that conditions the latent vectors parameters, k. Because the latent vectors do not factor independently anymore, mini-bated training can include the sequences but not the latent vectors, which limits the possible size of the data sets. Nonetheless, we show here that such a model is both computationally tractable and practically useful for realistic data sets concerning single protein families.
given phylogenetic tree, and an evolutionary model with latent
4 METHODS
4.1 THE DRAUPNIR MODEL
Figure I shows the corresponding graphical model. A summary of the variables, their dimensions and the notation is given in Tables and g in Appendix A.1
The pseudocode of the Draupnir model is given in Algorithm I
As inputs, we assume (a) a set of ny aligned sequences, each with length nu, organized in the ng x ny matrix, S, and (b) information on their phylogenetic tree in the form of their (ns + nA) x (ng + nA) patristic distance matrix, T.
The latent matrix Z of the model is a matrix with rows (one for each leaf sequence) and n2 columns, where nz is the size of the latent representation of the sequences. In all experiments, ny 30. Each column of Z (with ns elements) is sampled from univariate OU process on the phylogenetic tree, representing the evolution of a hidden trait underlying the sequences along the tree. Each row of Z corresponds to the latent vector of a standard VAE. However, unlike a standard VAE, the latent vectors do not factorize independently.
For each of the nz TOU processes, the parameters of the TOU process, corresponding to in Section 3.1 are sampled from a suitable prior distribution. The TOU process has three nz tree OU processes (one for each column of Z),
parameters: df. À and an- As need to sample ng sets of these three parameters.
For each of the three paramcters, ,,,A, On, we sample hyperparameter (01, 02 a3) from a half-normal distribution with scale parameter equal to one. These hyperparamelers serve as scale parameter for the half-normal priors over of, and an. Given the parameters of the TOU process obtained from the prior described above, an n. X ns covariance matrix can be calculated based on the patristic distance matrix of the leaves, (S,S), We need one such covariance matrix for each of the nz columns of the matrix of latent representations, Z. The clement I, with k. E 1.... ng, of covariance matrix h, with h E 1,... nz. is given by (Hadjipantelis et al. 2012 Jones & Moriarty 2013
Ch,KI 07,1 exp_eepp-110 (os(--Tka/Ma Tk.)/ + o o
(2)
4

---

Published as a conference paper at ICLR 2022
As decoder, bidirectional gated recurrent unit (GRU, Cho et al 2014 with length equal to the alignment length, nt. The imput at each position of the GRU for sequence Sk.: is concatenated vector, consisting of the latent vector kR representing sequence k and the BLOSUM embedding Eg,,, which is the result of applying a fully connected neural network,
N to the BLOSUM vector Vt, describing the amino acid preferences al position i in the MSA (see Section 4.2). For each of the nęş sequences and for each position i, the GRU states are mapped to a logit vector that specifies the probabilities the ne characters using another fully connected ncural network, ! The architecture of the networks is given in Appendix A.2
h E [I,....Z)
02
oCh
0=
x
Ah
2has
E (I....as)
01
On.h
Z,b
i (I....nL)
GRU/NN
Yki. f Ca.,
Sk.:
T(S,S)
u=00(n))
Figure 1: Draupnir graphical model. For notation and information variables and their dimensions see Tables [ and 2 in Appendix AA1 and A.2 Random variables are shown as ellipses, while deterministic quantities are shown as rounded boxes and observed random variables are shown as shaded ellipses. Parameters of priors and other given quantities are shown without boxes. The model contains three plates, respectively corresponding number of dimensions of the Ieaf sequence-specific latent vector Zk. (nz=30), the number of leaf sequences (ng) and the alignment length (nt). "cat" indicates the concatenation of two vectors. are the hyperparameters and on, Tf, À are the parameters of the TOU processes. Dn,,, the covariance matrix that used to sample component of the ng latent vectors from a multivariate Gaussian distribution (with mean 0(ns)). (SS.)) the patristic distance matrix containing the distances between the leaf sequences. Yk,; is the input vector for the GRC that produces the likelihood parameters for leaf sequence k; S..-- V represents the MSA fully connected neural network.
the
nı nc matrix of averaged BLOSUM vectors. NN denotes
4.2 BLOSUM EMBEDDINGS
A BLOSUM matrix B is an ne X substitution matrix used for sequence aligument, where each row contains the log-odds scores of replacing a given character with any of the other characters (Henikoff & Henikoff 1992p. Each position (column) in the MSA is represented by the weighted average of the BLOSUM vectors of the characters in that column (see Algorithm 2 The averaged BLOSUM vectors only need to be precomputed once. In the model, the BLOSUM vectors are processed into BLOSUM embeddings by neural network to provide position-specific information on the MSA, while the latent variables provide sequence-ppecific information (see Algorithm 2
4.3 MODEL IMPLEMENTATION AND TRAINING
Draupnir was implemented in the deep probabilistic programming language Pyro (Bingham et al. 2019) and trained using stochastic variational inference with Pyro's AutoDelfa guide by optimizing the ELBO (Kingına & Welling 2019). resulting in maximum posteriori (MAP) estimates for all parameters. We use Adam (Kingma & Ba 2014h as the optimizer using the default values. From the MAP estimates, we can sample the latent representations of the ancestral nodes (see Section 4.4 and equation 5 in the Appendix). These latent representations are then subsequently decoded to their respective ancestral sequences. We
5

---

Published as a conference paper at ICLR 2022
Algorithm 1 The Draupnir model
Require: Multiple sequence alignment S, patristic distance matrix T
for E 3] do 0j H1((
b Hyperpriors over the TOU process parameters
for h E [1, ...., Rz2] do
> Priors over the parameters of the nz TOU processes
ffh HNN(av) (i)
nnhh HN((a))
λb ~ HN((22)
for h E [1.... nz] do
D Kernels for the ng TOU processes
for k, E fl...., n do
Ch, t of, expl exp(-Th... (8.88) /2a) + oz Sx.
for h E [1.....,zz) do
Prior over treo-strucured latent matrix Z
Zi ~ MVN((0(s)) Chac)
for i E [I, do E NNN
& BLOSUM emboddings
for k E [L...., nsl do
Input vector Y for GRU
for i E (1...., nr. do YA,, f- cat Zk... E.,))
Concatenate sequence- and position-specific vectors
for & E (I,.... ns] do Ha:: GRUs (Y Aea)) for i E [I...., , do
Likelihood parameters (logits) from GRU
>> Bidirectional GRU states
Lk +- NN H (.) Sk,4 ~ Le.
Likelihood at position in sequence
Algorithm 2 Pre-computation of weighted averaged BLOSUM vectors Require: Multiple sequence alignment S
V <- ((12XR))
Initialize BLOSUM weighted average V
do for i E [l,..., n do
b Position in sequence alignment
for k (l,...,77 nss]
Index of leaf sequence
r r Sk
Character position i in lcaf sequence > Add BLOSUM vector corresponding to the amino acid
Vi Vi, + Br
VI. V ii:
Average
ng
also use custom guide to calculate variational posterior (Draupnir-variation) see Appendix A.6) Training details can be found in Appendix[A. [.4 All programs were executed on an Inte Xeon(R) Gold 6136 CPU @ 3.00GHz machine with Quadro RTX 6000 GPU.
4.4 INFERENCE OF THE ANCESTRAL SEQUENCES
In this section, for ease of nolation, let z Z,4..1 denote one of the h E {1,..., nz2 columns the latent representation matrix for both ancestral and leaf sequences, Z (A.s). First, note that can be partitioned latent representations the leaf and ancestral sequences, respectively. The prior p(z) multivariate Gaussian distribution with parameters,
(z) z (1), where (8) and (A) denote the

---

Published as a conference paper at ICLR 2022
10(5) p((A) (((n)) ) 2 )
E(Sss) E(1A)) E(1.))
E(SA)
) ,A= ( (111.5)
(A(8.)) 1(5,A)
A(A.S) 1 )
where A --1. The covariance matrices ( (AAA)), ((A.)) =\t (E(),))) are
respectively obtained from the distance matrices concerning distances within the leaves, within the ancestors and between the ancestors and the leaves, T(SS), T(A,4) and T(A.5), and the TOU process parameters (see Eq. 2
As p(z) is a multivariate Gaussian distribution, we can casily obtain the conditional distribution of 2(4) given 25() as follows (see Bishop (2006), page 689),
p(z) = MMW (z Σ) = P (2(4) (1(51) MDN (2(4) MAls: (Al4.2)) )1)
with
PAIS 14A1 (A(A.A)) -1 1(4.9) ( =
= -1alaa1))
1(4.9),
As MAIS corresponds to the MAP of 2 for any given values of (8) and the TOU process parameters when the ancestral sequences are not observed, the MAP estimate of the latent representation of the ancestral sequences is given by;
z (A) = -- -1((4 ((A.9),(5),M)
In addition, Gaussian approximation of the posterior of (A) based the MAP estimates is given by, (1)) ~ MMN (4))) ((())MAP (114.10))
(3)
The reconstructed ancestral sequences are subsequently obtained by applying the GRU decoder to the MAP estimates of their latent representations, as explained in Section 4.1
4.5 BENCHMARKING
In order assess the accuracy of the ASR we benchmark Draupnir (MAP and Marginal) against state-of-the-art phylogenetic methods. We selected methods that perform ASR using a given tree topology and given patristic distances, including one Bayesian method (PhyloBayes, Lartilot et al. (2013) and three maximum likelihood based methods (PAML,
Yang 2007): ; FastML, Ashkenazy et al. (2012); and IQTree, Nguyen et al. (2015) ) We apply the ASR methods to both protein sequences, and to their corresponding DNA sequences followed by subsequent translation to protein sequences. For Draupnir, we use the protein sequences, which is the harder problem. We use eleven data sets with different numbers of leaves (from 19 to 800), different alignment lengths (from 63 to 558) and with or without gaps (10 and 1 data set respectively). The data sets include eight simulated data sets generated using the software EvolveAGene (Hall 2016) and three data sets with experimentally determined ancestral sequences. Note that the simulated data sets were obtained from factored evolution models. We included large data set with 800 leaves. For Fig. 3/, using Equation 4 ii) samples from the marginal distribution (Draupnir-Mo Marginal in Fig. we report the average identity of 50 samples), using Equation 5 or fii) samples from the variational posterior using amortised guide (Draupnir- Variational, see Appendix A.6) 6 Details on the data sets and training can be found in Appendix A.3 and A.4 respectively.
prediction with Draupnir, we either use i) the most likely sequence (Draupuir-MAP
3
using Equation

---

Published as a conference paper at ICLR 2022
5 RESULTS
5.1 LATENT REPRESENTATIONS AND COEVOLUTION
In order to inspect the quality of the latent representations of the sequences, we use the 3-lactamase family with 32 leaf sequences. We visualize t-SNE projections (Van der Maaten & Hinton 2008) of the latent representations and compare the results with the structure of the phylogenetic tree (see Figure 2 The result indicates that the latent representations represent the structure of the tree and its different clades (subtrees) well, indicating that the TOU process performs well as an informative prior on evolutionary relationships. In Appendix A.6 we show how marginalizing over the latent representations allows Draupnir (marginal and variational) to model coevolution among sites.
Figure 2: Results for the 3-lactamase family with 32 leaves. Left: t-SNE projection of the latent representations of the ancestral and leaf nodes. Right: The phylogenetic tree. Both plots are coloured according clade membership.
5.2 BENCHMARKING RESULTS
In Figure 3 we compare the accuracy of Draupnir (MAP and marginal) with state-of-the thrt ASR methods by plotting the average percent identity between the ancestral sequences as reconstructed by Draupnir and the true sequences. The true ancestral sequences were either experimentally determined or simulated. The description and origin of the data sets can be found in Appendix A.3 Tables with benchmark results are shown in Appendix A.5
IQTree
PAML-CodeML
PhyloBayes
FastML
Draupnir
Draupnir marginal
100
DNA
Protein
Percentage of Identity
80
235
261
71
272
68
477
128
77
1558
99
233
261
71
272
63
477
128
77
+558
99
20
314
314
19
32 35*
50
71'
100
150
200
300
400
800
19*
32 35
50
71
100
150
200
300
400
800
Number of leaves
Figure 3: Comparison of the average percentage identity (y-axis) between predicted and true ancestral sequences for Draupnir (MAP and marginal) and ASR methods for data sets with different number of leaves (x-axis; experimental data sets are indicated with asterisk). Missing points indicate that the ASR method failed to produce results on the given hardware. The alignment size is shown the dotted lines. We compare with ASR methods using the DNA sequences (left) and the corresponding Protein sequences, subsequently translated to and Draupnir-
protein sequences (right). Tables with detailed results for Draupnir-mare MAP can be found in Appendix A.5 For the benchmarking settings see Appendix A.8
8

---

Published as a conference paper at ICLR 2022
5.3 ABLATION STUDIES
In the first ablation study, we investigate the influence of the BLOSUM embeddings by removing them as input to the GRU. Overall, the absence of the BLOSUM embeddings slows down convergence and sometimes make the learning process unstable, but ultimately does not strongly affect accuracy (see Figure
100
Percent identity
Percent idenity 20
Porcent identity
80
%
80
60
to
20
40
20
2000
4000
6000
8000
10000
12000
2000
4000
6000
8000
10000
12000
2000
4000
6000
8000
10000
12000
Number of epochs
Number of epocts
Number of epochs
(a) Leaves, stable case
Ancestors, stable case
Ancestors, unstable case.
Figure 5: BLOSUM embedding ablation study for the 800 leaves data set. For every 100 training epochs, the average percent identity and standard deviation are plotted for all leaves (training sct) or ancestors (test set), respectively. The results obtained with the BLOSUM embedding are shown in dark green (MAP) and light green (marginal, see Equation 3 The results without the BLOSUM embeddings are shown in pink (MAP) and purple (marginal).
In the second ablation study, we investigate the influence of the tree-structured prior by comparing with standard VAE with Gaussian prior. We do this by using diagonal unit covariance matrices for each of the nz columns the latent matrix Z. The rest of the model was identical. We then compare the latent representations obtained for the leaf nodes. The results (see Figure 7 indicate that the standard VAE is not capable of reconstructing the evolutionary relationships well: sequences belonging to the same clade often end up far apart in latent space. This indicates that the influence of the trce-structures prior is substantial. quantitative analysis of this ablation study can be found in Appendix
(a) Phylogenetic tree coloured by clade membership
(b) Standard VAE, leaves
(c) Leaves
(d) Leaves and ancestors
Figure 7: Bottom: t-SNE projections of the latent representations for the SRC-Kinase SH3 domain with 100 leaves, obtained from a standard VAE (left) and Draupnir-marginn (center and right), colored by clade. Note that only the latter model can be used to infer the latent representations of the ancestral sequences. Top: the corresponding tree.
9

---

Published as a conference paper at ICLR 2022
DISCUSSION AND FUTURE WORK
Draupnir demonstrates the potential value of incorporating evolutionary information and evolutionary models explicitly in deep generative models for representation learning of biological sequences. We point out that it is possible to extend the model with additional information beyond sequences, for example backbone angles describing protein structure (Golden et al. 2017) or measurements of protein stability. In future work, extending the model to genomic-size data can be done using inducing points for Gaussian processes, as explored in Jazbec et. al. (2021) and Vikram et al. (2019). The case of a latent phylogenet.cc tree can be addressed using a coalescent point process prior (Lambert & Stadler 2013 Vikram et al. 2019). Finallyy in the large data case, the current simple network architectures can be improved with more expressive compositions such as an MSA transformer (Rao et al. 2021 or deconvolutional model for sequences Hawkins-Hooker ct al., 2021). Finally, we point out that the learned parameters of the TOU processes might ofler interpretable information the evolutionary process.
ACKNOWLEDGMENTS
LSM acknowledges support from the Independent Research Fund Denmark under the grant "Resurrecting ancestral proteins in silico to understand how cancer drugs work". OR and ASA acknowledge support from the Independent Research Fund Denmark under the grant "Deep probabilistic programming for protein structure prediction". We thank Robert Schenck for technical support and contribution of computational resources. We thank the anonymous reviewers for their suggestions and comments.
AVAILABILITY
Draupnir be found httpp:///tttan.coo/yyyssbbeeeetr ASR installed as a python library.
and
REFERENCES
Naila O Alieva, Karen A Konzen, Steven F Field, Ella A Meleshkevitch, Marguerite F. Hunt, Victor Beltran-Ramirez, David Miller, Jörg Wiedenmann, Anya Salih, and Mikhail V Matz. Diversity and evolution of coral fluorescent proteins. PLoS ONE, 3(7):e2680, 2008.
Ethan Alley Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M Church. Unified rational protein engineering with sequence-based deep representation learning. Nature methods, 16(12):1315-1322 2019.
Haim Ashkenazy, Osnat Penn, Adi Moron-Faigenbor Ofir Cohen, Gina Cannarozzi, Oren Zomer, and Tal Pupko. FastML: a web server for probabilistic roconstruction of ancestral sequences. Nucleic acids research, (WI):W550 W584, 2012.
Ahmet Bakan, Anindita Dutta, Wenzhi Mao, Ying Liu, Chakra Chennubhotla, Timothy R Lezon, and Ivet Bahar. Evol and prody for bridging protein sequence evolution and structural dynamics. Bioinformatics, 30:2681-2683, 2014.
Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D Goodman. Pyro: Deep universal probabilistic programming. The journal of machine learning research, 20(1):973 978, 2019.
Christopher Bishop. Pattern recognition and machine learning. Springer, 2006.
Belinda SW Chang, Karolina Jönsson, Manija A Kazmi, Michael J Donoghue, and Thomas P Sakmar. Recreating a functional ancestral archosaur visual pigment. Molecular biology and evolution, (9):1483 1489, 2002.
10

---

Published as a conference paper at ICLR 2022
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409. 1259, 2014.
Megan F Cole and Eric Gaucher. Uuilizing natural diversity to evolve protein function: applications towards thermostabillty. Current opinion in chemical biology, (3):399 406, 2011.
Nicki Skafte Detlefsen, Soren Hauberg, and Wouter Boomsma. What is meaningful representation of protein sequences? arXiv preprint arXiv:2012.02677 2020.
Xinqiang Ding, Zhengting Zou, and Charles L Brooks III. Deciphering protcin evolution and fitness landscapes with latent. space models. Noture communiicalins,, (1):1 13, 2019.
Richard Durbin, Sean Eddy, Anders Krogh, and Graeme Mitchison. Biological sequence analysis: probubilistic models of proteins and nucleic acids. Cambridge university press, 1998.
Mohammed El-Kebir, Layla Oesper, Hannah -icld, and Benjamin J Raphacl. Reconstruction of clonal trees and tumor composition from multi-sample sequencing data. Bioinformalics, (12):i62 170, 2015.
Brian Gaschen, Jesse Taylor, Karina Yusim, Brian Foley, Feng Gao, Dorothy Lang, Vladimir Novitsky, Barton Haynes, Beatrice H Hahn, Tanmoy Bhattacharya, et al. Diversity considerations in HIV-1 vaccine selection. Science, (5577):2354-2388 2002.
Michacl Golden, Eduardo Garcfa-Portugués, Michael Sørensen, Kanti V Mardia, Thomas Hamehryck, and Jotun Hein. A generative angular model of protein structure evolution. Molecular biology and evolution, 34(8):2085 2100, 2017.
Joe G Greener, Lewis Moffat, and David T Jones. Design of metalloproteins and novel protein folds using variational autoencoders. Scientifie reports, pp. 12, 2018.
Pantelis Z Hadjipantelis, Nick S Jones, John Moriarty, David Springate, and Christopher G Knight. Ancestral inference from functional data: statistical methods and numerical examples. arXiv preprint arXiv: 208. 0628, 2012.
Barry G Hall. Comparison of the accuracies of several phylogenetic methods using protein and DNA sequences. Molecular biology and evolution, 22(3):792 802, 2005.
Barry G Hall. Effects of sequence diversity and recombination the accuracy of phylogenetic trees estimated by kSNP. Cladistics, (1):90 19: 2016.
Thomas F Hansen. Stabilizing selection and the comparative analysis of adaptation. Evolution, (1551113411135,, 1997.
Alex Hawkins- Hawkins-Hooker, Florence Depardieu, Sebastien Baur, Guillaume Couairon, Arthur Chen, and David Bikard. Generating functional protein variants with variational autoeneoders. PLoS computational biology, (2):e1008736, 2021.
Steven Henikoff and Jorja G Henikoff. Amino acid substitution matrices from protein blocks. Proceedings of the national academy of sciences, (22):10915- 10919, 1992.
Georg KA Hochberg and Joseph W Thornton. Reconstructing ancient proteins to understand the causes of structure and function. Annual Review of Biophysics, pp. 247- 269, 2017.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of machine learning research, 14(5), 2013.
Edwin Rodríguez Horta, Alejandro Lage- Cage-Castlllanos, and Roberto Mulet. Ancestral sequence reconstruction for co-evolutionary models. arXiv preprint arXiu:2108. 038000, 2021.
11

---

Published as a conference paper at ICLR 2022
Metod Jazbec, Matt Ashman, Vincent Fortuin, Michael Pearce, Stephan Mandt, and Gunnar Rätsch. Scalable Gaussian process variational autoencoders. In International Conference on Artificial Intelligence and Statistics, pp. 3511-3511 PMLR, 2021.
Nick Jones and John Moriarty. Evolutionary inference for function-valued traits: Gaussian process regression on phylogenies. Journal of the Royal Society Interface, 10(76):20120116 2013.
Jeffrey B Joy, Richard H Liang, Rosemary M McCloskey, T Nguyen, and Art. FY Poon. Ancestral reconstruction. PLoS computational biology, 12(7) 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv
preprint or arxiv: 14 11266990, 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint arxiv: 1oXx:111122 611,6114, 2013.
Diederik P Kingma and Max Welling. An introduction to variational autoencoders.
Poundations and trends in machine learning, 12:307-300 -92, 2019.
Jeffrey M Koshi and Richard A Goldstein. Probabilistic reconstruction of ancestral protein sequences. Journal of molecular evolution, 42(2):313-33 320, 1996.
Amaury Lambert and Tanja Stadler. death models and coalescent point processes: The shape and probability of reconstructed phylogenies. Theorelionl population biology, 90:113 128, 2013.
Nicolas Lartillot. A phylogenetic Kalman filter for ancestral trait reconstruction using molecular data. Bioinformatics, pp. 488 2984996,440 2014.
Nicolas Lartillot, Nicolas Rodrigue, Daniel Stubbs, and Jacques Richer. PhyloBayes MPI: Phylogenetic reconstruction with infinite mixtures of profiles in parallel environment. Systematic biology, (4):611-614) 2013.
Emile Mathieu, Charline Le Lan, Chris Maddison, Ryota Tomioka, and Yee Whye Teh. Continuous hierarchical representations with Poincaré variational auto-encoders. arXiv preprint arx1: 060033, 2019.
Jaina Mistry, Sara Chuguransky, Lowri Williams, Matloob Qureshi, Gustavo A Salazar, Erik LL Sonnhammer, Silvio CE Tosatto, Lisanna Paladin, Shriya Raj, Lorna Richardson, et al. Pfam: The protein families database in 2021. Nucleic Acids Research, 49(D1):D412 D419, 2021.
Faruck Morcos, Andrea Pagnani, Bryan Lunt, Arianna Bertolino, Debora S Marks, Chris Sander, Riccardo Zecchina, José Onuchic, Terence Hwa, and Martin Weigt. Direct- coupling analysis of residue coevolution captures native contacts across many protein families. Proceedings of the National Academy of Sciences, 108:E1293 E1301, 2011.
Lam-Tung Nguyen, Heiko A Schmidt, Arndt Von Haeseler, and Bui Quang Minh. IQ-TREE: fast and effective stochastic algorithm for estimating pexiimum-lkkellhoo phy hoggeeiess. Molecular biology and evolution, 32(1):268-22 274, 2015.
Linus Pauling, Emile Zuckerkandl, T Henriksen, and Lövstad. Chemical paleogenetics. Acta chemica Scandinarica, 17:89-S16, 1963.
Ryan N Randall, Caclan E Radford, Kelsey Roof, Divya K Natarajan, and Eric A Gaucher. An experimental phylogeny to benchmark ancestral sequence reconstruction. Nature communications, 7(1):1-6, 2016.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with TAPE. Advances in neural information processing systems, 32:9689, 2019.
12

---

Published as a conference paper at ICLR 2022
Roshan Rao, Jason Liu, Robert Verkuil, Joshua Meier, John F Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA transformer. bioRrio, 2021.
Adam J Riesselman, John B Ingraham, and Debora S Marks. Deep generative models of genetic variation capture the effects of mutations. Nature methods, 15(10):816 822, 2018.
Avery GA Selberg. Eric Gaucher, and David A Liberles. Ancestral sequence reconstruction: From chemical palcogenetics to maximum likelihood algorithms and beyond. Journal of Molecular Evolution, (3):157-164, 2021.
Matthew Spence, Joe A Kaczmarski, Jake W Saunders, and Colin Jackson. Ancestral sequence reconstruction for protein engincers. Current Opinion in Structural Biology, pp. 131-141, 2021.
Max R Tolkoff, Michael E Alfaro, Guy Baele, Philippe Lemey, and Marc A Suchard. Phylogenetic factor analysis. Systematic biology, (3):384 399 2018.
Laurens Van der Maaten and GeofTrey Hinton. Visualizing data using I-SNE. Journal of machine learning research, (11), 2008.
Sharad Vikram, Matthew D Hoffman, and Matthew J Johnson. The LORACs prior for VAEs: Letting the trees speak for the data. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2292-3301. PMLR, 2019.
C Wilson, RV Agalonov, M Hoemberger, Kutter, A Zorba, J Halpin, Buosi, R Otten, D Waterman, DL Theobald, et al. Using ancient protein kinases to unravel a modern cancer drug's mechanism. Science, 347(6224):882-888) 2015.
Ziheng Yang. PAML 4: Phylogenetic analysis by maximum likelihood. Molecular biology and evolution, 24(8):1586 1591, 2007.
Ziheng Yang, Sudhir Kumar, and Masatoshi Nei. A new method of inference of ancestral nucleotide and amino acid sequences. Genetics, 141 14)666064166650 1995.
13

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 NOTATION AND VARIABLES
Name
Description
ng
Number of TOU processes, length of latent vector (30)
Alignment length
ng
Number of leaf sequences
Number of ancestral sequences (n. ns
DC
Number of character types
S
Sequence alignment matrix of leaf sequences
Z Z
Matrix of latent representations of the leaf sequences
Note: We use for notational convenience where possible.
Z(4)
Matrix of latent representations of the ancestral sequences Matrix of latent representations of the leaf and ancestral sequences
Z(4,5)
T
Patristic distance matrix (given)
((5.))
Patristic distance submatrix of distances between leaf sequences (given)
Patristic distance submatrix of distances bet ancest.al sequences (given)
((44)
(4.8)
Patristic distance submatrix of distances between leaf and ancestral sequences (given)
B
BLOSUM matrix (given)
a
Vector of OU hyperprior paramcters
Vector of intensities of inherited variation (TOU process)
Vector of intensities of specific variation (TOU process)
X
Vector of characteristic lenght-scales (TOU process)
C
Tensor of nz TOU process covariance matrices
Matrix of weighted BLOSUM vectors Matrix of BLOSUM embeddings in the model
E
Tensor of BLOSUM embeddings the guide Input tensor to GRU State tensor of GRU Tensor of logits of the sequence characters Parameters of the neural networks and the GRU d-dimensional vector of zeros
Y
L
(4d)
O (mxn)
(m di-dimensional matrix of zeros Gated recurrent unit
GRU
NN
Fully connected neural network Half-normal distribution
HN
MVN
Multivariate Gaussian distribution
cat
Concatenation of two vectors.
Table 1: Variables and notation used for the Draupnir model.
14

---

Published as a conference paper at ICLR 2022
Name
Dimensions
3
Øt. on X
nz
T
(ns + nA) X (ns + nAA)
T(4,))
nA X 7A
(SSS)
US X TLS
T(4.9)
TA X ng
C
nz TIS 115
Z =
ng X T.Z
Z(4)
AA X T2
((A,S)
(ns + TAA) X 11%
Y
ng X n X (nz + nc)
AL X TC
ng X T. X TCC
H
ng RL 60
ng X nL X TC
B
NC TCC
S
TS X TLT.
V
nL X nc
Table 2: Dimensions of variables
A.2 DRAUPNIR SETTINGS
The Draupnir model contains three neural networks (see Algorithm Figure and Figure fully connected network, NN that maps the pre- computed BLOSUM vectors, V, to BLOSUM embeddings, E; bidirectional GRU, (Cho et al. 2014) with a single layer that takes as input the BLOSUM embeddings and the latent vector of the k-th leaf sequence, Zk, and second fully connected network, NN that maps the GRU, states to the logit vectors of the sequence characters. The dimensionality of the state of the bidirectional GRU, is 2 x 60.
Neural network architecture
In the guide, we re-use the neural network architectures described above: NN and GRU, are identical NN(2) and GRU,, respectively, except that the output size of NN is (2xnz) instead of nc; corresponding to the length the mean vector and standard deviation vector of the latent representation. NN is identical to NN
2x60
21x50
1c3
120x60
fc1
50x21
Fe4
60x21
fx2
LogSoftmax
Figure 8: Architectures and dimensions of the neural networks used in Draupnir. Left: Architecture of the bidirectional GRU, (red) and NN(2). Right: Architecture of NN!!) "fe" indicates fully connected layer.
15

---

Published as a conference paper at ICLR 2022
Additional Draupnir settings We chose BLOSUM62 the base substitution matrix, except for the CFP datasets (71 and 35 leaves) where we use PAM70 due to the presence of special amino acids.
A.3 DATA SETS
Descriptions of the eleven data sets used for benchmarking and the leaves-only data set used in the co-evolutionary analysis. All data sets contain insertions and deletions (gaps), except the one in italic (top), which only contains substitutions. Data sets labelled with asterisk only contain the sequence of the root node.
Table 3:
Number
liggmeent
leaves
length
Datasets with experimentally detertmined ancestral sequences
Rondall's Comal fuuresond proteino (CFP) benchnark
225
Randall 2016
"Coral Hhooeeoent proteins (CFP) Faviine aubodade
35
261
Signo yoor tour sequeence from A Nlwves et al. allow roor, node sqquene from
2006
*Coral thressrent proteins (CPP) subcladee
71
272
EvolveAGeued
Hal 2005
simulations
Simulation (-Lactomase
GenBank pecossion
Stmulation SRC-Kinase SEE3 domain
100
G3
GenBani PC 11-1-7
Sumulation Sartuin
150
477
NBCI OL 1ES-111111
Simulation SRC-Kinase SH3 dlomain
200
128
Gin DETRE 1720
Simulation PIGBOS
300
77
NBCI
Ct 17S91851
Simnhstion SRC Kinase SEES domain
800
99
GenBan
PF0000
125
138
PEAM famuly Do. PF
16

---

Published as a conference paper at ICLR 2022
A.4 TRAINING
Table 4: Training settings and running times (Iutel(R) (R) Gold 6136 CPU 3.00GHz, Quadro RTX 6000 GPU). On the largest dataset (800 leaves), make use of Reduce plateau learning scheduler combined with plating to further improve the accuracy results.
Learning rate scheduler
Dataset
Epochs
Plate
size
Model parameters
Running times
Datazels with experimeutally determined ancestral sequences
Randall's Coral
(CFP) bendhmark,
53 h min
Cord Uluerssenn proteins (CFP) Faviiux subelade, 35 leaves Cord Blerrsonnt proteins (CFP) clrde, 71 Inaer
23000
No
55181 52535
23000
71
No
stmulations
Simmlation A-Lactaneee, Simolation SRC SRC-Kius Kinase SIB domaan, 100 leaves
15000
No
min 2 main
21000
100
No
58685
Simalation Sirtuin, 150 laaves
20000
150
No
55985
la mit 23
Simolation SRC Kitase SH3 domain, leaves
22000
200
No
57085
49 min
Simolation PIGBOS,
Simmlution Insalin,
18400
63485
Simolation SRC-Kinase SIB domain, 800 leaves
25000
50
Yes
141005
30 miin
only
PF00400, leaves Marginal
23000
125
No
55205
52 min
Variational
17

---

Published as a conference paper at ICLR 2022
A.5 BENCHMARKING TABLES
Table 5: Benchmarking results using protein sequences.. The table shows the means and the standard deviation of the percent identity for all the predicted ancestral sequences and their corresponding truc soquences. In the case of Draupnir-MAP, the means and standard deviations are calculated using the most likely sequence of each ancestral node. In the case of Draupnir-margnnal samples, they are calculated using 50 samples per ancestral node. "Not available" indicates the ASR method failed to produce results for that data set on the given hardware. The results for the standard ASR methods were obtained using the amino acid sequences.
Nariio of lrone
Aligeneet lergii
Derqeie MAP
samppl
Prykilope
Fus
Rushl'ss
Cimi Froneet pottris (CIPP) dlate
71
373
T1 12-
71-88-100 06
30 111010
61:244855
71 72:09115
(0:9110
Cord Exesed
seledada
211
76.3181334 14
Steokrio Tartursso
S0.1316
67 (7.336.33
5310-500
Stankss caxbasis 91
50
T1
807725422
4:340:1:44 21
10-4122.55
06:123100
01542133
Stanksso Sanda SutT. L00
190
427
$9.00:199
T7.7622.55
12.5757.38
11:3121.3
Not ooddale
191021.0
Stankrsn SEC Krose Ammonn, 309
00.55-
300-7:99
1100- 68
01 01:94:544
41-33311
Storktion borfin Factie Tie
130
538
06 16
70.41-9111 16
16. 765.37
25.31-1174
Vit costale
35 15:00-1 01
Stankkso SRU--kzee 3855 dennin, SCD
931
56.3605083
Not sostadds
Not sosBalds
08/013355
Table 6: Benchmarking results using DNA sequences. The table is similar to Table 5 but the reconstructions for the standard ASR methods were obtained using DNA instead of amino acid sequences. The DNA sequences were subsequently translated into protein sequences before comparison.
Nanion of
Alagazant laplis
sample
PAME-OokMlL
Feabil
Corl Lnesest prstauu (CFP)
T1
272
764118107
14.3521 08
00.5555555
54.0220775
T4:461.00
SL.002)29
Carl Fraeweet pestels (CTT) Fastia relvcate
351
1124
56.95
T-67111
46751130
26
Beo-lattraes
Sterkeinn Cayharls, 5)
50
71
17:31-13
188--35 1/
96 06311633
01:05:100
Sunkkso SBC-Rinee daniin
00.2722.52
Sende SIET Es
195
437
8530-19
15:347255
00.511204
Stankaa SitC-Siasee 88S deanin 250
230
125
86-3225344
80.5325.33 63
177228.17
11/0122:33
444821.3
01/22171
regilstte
Stanksso landis Eactie akk
433
538
75.5122-15
28.24
2--1222:22
Not condald
24.80210
Stenkesn SRC Krowe S11 Amniin sED
350
91 07-13
scstalde
Via ausdall
19:6-110
A.6 CORVOLUTION ANALYSIS
Conventional ASR methods use models that assume factorized evolution (Horta et al. 2021 that is, they assume that each site evolves independently of all other sites in the sequence. In reality, some sites are coupled in evolution, resulting in dependencies between sites. This phenomenon is called epistasis (Hochberg & Thornton 2017). Draupnir is model that goes beyond factorized evolution, and thus potentially models coevolving sites. Here, we analyze to what extend this is indeed the case.
In order to evaluate modelling of coevolution, we make use of direct coupling analysis (DCA) (Morcos et al., 2011). DCA identifics cocvolving pairs of residues that directly influcnce each other by calculating quantity called Direct Information (DI), which is obtained by fitting a Markov random field. We calculated the DI using ProDy (iBakan et al. 2014)). As data set, used 125 sequences from the WB40 domain of the PF00400 lamily Trom the PFAM data base Mistry et al. (021)
The DIs of the leaf sequences sampled from Draupnir at the root node in three different ways (see below). coevolution is at least partially modclled, the DIs of the leaf sequences will be similar (but not completely identical) to the DIs of the sampled sequences at the root node. We sample sequences from Draupnir using three different methods.
as ground truth and are compared with sequences
The first method (MAP) simply uses the MAP estimates of the probability vectors at each position:
18

---

Published as a conference paper at ICLR 2022
a II P a ((MA)))
(4)
where a is an ancestral sequence and 0 (APP) is the MAP estimate of the amino acid probability vector at position for that sequence. This baseline results in independent sites. Picking the most likely amino acid at each position according to the above expression corresponds Draupnir -MAP in Section 5.2
The second method (Draunnir-marginal) makes use of the MAP estimates of the leaf representations but marginalizes over the ancestral representations (see Section 4.4h:
II p
0
(2()))
p(z(a) Z (MA)) dzz
(5)
where the probability vectors 0, are obtained from GRU applied to the latent representation z(a) of the ancestral sequence (see Algorithm L The last factor involves conditional Gaussian distributions (see Equation and Section 4.4 Integrating over z(a) introduces correlations along the sequence. Therefore, this method is in principle capable to model some coevolution as we integrate over the latent representations of the ancestors (while using a point estimate for the latent representations of the leaves). This corresponds to the "Draupnir marginal samples" Section 5.2
In the third method (Variational), make use of guide (Z S)) to obtain variational posterior over the latent representations Z (see Algorithm 3):
/
(z. (()1))1) p(z(a) Z)go S)dz dZ.
II
(6)
In this case, we marginalize over the latent representations of both leaves and ancestors. This method should capture the coevolutionary signal a greater extent than the second method.
The results are shown in Figure As expected the third method, (Variational) does best, while the first method (MAP) does worst. The Variational method improves considerably upon the Marginal method, indicating that replacing the MAP estimate with variational distribution for the latent representations of the leaves has significant. impact. The results indicate that Draupnir indeed to significant extent can capture coevolutionary information.
19

---

Published as a conference paper at ICLR 2022
Algorithm 3 Architecture of the variational guide, (2 S). We use point estimates for the hyperparameters and parameters of the TOU processes, and multivariate diagonal Gaussian distributions for the latent representations. An asterisk indicates the value of point estimate, 6(.) is the Dirac delta function.
Require: Multiple sequence alignment S
for jE (1,2,3) do 0j a)~5(a7)
b Point estimates of TOU hyperprior parameters
for h E {l....
nz} do
> Point estimates of the parameters of the nz TOU processes
Jsh _(0),1)
Onh (() An ~- 8(1,)
for k E ,I...., ns do
ng| for i E [l..... ,.. n2] do
r r Sk
b Amino acid at position i in leaf sequence BLOSUM embedding
(1) F FM. (N & (B,))
for l E [1..... ng] do
H Ha GRU 2 k
D Bidirectional GRU states Mean and standard deviations of Z
c + NN H ZR MVN (m, diag(c))
Leaves
MAP
[1]
40
1:11
40
100
120
120
Marginal
Variational
0.1
60
40
80
00
0.0
40
80
120
40
120
Figure 9: DI values for all position pairs of the WB40 data set obtained from the leaf sequences (upper left) and sequences sampled the root node (other plots) using the MAP, Marginal and Variational methods. We sampled 125 root sequences, which is equal to the number of leaves. The corrclation coefficients between the DIs of the leaves and the sampled sequences are 0.05 (MAP), 0.66 (Marginal) and 0.79 (Variational).
20

---

Published as a conference paper at ICLR 2022
A.7 QUANTITATIVE ANALYSIS OF THE LATENT SPACE REPRESENTATIONS
In order to compare the standard VAE with Draupnir in a quantitative way (sce Figure analyze the correlation between (a) the Euclidean distances between the latent representations of the leaves and (b) the corresponding branch lengths in the phylogenetic tree for both models. The results shown in Figure T0
Standard VAE
Draupnir marginal
Euclidean distance
Euclidean distance
Branch length
Branch length
Figure 10: Comparison of the Euclidean distances between the latent representations of the leaves (y-axis) and the corresponding branch lengths the phylogenetic tree (x-axis). We use the same color scheme in Figure We traverse the tree in level order and assign the colour of the clade of the first leaf. (Left) The standard VAE. The correlation coefficient is 0.79; the Spearman correlation coefficient is 0.85. (Right) Draupnir. The correlation coefficient 0.91; the Spearman correlation coefficient is 0.94.
21

---

Published as a conference paper at ICLR 2022
A.8
BENCHMARK SETTINGS
PAML-CodeML settings PAML-CodeML (provided by Biopython 1.78) was used with
the following settings:
verbose
(includes detailed information of the posterior probabilities per node) (utilize given tree) (amino acids)
runmode
seqtype
2
clock
(no molecular clock, genes are evolving at different rates)
aaDist
getSE
O
RateAncestor
aaRateFile
WAG
method
model
2 (more dn/ds (selection coelficient) ratios per branch)
fixaalpha
(estimate gamma)
alpha
0.5 (initial alpha value for gamma distribution)
fix blenghhh
(keep given branch lengths)
PhyloBayes 4.1 settings
pd -s -T treefile -cat -gtr -d alignmentfile chainname
In the case of PhyloBayes, as recommended, we run 2 Markov chains until the convergence criteria are met. The recommended convergence criteria are minimum effective sample size above 300 and mar diff among the chains below 0.1. Both for evaluation of the convergence of the chains and for sampling the ancestral sequences we use 100 samples.
FastML v3.11 settings
perl fastml
MSA-file aligumentfile seqType aa nue SubMatrix WAG GTR indelReconstrructon ML -Tree treefile
IQ-Tree v2.0.3 settings
For IQ-Tree (multicore version 2.0.3), model choice and settings are automatically optimized by using the options "-m TEST" and "-nt AUTO".
iqtree - aligumentfile -m TEST -asr -le treefile -nt AUTO
22


Reviews:
[{"review_id": "y12VJ2tcSRF", "reviewer": "Reviewer_mia3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n- I think the authors approach is statistically well grounded and applied to the appropriate problem.\n\n- The results in Figure 3 are quite convincing: This approach is state-of-art for ancestral sequence reconstruction.\n\n- A5 Table 5 and Table 6 are quite convincing. I think it is interesting that some of the other approaches can't even run with the size of the multiple sequence alignment and sequences!\n\nCons:\n\n- I think there could be better metrics to assess the structure of the latent variables. Could the pairwise distance of latent variables be compared to the distance when fitting a proper phylogenetic tree? Otherwise this comparison is entirely qualitative with tSNE.\n\nNeutral:\n\n- Can the BLOSUM parameters be estimated while fitting the model instead of being precomputed (i.e. they are free parameters)? If there is ablation (Figure 5), does the GRU still learn the covariance of the utilization of amino acids at that position?\n\n- Could you better describe the approach to sample sequences from the model, as discussed in section 4.5? It would be interesting to use something like beam search to generate even better sequences from the model.\n\n- It would be interesting to compare a direct generative approach to trees to a Gaussian Process Variational Autoencoder (https://arxiv.org/abs/1810.11738) with kernels, in which the kernel encodes distance  in a precomputed tree. \n\n- Since these sequences are aligned (thus all the same length), why does a GRU decoder need to be used? Why can't an MLP be used to decode?", "summary_of_the_paper": "The authors propose replacing the diagonal gaussian prior in a variational autoencoder with an Ornstein-Uhlenbeck process for fitting a generative model to biological sequences. They they apply their approach to ancestral sequence reconstruction.", "summary_of_the_review": "The authors incorporate biologically-motivated priors into deep models for biological sequences. Their approach is statistically well grounded, and they apply their model to a relevant, difficult task in biology and improve upon it.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "l7s0Ii9R3q", "reviewer": "Reviewer_Exob", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "positives:\n\n1) The motivation of the paper is clear and interesting: the standard methods for phylogenetic inference typically assume independent evolution of the characters in the sequence and the paper proposes to use VAEs in combination with the OU process in the latent dimension to capture possible higher order interactions (epistatsis) in the sequence evolution using a (given) tree structure.\n\n2) The description of the steps and derivations are clear and the approach seems novel and interesting.\n\nnegatives:\n\n1) However, the paper falls short in the results section demonstrating the claims and power of the proposed method: It is hard to see the improved performance in terms of accuracy of the method from Figures 2/3/4 given the gap. Adding some standard deviation bars will help a lot in demonstrating this claim. Also seems that the method is mostly outperformed by FastM: the paper would benefit from a discussion on the method of choice for small sequences; is there any reason to pick the proposed method over AutoML when both run. \n\n2) A very important ablation study has been done in Figure 7 to asses the importance of the added tree structure in the latent space, however, the result is purely quantitative and relies on TSNE embedding. A more quantitive analysis similar to Figure 3 would improve the paper.\n\n4) The paper alludes to computational benefits of the method but never completely delve into the computational aspect of the work. If the computational side is really the selling point I suggest adding plots on the time complexity and discussion on why the propsoed method achieves a better scaling compared to others.\n\n5) I think beyond the usual metrics (accuracy, scaling, and time), the paper can investigate more on its strength which is the ability to generate samples informed by the evolutionary process. Therefore, it can be the case that the latent dimension learned in the proposed VAEs provide some evolutionary information beyond the usual phylogenetic trees. An interesting one that he paper is motivated by is the presence of possible epistasis that the conventional methods are missing. Showing such results helps the paper excel the standard methods beyond the accuracy metrics. This is sth that the authors mention in the last sentence of the paper but I think can be extended more. \n\nquestions:\n\nThe algorithm needs the information of the phylogenetic tree as input. How easy is it to construct the tree? and how stable is the proposed generative model to the possible variations in the construction of the phylogenetic tree?\n\n\n\n", "summary_of_the_paper": "The paper develops a new generative model for sequences using Variational AutoEncoders (VAEs) with the key difference that the latent dimension is modeled as a tree-structured Ornstein-Uhlenbeck (OU) process which captures the phylogenetic tree of the sequences.", "summary_of_the_review": "Overall, the paper tries to address an important and interesting problem in learning generative models of homologous sequences informed by their phylogenetic structure. However, the results in particular are not convincing enough to grant acceptance. The paper can be improved by adding evidences of some biological insights gleaned from the model re the evolutionary process behind the sequence. Alternative the paper can focus on the computation aspect and expand on the reasons behind scaling and the extend it can tolerate larger data sets.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "5z7gDya2Ix-", "reviewer": "Reviewer_1Min", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall this is an exciting paper and the field should welcome more work along these lines. The authors have put a lot of work into developing their model and the presentation is quite clear. My main feedback is concerned with the model evaluations and how they relate to overall claims of the paper.\n\nStrengths:\n- The technical novelty of relaxing independent sites is interesting and important. Careful understanding of when it is useful to relax this assumption and how to do it best could yield scientific insights and improve models.\n- The use of a tree-structured OU process over latent space is novel and natural for this problem setting.\n- The authors compare to a range of baselines that seem relevant and appear to have run baselines correctly (although this is important to confirm).\n- The exposition of the model itself is easy to follow and well-written.\n\nWeaknesses:\nThe main weakness of the paper is that the evaluations are not systematic enough to back up certain core claims of the paper. Please note that some of the suggestions I make below are *not* feasible in a rebuttal period, so I offer them for consideration in followup work or future iterations.\n\nScale - The authors claim in the abstract that their method \"scale[s] to larger data sets.\" I see a few key issues here\n1. The authors do not explain why they can't run methods on larger trees. They simply say \"most ASR methods stop working for data sets with more than 200 leaves.\" Is this because of excessive compute time, memory constraints, software failing, etc.? I can't find any information on hardware used for these experiments either, which is essential to claims about scale. \n2. FastML is the best-performing method for nearly all instances where it's run. The authors should subsample large trees to 200 leaves and compare their method trained on the whole tree to FastML on a subset of the tree. The algorithm \"prune down over-represented clades and run FastML\" is a reasonable one for trees with lots of leaves and an important baseline.\n\nRepresentation Quality - In the abstract the authors state that \"Our results and ablation studies indicate that the explicit representation of evolution using a suitable tree-structured prior has the potential to improve representation learning of biological sequences considerably\"\n1. The main evidence I see for this claim is the tSNE clustering plots in figures 2 and 7. These plots are excellent evidence that the tree OU is implemented well, but this evidence alone doesn't suggest major potential for improvement.\n2. A direct comparison to other representation learning methods would ground this claim substantially. Without any comparison to other sequence representations, it's hard to assess. For example, the set of MSAs + DMS studies used by DeepSequence could be turned into trees + DMS studies. Then the correlation of DeepSequence ELBOs and Draupnir ELBOs to fitness could be compared.\n\nIndependent Sites - The authors state in Section 2.2 \"The aim of this work is to go beyond the assumption of independent, factorized evolution.\" This is an exciting aim but I feel that it needs to be evaluated more carefully.\n1. The authors should provide an ablation study of their sequence likelihood where they use a position-wise MLP instead of a GRU. The ASR performance of this model to Draupnir would give a sense for the role of independent sites assumptions in performance. (Note: I do not believe that model is an independent sites model because you still integrate out the latent z's)\n2. For the benchmarks offered, we are given no sense of how violated independent sites is in these protein families. One very helpful baseline would be comparing FastML ASR performance on a full tree versus a tree where known coevolving positions are filtered. Showing that Draupnir has a smaller delta in this setting would provide clearer evidence that independent sites is being relaxed in a useful way.\n\nGeneral Suggestion: Something to consider for future work. It is common for papers that train VAEs on MSAs to claim they have implicitly learned phylogenetic signal in a protein family. This paper has an excellent opportunity to evaluate that claim by comparing models which do not use phylogeny to one which explicitly does!\n\nThings that would improve my score:\n- Reporting results with sequence likelihoods that factorize over positions\n- Running FastML on filtered versions of large trees for comparison\n- Direct comparison of Draupnir to other representations of protein or DNA sequences\n- Taking a family where independent sites is known to be violated and doing a case study on baselines vs Draupnir", "summary_of_the_paper": "The authors introduce a VAE for modeling individual protein families that incorporates phylogenetic trees through an OU process on latent space. They also use a sequence likelihood which does not factorize over positions. The authors claim these two advances represent a more expressive and efficient model of protein evolution and apply it to ancestral sequence reconstruction.", "summary_of_the_review": "This paper is timely and addressing a very important area of work for the protein-ml field. The model is thoughtfully implemented and presented, but the evaluations do not enable strong conclusions to be reached yet. More careful ablations and comparisons are needed to understand implications of this work for scalability, representation learning, and evolutionary modeling.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper describes a VAE model for individual protein families that uses phylogenetic trees through an Ornstein-Uhlenbeck process on latent space. They also use a sequence likelihood which does not factorize over positions. The paper claims these two advances represent a more expressive and efficient model of protein evolution and apply it to ancestral sequence reconstruction.

Strengths:

- The technical novelty of relaxing independent sites is interesting and important
- The use of a tree-structured OU process over latent space is novel and natural for this problem setting
- The exposition of the model itself is easy to follow and well-written
- A statistically well grounded approach

Weaknesses:

- The evaluations do not enable strong conclusions to be reached yet. More careful ablations and comparisons are needed to understand implications of this work for scalability, representation learning, and evolutionary modeling
- Lack of computational cost details

A majority of reviewers voted with high confidence for acceptance. The only reviewer voting for rejection did not respond to the author's rebuttal and did not give strong arguments for rejection. The authors are encouraged to address the reviwers' comments and improve the paper for the camera ready version.
Decision: Accept (Poster)