Evaluating probabilistic programming and fast variational Bayesian inference in phylogenetics

Attempts to accelerate MCMC

The tree topology is arguably the most difficult part of a phylogenetic model to estimate. The number of possible topologies for binary trees grows super-exponentially in the number of sequences (tips). As a result, sophisticated MCMC proposals are required to efficiently explore the enormous space of this parameter. Typically, MCMC samplers propose a new topology using simple rearrangements of the current topology such as nearest neighbor interchange and subtree-prune-regraft operations. However these relative small, local moves through the state space can lead to inadequate sampling, especially when the posterior is multimodal (Whidden & Matsen, 2015). Several advances in tree topology sampling have been proposed, such as the use of a conditional clade probability distribution to guide proposals (Höhna & Drummond, 2011). Another recently proposed approach involved an extension of Hamiltonian Monte Carlo, which uses gradient information to guide proposals, to work in the discrete state space of tree topologies (Dinh et al., 2017).

When it comes to continuous parameters, the simplest proposals update a single parameter at a time (e.g. via a multiplier or sliding window proposals). This comes at the expense of ignoring correlation among parameters. Building upon research on adaptive MCMC (Haario, Saksman & Tamminen, 2001; Roberts & Rosenthal, 2009), recent efforts have implemented multivariate proposals to update blocks of continuous phylogenetic model parameters (Baele et al., 2017). While Baele et al. (2017) showed that this proposal is statistically more efficient than standard proposals in terms of effective sampling size per unit of time, this proposal is currently not available for updating branch lengths in the BEAST packages (Suchard et al., 2018; Bouckaert et al., 2019). Finally, in order to increase the efficiency of an MCMC sampler, Aberer, Stamatakis & Ronquist (2015) developed an independence sampler to update branch lengths in an informed way. Unfortunately, none of these approaches provide the magnitude of improvement required to carry out Bayesian phylogenetic inference on the datasets of 1,000s of sequences which have now become commonplace.

Variational inference: an alternative to MCMC

One alternative to MCMC that has been proposed for Bayesian inference of model parameters is variational Bayes (VB) (Jordan et al., 1999; Wainwright & Jordan, 2008). Like MCMC, VB is a technique used to approximate intractable integrals. The main idea behind variational inference is to transform the posterior approximation of an intractable model into an optimization problem using a family of approximate densities that are tractable. The aim is to find the member of that family with the minimum Kullback–Leibler (KL) divergence to the posterior distribution of interest. Although variational inference does not provide the guarantee that MCMC would of generating samples from the true posterior distribution, variational inference tends to be much faster than MCMC since it relies on fast optimization methods. It is common for VB inference to require 100-fold less compute time than MCMC to approximate the same posterior distribution, although the quality of the approximation may not be as good. Variational inference has become popular in machine learning, as evidenced by the large number of software libraries implementing it (e.g. Stan (Carpenter et al., 2017), PyMC (Salvatier, Wiecki & Fonnesbeck, 2016), TensorFlow (Abadi et al., 2015), Edward (Tran et al., 2016)).

To date, VB has received relatively little attention in the field of phylogenetics, possibly due to the presence of discrete model components such as the tree topology which can not be inferred using standard VB algorithms. However, in recent years interest in VB for phylogenetics has begun to grow. Using the Jukes–Cantor (JC69) nucleotide substitution model and fixed topologies, Fourment et al. (2018b) employed VB to approximate the marginal likelihood of fixed topologies. Zhang & Matsen (2019) have recently shown that VB is a superior method to infer phylogenetic tree posteriors in the classical case of unrooted trees with branch lengths and the JC69 substitution model. Dang & Kishino (2019) used variational inference to approximate the CAT-Poisson model for amino acid data (Lartillot & Philippe, 2004). In that work the authors used a partial Gibbs sampling strategy to update the topology and therefore only the continuous parameters were approximated with VB.

In the present work, we describe implementations of several phylogenetic models for nucleotide datasets using the Stan language (Carpenter et al., 2017). These models are substantially richer than those described in previous studies of variational inference for phylogenetics as they include more general nucleotide substitution models, such as the general time reversible (GTR) substitution model, and rate heterogeneity across sites. We also implemented molecular clock and coalescent models to infer the substitution rate and divergence times of heterochronous and homochronous sequence data. As with previous studies, we consider the topology to be fixed (Fourment et al., 2018b). Despite this limitation, implementing the type of model we describe here under a fixed tree is a useful step forward toward fast Bayesian inference of complex phylogenetic models, as it helps us understand the quality and speed of posterior approximation that can be achieved using a generic modeling language like Stan. We compared the performance of the variational approximations to that obtained with the widely used MCMC-based phylogenetic software package BEAST, run on a fixed tree topology (Suchard et al., 2018; Bouckaert et al., 2019). Because our models are implemented in the Stan language, we are able to carry out inference using any of the inference engines available in Stan, and so we used the No-U-Turn Sampler (NUTS) (Hoffman & Gelman, 2014) to further validate our phylogenetic model implementations. Finally, we compared the mean-field approximations obtained with phylostan to those obtained with a highly specialized implementation of mean-field variational inference in the C language software physher (Fourment & Holmes, 2014; Fourment et al., 2018b) available from https://www.github.com/4ment/physher.

Variational inference

The main idea behind variational inference is to transform posterior approximation into an optimization problem using a family of approximate densities. The aim is to find the member of that family with the minimum KL divergence to the posterior distribution of interest: $${\text{ϕ}}^{\ast }={\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}_{\text{ϕ}\in \text{Φ}}\mathrm{K}\mathrm{L}(q(\text{θ};\text{ϕ})\parallel \mathrm{p}(\text{θ}|\mathrm{D},\tau )),$$where q(θ; ϕ) is the variational distribution parameterized by a vector ϕ ∈ Φ, while θ are the model parameters (e.g. branch lengths), D are the aligned sequence data, τ is a fixed tree topology, and KL is defined as

$$\mathrm{K}\mathrm{L}(q\parallel p)={\int }_{\text{θ}}q(\text{θ};\text{ϕ})\log {\displaystyle \frac{\mathrm{q}(\text{θ};\text{ϕ})}{\mathrm{p}(\text{θ}|\mathrm{D},\tau )}.}$$

To minimize the KL divergence, we first rewrite the KL equation: $$\begin{array}{rcl}\text{KL}(q(\text{θ};\text{ϕ})\Vert p(\text{θ}\mid D,\tau )& =& \mathbb{E}[\log q(\text{θ};\text{ϕ})]-\mathbb{E}[\log p(\text{θ}\mid D,\tau )]\\ & =& \mathbb{E}[\log q(\text{θ};\text{ϕ})]-\mathbb{E}[\log p(\text{θ},D\mid \tau )]+\log p(D\mid \tau ),\end{array}$$ where the expectations are taken with respect to the variational distribution q. The third term log p(D | τ) on the right hand side of the last equality is a constant with respect to the variational distribution so it can be ignored for the purpose of the minimization. After switching the sign of the other two terms, the minimization problem can be framed as a maximization problem of the evidence lower bound (ELBO) function

The ELBO is easier to calculate than the KL divergence as it does not involve computing the intractable posterior nnormalization term p(D | τ). We entertained two classes of variational distributions: mean-field and full-rank Gaussian distributions. The mean-field approximation assumes a complete factorization of the distribution over each of the N parameters of the model (e.g. branch lengths, GTR parameters) and each factor is governed by its own variational parameters ϕ_i: $$q({\text{θ}}_1,\dots ,{\text{θ}}_N;\text{ϕ})=\prod _{i=1}^{N}q({\text{θ}}_i;{\text{ϕ}}_i),$$where q(θ_i; ϕ_i) is a Gaussian density and ϕ_i = (μ_i, σ_i).

The full-rank approximation is $$q({\text{θ}}_1,\dots ,{\text{θ}}_N;\text{ϕ})=\mathcal{N}(\text{θ};\text{ϕ}),$$where the term ϕ = (μ, Σ) concatenates the mean vector μ and covariance matrix Σ of a multivariate Gaussian distribution. The mean vector and the covariance matrix contribute respectively N and N(N + 1)/2 parameters to the full-rank variational model.

When provided with a model definition via phylostan, the Stan software is able to estimate the variational parameters by applying stochastic gradient ascent in a black box approach (Ranganath, Gerrish & Blei, 2014; Kucukelbir et al., 2015).

In phylogenetics, parameters are usually constrained, for example, branch lengths and substitution rates are non-negative while nucleotide frequencies sum to one. Since the support of the Gaussian distribution is ℝ, Stan automatically transforms constrained parameters into unconstrained variables. For example, branch lengths are log-transformed such that they live in the real coordinate space. Stan maintains a set of transformations for common constraints (e.g. lower and/or upper bound constraints) and their corresponding Jacobians. Further information on the library of transformations can be found in the online Stan manual https://mc-stan.org/docs/2_18/reference-manual/variable-transforms-chapter.html.

Models with clocks

The simplest phylogenetic models do not distinguish between evolutionary rate and time, and for those models the only constraint on the branch lengths of a tree is that they must be non-negative. Since our models assume that branches are independent, the non-negativity is easily accommodated in a Stan model via Stan’s application of a log transform on the branch length parameters.

On the other hand, trees in molecular clock models are constrained so that parent nodes must be older than any of their descendent nodes. Here we consider that time is going backwards: the earliest taxon was sampled at time 0 and the age of every internal node is greater than 0. In the case of homochronous sequence data (i.e. every sequence was collected at the same time), the age of an internal node can be reparameterized as a proportion of the age of its parent, except for the root age. For each internal node i (excluding the root), the height h_i of the node is reparameterized as h_i = p_i h_pa(i) where the subscript pa(i) denotes the parent of node i and 0 < p_i < 1. The height of the root is constrained to be non-negative.

In the case of heterochronous sequence data, the age of the sequence data at the tips of the tree must to be taken into account. For each internal node i (excluding the root), the height h_i of the node can be reparameterized as h_i = h_d(i) + p_i (h_pa(i) − h_d(i)) where the subscript d(i) denotes the earliest taxon in the set of descendant nodes of node i. The height of the root is now constrained to be greater than the earliest taxa.

The reparameterized model consists of a vector of proportions $p=(p_1,\dots ,p_{N-2})$, which Stan automatically transforms using a log-odds transform.

This reparametrization of the node heights is also used in the PAML software package (Yang, 2007). In our case, the transformation of the node ages requires an adjustment to the joint density. Specifically, in the multivariate case, the density is scaled by the absolute value of the Jacobian of the inverse of the transforms. The Jacobian of homochronous data is a triangular matrix for which the determinant is ${\prod }_{i=1}^{N-2}{h}_{pa(i)}$. For heterochronous data, the Jacobian is also triangular for which the determinant is ${\prod }_{i=1}^{N-2}{h}_{pa(i)}-h_{d(i)}$.

Under these reparameterizations, different ranked topologies will be evaluated during the gradient ascent optimization of the variational parameters.

Rate heterogeneity among sites

The gamma distribution is commonly used to model rate heterogeneity across sites in phylogenetic models (Yang, 1994, 1996). However, it is currently not possible to use the quantile approximation of the gamma distribution proposed by Yang (1994) since the Stan language does not provide inverse cumulative density functions (CDF). Instead we opted to use the Weibull distribution to describe rate heterogeneity across sites. The Weibull distribution has the closed-form inverse CDF F⁻¹(p | λ, k) = λ ( − ln(1 − p))^1/k. Fixing the scale parameter λ to 1, the shape of the Weibull distribution is determined by the shape parameter k. When k ≤ 1 the Weibull distribution is skewed, suggesting strong rate heterogeneity. When k > 1 the distribution is bell shaped, suggesting no or low rate heterogeneity. Since the mean of the Weibull distribution must be equal to 1 in order to preserve branch length interpretation (i.e. expected number of substitutions per site) we rescale the rates $r_1\dots r_K$ such that ${\sum }_{i=1}^K{p}_i{r}_i=1$. Such rescaling is also necessary when a proportion of invariant sites is included in the model.

We also implemented rate heterogeneity across sites using a general discrete distribution with 2K − 2 parameters, by specifying K rates $0\le r_1<r_2\dots <r_K$ and probabilities p_i = Pr{r = r_i} with ${\sum }_{i=1}^K{p}_i{r}_i=1$ and ${\sum }_{i=1}^K{p}_i=1$ (Kosakovsky Pond & Frost, 2005). The values of r_i are constrained to be ordered from smallest to largest merely for the sake of model identifiability in Stan. This distribution is parameter rich compared to the discretized Weibull distribution scheme and as a result it may be more difficult to fit to data.

Datasets and validation

To evaluate the accuracy of variational inference we applied phylostan to three empirical datasets and compared the posterior approximations to those obtained with BEAST2 (Bouckaert et al., 2019), BEAST (Suchard et al., 2018), or MrBayes (Ronquist et al., 2012). The first dataset is made of 69 human influenza virus haemagglutinin nucleotide sequences isolated between 1981 and 1998. The dataset was analyzed under the Hasegawa, Kishino and Yano (HKY) substitution model allowing for rate variation among sites (Weibull distribution with four rate classes). We used a strict clock and a constant size coalescent prior on the tree. The topology used in the phylostan and BEAST2 analyses was drawn randomly from a sample of trees generated by a preliminary analysis with BEAST2 without topological constraints.

The second dataset comprises 63 RNA sequences of type 4 from the E1 region of the hepatitis C virus (HCV) genome that were isolated in 1993. As in previous studies (Pybus et al., 2001; Faulkner et al., 2018), the substitution rate was fixed to 7.9 × 10⁻⁴ substitutions per site per year. We used the GTR substitution model with Weibull distributed rate heterogeneity (four categories). We used the Bayesian skyride tree prior and a Gaussian Markov random field prior for the effective population size trajectory (Minin, Bloomquist & Suchard, 2008). The topology used in the phylostan and BEAST (Suchard et al., 2018) analyses was also drawn randomly from a sample of trees generated by a preliminary analysis with BEAST without topological constraints. We assigned a gamma prior with rate and scale equal to 0.005 on the precision parameter.

The last dataset consists of twenty-seven 18S rRNA sequences of length 1949. This dataset, commonly referred to as DS1, has been studied several times and has become a de facto standard dataset for evaluating MCMC methods (Hedges, Moberg & Maxson, 1990; Lakner et al., 2008; Fourment et al., 2018b; Whidden and Matsen IV, 2015; Whidden et al., 2018). DS1 was analyzed under the Jukes Cantor substitution model (JC69) without a clock. In (Fourment et al., 2018b), the authors used very long runs of MrBayes (Ronquist et al., 2012) to estimate the posterior distribution of the tree topologies (see corresponding paper for more details). Akin to that study, we focused on the 42 tree topologies contained in the 95% credible set and renormalized their posterior probabilities so that they sum to one. For each topology we then estimated its ELBO under the mean-field approximation using either Stan or physher (Fourment & Holmes, 2014), and converted these approximations of the marginal likelihood to a posterior probability. Specifically, we use the ELBO estimate ${\hat{p}}_{\mathrm{E}\mathrm{L}\mathrm{B}\mathrm{O}}(D|\tau ):={max}_{\text{ϕ}\in \text{Φ}}\mathrm{E}\mathrm{L}\mathrm{B}\mathrm{O}(\text{ϕ})$ to approximate the marginal likelihood of a topology. Finally, we compare the ELBO-based posterior estimates to the MCMC-based estimates from MrBayes, which we consider as the ground truth. We placed independent exponential priors on branch lengths with a prior expectation of 0.1 substitutions.

In every analysis the gradient of the ELBO is evaluated with one sample and the ELBO is computed using 100 samples. While we used the default parameters provided by pystan, we used a smaller relative tolerance of 0.001 for every analysis as suggested by Yao et al. (2018).