Hamiltonian zigzag speeds up large-scale learning of direct effects among mixed-type biological traits

Inferring correlation among mixed-type biological traits while controlling for the evolutionary relationship among taxa is of great scientific interest yet remains computationally challenging. The recently developed phylogenetic multivariate probit model accommodates binary and continuous traits by assuming latent parameters underlying binary traits. The most expensive inference step is to sample the latent parameters from their conditional posterior that is a high dimensional truncated normal. The current best approach uses the bouncy particle sampler (BPS) optimized with a linear-order gradient evaluation method that employs a dynamic programming strategy on the directed acyclic structure of the phylogeny. Despite its significant improvement upon previous methods, with increasing sample sizes BPS encounters difficulty in exploring the parameter space and fails to provide reliable estimates for the across-trait partial correlation that describes the direct effects among traits. We develop a new inference scheme that highlights Zigzag Hamiltonian Monte Carlo (Zigzag-HMC), a variant of traditional HMC that uses Laplace momentum. Zigzag-HMC can utilize the same gradient evaluation method that speeds up BPS, yet it is much more efficient. We further improve the efficiency by jointly updating the latent parameters and correlation elements using a differential operator splitting technique. In an application exploring HIV-1 evolution that requires joint sampling from a 11,235-dimensional truncated normal and a 24-dimensional covariance matrix, our method yields a 5× speedup compared to BPS and makes it possible to estimate the direct effects among important viral mutations and virulence. We also extend the phylogenetic probit model to categorical traits for broader applicability, and demonstrate its use to study Aquilegia flower and pollinator co-evolution. 1 ar X iv :2 20 1. 07 29 1v 1 [ st at .M E ] 1 8 Ja n 20 22


Introduction
An essential goal in evolutionary biology is to understand the association among multiple traits observed on biological samples, or taxa, that are organisms ranging from animals and plants to microorganisms including important pathogens such as human immunodeficiency virus (HIV) and influenza. All taxa are implicitly correlated through their shared evolutionary history often described with a reconstructed phylogenetic tree. The tree tips correspond to the taxa and the internal nodes are their unobserved ancestors. Inferring across-trait covariation requires a highly-structured model that can explicitly describe the tree structure and adjust for across-taxa covariation. With a tree structure embedded, phylogenetic models are computationally challenging because we need to integrate out the unobserved internal node traits and also account for the uncertainties in tree estimation. The computational burden increases when the number of taxa and traits is large, and when traits are mixed with continuous and discrete quantities. The recently developed the phylogenetic multivariate probit model along with an efficient inference scheme  that achieves order-of-magnitude efficiency gain compared to the previous best approach (Cybis et al., 2015) and provides a promising tool to learn correlation among mixed-type traits at scale. Zhang et al. (2021) demonstrate their method on a data set with N = 535 HIV viruses and P = 24 traits that requires sampling from a truncated normal distribution with more than 11,000 dimensions. In this work, we overcome several shortcomings of Zhang et al. (2021) and largely push forward the inference performance to solve more challenging problems, and more importantly, to learn across-trait direct effects that bring insights about the potential causal pathways.
To jointly model mixed-type traits, the phylogenetic probit model assumes continuous latent parameters underlying discrete traits and the latent parameters follow a Brownian diffusion along the tree (Felsenstein, 1985;Cybis et al., 2015;Zhang et al., 2021). Assuming latent processes is a common strategy for modeling mixed-type data and it finds uses across various fields (Fedorov et al., 2012;Schliep and Hoeting, 2013;Irvine et al., 2016;Pourmohamad et al., 2016;Clark et al., 2017). For a data set with N taxa and P continuous or binary traits, Bayesian inference for the phylogenetic probit model involves repeatedly sampling latent parameters from their conditional posterior, an N P dimensional truncated normal distribution. For this task, Zhang et al. (2021) develop a bouncy particle sampler (BPS) (Bouchard-Côté et al., 2018) augmented with an efficient dynamic programming approach that speeds up the most expensive step in the BPS implementation. Two remaining limitations of BPS motivate this present work. First, BPS can have a near-reducible behavior without frequent velocity refreshment, but such refreshment largely costs efficiency and may lead to a "random-walk" behavior (Bouchard-Côté et al., 2018;Fearnhead et al., 2018;Neal, 2011;Andrieu and Livingstone, 2019). Second, in the phylogenetic probit model the latent parameters X and the trait correlation C are highly correlated by model assumption, so a joint update is more efficient. BPS, however, does not allow such joint sampling  use a Hamiltonian Monte Carlo sampler (HMC) (Neal, 2011) for C and update the two sets of parameters alternately within a random-scan Gibbs scheme (Liu et al., 1995).
Our solution to these issues highlights a state-of-the-art Markov chain Monte Carlo (MCMC) method called Zigzag-HMC (Nishimura et al., 2020) that overcomes the two limitations of BPS. Compared to BPS, Zigzag-HMC better explores the parameter space of a high-dimensional truncated normal (Section 3.1). Also, Zigzag-HMC enables a joint update of X and C through differential operator splitting (Strang, 1968;Nishimura et al., 2020), generalizing the previously proposed split HMC framework based on Hamiltonian splitting (Neal, 2011;Shahbaba et al., 2014). Zigzag-HMC takes advantage of the same linear-order in sample size gradient evaluation method in Zhang et al. (2021). The new inference scheme largely improves the mixing of elements in C and provides us with reliable estimates of across-trait partial correlation that describes the direct effect between any two traits, without confounding from other traits in the model. As seen in our applications, these direct effects help to understand the potential causal pathways of the biological processes.
We apply our method on two real-world examples. First, we re-evaluate the HIV evolution application in Zhang et al. (2021) and identify HIV-1 gag immune-escape mutations that directly affect the virulence. Our findings closely match with the experimental literature and also indicate a potentially general pattern in the immune escape mechanism of HIV. Second, we investigate how floral traits of Aquilegia flower attract different pollinators, for which we generalize the phylogenetic probit model to accommodate the categorical pollinator trait.

Mixed-type trait evolution
We describe biological trait evolution with the phylogenetic multivariate probit model following Zhang et al. (2021) and extend it to categorical traits as in Cybis et al. (2015). Consider N taxa on a phylogenetic tree F = (V, t) that is a directed, bifurcating acyclic graph. The tree is either known or informed by molecular sequence alignment S .
The node set V of size 2N − 1 contains N tip nodes, N − 2 internal nodes, and one root node. The branch lengths t = (t 1 , . . . , t 2N −2 ) denote the child-parent distance in real time.
We observe P traits of mixed-type for each taxon. The trait data Y = {y ij } = Y cont , Y disc partition as Y cont , an N × P cont matrix of continuous traits and Y disc , an N × P disc matrix of discrete ones. For each node i in F , we assume a d-dimensional latent parameter X i ∈ R d , i = 1, . . . , 2N − 1, where d = P cont + P disc j=1 (k j − 1) and k j is the number of classes for the jth discrete trait. To relate the latent parameters to observed traits, we assume a threshold model for binary traits and a choice model for traits with more than two classes, and continuous traits are directly observed. For a categorical trait y ij , the possible classes are {c 1 , . . . , c k j } with the reference class being c 1 . Multiple latent parameters x i,j , . . . , x i,j +k j −2 decide the value of y ij . We summarize the mapping from X to Y as where sign(x ij ) returns the value 1 on positive values and -1 on negative values, and sup(·) is the supremum. This data augmentation strategy is a common choice to model categorical data (Albert and Chib, 1993).
The latent parameters follow a multivariate Brownian diffusion process along F such that X i distributes as a multivariate normal (MVN) where X pa(i) is the parent node value and the d × d covariance matrix Ω describes the across-trait association. Assuming a conjugate root prior X 2N −1 ∼ N (µ 0 , ω −1 Ω) with prior mean µ 0 and prior sample size ω, we can analytically integrate out latent parameters on all internal nodes. Marginally, then, the N × d tip latent parameters X have the matrix normal distribution where M = (µ 0 , . . . , µ 0 ) T is an N × d mean matrix and the across-taxa covariance matrix Pybus et al., 2012). The tree F determines the diffusion matrix V(F ) and ω −1 J comes from the integrated-out tree root prior, where J is an all-one N × N matrix. The augmented likelihood of X and Y factorizes as where p(Y | X) = 1 if X are consistent with Y according to Equation (1) and 0 otherwise.
To ensure that the model is parameter-identifiable and also allow a non-informative prior on trait correlations, we decompose Ω as DCD such that C is the d × d correlation matrix and D is a diagonal matrix with marginal standard deviations .

A novel inference scheme
We sample from the joint posterior to learn the across-trait correlation C where we drop the dependence on hyper-parameters (Υ, µ 0 , ω) to ease notation. Zhang et al. (2021) use a random-scan Gibbs (Liu et al., 1995) scheme to alternately update X, (C, D) and F from their full conditionals . They sample X from an N ddimensional truncated normal distribution with BPS and deploy the standard HMC based on Gaussian momentum (Hoffman and Gelman, 2014) to update (C, D). Instead, we simulate the joint Hamiltonian dynamics on (X, C, D) by combining the novel Hamiltonian zigzag dynamics on X  and the traditional Hamiltonian dynamics on (C, D).
This strategy enables an efficient joint update of the two highly-correlated parameters. We first describe how Zigzag-HMC samples X from a truncated normal, and then detail the joint update of (X, C, D), followed by an auto-tuning of the algorithm.

Zigzag-HMC for truncated MVNs
We start by outlining the main ideas behind HMC (Neal, 2011) and then describe Zigzag-HMC as a version of HMC based on Hamiltonian zigzag dynamics (Nishimura et al., 2020. In order to sample a d-dimensional parameter x = [x 1 , . . . , x d ] from the target distribution π(x), HMC introduces an auxiliary momentum p = [p 1 , . . . , p d ] ∈ R d and samples from the product density π(x, p) = π(x)π(p) by approximating Hamiltonian dynamics described by the differential equation where U (x) = − log π(x) and K(p) = − log π(p) are the potential and kinetic energy.
In each HMC iteration, we first draw p from its marginal distribution π (p) ∼ N (0, I), a standard Gaussian, then approximate (6) from time t = 0 to t = τ by L = τ / steps of the leapfrog update with stepsize : The end state is a valid Metropolis proposal and one accept or reject it according to the standard acceptance probability formula (Metropolis et al., 1953;Hastings, 1970).
Zigzag-HMC differs from HMC in having a Laplace momentum π (p) ∝ i exp (−|p i |) , i = 1, . . . , d. The Hamiltonian differential equations now become and the velocity v := dx/ dt ∈ {±1} d depends only on the sign of p and thus remains constant until one of p i 's undergoes a sign change (an "event"). To understand how the Hamiltonian zigzag dynamics (8) evolve over time, one must investigate when such events happen. Before moving to the truncated MVN, we first review the event time calculation for a general π(x) following Nishimura et al. (2021). Let τ (k) be the kth event time and follows a piecewise linear path and the dynamics evolve as and Therefore we can derive the (k + 1)th event time and the dimension causing this event is i * = argmin i t i . At the moment of τ (k+1) , the i * th component of velocity has a sign change Then the dynamics continue for the next interval τ (k+1) , τ (k+2) . Now we consider simulating the Hamiltonian zigzag dynamics for a d-dimensional truncated MVN formed as where y ∈ R P is the mixed-type data and map(·) is the mapping from latent parameters x to y as in Equation (1). Here x ∈ R d and d ≥ P . It turns out we can analytically simulate the dynamics for a truncated MVN . We handle the constraint map(x) = y with the technique in Neal (2011) where the constraint boundaries are seen as "hard walls". Essentially, the trajectory of Hamiltonian zigzag dynamics would "bounce" against a boundary when reaching it. To distinguish different types of events, we define gradient events arising from solutions of Equation (11), binary events arising from hitting binary data boundaries, and categorical events arising from hitting categorical data boundaries. Starting from a state (x, v, p), we can calculate the gradient event times by solving d quadratic Then the gradient event time t g is the minimum among all positive roots from Equation (14) and the gradient event happens at the corresponding dimension. We refer interested readers to Nishimura et al. (2021) for a detailed derivation.
Next we focus on the binary and categorical events. We partition x into three sets: x i is for categorical data}. Since latent parameters in S cont are fixed, we "mask" them out with the same method in Zhang et al. (2021). Starting from a state (x, v, p), a binary event happens at time t b when the trajectory first reaches a binary boundary at dimension i b Note here we only need to check the dimensions with We turn to the categorical events. Let a categorical trait y j obtain class c k out of n possible classes, and x 1 , x 2 , . . . , x n−1 the underlying latent parameters. Equation (1) specifies the boundary constraints. If k = 1, the n − 1 latent parameters must be all negative, which poses the same constraint as if they were for n − 1 binary traits, therefore we can solve the event time using Equation (15). If k > 1, we must check when and which two dimensions first violate the order constraint x k−1 = sup(x 1 , . . . , x n−1 ) > 0. Let the current state be (x, v, p) and t j c is the categorical event time, we have when x ic reaches x k−1 and violates the constraint. To identify i c we only need to check is for a single y j and we need to consider all categorical data to find the actual categorical event time t c = min j t j c .
Now we describe the dynamics simulation with all three event types included, starting from a state (x, v, p) that satisfies all truncation constraints: 1. Solve t g , t b , t c using Equations (14), (15) and (16) respectively.
2. Determine the actual (first) event time t = min{t g , t b , t c } and update x and p as in Equations (9) and (10) for a duration of t.
3. Make instantaneous velocity and momentum sign flips according to the rules of the actual event type, then go back to Step 1.
Based on the above discussion, Algorithm 1 describes one iteration of Zigzag-HMC on truncated MVNs where we simulate the Hamiltonian zigzag dynamic for a pre-specified duration t total . For a truncated MVN arising from the phylogenetic probit model, we adopt the dynamic programming strategy of Zhang et al. (2021) to speed up the most expensive gradient evaluation step in line 3 and reduce its cost from O(N 2 d + N d 2 ) to O(N d 2 ). In brief this strategy avoids explicitly inverting Υ by recursively traversing the tree (Pybus et al., 2012) to obtain N conditional densities that directly translate to the desired gradient.

Jointly update latent parameters and across-trait covariance
We now turn to the joint update of the N × d latent parameters and d × d across trait covariance. As described in the previous section, we can efficiently simulate Hamiltonian zigzag on the truncated normal conditional of X. The covariance components C and D has no such special structure, so we can deploy standard Hamiltonian dynamics based on Gaussian momentum after transformation. We then combine the two dynamics via differential operator Algorithm 1 Zigzag-HMC for multivariate truncated normal distributions find binary boundary event time find categorical boundary event time, n c = number of categorical traits 8: if a gradient event happens at i g then else if a binary boundary event happens at i b then else if a categorical boundary event happens at i c1 , i c2 then splitting (Strang, 1968;Nishimura et al., 2020) to approximate the joint dynamics of Laplace-Gauss mixed momenta.
We denote the two concatenated sets of parameters X and (C, D) as x = (x G , x L ) with momentum p = (p G , p L ), where indices G and L refer to Gaussian or Laplace momentum. The joint sampler updates (x G , p G ) first, then (x L , p L ), followed by another update of (x G , p G ). Algorithm 2 describes the process where we analytically simulate the Hamiltonian zigzag dynamics for (x L , p L ), and utilize the standard leapfrog method (7) to approximate the Hamiltonian dynamics for (x G , p G ). The function SPLITHMC returns a Metropolis proposal that we accept or reject following the standard acceptance probability formula (Metropolis et al., 1953;Hastings, 1970). Because x G and x L can have very different scales, we incorporate a tuning parameter, the step size ratio r, to allow different step sizes for the two samplers. Appendix A provides an empirical method to automatically tune r. Within our random-scan Gibbs scheme, each MCMC iteration to update (X, C, D) calls the joint sampler in Algorithm 2 m times before scanning to a different set of model parameters. In practice, the overall sampling efficiency largely depends on m, the step size and the step size ratio r, so it is preferable to auto-tune all of them. We utilize the No-U-Turn algorithm to automatically decide m and adapt the step size with primal-dual averaging to achieve an optimal acceptance rate (Hoffman and Gelman, 2014).

Zigzag-HMC better explores the energy space than BPS
In our hands, BPS tends to generate highly auto-correlated samples from its target distribution π(x). In other words, it moves slowly in the energy space − log π(x), even when the marginal dimensions all appear to demonstrate good mixing. We apply BPS and Zigzag-HMC to a 256-dimensional standard normal target and Zigzag-HMC returns a clear win in the mixing of joint density (Figure 1). The inefficiency in sampling log π(x) is less a problem Algorithm 2 A joint update of parameters with Laplace-Gauss mixed momenta

12: end function
if one only needs to sample from a truncated normal with a fixed covariance matrix. But we are also updating the covariance, and such an inefficiency in the energy space harms the sampling efficiency for the interested parameters (Section 3.2). While a formal theoretical analysis is out of the scope of this work, we provide an ad hoc analysis to examine BPS's slow movement in energy space. Assume the d-dimensional parameter at the tth MCMC iteration is x(t) = {x 1 (t), . . . , x d (t)} ∈ R d , t = 1, . . . , T , with T being the total number of iterations. For a standard normal, its log density log π(x) ∝ d i x 2 i , and a high auto-correlation suggests log π(x) changes little between successive iterations, that is, the squared jumping distances , t = 0, . . . , T − 1 are small. We then decompose J D into two components where J 1 measures the squared travel distance in all dimensions and J 2 can be seen as the empirical covariance of x 2 j (t + 1) − x 2 j (t) and x 2 k (t + 1) − x 2 k (t) for j = k ∈ {1, . . . , T }. We compare J D , J 1 and J 2 between BPS and Zigzag-HMC on a 256-dimensional standard normal distribution (Table 1). Clearly, BPS yields a much lower J D than Zigzag-HMC because its J 2 is largely negative, suggesting strong negative correlation among the marginal travel distance in different dimensions. Table 1: Squared jumping distance (J D ) of log π(x) sampled by the bouncy particle sampler (BPS) and Zigzag Hamiltonian Monte Carlo (Zigzag-HMC). We report posterior expected J 1 and J 2 in their means and standard deviations (SD) across ten independent simulations with T = 2000 samples. To avoid reducible behavior, we include a Poisson velocity refreshment for BPS and set the refreshment rate to an optimal value 1.4

Efficiency gain from the new inference scheme
We demonstrate that Zigzag-HMC and the joint update of latent parameters X and the covariance matrix Ω largely improve the inference efficiency. Table 2 compares the performance of four sampling schemes on the HIV immune escape example with N = 535, P disc = 21, P cont = 3 (described in more detail later). We choose our efficiency criterion to be the per run-time, effective sample size (ESS) for the across-trait partial correlation R = {r ij } that is of the most scientific interest. After sampling Ω, we transform it to R through BPS and Zigzag-HMC only update X and we use an HMC-NUTS transition kernel for the Ω elements. SplitHMC employs the joint update of X and Ω described in Section 2.2.2, and SplitHMC-NUTS employs NUTS to automatically decide step size and the number of steps. For a fair comparison, we set the same t total for BPS and Zigzag-HMC, and supply SplitHMC with the optimal step size learned by SplitHMC-NUTS. We run every MCMC chain until the minimal ESS across all r ij is above 100 and we discard the first 20% of the samples as burn-in. We run all simulations on a node equipped with AMD EPYC 7642 server processors. As reported in Table 2, BPS loses to all the other three samplers and SplitHMC performs the best, yielding a 5× speed-up. SplitHMC-NUTS has a slightly lower efficiency than SplitHMC mainly because the No-U-Turn algorithm takes extra steps to decide when the trajectory would make a "U-turn". In practice we recommend using the tuning-free SplitHMC-NUTS.

HIV immune escape
We revisit the HIV evolution application in Zhang et al. (2021) where a main scientific focus lies on the association between HIV-1 immune escape mutations and virulence, the pathogen's ability to cause disease. The human leukocyte antigen (HLA) system plays an important role in the immune response against HIV-1 and is predictive of the disease course.
Through its rapid evolution, HIV-1 can acquire mutations that help to escape HLA-mediated immune response, but the escape mutations may reduce its fitness and virulence (Nomura et al., 2013;Payne et al., 2014). Zhang et al. (2021) identify HLA escape mutations asso- In other words, holding one of CD4 and RC as constant, the other does not affect VL, suggesting that RC increases VL via reducing CD4. The fact that RC is not found to have a direct impact on VL may be explained by the strong modulatory role of immune system on VL. Only when viruses with higher RC also lead to more immune damage, as reflected in the CD4 count, higher VL may be observed as a consequence of less suppression of viral replication. As such, our findings are in line with the demonstration that viral RC impacts HIV-1 immunopathogenesis independent of VL (Claiborne et al., 2015).
The partial correlation also helps to decipher epistatic interactions and how the escape mutations and potential compensatory mutations affect HIV-1 virulence. For example, we find a strong positive partial correlation between T186X and T190X. Studies have shown that T186X is highly associated with reduced VL (Huang et al., 2011;Wright et al., 2010) and it requires T190I to partly compensate for this impaired fitness so the virus stays replication competent (Wright et al., 2012). The negative direct effect between T186X and RC and the positive direct effect between T190I and RC are consistent with this experimental observation. In contrast, because of the strong positive association between T186X and T190, they are both suggested to have a negative association with RC, so the marginal association alone is sufficient to identify which effect is real. Another pair of mutations that potentially shows a similar interaction is H28X and M30X, which have a positive and negative partial correlation with VL, respectively. These mutations have indeed been observed to co-occur in gag epitopes from longitudinally followed-up patients (Olusola et al., 2020). Figure 2b keeps all the other compensatory mutation pairs in Figure 2a such as A146X-I147X and A163X-S165X that find confirmation in experimental studies (Troyer et al., 2009;Crawford et al., 2007).
More generally, when considering the viral trait RC and the infection trait VL, for which their variation are to a considerable extent attributable to viral genetic variation (Blanquart et al., 2017), we reveal an intriguing pattern. As in Figure 3, when two escape mutations impair virulence, and there is a direct effect between them, it is always negative. When two mutations have opposing effects on these virulence traits, the direct effect between them (if present) is positive. For example, T186X and I61X both have a negative impact on RC and the negative effect between them suggests that their additive, or even potentially synergistic, impact on RC is inhibited. Moreover, they appear to benefit from a compensatory mutation, T190X, which has been corroborated for the T186X-T190X pair at least as reported above.
Also for VL, the direct effect between mutations that both have a negative impact on this virulence trait is consistently negative. Several of these individual mutations may benefit from H28X as a compensatory mutation, as indicated by the positive effect between pairs that include this mutation, and as suggested above for H28X -M30X. This illustrates the extent to which escape mutations may have a negative impact on virulence and the need to evolve compensatory mutations to restore it. We note that our analysis is not designed to recover compensatory mutations at great length as we restrict it to a limited set of known escape mutations, while mutations on many other sites may be compensatory. In fact, our analysis suggests that some of the considered mutations may be implicated in immune escape due to their compensatory effect rather than a direct escape benefit.

Aquilegia flower and pollinator co-evolution
Reproductive isolation allows two groups of organisms to evolve separately, eventually forming new species. In plants, pollinator plays an important role in reproductive isolation (Lowry et al., 2008). We exam the relationship between floral phenotypes and the three main pollinators for columbine genus Aquilegia that are bumblebees, hummingbirds, and hawk moths (Whittall and Hodges, 2007). The pollinator is a categorical trait with three classes and we choose bumblebee with the shortest tongue as the reference class. Figure 4 is the across-trait correlation and partial correlation. Compared to a similar analysis on the same data set that only looks at correlation or marginal association (Cybis et al., 2015), partial correlation controls confounding and indicates the direct effect between pollinators and floral phenotypes that can bring new insights.
For example, we observe a positive marginal association between hawk moth pollinator and the spur length, but no direct effect between them. The marginal association matches with the observation that flowers with long spur length have pollinators with long tongues (Whittall and Hodges, 2007;Rosas-Guerrero et al., 2014). The absence of a direct effect makes intuitive sense because hawk moth's long tongue is not likely to stop them from visiting a flower with short spurs when the other floral traits are held constant. In fact, researchers observe that shortening the nectar spurs does not affect hawk moth visitation (Fulton and Hodges, 1999). Similarly, the positive partial correlation between orientation and hawk moth also finds experimental support. The orientation trait is the angle of flower axis relative to gravity, in the range of (0, 180). A small orientation value means a pendent flower and a large value for a more upright flower (Hodges et al., 2002). Due to their different morphologies, hawk moth prefer upright flowers while hummingbird tend to visit pendent ones. Making the naturally pendent Aquilegia formosa flowers upright increases hawkmoth visitation (Hodges et al., 2002). These results suggest that partial correlation may have predictive power for results from carefully designed experiments with controlled variables.

Discussion
Learning how different biological traits interact with each other from many evolutionarily related taxa is a long-standing problem of great interest, since it helps us to understand various phenomena in evolution. We present a scalable solution for this purpose and it largely improves the inference efficiency compared to the previous best approaches (Cybis et al., 2015;Zhang et al., 2021). Our new inference scheme allows learning the across-trait direct effect that is more informative than the previous marginal association. Our approach can provides reliable estimates of the across-trait partial correlation on large problems while the BPS method previously employed suffers from the computational complexity. In the large-scale application on the viral evolution of HIV-1, the improved efficiency allows us to infer direct effects among traits of scientific interest and therefore investigate some of the most important molecular mechanisms under the disease. In addition, our approach is tuning free that the most influential tuning parameters automatically adapts to the specific data set. We also extend the phylogenetic probit model to include categorical traits and illustrate its use in examining the co-evolution of Aquilegia flower and pollinators.
The novelty of our approach lies in two aspects. First, we highlight one cutting-edge MCMC method, Zigzag-HMC (Nishimura et al., 2020) to tackle the most challenging computational task of sampling from a high-dimensional truncated normal distribution. Zigzag-HMC is more efficient than the previous best approach that uses the BPS (Section 3.2).
It is worth mentioning that another closely related sampler, the Markovian zigzag sampler (Bierkens et al., 2019), or MZZ, also applies in this task but with lower efficiency than Zigzag-HMC . Zigzag-HMC, BPS, and MZZ are all state-of-the-art MCMC algorithms. Zigzag-HMC is a recent and less explored version of HMC, and BPS and MZZ are the two main methods based on piecewise deterministic Markov processes that have attracted growing interest in recent years (Fearnhead et al., 2018;Dunson and Johndrow, 2020). Intriguingly, the most expensive step of all three samplers is to obtain the log-density gradient, and the same linear-order gradient evaluation method  largely speeds it up. Second, we utilize differential operator splitting to jointly update two sets of parameters X and Ω that are correlated. This further improves efficiency and therefore allows us to obtain reliable estimates of the direct effect among traits. As in our applications, the direct effect is likely to better describe the trait interactions than the marginal association.
We now consider limitations of this work and also the future directions to which they point. First, the phylogenetic probit model does not currently accommodate a directional effect among traits since it only describes pair-wise and symmetric correlations. However, the real biological processes are often not symmetric but directional, where it is common that one reaction may trigger another but not the opposite way. A model allowing directed paths is preferable since it better describes the complicated causal network among multiple traits. Graphical models with directed edges (Lauritzen, 1996) are commonly used to learn molecular pathways (Neapolitan et al., 2014;Benedetti et al., 2017), but challenges remained to integrate these methods with a large and random phylogenetic tree. One may also try constructing a continuous-time Markov chain to describe how discrete traits evolve (Pagel, 1994;O'Meara, 2012), but with P binary traits the transition rate matrix grows to the astronomical size 2 P . Second, though our method achieves the current best inference efficiency under the phylogenetic probit model, there is still room for improvement. One potential solution is to de-correlate some latent parameters by grouping them into independent factors using phylogenetic factor analysis (Tolkoff et al., 2018;Hassler et al., 2021). Also, we can consider a logistic or softmax function to map latent parameters to the probablity of a discrete trait. This avoids the hard truncations in the probit model but also adds another layer of noise. It requires substantial effort to develop an approach that overcomes above limitations and also allows an efficient inference at the scale of applications in this work.

Acknowledgments
We A Auto-tuning of r We describe a simple heuristic to auto-tune the step size ratio r on the fly. Let Σ G and Σ L be the covariance matrices for x G and x L respectively, then their minimal eigenvalues λ min,G and λ min,L describe the variance magnitude in the most constrained direction. Intuitively, for both HMC and Zigzag-HMC, the step size should be proportional to the diameter of this most constrained density region, which is λ min,G or λ min,L . Therefore we propose a choice of r = √ λ min,L √ λ min,G , assuming the two types of momenta lead to similar travel distance during one unit time. It is straightforward to check this assumption. At stationarity, HMC has a velocity v G ∼ N (0, I), so its velocity along any unit vector u would be distributed as v G , u ∼ N (0, 1), and the travel distance E| v G , u | = 2/π. For Zigzag-HMC, as v L , u does not follow a simple distribution, we estimate E| v L , u | by Monte Carlo simulation and it turns out to be ≈ 0.8, close to 2/π.
We test this intuitive choice of r on a subset of the HIV data in Zhang et al. (2021) with 535 taxa, 5 binary and 3 continuous traits. We calculate the optimal r = √ λ min,L √ λ min,G ≈ 2.5 with Σ G and Σ L estimated from the MCMC samples. Clearly, r has a significant impact on the efficiency as a very small or large r leads to lower ESS (Table 3). Also, an r in the order of our optimal value generates the best result, so we recommend this on-the-fly automatic tuning r = √ λ min,L √ λ min,G (Table 3).