Ancestral protein sequence reconstruction using a tree-structured Ornstein-Uhlenbeck variational autoencoder
Year: 2021

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 genomicscale data sets and the case of a latent phylogenetic tree.

heading: 1 Introduction
text: 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 (VAEs) (Kingma & Welling, 2013; 2019) have been especially useful for this purpose (Riesselman et al., 2018; Greener et al., 2018). However, current models do not take evolutionary information fully into account, i.e., by relating the sequences belonging to a protein family in a phylogenetic tree and incorporating parameterized evolutionary models (Durbin et al., 1998). To address this problem, we replace the standard multivariate Gaussian prior of a conventional VAE with a tree-structured prior that takes into account a given evolutionary tree. We propose a prior based on the Ornstein-Uhlenbeck Gaussian process on a tree (Hansen, 1997; Jones & Moriarty, 2013). We apply the model to a classic problem in phylogenetics, namely the inference of ancestral sequences. Ancestral sequence reconstruction (ASR), i.e., 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 et al., 2021), modeling tumour evolution (El-Kebir et al., 2015), evaluating virus diversity and vaccine design (Gaschen et al., 2002), understanding drug mechanisms (Wilson et al., 2015) and reconstructing ancient proteins in vitro (Chang et al., 2002; Wilson et al., 2015; Hochberg & Thornton, 2017). As input, we assume a set of nS known, aligned leaf sequences and their phylogenetic tree. The task we want to address is the inference of the nA = nS−1 unknown, ancestral sequences (Joy et al., 2016). We show that our probabilistic model, called Draupnir, is about on par with or better than the accuracy of established ASR methods for a standard experimentallyderived data set (Alieva et al., 2008; Randall et al., 2016) and several simulated data sets. The paper is organised as follows. In Background, we briefly discuss evolution of biological sequences, ancestral sequence reconstruction and the tree-structured Ornstein-Uhlenbeck process. In Related Work, we 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-art 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 a latent phylogenetic tree.

heading: 2 Background

heading: 2.1 Protein sequences and evolution
text: Biological molecules such as proteins and nucleic acids (DNA, RNA) can be characterised by sequences of characters from an alphabet of size nC , where typically nC = 21 for proteins and nC = 5 for nucleic acids (Durbin et al., 1998). These alphabets include one character that represents a gap, which is useful in aligning related sequences in a 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 as nS leaf nodes in a binary tree or phylogeny, where the nA 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 a (nS +nA)× (nS +nA) matrix containing the pairwise distances between all nodes, called the patristic distance matrix, T. In the context of biological sequences, the field of phylogenetics 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 of methods based on heuristics (such as maximum parsimony) or probabilistic, evolutionary models (Joy et al., 2016).

heading: 2.2 Ancestral protein reconstruction
text: The ASR problem amounts to inferring the composition of the nA ancestral sequences from the nS extant sequences, making use of a tractable model of evolution (Joy et al., 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 in the sequence, which is a computationally convenient but unrealistic assumption. For example, in proteins, amino acids are involved in an intricate 3-dimensional network of interactions that can 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. Nonetheless, 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 to go beyond the assumption of independent, factorized evolution by using a 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.

heading: 2.3 The Ornstein-Uhlenbeck process on a phylogenetic tree
text: 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-Uhlenbeck (OU) process (Hansen, 1997; Jones & Moriarty, 2013). As the OU process is a Gaussian process, it has a Gaussian equilibrium distribution, as well as Gaussian marginal distributions. We use an OU process on a 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 to be zero, the TOU process has three parameters: the variation unattributable to the phylogeny or the intensity of specific variation σn, the characteristic length scale of the evolutionary dynamics, λ and the intensity of inherited variation, σf . The covariance function for the corresponding multivariate Gaussian distribution is then given by (Hadjipantelis et al., 2012; Jones & Moriarty, 2013),
Σk,l = σ2f exp ( −Tk,l λ ) + σ2nδk,l (1)
where Tk,l is the patristic distance between nodes k and l in the tree, and the Kronecker delta δk,l = 1 if k = l, and 0 otherwise. The TOU process and related diffusive processes on trees are well-established 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 phylogenetic factor analysis, in which a latent variable under the control of a small number independent univariate Brownian diffusion processes is related to observed traits through a loading matrix; Horta et 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.

heading: 3 Related work

heading: 3.1 Representation learning of biological sequences
text: A VAE (Kingma & Welling, 2013; 2019) is a probabilistic, generative model featuring latent vectors or representations, {z}Nn=1, that are independently sampled from a prior distribution, zn ∼ π(zn). The latent vectors are passed to a neural network (the decoder) with parameters θ, leading to a likelihood, xn ∼ pθ(xn | NNθ(zn)), for the data, {x}Nn=1. The prior is typically a standard multivariate Gaussian distribution, but other priors have been used, such as distributions on the Poincaré ball to recover hierarchical structures (Mathieu et al., 2019). The posterior distribution p(zn | xn) is intractable, but can be approximated with a variational distribution or guide, qφ(zn | NNφ(xn)), involving a second neural network (the encoder). Point estimates of the parameters θ and φ are obtained by maximizing the evidence lower bound (ELBO),
Lθ,φ(x) = Eq [ log ( pθ(x, z) qφ(z | x) )] ,
using stochastic gradient ascent (Hoffman et al., 2013). VAEs are increasingly used for representation learning of biological sequences (Detlefsen et 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 a VAE can capture evolutionary relationships between sequences. The above models do not represent the phylogenetic 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 et al., 2020). Hawkins-Hooker et al. (2021) use a VAE with a convolutional encoder and decoder, combining upsampling and autoregression, without relying on a MSA. The above models assume that the latent vectors factor independently, which is computationally convenient but unrealistic if the sequences are related to each other in a phylogeny. A more realistic approach thus uses a prior π({z}Nn=1 | τ, κ)π(κ) that conditions the latent vectors on a given phylogenetic tree, τ , and an evolutionary model with latent parameters, κ. Because the latent vectors do not factor independently anymore, mini-batch 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.

heading: 4 Methods

heading: 4.1 The Draupnir model
text: The pseudocode of the Draupnir model is given in Algorithm 1; Figure 1 shows the corresponding graphical model. A summary of the variables, their dimensions and the notation is given in Tables 1 and 2 in Appendix A.1. As inputs, we assume (a) a set of nS aligned sequences, each with length nL, organized in the nS × nL matrix, S, and (b) information on their phylogenetic tree in the form of their (nS + nA)× (nS + nA) patristic distance matrix, T. The latent matrix Z of the model is a matrix with nS rows (one for each leaf sequence) and nZ columns, where nZ is the size of the latent representation of the sequences. In all experiments, nZ = 30. Each column of Z (with nS elements) is sampled from a 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 parameters: σf , λ and σn. As we use nZ tree OU processes (one for each column of Z), we need to sample nZ sets of these three parameters. For each of the three parameters, σf , λ, σn, we sample a hyperparameter (α1, α2, α3) from a half-normal distribution with scale parameter equal to one. These hyperparameters serve as scale parameter for the half-normal priors over σf , λ and σn. Given the parameters of the TOU process obtained from the prior described above, an nS × nS covariance matrix can be calculated based on the patristic distance matrix of the leaves, T(S,S). We need one such covariance matrix for each of the nZ columns of the matrix of latent representations, Z. The element k, l, with k, l ∈ 1, ..., nS , of covariance matrix h, with h ∈ 1, ..., nZ , is given by (Hadjipantelis et al., 2012; Jones & Moriarty, 2013),
Ch,k,l = σ2f,h exp(−Tk,l/λh) + σ2n,hδk,l. (2)
As decoder, we use a bidirectional gated recurrent unit (GRU, Cho et al. (2014)) with length equal to the alignment length, nL. The input at each position i of the GRU for sequence Sk,: is a concatenated vector, consisting of the latent vector Zk,: representing sequence k, and the BLOSUM embedding Ei,:, which is the result of applying a fully connected neural network, NN(1)θ , to the BLOSUM vector Vi,: describing the amino acid preferences at position i in the MSA (see Section 4.2). For each of the nS sequences and for each position i, the GRU states are mapped to a logit vector that specifies the probabilities of the nC characters using another fully connected neural network, NN(2)θ . The architecture of the networks is shown in Appendix A.2.

heading: 4.2 BLOSUM embeddings
text: A BLOSUM matrix B is an nC×nC substitution matrix used for sequence alignment, where each row contains the log-odds scores of replacing a given character with any of the other characters (Henikoff & Henikoff, 1992). Each position (column) in the MSA is represented by the 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 a neural network to provide position-specific information on the MSA, while the latent variables provide sequence-specific information (see Algorithm 2).

heading: 4.3 Model implementation and training
text: Draupnir was implemented in the deep probabilistic programming language Pyro (Bingham et al., 2019) and trained using stochastic variational inference with Pyro’s AutoDelta guide, optimizing the ELBO (Kingma & Welling, 2019) and resulting in maximum a posteriori (MAP) estimates for all parameters. We use Adam (Kingma & Ba, 2014) 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). These latent representations are then subsequently decoded to their respective ancestral sequences. We also use a custom guide to calculate a variational posterior (see Appendix A.6). Training details can be found in Appendix A.4. All programs were executed on an Intel(R) Xeon(R) Gold 6136 CPU @ 3.00GHz machine with a Quadro RTX 6000 GPU.

heading: 4.4 Inference of the ancestral sequences
text: In this section, for ease of notation, let z ≡ Z(A,S):,h denote one of the h ∈ {1, . . . , nz} columns of the latent representation matrix for both ancestral and leaf sequences, Z(A,S). First, note that z can be partitioned as z = (z(S), z(A)), where z(S) and z(A) denote the latent representations of the leaf and ancestral sequences, respectively. The prior p(z) is a multivariate Gaussian distribution with parameters,
µ = ( µ(S)
µ(A)
) = (
0(nZ ) 0(nA)
) ,Σ = ( Σ(S,S) Σ(S,A) Σ(A,S) Σ(A,A) ) , Λ = ( Λ(S,S) Λ(S,A) Λ(A,S) Λ(A,S) ) ,
where Λ ≡ Σ−1. The covariance matrices Σ(S,S),Σ(A,A),Σ(A,S) ≡ ( Σ(S,A) )T are respectively obtained from the distance matrices concerning distances within the leaves, within the ancestors and between the ancestors and the leaves, T(S,S),T(A,A) and T(A,S), and the TOU process parameters (see Eq. 2).
Algorithm 1 The Draupnir model Require: Multiple sequence alignment S, patristic distance matrix T
for j ∈ {1, 2, 3} do . Hyperpriors over the TOU process parameters αj ∼ HN (1)
for h ∈ {1, ..., nZ} do . Priors over the parameters of the nZ TOU processes σf,h ∼ HN (α0) σn,h ∼ HN (α1) λh ∼ HN (α2)
for h ∈ {1, ..., nZ} do . Kernels for the nZ TOU processes for k, l ∈ {1, ..., nS} do
Ch,k,l ← σ2f,h exp(−T (S,S) k,l /λh) + σ2n,hδk,l
for h ∈ [1, . . . , nZ ] do . Prior over tree-strucured latent matrix Z Z:,h ∼MVN (0(nS),Ch,:,:)
for i ∈ [1, . . . , nL] do . BLOSUM embeddings Ei,: ← NN(1)θ (Vi,:)
for k ∈ [1, . . . , nS ] do . Input vector Y for GRU for i ∈ [1, . . . , nL] do
Yk,i,: ← cat(Zk,:,Ei,:) . Concatenate sequence- and position-specific vectors
for k ∈ [1, . . . , nS ] do . Likelihood parameters (logits) L from GRU Hk,:,: ← GRUθ(Yk,:,:) . Bidirectional GRU states for i ∈ [1, . . . , nL] do
Lk,i,: ← NN(2)θ (Hk,i,:) Sk,i ∼ Categorical(Lk,i,:) . Likelihood at position i in sequence k
Algorithm 2 Pre-computation of averaged BLOSUM vectors Require: Multiple sequence alignment S
V← 0(nL×nc) . Initialize BLOSUM weighted average V for i ∈ [1, . . . , nL] do . Position in sequence alignment
for k ∈ [1, . . . , nS ] do . Index of leaf sequence r ← Sk,i . Character at position i in leaf sequence k Vi,: ← Vi,: + Br,: . Add BLOSUM vector corresponding to the amino acid Vi,: ← 1nS Vi,: . Average
As p(z) is a multivariate Gaussian distribution, we can easily obtain the conditional distribution of z(A) given z(S) as follows (see Bishop (2006), page 689),
p(z) =MVN (z | Σ)⇒ p ( z(A) | z(S) ) =MVN ( z(A) | µA|S , ( Λ(A,A) )−1) ,
with
µA|S = µA − ( Λ(A,A) )−1 Λ(A,S) ( z(S) − µ(S) ) = − ( Λ(A,A) )−1 Λ(A,S)z(S).
As µA|S corresponds to the MAP of z(A) for any given values of z(S) 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),MAP = − ( Λ(A,A) )−1 Λ(A,S)z(S),MAP.
In addition, a Gaussian approximation of the posterior of z(A) based on the MAP estimates is given by,
z(A) ∼MVN ( z(A) | z(A),MAP, ( Λ(A,A) )−1) . (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.

heading: 4.5 Benchmarking
text: In order to assess the accuracy of the ASR we benchmark Draupnir 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, Lartillot 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 nine data sets with different numbers of leaves (from 19 to 800), different alignment lengths (from 63 to 477) and with or without gaps (8 and 1, respectively). The data sets include six 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 a large data set with 800 leaves. For prediction with Draupnir, we either use the most likely sequence, sample sequences from the MAP estimates using Equation 3 (we report the average identity of 50 samples)) or sample sequences from the variational posterior using an amortised guide (Appendix A.6). Details on the data sets and training can be found in Appendix A.3 and A.4, respectively.

heading: 5 Results

heading: 5.1 Latent representations and coevolution
text: In order to inspect the quality of the latent representations of the ancestral and leaf sequences, we use the β-lactamase family with 32 leaf sequences. We visualize t-SNE projections (Van der Maaten & Hinton, 2008) 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 to model coevolution among sites.

heading: 5.2 Benchmarking Results
text: In Figure 3, we compare the accuracy of Draupnir with state-of-the-art 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.

heading: 5.3 Ablation studies
text: 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 5).
In the second ablation study, we investigate the influence of the tree-structured prior by comparing with a standard VAE with a Gaussian prior. We do this by using diagonal unit covariance matrices for each of the nZ columns of 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 tree-structures prior is substantial. A quantitative analysis of this ablation study can be found in Appendix A.8.

heading: 6 Discussion and future work
text: 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 phylogenetic tree can be addressed using a coalescent point process prior (Lambert & Stadler, 2013; Vikram et al., 2019). Finally, 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 a deconvolutional model for sequences (Hawkins-Hooker et al., 2021). Finally, we point out that the learned parameters of the TOU processes might offer interpretable information on the evolutionary process.

heading: A.1 Notation and variables

heading: A Appendix

heading: A.2 Draupnir settings
text: Neural network architecture The Draupnir model contains three neural networks (see Algorithm 1, Figure 1 and Figure 8): a fully connected network, NN(1)θ , that maps the precomputed BLOSUM vectors, V, to BLOSUM embeddings, E; a 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 a second fully connected network, NN(2)θ , that maps the GRU states to the logit vectors of the sequence characters. The dimensionality of the state of the bidirectional GRU is 2× 60.
In the guide, we re-use the neural network architectures described above: NN(2)φ and GRUφ are identical to NN(2)θ and GRUθ, respectively, except that the output size of NN (2) φ is (2×nZ) instead of nC , corresponding to the length of the mean vector and standard deviation vector of the latent representation. NN(1)φ is identical to NN (1) θ .
Additional Draupnir settings We chose BLOSUM62 as 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.

heading: A.3 Data sets

heading: A.4 Training

heading: A.5 Benchmarking tables

heading: A.6 Coevolution analysis
text: 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 a 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 identifies coevolving pairs of residues that directly influence each other by calculating a quantity called Direct Information (DI), which is obtained by fitting a Markov random field. We calculated the DI using ProDy (Bakan et al. (2014)). As data set, we used 125 sequences from the WB40 domain of the PF00400 family from the PFAM data base. The DIs of the leaf sequences serve as ground truth and are compared with sequences sampled from Draupnir at the root node in three different ways (see below). If coevolution is at least partially modelled, 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. The first method (MAP) simply uses the MAP estimates of the probability vectors at each position:
a ∼ nL∏ i=1 p ( ai | θ(MAP)i ) ,
where a is an ancestral sequence and θ(MAP)i is the MAP estimate of the amino acid probability vector at position i for that sequence. This baseline results in independent
sites. Picking the most likely amino acid at each position according to the above expression corresponds to “Draupnir MAP” in Section 5.2. The second method (Marginal) makes use of the MAP estimates of the leaf representations but marginalizes over the ancestral representations (see Section 4.4):
a ∼ ∫ (nL∏ i=1 p ( ai | [ θ ( z(a) )] i )) p(z(a) | Z(MAP))dz(a),
where the probability vectors θi are obtained from a GRU applied to the latent representation z(a) of the ancestral sequence a (see Algorithm 1). The last factor involves conditional Gaussian distributions (see Equation 3 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 “Draupnir samples” in Section 5.2. In the third method (Variational), we make use of a guide qφ(Z | S) to obtain a variational posterior over the latent representations Z (see Algorithm 3):
a ∼ ∫ (nL∏ i=1 p ( ai | [ θ ( z(a) )] i )) p(z(a) | Z)qφ(Z | S)dz(a)dZ.
In this case, we marginalize over the latent representations of both leaves and ancestors. This method should capture the coevolutionary signal to a greater extent than the second method. The results are shown in Figure 9. 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 a variational distribution for the latent representations of the leaves has significant impact. The results indicate that Draupnir indeed to a significant extent can capture coevolutionary information.
Algorithm 3 Architecture of the guide. 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 a point estimate. δ(.) is the Dirac delta function. Require: Multiple sequence alignment S
for j ∈ {1, 2, 3} do . Point estimates of TOU hyperprior parameters αj ∼ δ(α∗j )
for h ∈ {1, ..., nZ} do . Point estimates of the parameters of the nZ TOU processes σf,h ∼ δ(σ∗f,h) σn,h ∼ δ(σ∗n,h) λh ∼ δ(λ∗h)
for k ∈ [1, . . . , nS ] do for i ∈ [1, . . . , nL] do
r ← Sk,i . Amino acid at position i in leaf sequence k Fk,i,: ← NN(1)φ (Br,:) . BLOSUM embedding
for k ∈ [1, . . . , nS ] do Hk,:,: ← GRUφ(Fk,:,:) . Bidirectional GRU states m, c← NN(2)φ (Hk,−1,:) . Mean and standard deviations of Zk,: Zk,: ∼MVN (m,diag(c))
A.7 Quantitative analysis of the latent space representations
In order to compare the standard VAE with Draupnir in a quantitative way (see Figure 7), we 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 are shown in Figure 10.

heading: A.8 Benchmark settings
text: PAML-CodeML settings PAML-CodeML (provided by Biopython 1.78) was used with the following settings:
verbose 2 (includes detailed information of the posterior probabilities per node) runmode 0 (utilize given tree) seqtype 2 (amino acids) clock 0 (no molecular clock, genes are evolving at different rates) aaDist 0 getSE 0 RateAncestor 1 aaRateFile WAG method 1 model 2 (more dn/ds (selection coefficient) ratios per branch) fix alpha 0 (estimate gamma) alpha 0.5 (initial alpha value for gamma distribution) fix blength 0 (keep given branch lengths)
PhyloBayes 4.1 settings
pd -s -f -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 max 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.

heading: FastML v3.11 settings
text: perl fastml –MSA-file alignmentfile –seqType aa or nuc –SubMatrix WAG or GTR –indelReconstruction ML –Tree treefile

heading: IQ-Tree v2.0.3 settings
text: 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 -s alignmentfile -m TEST -asr -te treefile -nt AUTO

Review summaries:
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.