Zig-zag sampling for discrete structures and non-reversible phylogenetic MCMC

The zig-zag process is a non-reversible, piecewise deterministic Markov process introduced by [BR17, BFR19b] for continuous-time MCMC. It has several advantages over reversible methods such as Metropolis–Hastings [Has70] and Gibbs sampling [GS90]: it avoids diffusive backtracking which slows their mixing, and is rejection-free so that no computation is wasted on rejected moves.

In brief, the generator of the zig-zag process $(\mathbf{x}_{t},\mathbf{v}_{t})_{t\geq 0}$ targeting the probability density (with respect to the product of the Lebesgue and counting measures) $\tilde{\pi}(\mathbf{x},\mathbf{v}):=\pi(\mathbf{x})/2^{d}$ on a state space $X\times\{-1,1\}^{d}\subseteq\mathbb{R}^{d}\times\{-1,1\}^{d}$ is

where $\partial_{i}$ is the derivative with respect to $x_{i}$, and $F_{i}$ flips the sign of $v_{i}$. The flip rates

with $(x)^{+}:=\max\{x,0\}$, ensure that $(\mathbf{x}_{t},\mathbf{v}_{t})_{t\geq 0}$ leaves $\tilde{\pi}(\mathbf{x},\mathbf{v})$ invariant [BFR19b, Theorem 2.2]. In words, the coordinates of $\mathbf{x}_{t}$ move at constant velocities $\mathbf{v}_{t}$ until a flip at coordinate $i$, when the corresponding velocity changes sign.

To date, the zig-zag processes have been applied to targets such as the Curie–Weiss model [BR17] and logistic regression [BFR19b], whose state spaces have simple geometric structures with natural notions of direction and velocity. Discrete variables (other than the velocities) have been restricted to cases with simple partial orders, such as model selection [CFS20, GD21]. We construct a zig-zag process on a general hybrid state space with both continuous and discrete coordinates, without imposing any structure on discrete coordinates. This is done by introducing a separate space of continuous variables for each value of the discrete variable, introducing boundaries into the continuous spaces, and updating the discrete variable when boundaries are hit. The strategy takes advantage of the full generality of piecewise-deterministic Markov processes [Dav93, Section 24], and is resembles similar work for Hamiltonian Monte Carlo (HMC) [DBZM17, NDL20]. Our method can also been seen of as a generalization of the zig-zag process on a restricted domain [BBCD⁺18] to a union of many restricted sub-domains, with jumps between sub-domains at boundary hitting events. We illustrate our method with an application to the coalescent [Kin82]: a tree-valued target with continuous branch lengths, discrete tree topologies with no natural partial order, and no canonical geometric structure.

The coalescent examples illustrate the need for methods which are implementable on complex state spaces. They are also of interest because existing MCMC algorithms for coalescents tend to mix slowly. The key difficulty lies in designing Metropolis–Hastings proposal distributions which combine efficient exploration with a high acceptance rate [MV05, HDPD08, LVH⁺08]. Workarounds consist of empirical searches for efficient proposals [HD12, ASR16] or Metropolis-coupled MCMC [Gey92]. The former does not scale to problems for which empirical optimization is infeasible. The latter helps mixing between modes, but does not overcome low acceptance rates or the backtracking behavior of reversible MCMC.

The zig-zag process has some similarities with HMC [Nea10], which augments the state space with momentum and uses Hamiltonian dynamics to propose large steps which are accepted with high probability, though they are not rejection-free. Like the zig-zag process, HMC requires gradient information and a suitable geometric embedding of the target. [DBZM17] provided those for the coalescent using an orthant complex construction of phylogenetic tree space [BHV01]. Our examples differ from the method of [DBZM17] in three ways. Firstly, we replace the embedding of [BHV01] with $\tau$-space [GD16], which is better suited to ultrametric trees. Secondly, the zig-zag process is readily implementable on $\tau$-space via Poisson thinning without a numerical integrator such as leap-prog [DBZM17, Algorithm 1]. Finally, the zig-zag process has simple boundary behavior between orthants and does not require boundary smoothing [DBZM17, Section 3.3], chiefly because discontinuous gradients are easier to handle on continuous rather than discretized paths.

The rest of the paper is structured as follows. Section 2 defines the zig-zag algorithm with discrete and continuous variables, and proves that it has the desired invariant distribution. Section 3 recalls the coalescent and the $\tau$-space embedding. In Sections 4 and 5 we recall the popular infinite and finite sites models of mutation, derive zig-zag processes for each, and demonstrate their performance via simulation studies. Section 6 concludes with a discussion. The algorithms and data sets used in the simulation studies are available at https://github.com/JereKoskela/tree-zig-zag.

The definition of our zig-zag process follows [Dav93, Section 24]. Let $\mathbb{F}$ be a countable set. For each $m\in\mathbb{F}$, let $\Omega_{m}^{o}$ be an open subset of $\mathbb{R}^{d}$, $\overline{\Omega_{m}^{o}}$ be its closure, and $\partial\Omega_{m}^{*}:=\overline{\Omega_{m}^{o}}\setminus\Omega_{m}^{o}$ be its boundary. We assume that $\{\partial\Omega_{m}^{*}\}_{m\in\mathbb{F}}$ are piecewise Lipschitz and denote by $\partial\Omega_{m}$ the restriction of $\partial\Omega_{m}^{*}$ to non-corner points. Let $\Omega^{o}:=\cup_{m\in\mathbb{F}}\Omega_{m}^{o}=\{(m,\mathbf{x}):m\in\mathbb{F},\mathbf{x}\in\Omega_{m}^{o}\}$ and $\partial\Omega:=\cup_{m\in\mathbb{F}}\partial\Omega_{m}$. For a point $(m,\mathbf{x})\in\partial\Omega_{m}$, let $n(m,\mathbf{x})$ be the outward unit normal, and let $\Gamma^{\pm}(m,\mathbf{x}):=\{\mathbf{v}\in\{-1,1\}^{d}:\pm(\mathbf{v}\cdot n(m,\mathbf{x}))>0\}$ be the sets of velocities with which a zig-zag process can exit (+) or enter (-) $\Omega_{m}^{o}$ at $\mathbf{x}$. Zig-zag dynamics imply $\mathbf{v}\in\Gamma^{+}(m,\mathbf{x})\Leftrightarrow-\mathbf{v}\in\Gamma^{-}(m,\mathbf{x})$. We also define $\Gamma^{\pm}(\partial\Omega):=\cup_{(m,\mathbf{x})\in\partial\Omega}(\{(m,\mathbf{x})\}\times\Gamma^{\pm}(m,\mathbf{x}))$ and $\Omega^{*}:=(\Omega^{o}\times\{-1,1\}^{d})\cup\Gamma^{-}(\partial\Omega)$. Integrals over $\Omega^{o}$ and $\partial\Omega$, or subsets thereof, are taken to incorporate discrete sums in the $m\in\mathbb{F}$ coordinate.

The zig-zag process $(m_{t},\mathbf{x}_{t},\mathbf{v}_{t})_{t\geq 0}$ is defined on $\Omega^{*}$, with target $\tilde{\pi}(m,\mathbf{x},\mathbf{v}):=\pi(m,\mathbf{x})/2^{d}$ on $\Omega^{*}\cup\Gamma^{+}(\partial\Omega)$ for a given density $\pi(m,\mathbf{x})$. At $(m,\mathbf{x},\mathbf{v})\in\Omega^{*}$, the process moves with velocity $\mathbf{v}$ and each coordinate $v_{i}$ flips at rate $\lambda(m,\mathbf{x},\mathbf{v})$, defined as in (1) since $m$ is fixed between boundary hitting events. When $(m,\mathbf{x},\mathbf{v})\in\Gamma^{+}(\partial\Omega)$, the process jumps according to a Markov kernel $Q:\Gamma^{+}(\partial\Omega)\mapsto\mathcal{P}(\Gamma^{-}(\partial\Omega))$, where $\mathcal{P}(A)$ denotes the set of probability measures on $(A,\mathcal{B}(A))$. We assume that $Q$ and $\tilde{\pi}$ satisfy the skew-detailed balance condition

for any $(m,\mathbf{x},\mathbf{v})\in\Gamma^{+}(\partial\Omega)$ and $(j,\mathbf{y},\mathbf{w})\in\Gamma^{-}(\partial\Omega)$, as well as

for any $(\mathbf{x},\mathbf{v})\in\Gamma^{+}(\partial\Omega)$, and exclude jumps to paths pointing into corners by assuming

where $T_{\Gamma^{+}(\partial\Omega)}(j,\mathbf{y},\mathbf{w})$ is the time the line $(j,\mathbf{y}+t\mathbf{w})_{t\geq 0}$ hits $\Gamma^{+}(\partial\Omega)$. We will abuse notation and use $\tilde{\pi}(m,\mathbf{x},\mathbf{v})\mathrm{d}\mathbf{x}$ as the target density and Lebesgue measure on $\Omega^{o}$, and as their restrictions to the surface $\partial\Omega$, on which $\tilde{\pi}$ is not a probability density.

By [Dav93, Theorem 26.14], the zig-zag process with dynamics defined above is a piecewise-deterministic Markov process with extended generator

whose domain $D(L)$ consists of measurable functions $f(m,\mathbf{x},\mathbf{v})$ on $\Omega^{*}$ satisfying:

1. 1.

   For each $(m,\mathbf{x},\mathbf{v})\in\Gamma^{+}(\partial\Omega)$, the limit $\lim_{t\to 0}f(m,\mathbf{x}-t\mathbf{v},\mathbf{v})=:f(m,\mathbf{x},\mathbf{v})$ exists.

2. 2.

   For each $(m,\mathbf{x},\mathbf{v})\in\Omega^{*}$, the function $t\mapsto f(m,\mathbf{x}+t\mathbf{v},\mathbf{v})$ is absolutely continuous on $t\in[0,T_{\Gamma^{+}(\partial\Omega)}(m,\mathbf{x},\mathbf{v}))$.

3. 3.

   For $(m,\mathbf{x},\mathbf{v})\in\Gamma^{+}(\partial\Omega)$,

   $$f(m,\mathbf{x},\mathbf{v})=\int_{(j,\mathbf{y})\in\partial\Omega}\sum_{\mathbf{w}\in\Gamma^{-}(\mathbf{y})}f(j,\mathbf{y},\mathbf{w})Q(m,\mathbf{x},\mathbf{v};j,\mathrm{d}\mathbf{y},\mathbf{w}).$$

   (6)

4. 4.

   The random variable $\sum_{k=1}^{n}f(m_{T_{k}},\mathbf{x}_{T_{k}},\mathbf{v}_{T_{k}})-f(m_{T_{k}},\mathbf{x}_{T_{k}-},\mathbf{v}_{T_{k}-})$ is integrable for each $n\in\mathbb{N}$, where $\{T_{k}\}_{k\geq 1}$ are the successive jump times (both velocity flips and jumps due to hitting a boundary) of $(m_{t},\mathbf{x}_{t},\mathbf{v}_{t})_{t\geq 0}$.

For a set $A$, let $B(A)$ and $C^{1}(A)$ be the respective spaces of bounded and continuously differentiable functions on $A$. Let $C_{b}^{1}(A):=B(A)\cap C^{1}(A)$. For $t>0$, we define

We conclude this section with a pseudocode specification of our zig-zag algorithm.

An ultrametric binary tree with $n$ labeled leaves is a rooted binary tree in which each leaf is an equal graph distance away from the root. We follow [GD16] and encode such a tree with leaf labels $\{1,\ldots,n\}$ as the pair $(E_{n},\mathbf{t}_{n})$, where $E_{n}$ is a ranked topology and $\mathbf{t}_{n}\in(0,\infty)^{n-1}$. The continuous variables $\mathbf{t}_{n}$ encode times between mergers. The time from the leaves to the first merger is $t_{1}$, and subsequent $t_{i}$ variables are times between successive mergers. The ranked topology $E_{n}$ is an $(n-1)$-tuple of pairs of labels, where the $i$th pair specifies the two child nodes of the $i$th merger. Non-leaf nodes are labeled by the leaves they subtend. For example, the ranked topology

encodes the four leaf caterpillar tree depicted in Figures 3 and 7 with nodes labeled left to right. We order the two entries of $E_{n,i}$ by their least element for definiteness.

The coalescent [Kin82] is a seminal model for the genetic ancestry of samples from large populations. Under the coalescent, a tree $(E_{n},\mathbf{t}_{n})$ has probability density

which arises as the law of a tree constructed by starting a lineage from each leaf, and merging each pair of lineages at rate 1 until the most recent common ancestor (MRCA) is reached. The success of the coalescent is due to robustness: distributions of ancestries of a large class of individual-based models converge to the coalescent in the infinite population limit under suitable rescalings of time. For details, see e.g. [Wak09].

We define swap operators $s_{i}$ for $i\in\{2,\ldots,n-2\}:E_{n,i-1}\notin E_{n,i}$ via

In words, $s_{i}$ swaps the order of the $(i-1)$th and $i$th mergers. We also define pivot operators $p_{i}^{\downarrow}$ and $p_{i}^{\uparrow}$ for $i\in\{2,\ldots,n-1\}:E_{n,i-1}\in E_{n,i}$ as

where $E_{n,i}^{\uparrow}$ (resp. $E_{n,i}^{\downarrow}$) is the entry of $E_{n,i}$ with the higher (resp. lower) least element, and $E_{n,i}^{s}$ is the sibling: the entry of $E_{n,i}$ that is not $E_{n,i-1}$. Pivot $p_{i}^{\uparrow}$ is defined by interchanging $\downarrow$ and $\uparrow$ in (10). The pivots are the two nearest neighbor interchanges between the $i$th merger and the merger involving its nearest child. Figure 1 illustrates all three operators.

Next we describe $\tau$-space, which gives a geometric structure to the set of pairs $(E_{n},\mathbf{t}_{n})$. For fixed $E_{n}$, the space of $\mathbf{t}_{n}$-vectors is the orthant $[0,\infty)^{n-1}$. Each boundary point with $t_{i}=0$ corresponds to a degenerate tree in which one of three things happens:

1. 1.

   The two leaves of $E_{n,1}$ merge at time 0 if $i=1$.

2. 2.

   There are two simultaneous mergers if $E_{n,i-1}\notin E_{n,i}$.

3. 3.

   There is a simultaneous merger of three lineages if $E_{n,i-1}\in E_{n,i}$.

Type 1 boundaries are boundaries of the whole $\tau$-space. Type 2 boundaries separate orthants corresponding to two ranked topologies separated by an $s_{i}$-step. Trajectories crossing the boundary move from one ranked topology to the other. Type 3 boundaries separate the three orthants which resolve the triple merger into two binary mergers, which differ by a $p_{i}^{\downarrow}$ or $p_{i}^{\uparrow}$-step. A trajectory that crosses the boundary visits a tree with the triple merger. Figure 2 depicts the $\tau$-space with three leaves, and a type 2 boundary in a general $\tau$-space. An example with five leaves is depicted in [GD16, Figure 2], and the $\mathbf{t}_{n}$ variables are illustrated in Figures 3 and 7 in Sections 4 and 5.

We will use $\tau$-space to construct zig-zag processes whose state spaces consist of tree topologies, branch lengths, and a scalar parameter introduced in the next section. The discrete variables $\mathbb{F}$ will be the ranked topologies, with the boundary crossings described above defining the boundary jump kernel $Q$. For more on $\tau$-space, e.g. existence and uniqueness of geodesics and Fréchet means, we refer to [GD16].

The infinite sites model [Wat75] connects the coalescent tree to DNA sequence data by associating the MRCA with the unit interval $(0,1)$. Mutations with independent, $U(0,1)$-distributed locations accrue along branches of the tree (with branch lengths as specified by $\mathbf{t}_{n}$) at rate $\theta/2$. The type of a leaf consist of the mutations along branches separating it from the MRCA. We denote the resulting list of types of the leaves by $D_{n}$. A realization of a coalescent tree and associated $D_{n}$ is shown in Figure 3. The typical task is to sample from the conditional law $(E_{n},\mathbf{t}_{n},\theta)|D_{n}$ corresponding to observing $D_{n}$, but not the tree which gave rise to it.

To handle mutations, we define $F_{n}$ as the rooted graphical tree with $2n-1$ nodes, the first $n$ of which are leaves labeled $1,\ldots,n$, while the remaining $n-1$ are labeled as in $E_{n}$. Edges connect children to their parents, and edge lengths are determined by $\mathbf{t}_{n}$. For an edge $\gamma\in F_{n}$, we denote by $c_{\gamma}$ and $p_{\gamma}$ the respective labels of the child and parent nodes of $\gamma$, by $m_{\gamma}$ the number of mutations on $\gamma$, and by $l_{\gamma}:=\sum_{t_{j}\in\gamma}t_{j}$ the edge length, where we write $t_{i}\in\gamma$ if $t_{i}$ contributes to the length of $\gamma$ in which case we say $\gamma$ spans $t_{i}$.

Given a prior density $\pi_{0}(\theta)$ for $\theta$, the posterior distribution of $(E_{n},\mathbf{t}_{n},\theta)|D_{n}$ follows from (9), the fact that the number of mutations on branch $\gamma\in F_{n}$ is Poisson($\theta l_{\gamma}/2$)-distributed given $l_{\gamma}$, and that mutations on distinct branches are independent. In particular,

provided $E_{n}$ is consistent $D_{n}$, and $\pi=0$ otherwise. This distribution can be sampled using a zig-zag algorithm by taking $\mathbb{F}$ to be the set of ranked topologies on $n$ leaves which are consistent with $D_{n}$, as well as $\Omega_{E_{n}}^{o}:=\{(\mathbf{t}_{n},\theta)\in(0,\infty)^{n}\}$ and $\partial\Omega_{E_{n}}:=\{(\mathbf{t}_{n},\theta)\in[0,\infty)^{n}:t_{i}=0\text{ for one }i\in\{1,\ldots,n-1\}\text{ or }\theta=0\}$ for each $E_{n}\in\mathbb{F}$. In the boundary classification of Section 3, $\theta=0$ is another type 1 boundary. For $(\mathbf{t}_{n},\theta)\in\partial\Omega$, we define $k(\mathbf{t}_{n},\theta)$ as the index of the zero variable, taken to be $n$ in the case of $\theta$.

We augment the state space with $n$ zig-zag velocities $\mathbf{v}_{n}$, of which $(v_{1},\ldots,v_{n-1})$ drive $\mathbf{t}_{n}$ and $v_{n}$ drives $\theta$. For $\gamma\in F_{n}$, we also define $v_{\gamma}:=\sum_{j:t_{j}\in\gamma}v_{j}$ as the rate of change of $l_{\gamma}$. The boundary kernel $Q$ is defined separately on each boundary type:

At a type 1 boundary the process reflects back into $\Omega_{E_{n}}^{o}$ via a velocity flip. At a type 2 boundary it undergoes an $s_{i}$-step and a velocity flip to pass through the boundary. For type 3 boundaries it chooses an adjacent orthant uniformly at random.

In the interiors of orthants, velocity flip rates are

Simulating holding times with these rates is difficult due to the time intervals during which they vanish. One strategy is Poisson thinning via dominating rates consisting of only those terms in the round brackets in (13) and (14) whose sign matches that of the corresponding velocity $v_{i}$, but these can result in loose bounds and inefficient algorithms. Instead, we define $t_{n}:=\theta$ for brevity, and for $i\in\{1,\ldots,n-1\}$ define $\gamma(E_{n},i):=\operatorname{argmin}\{l_{\gamma}:p_{\gamma}=E_{n,i}\}$ as the shorter child branch from parent node $E_{n,i}$. For $i\in\{1,n\}$ we adopt the short-hands

and finally define the time localization $T\equiv T(E_{n},\mathbf{t}_{n},\theta;\mathbf{v}_{n})$ as

where $\min_{\emptyset}=\infty$, and $K\gg 0$ is a maximum increment for the case when all velocities are positive. The indicator functions in the denominator pick out boundaries where (11) vanishes: type 1 or 3 boundaries corresponding to length 0 branches which carry at least one mutation, and the $\theta=0$ boundary if there is at least one mutation in total. The parameter $c>0$ ensures that, when the current process time is $t$, such boundaries cannot be hit on $[t,t+T]$, and at most one other boundary can be reached. We found $c=4$ gave good performance across our tests in Sections 4 and 5. A larger value results in tighter bounds on (13) and (14), but wastes more computation as $t+T$ is hit more often before an accepted velocity flip.

On time interval $[t,t+T]$, flip rates (13) and (14) are bounded above by constant rates

where, for each $\lambda_{i}^{*}$, $i\in\{1,\ldots,n-1,\theta\}$, $(x)^{\pm}:=(x)^{+}$ if $v_{i}>0$ and $(x)^{\pm}:=(x)^{-}:=\min\{x,0\}$ if $v_{i}<0$. Algorithms 2 and 3 in the appendix give pseudocode for simulating holding times, velocity flips, and boundary crossings as outlined above.

We compared the zig-zag process to a Metropolis–Hastings algorithm by reanalyzing the data of [WFDP91] with $n=55$, 14 distinct types, and 18 mutations. We used the improper prior $\pi_{0}(\theta)\propto\mathds{1}_{(0,\infty)}(\theta)$, set $v_{i}=\pm 2/[(n+1-i)(n-i)]$, and set $v_{n}$ from trial runs to cross the $\theta$-mode in unit time. The appendix details the Metropolis–Hastings algorithm, and other tuning parameters in Table 3. We compared both methods to a hybrid combining zig-zag dynamics with continuous time Metropolis–Hastings moves at rate $\kappa=10$. Performance was insensitive to $\kappa$ provided it was not extreme: small values resemble a zig-zag process, while large values resemble Metropolis–Hastings.

Figure 4 shows that the zig-zag and hybrid methods mix visibly better than Metropolis–Hastings over the latent tree, as measured by the tree height $H_{n}:=t_{1}+\ldots+t_{n-1}$. However, they are not as effective at exploring the upper tail of the $\theta$-marginal, likely because they do not stay in regions of short trees for long enough for $\theta$ to increase into the tail.

To assess scaling, we simulated two data sets: one of size $n=550$ with mutation rate $\theta=5.5$ (the approximate posterior mean in Figure 4) and one with $n=55$ and $\theta=55$, which models a segment of DNA 10 times longer. Figures 5 and 6 demonstrate that the zig-zag and hybrid processes scale far better than Metropolis–Hastings, particularly when $\theta=55$. Estimates in Table 1 quantify the improvement to 1–3 orders of magnitude.

The finite sites model [JC69] is more detailed than the infinite sites model, but has greater computational cost. Consider a finite set of sites $S$ with a finite number of possible types $H$ per site; for example $H=\{0,1\}$ or $H=\{A,T,C,G\}$. Mutations occur along branches of the coalescent tree at each site with rate $\theta/(2|S|)$, and the type of a mutant child is drawn from stochastic matrix $P$ with unique stationary distribution $\nu$. We denote the transition matrix of the $H$-valued compound Poisson process with jump rate $\theta/(2|S|)$ and jump transition matrix $P$ by $(Q_{hg}^{\theta}(t))_{h,g\in H;t\geq 0}$. Figure 7 depicts a realization of the finite sites coalescent.

As in Section 4, we denote the configuration of types at the leaves by $D_{n}$, and seek to sample from the posterior $\pi(E_{n},\mathbf{t}_{n},\theta|D_{n})$, which can be written as a sum over the types of internal nodes:

where $h(s;\eta)\in H$ is the type at site $s\in S$ on the node with label $\eta\in E_{n}$, and $\gamma\in F_{n}:|c_{\gamma}|>1$ denotes edges which do not end in a leaf. The target (16) can be evaluated efficiently using the pruning algorithm of [Fel81].

The posterior (16) can be sampled using zig-zag dynamics with the same construction as in Section 4. Velocity flip rates can be written in terms of branch-specific gradients as