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

1 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 ( 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 . 2 Background . 2.1 Protein sequences and evolution . 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 ) . 2.2 Ancestral protein reconstruction . 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 . 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-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 . 3 Related work . 3.1 Representation learning of biological sequences . 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 .

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