Penn

# Leaping through tree space: continuous phylogenetic inference for rooted and
unrooted trees

Matthew J Penn  Neil Scheidwasser  Joseph Penn  Department of Physics,
University of Oxford, Oxford, United Kingdom  Christl A Donnelly  Department
of Statistics, University of Oxford, Oxford, United Kingdom  Pandemic Sciences
Institute, University of Oxford, Oxford, United Kingdom  MRC Centre for Global
Infectious Disease Analysis, Department of Infectious Disease Epidemiology,
School of Public Health, Faculty of Medicine, Imperial College London, London,
United Kingdom  David A Duchêne  Samir Bhatt  Section of Epidemiology,
University of Copenhagen, Copenhagen, Denmark

###### Abstract

Phylogenetics is now fundamental in life sciences, providing insights into the
earliest branches of life and the origins and spread of epidemics. However,
finding suitable phylogenies from the vast space of possible trees remains
challenging. To address this problem, for the first time, we perform both tree
exploration and inference in a continuous space where the computation of
gradients is possible. This continuous relaxation allows for major leaps
across tree space in both rooted and unrooted trees, and is less susceptible
to convergence to local minima. Our approach outperforms the current best
methods for inference on unrooted trees and, in simulation, accurately infers
the tree and root in ultrametric cases. The approach is effective in cases of
empirical data with negligible amounts of data, which we demonstrate on the
phylogeny of jawed vertebrates. Indeed, only a few genes with an ultrametric
signal were generally sufficient for resolving the major lineages of
vertebrates. Optimisation is possible via automatic differentiation and our
method presents an effective way forwards for exploring the most difficult,
data-deficient phylogenetic questions.

###### keywords:

Keywords: phylogenetic inference | balanced minimum evolution | gradient descent | distance matrix

s.bhatt and imperial.ac.uk

## Significance statement

Phylogenetics is vital in life sciences, revealing early life origins and
epidemic dynamics. A central challenge in inferring the best tree from a set
of data is that exploring the vastness of tree space is computationally hard.
Here we recast the exploration and inference problem from a discrete to a
continuous space. Our method performs well on unrooted trees and can
accurately infer both the tree and the root in ultrametric cases, but at a
substantially increased computational cost. Our approach represents a shift in
the methodology to explore tree space, and opens the possibilities of new
efficient forms of inference.

## Introduction

Phylogenetic inference, the task of reconstructing the evolutionary
relationships across taxonomic units given observational data, has a wide
range of theoretical and practical applications in biology, such as evolution
[13, 33, 76], conservation [87] and epidemiology [40, 27], but also in
comparative linguistics [65] and cultural anthropology [16, 72]. In
particular, the COVID-19 pandemic has catalysed the development of efficient
phylogenetic tools and methods to better understand the virus’ origin, spread,
and evolution [62, 77, 3, 94, 106, 108, 19]. For biological problems, tree
inference is primarily informed by molecular sequence data (i.e., nucleotide
or amino acid sequences), for which an extensive body of literature exists
[93, 112, 114]. Other types of biological data such as morphology [60],
fossils [71], and auditory communication in animals [1] can also be used as
input.

Two key parameters considered when inferring a phylogenetic tree include the
_topology_ , the branching pattern that specifies the evolutionary
relationships between operational taxonomic units, and _branch lengths_ , the
amount of evolutionary divergence that occurred between the branching events.
A substantial amount of research has been conducted on how to parameterise
branch lengths [9, 24], especially through the use of various molecular clocks
[120]. Similarly, although to a lesser degree, progress has been made on
methods for efficient exploration of the space of tree topologies[99], which
is fundamentally challenging due to its combinatorial complexity. Indeed, for
$n$ taxa, there are $(2n-3)!!$ possible rooted tree arrangements, where $n!!$
denotes the semifactorial of $n$ — even a small dataset of ten taxa can be
enumerated by 34 million unique rooted trees. Moreover, finding the global
optimal tree is NP-hard for all major optimality criteria (e.g., maximum
parsimony [35], minimum evolution [18], maximum likelihood [85]). Methods such
as linear programming [12] or branch and bound [44] can provide exact
solutions, but are practically limited to problems with $\approx$ 15 or fewer
taxa. To overcome these challenges, the overwhelming majority of state-of-the-
art software (e.g., MrBayes [47], PAUP [110], BEAST [25], PAML [113],
RAxML(-NG) [99, 54], FastME [61], IQ-TREE [75, 70]) rely on hand-engineered
search heuristics to perform tree topology optimisation or Bayesian analysis.
These are traditionally based on subtree pruning and regrafting (SPR) and tree
bisection and reconnection (TBR) operations, which have empirically been shown
to be the best available methods for exploring tree topology space [99, 79].

However, such methods still have limitations. First, hill climbing using
heuristic approaches necessitates multiple evaluations of the objective
function to pick the best move. While these heuristics are still polynomial,
exhaustive exploration of single SPR operations is quadratic in complexity,
and paired operations (two sequential SPR changes) are quartic. Second, all
the aforementioned tree arrangements are prone to being trapped in local
optima and even if a global optimum is found, terraces of trees with identical
quality exist [92]. The challenge of exploring tree space is exacerbated when
concatenating multiple genes in supermatrices [86, 83, 15] or when using
genomic-scale datasets which require extensive computational resources.

To address these shortcomings, we propose GradME, a new direction for tree
topology inference which expands the problem space using a continuous rather
than discrete parameterisation of a phylogenetic tree. Generally, aside from
considering metrics (e.g., distances in tree space) [98, 23, 7, 15],
performing topological search in a continuous tree space has rarely been
explored (for recent work in hyperbolic spaces, see [66, 111, 64, 69]).
Furthermore, very few approaches have made use of gradient-based tree
proposals [66, 118, 74, 23]. Although maximum likelihood and Bayesian
inference criteria are more popular and generally considered state-of-the-art
[47, 99, 25, 70, 110, 4], the GradME framework optimises tree topology under a
balanced minimum evolution (BME) criterion [80, 20] using distance matrices as
an input. This criterion is well-principled [52] but generally performs worse
than likelihood-based [82, 114, 61]. However, the framing of the minimum
evolution criterion [52, 89] has been proven to be statistically consistent
[21, 33] and has repeatedly shown good (although not state-of-the-art)
performance in various settings [56, 57, 38, 61].

To better explore the space of possible trees, we expand the space over which
we need to search. Our novel vector representation of a phylogenetic tree,
Phylo2Vec [81], has a natural continuous extension, allowing us to improve the
ability to search parts of this space. Appealing to a common analogy that
casts the optimal tree search problem as finding a needle in a haystack, our
approach observes a much bigger haystack, but the hay is in very large
bundles, many of whom have a needle, and for these bundles we have access to a
(weak) magnet. Providing details to this analogy, the size of the usual
phylogenetic haystack with $n$ taxa is $(2n-3)!!$ [29], while we search a much
larger haystack of size $(n!)^{2}$. There are $n!$ bundles in this larger
haystack, each of which contains $n!$ trees, but for any tree, $2^{n-1}$
bundles will contain that tree. Although the proportion of bundles containing
a needle shrinks exponentially, we propose a novel approach (Queue Shuffle)
that chooses bundles that should be closer to one with a needle. For any given
bundle, we also introduce a continuous objective function that can be
efficiently traversed using gradient descent approaches (the weak magnet)
developed for large-scale machine learning problems [53, 63, 96]. This
continuous objective facilitates enormous changes to tree topology in a single
step in a direction that improves the objective function. After searching any
given bundle using the continuous objective, we use Queue Shuffle, which
improves the switch towards the next bundle to search. This counterintuitive
approach offers a new addition to the existing heuristic methods used for
topological inference, outperforming the current state-of-the-art but as
currently stands with a larger overall complexity comparable to other minimum
evolution approaches such as Neighbour Joining.

## Results

#### Tree traversal in continuous space

For any choice of label ordering, our approach admits a continuous gradient
across $n!$ trees for $n$ leaves. This gradient, which can be obtained readily
via automatic differentiation, can rapidly traverse tree space to find trees
with a close to optimal objective value. Fig. 1(a) shows a single gradient
step for the small Primates dataset [78, 113]. Simply subtracting the gradient
from a random initial tree, followed by softmax activation, results in an
almost discrete $W$ which corresponds to the best BME tree for a given
substitution model. Note that if more gradient steps were taken, the $W$
matrix would quickly become discrete (from Lemma 0.8). The jump taken
corresponds to six subtree-prune and regraft moves [78].

For larger alignments such as the popular Eutherian dataset [97], a single
gradient step can result in 14 to 18 SPR moves. While the number of SPR moves
achieved is large, this is achieved with a substantial increase in overall
computational complexity when compared to FastME. We note that the gradient
step size is dependent on the data and, as expected, greatly reduces as we
approach an optimum.

(a) Dataset: Primates

(b) Dataset: Simulated (20 taxa, 100k sites)

(c) Dataset: Jawed

Figure 1: Results on empirical data (a) Starting from a random tree, represented by an $n\times n$ stochastic matrix, we compute the continuous gradient, apply softmax activation and increment the original matrix. In a single step, our gradient finds the correct tree at a distance of 6 subtree-prune and regraft moves from the random starting tree. (b) Simulating ultrametric trees of 20 taxa and 100,000 sites under an LG model of protein evolution. We add random uniform noise to all branch lengths to simulate departures from ultrametricity. Compared to the true tree via Robinson-Foulds distance, light blue bars are midpoint rooting the best FastME tree and dark blue bars are the inferred root from our approach. (c) Phylogenies for jawed vertebrates, where the number of genes (hence sites) are reduced to be more clocklike. Normalised Robinson-Foulds distance are shown between the best ASTRAL [117] tree, the best unrooted FastME tree which has been midpoint rooted (light blue) and our inferred rooting algorithm (dark blue). Performance for FastME reduces when the number of sites is small. Table 1: Balanced minimum evolution loss scores for 11 phylogenetic benchmark datasets. Lower is better. Scores from BioNJ and FastME were obtained following the implementations in ape [78] using the same distance matrix as GradME. The distance matrix was estimated from a GTR+$\Gamma$ model via maximum likelihood [112]. Our GradME approach always starts from a uniform tree distribution (every tree is equiprobable) with a random taxon ordering (optimised by Queue Shuffle). The best performing approaches for each dataset are denoted in bold. GradME either equalled or performed better than FastME. The topological accuracy, measured as one minus the Robinsons-Foulds distance is shown between GradME and FastME and GradME and a maximum likelihood gold standard from IQ-TREE also using a GTR+$\Gamma$ model. Dataset | BioNJ | FastME | GradME |  Topological accuracy between GradME and FastME |  Topological accuracy between IQ-TREE and the best distance tree
---|---|---|---|---|---
DS1 | 0.3118613 | 0.3101232 | 0.3101232 |  1.00 |  0.54
DS2 | 3.725205 | 3.7239944 | 3.7239944 |  1.00 |  0.77
DS3 | 8.0115913 | 8.0075588 | 8.0075588 |  1.00 |  0.97
DS4 | 2.2528503 | 2.2447615 | 2.2447615 |  1.00 |  0.68
DS5 | 6.3077156 | 6.2606057 | 6.2606057 |  1.00 |  0.70
DS6 | 0.6249236 | 0.6228563 | 0.6219367 |  0.87 |  0.67
DS7 | 9.9174641 | 9.882138 | 9.882138 |  1.00 |  0.91
DS8 | 1.337924 | 1.3252984 | 1.3252984 |  1.00 |  0.82
DS9 | 0.3788481 | 0.3788481 | 0.3788481 |  1.00 |  0.66
DS10 | 1.1286037 | 1.1247627 | 1.1247627 |  1.00 |  0.78
DS11 | 1.313921 | 1.3096422 | 1.3096415 |  0.88 |  0.53

#### A comparison to benchmark phylogenetic data sets

Table 1 presents a comparison of GradME with neighbour joining (BioNJ) and
FastME (subtree-prune and regraft version) over 11 popular phylogenetic
benchmark datasets [109]. Both neighbour joining and FastME are only able to
infer a minimum length unrooted tree, and therefore we compare estimates only
on unrooted trees. We always initialise our algorithm with a uniform,
equiprobable tree, where the starting taxon labelling is random and optimised
using Queue Shuffle. We estimate tree using distances from a GTR+$\Gamma$
model estimated via maximum likelihood (see Appendix 0.9 for details). As
expected, FastME consistently outperforms BioNJ, with lower BME loss on all
alignments. On the other hand, GradME always achieves a better or equal loss
compared to FastME. We observe similar results when using different
substitution models (e.g., F81). In the two examples where GradME does better
than FastME, the topological accuracy, measured by one minus the Robinson-
Foulds distance [84] is close to 0.9, suggesting FastME has converged to a
similar tree. We note that FastME’s performance is generally worse when using
the nearest neighbour interchange heuristic (instead of the SPR-based
heuristic). When compared to a maximum likelihood gold standard (IQ-TREE
[70]), the best distance method does not recover the same tree as that from
maximum likelihood, but in some cases, is very close (e.g., DS3 and DS7).
Finally, we note that while GradME outperforms FastME, it is orders of
magnitude slower and in most of the data sets FastME finds the same optimal
tree as GradME.

#### Rooting ultrametric trees

Despite being applicable to the unrooted problem, our approach, at its core,
works with rooted trees. As previously discussed, if we assume the existence
of a distant outgroup, then the balanced minimum evolution objective can be
used to optimise a rooted phylogenetic tree. In Appendix 0.1, we show that,
given an ultrametric unrooted tree, the optimal rooting maximises a heuristic
for the root-to-tip distance in the tree. Equivalently, the optimal rooting
ensures that the root is estimated to be the maximal possible distance back in
time. This is not an immediately biologically plausible objective for the
root. Indeed, the cornerstone of balanced minimum evolution is finding the
tree of minimum length, and it hence seems counter-intuitive to require the
root that is the maximum distance backwards in time (though this does in fact
minimise the tree length). However, our assumption of a distant ancestor means
that the root of our tree must be the point that is furthest backwards in
time. In particular, this means that the evolutionary direction needs to be
away from the root. By setting our root such that the root-to-tip distance is
maximised, we ensure that the root satisfies this constraint. However, this
property does not hold for trees that are not ultrametric - in these cases,
the root will be drawn towards branches with higher mutation rates.

While this property only holds for ultrametric trees, our approach still works
well for near clock-like trees. As an experiment, we draw small (20 taxa)
random ultrametric phylogenies with a total length of one, and simulate
100,000-residue protein sequences [78] down these trees under an LG [60] model
of protein evolution, assuming random uniform amino acid base frequencies. In
the ultrametric cases, all taxa are equidistant to the root, which corresponds
to a strict molecular clock. We add uniform noise to all branch lengths to
simulate departure from a strict clock. Fig. 1(b) shows the Robinson-Foulds
[84] distance from the true tree to the _midpoint-rooted_ best _unrooted_
FastME tree (when SPR moves were used by FastME), and the distance to our
inferred rooted tree. We see that when the tree is ultrametric, or close to
ultrametric, our approach recovers the correct rooted tree. As expected, an
increase in noise leads to a decrease in topological accuracy, although our
approach still performs substantially better than midpoint rooting. We note
that uniform noise is unlikely to be biologically realistic. Instead,
deviations from a strict clock are more likely to be heterogeneous in certain
clades or internal branches. However, for small departures, we believe our
algorithm to reliably infer the correct tree and root simultaneously.

We implement our rooting algorithm on the popular mammal data from [97]. We
infer a rooted tree via Queue Shuffle and also midpoint root the best FastME
tree. Both trees, unrooted, have the same balanced minimum evolution loss, but
our rooted loss is less than the FastME midpoint rooted loss. Our rooted tree
correctly identifies Gallus gallus (red junglefowl) as the outgroup, while
midpoint rooting pairs Gallus gallus with Ornithorhynchus Anatinus (platypus)
(see Fig. S2 for the rooted phylogenies).

Figure 2: Phylogenetic inferences of the jawed vertebrates’ phylogeny using
the two most ultrametric loci from a data set of 99 taxa and 4593 genes [50].
(a) Inference using our approach leads to high accuracy in identifying the
root and all major jawed vertebrate taxa. Note that, we do not estimate branch
lengths, but only topology via balanced minimum evolution (b) inference using
FastME and midpoint rooting leads to widespread error, primarily and
critically near the root of the process.

#### Rooting the phylogeny of all jawed vertebrates

To perform a more detailed evaluation of our framework, we tested GradME’s
robustness for topological inference by finding the root of the large jawed
vertebrates dataset from [50] with 99 taxa and 4593 genes. Given the reliance
of our method on ultrametric data for inference of the root, we first made a
fast measure of the ultrametricity of each gene-tree. To do this, we inferred
the phylogeny of each gene using GradME, followed by midpoint rooting. The
coefficient of variation in root-to-tip lengths was taken as a measure of
ultrametricity. We then concatenated ranked genes into supermatrices including
decreasing numbers of genes, and examined the performance of GradME with
midpoint rooting against our method. All inferences were performed using the
LG amino-acid substitution model to maintain simplicity. We placed special
focus on our ability to use small portions of data for recovering the main
groupings of vertebrates; these key groupings include the root separating
cartilaginous (Chondrichthyes) versus boned vertebrates, ray-finned fishes
(Actinopterygii), and the major groups of tetrapods and amniotes (amphibians,
mammals, archosaurs, turtles, and lizards and relatives).

A small numbers of genes with an ultrametric signal were generally sufficient
for resolving many of the major lineages of vertebrates using both midpoint
rooting and our approach (Fig. 1(c)). For larger numbers of genes, midpoint
rooting and our approach are broadly similar. However, at the smallest numbers
of genes (0.05%, 2 genes), midpoint rooting was unable to recover many of the
early relationships among vertebrates, such as the root, monophyly of
cartilaginous fishes, ray-finned fishes, Tetrapoda, or the mammals. Even small
amounts of data (1460 amino-acids of 1,964,439; 0.07% of the original data)
were sufficient for GradME to resolve the root as well as every major grouping
of jawed vertebrates accurately (Fig. 2). The only exception was the
controversial position of the Coelacanth, which was found to be sister of
Dipnoi (lungfish) rather than the more widely accepted position as sister of
Dipnoi plus Tetrapoda. While this remarkable performance under the simple LG
model is in part attributed to the informative nature of highly ultrametric
genes, our tree topology demonstrates the superiority of our approach in
accuracy and efficiency over other fast methods in phylogenetics.

## Discussion

We have introduced a new approach for exploring the vastness of tree space.
Counterintuitively, our approach explores a much bigger space than the space
of possible trees, but this larger space allows for new ways to find the best
tree. The key to our method’s success lies in transforming the phylogenetic
tree search problem from a discrete to a continuous one, allowing us to
achieve superior performance. To our knowledge, this is the first time a
continuous, differentiable objective function for the inference of tree
topology has been proposed, and it opens new possibilities for phylogenetic
inference. Bayesian phylogenetics can be regarded as the most robust framework
for inferring phylogenies, but has been to-date limited by the poor ability of
random walk Metropolis-Hastings algorithms to explore tree space [6]. More
efficient Hamiltonian Monte Carlo samplers have been proposed [23] to tackle
this problem, and our framework presents a new avenue to jointly explore
topology and branch lengths with efficient samplers. A remaining limitation of
our approach is the need to shuffle labels to fully explore the space of all
possible trees, and while the approach we use, Queue Shuffle, is
mathematically and practically powerful, this step is still discrete. The
possibility of permutation distributions such as the Gumbel-Sinkhorn
distribution could allow for a fully differentiable algorithm. Finally, the
complexity of our approach is $\mathcal{O}(n^{5})$, which easily allows for
large phylogenies up to a thousand, but not tens of thousands. However,
computation on GPUs or TPUs in parallel can facilitate computational
tractability.

A major benefit of our approach is that it naturally enables the estimation of
the root node, which has been a long matter of interest in the biological
sciences [46, 105, 73]. For genes where a strict clock is a reasonable
assumption, our method of traversing tree space in large steps reliably
estimates both the correct tree topology and the root. Our approach will
likely be useful in settings where genetic sequences are contemporaneous and
time for measurable evolution is short, such as early epidemics or nosocomial
settings. However, as we showed analytically, our approach will have reduced
performance when considering rate heterogeneity and departures from a strict
clock.

Tests on the relationships among jawed vertebrates demonstrate that even
minimal amounts of data can be sufficient for our method to reach high
accuracy in topology and root estimates. These results are consistent with
previous work on large amounts of genome-scale data showing that clocklike
loci to be the most suitable for phylogenetic inference [107]. Furthermore,
our approach is effective with negligible amounts of data – where other
methods are ineffective – making it a powerful addition to the existing
toolkit for addressing recalcitrant questions of the tree of life.

Our approach is based on the minimum evolution principle, which has repeatedly
shown to produce fast and accurate inference. Nonetheless, an interesting area
for further study is to extend the continuous path length formulation to
approximations of traditional phylogenetic likelihoods [31]. This would be
particularly beneficial for implementation in Bayesian inference, since tree
topology inference is a major obstacle to large hierarchical models [100, 26].
Our method is therefore a step towards more efficient sampling of the complex
posterior distributions over tree topology.

## Methods

In the following, we describe GradME, a distance-based method for continuous
phylogenetic inference of rooted and unrooted trees using gradient descent.
The framework can be divided into three components: i) a continuous tree
representation based on Phylo2Vec [81], a bijective integer representation of
phylogenetic trees; ii) gradient-based optimisation using a continuous version
of the balanced minimum evolution criterion [80, 20], iii) Queue Shuffle, a
method to shuffle the integer-to-taxon mapping underlying Phylo2Vec for full
tree space exploration. The overall approach works for both rooted and
unrooted trees.

#### Balanced minimum evolution

Popular objective functions to infer the optimal tree from phylogenetic data
include maximum parsimony [34], maximum likelihood [31] and minimum evolution
[90]. Maximum likelihood and minimum evolution are provably statistically
consistent [21, 33], whereas maximum parsimony can be inconsistent under
certain conditions [28]. For small to moderate sized phylogenies, methods
based on maximum likelihood (and Bayesian extensions) are generally considered
state-of-the-art [99, 25, 70, 110]. However, approaches based on minimum
evolution (ME) have also shown to yield adequate performance [56, 38, 61, 57,
85]. The first introductions of the minimum evolution (ME) paradigm [52, 89]
sought to express evolutionary relationships through dissimilarity. They
proved that, given unbiased estimates of the true evolutionary distances, the
true phylogeny has an expected length shorter than any other possible
phylogeny – thereby establishing the principled ME criterion. Currently, the
best performing ME approach is that of balanced minimum evolution (BME) [80,
20], with FastME [61] being a popular software implementation. Its objective
function can be written as:

$\mathcal{L}(T)=\sum_{i,j}D_{ij}2^{-e_{ij}}$ (1)

where $D_{ij}$ denotes a distance (e.g., based on molecular sequence data)
between two taxa $i$ and $j$ and $e_{ij}$ the number of branches in the path
between taxa $i$ and $j$ (the path length [95]). This objective can be
computed in a numerically stable fashion using the log-sum-exp trick (see
Appendix 0.11 for an example snippet). A widely used approach to estimate the
optimal tree greedily [21, 38] is the neighbour joining (NJ) method [90]. When
neighbour joining is based on an additive distance measure, it reconstructs a
unique tree, but still performs well with near-additive trees [2] and under
small perturbations in the data [67]. However, despite these highly favourable
properties, further heuristic optimisation on a NJ tree using SPR moves have
proven to be even more accurate [61]. Once a tree topology is found, quadratic
algorithms exist for estimating the branch lengths [68] as well as efficient
approaches for molecular clock dating [104].

#### Balanced minimum evolution for rooted trees

Inference using BME is always restricted to unrooted trees [95, 11] with
rooting chosen after inference through heuristics (e.g., midpoint rooting) or
via a molecular clock (e.g., for serially sampled data). However, it is often
of interest to find the optimal rooted tree for a set of taxa, as this
provides extra biological context (e.g., to represent evolutionary paths).

In an unrooted tree, the BME objective function (Eq. 1) provides an efficient
way of calculating the total length of a tree where the branch lengths are the
least squares estimators for approximating each $D_{ij}$ with the distance
from nodes $i$ to $j$ in the tree. However, this result does not hold in a
rooted tree, as the addition of a root changes many of the path lengths. To
remedy this, we consider adding a “root taxon” to the tree by joining it to
the root node as taxon $n$. If the tree is roughly ultrametric, then we expect

$D_{ni}\approx D^{*}\quad\forall i\neq n$ (2)

where $D^{*}$ is the (assumed constant) root-to-taxa distance. Of course, we
do not know the sequence of the root but, as we will show, the value of
$D^{*}$ is unimportant — it is instead simply important that it is independent
of $i$. Adding this root taxon as a leaf node transforms the tree from being
rooted to being unrooted, where standard BME can be used. From this
assumption, we prove two lemmas to ensure the framework’s validity, showing
that the optimal unrooted tree is obtained when the variation in the root-to-
taxa distance is sufficiently small (Lemma 0.2), and subsequently that, in all
cases, the optimal rooting for an unrooted tree solves a biologically
plausible optimisation problem (Lemma 0.4).

First, Lemma 0.2 shows that, if

$|D_{ni}-D^{*}|<\delta\quad\forall i\neq n$ (3)

then, using $e_{ij}^{u}$ and $e_{ij}^{r}$ to denote path lengths in the
unrooted tree ($u$) containing taxon $n$ and the rooted tree ($r$) formed by
removing taxon $n$,

$\bigg{|}\sum_{i=0}^{n}\sum_{j=0}^{n}D_{ij}2^{-e^{u}_{ij}}-\sum_{i=0}^{n-1}\sum_{j=1}^{n-1}D_{ij}2^{-e^{r}_{ij}}-D^{*}\bigg{|}\leq\delta$
(4)

where $\delta$ denotes a small number. Hence, the difference between the
rooted and unrooted objective functions is approximately equal to the
constant, $D^{*}$. Thus, for sufficiently small $\delta$, by the discreteness
of tree space, we can see that, if it is unique, the optimal unrooted tree
under the rooted objective (using $e_{ij}^{r}$) will be the same as that under
the unrooted objective (using $e_{ij}^{u}$), when the root taxon is used as an
additional leaf.

Subsequently, Lemma 0.4 shows that, for the correct unrooted tree, the BME-
optimal rooting maximizes a simple heuristic (defined in Definition 0.3) for
the root-to-tip distance. Equivalently, the optimal rooting ensures that the
root is estimated to be the maximal possible distance back in time.

This is not an immediately biologically plausible objective for the root.
Indeed, the cornerstone of BME is that we want the tree of minimum length, and
it hence seems counter-intuitive to require the root that is the maximum
distance backwards in time (though, by Lemma 0.4, this does create the minimum
length tree). However, the root of our tree must be the point that is furthest
backwards in time. In particular, this means that the evolutionary direction
needs to be away from the root. By setting our root such that the root-to-tip
distance is maximised, we ensure that the root satisfies this constraint.

This method shares a similar motivation with midpoint rooting, which also
seeks to maximise a heuristic for root-to-tip distance. However, the heuristic
used in midpoint rooting uses only the two taxa which are furthest apart,
while our rooting method uses distances from all taxa. Thus, we expect our
method to be more robust, particularly as large inter-taxa distances are
difficult to estimate, meaning that the additional information used in our
heuristic should help to reduce errors. This is evidenced in Fig. S2, where
midpoint rooting leads to incorrect root placement.

Nonetheless, this property will not hold if the tree is not ultrametric. If
taxa evolve at different rates at different times throughout the tree, then
the root will be drawn towards taxa with high evolutionary rates. Thus,
caution must be used when applying our rooted algorithm to such trees,
although the unrooted algorithm will still give a correct unrooted tree. In
this case, it may be best to find the optimal unrooted tree topology and then
solve the rooting problem for this tree, rather than finding the optimal
rooted tree, as the former will reduce the skewing effect of the heterogeneity
in evolutionary rates. Because of this, the algorithm introduced in this paper
has the flexibility to find the optimal unrooted or rooted tree.

Figure 3: An example of the left-to-right construction of the ordered tree
$v=[0,0,0,2]$.

#### An ordered bijection to tree space

Previously, we introduced Phylo2Vec [81], a novel bijection between the space
of phylogenetic trees and a space of integer vectors. In contrast to other
bijections such as permutation matchings [22], changes in Phylo2Vec correspond
to smooth changes in the tree space, e.g., single changes in a Phylo2Vec
vector correspond to a limited set of SPR changes. On the other hand, Prüfer
codes [14] form a bijection to the space of all $m$-ary trees, meaning that
there is no guarantee to sample binary trees from random Prüfer sequences.

Here, we focus on the notion of _ordered trees_ from [81], where it is
possible to construct a tree from its vector in linear time. An ordered tree
can be thought of as a birth process, such that when a birth occurs, the
original node continues to live and retains its label, while the new node
receives an incremented label. Accordingly, we introduce an equivalent but
more intuitive tree construction process for these ordered trees (see Fig. 3
for an example). We begin with two leaf nodes and two edges labelled 0 and 1.
We then append nodes by joining them, _in order_ , to edges connecting leaf
nodes to the tree. This tree construction process can be summarised by a
single vector $\bm{v}$, with $v_{0}=v_{1}=0$ and, for $m\geq 2$, $v_{m}$ being
the label of the edge to which node $m$ is appended. Each index in $v_{m}$ is
subject to the simple constraint:

$\displaystyle v_{0}=v_{1}=0\quad\text{and}$ (5) $\displaystyle
v_{m}\in\\{0,1,...,m-1\\}\quad\forall\quad m\geq 2$

which is equivalent to the definition of ordered trees in [81]. Intuitively, a
tree is ordered if, starting with a branch connecting the root to taxon 0, the
taxa can be added in order of their label by appending a new branch to a
terminal branch of the existing tree (i.e. a branch connecting a leaf node to
the rest of the tree). Thus, the ordering of the taxa is in some sense
“natural” for each possible ordered tree. It is proved in Lemma 0.6 that for
ordered $\bm{v}$, the algorithm presented above and the more general Phylo2Vec
algorithm in [81] produce equivalent trees.

Note that, for a fixed integer-taxon labelling, the number of ordered trees is
a subset of the number of possible trees. We discuss an efficient method to
remedy this problem and explore all tree space called the Queue Shuffle.

#### A continuous representation of a tree

We introduce continuous, probabilistic, representation of trees using a square
matrix $W$ which gives the distribution of a random ordered vector $\bm{v}$
with independent entries such that $W_{ij}=\mathbb{P}(v_{i}=j)$. Given Eq. 5,
$W$ is a lower-triangular, stochastic matrix (row sums to 1). Thus, $W$ can
probabilistically represent any _ordered_ phylogenetic tree (a space of
$(n-1)!$ trees). A simple approach to determining the most likely single tree
from $W$ is to take the column-wise argmax, yielding a single tree $\bm{v}$.

#### Gradient-based optimisation using the BME criterion

Using this continuous representation, we can find the optimal ordered tree.
Defining $f(\bm{v})$ to be the BME objective function for the tree generated
by an ordered vector $\bm{v}$, we then create a continuous objective, $F(W)$
by $F(W)=\mathbb{E}\big{[}f(\bm{v})\big{]}$.

The calculation of $F(W)$ follows our new method of constructing ordered trees
from the vector $\bm{v}$. For a fixed (randomly chosen) tree, we define
$e_{ij}^{k}$ to be the path length between nodes $i$ and $j$ when nodes
$0,1,...,k-1$ have been added to the tree (for $i,j<k$). Note that to find the
rooted objective function, we initialise with $e_{10}^{2}=e_{01}^{2}=2$, while
to find the unrooted objective, we initialise with $e_{10}^{2}=e_{01}^{2}=1$
(while the Phylo2Vec representation is an inherently rooted representation,
this unrooted objective finds the tree length if the root were removed from
the random rooted tree with distribution given by $W$). This is because, in a
tree where the only leaf nodes are $0$ and $1$, these nodes are a path length
of 2 apart if there is also a root (as the root is on this path) while
otherwise, they are a path length of 1 apart.

If node $k$ is appended to the edge joining either node $i$ or $j$ to the
tree, then $e_{ij}^{k+1}=e_{ij}^{k}+1$; otherwise, $e_{ij}^{k+1}=e_{ij}^{k}$.
Similarly, if node $k$ is appended to the edge joining node $i$ to the tree,
then $e_{ik}^{k+1}=2$ and otherwise, if node $k$ is appended to the edge
joining node $x$ to the tree, then $e_{ik}^{k+1}=e_{ix}^{k}+1$. Thus, using
$V_{k}$ to denote the random value of $v_{k}$, we can write

$\displaystyle e^{k+1}_{ij}=e^{k}_{ij}+G^{k}_{ij}(V_{k})$ (6)
$\displaystyle\Rightarrow
e^{n-1}_{ij}=\sum_{k=2}^{n-2}G^{k}_{ij}(V_{k})+e^{2}_{ij}$ (7)

for some functions $G^{k}_{ij}$ which are derived explicitly in Lemma 0.7.
Importantly, each term in the sum is independent, and hence, in Lemma 0.7, a
closed iterative system for the quantities
$E_{ij}^{k}=\mathbb{E}(2^{-e_{ij}^{k}})$ can be calculated for $i<j$ as

$E_{ij}^{k+1}=\left\\{\begin{matrix}E^{k}_{ij}\bigg{[}1-\frac{1}{2}(W_{ki}+W_{kj})\bigg{]}&\text{if
$i<j<k$}\\\ \bigg{[}\frac{1}{2}\sum_{x\neq
i}E_{ix}^{k}W_{kx}\bigg{]}+\frac{1}{4}W_{ki}&\text{if
$i<k$}\end{matrix}\right.$ (8)

with the remaining values for $i>j$ following by symmetry.

The objective function is a polynomial function of the entries of $W$ and is
linear in each fixed entry (that is, the diagonal entries of $\nabla^{2}F$ are
zero). Thus, by Lemma 0.8, there is always a minimum at a “discrete tree”
(that is, at a matrix $W$ where for each row, one value is 1 and all the
others are 0). Moreover, this simple form makes it easy to differentiate $F$
analytically, numerically or automatically. Using state-of-the-art automatic
differentiation [8], gradient descent can be used to efficiently minimise $F$
and find the optimal ordered tree.

There may also be minima at non-discrete trees if multiple trees share the
same, optimal, objective value. In our rooted optimisation, this is highly
unlikely (as two topologically different trees having equal objectives places
a dimension 1 condition on the distance matrix $D$, meaning the set of
distance matrices for which this happens has measure 0 in the set of possible
distance matrices), but when we use our unrooted algorithm, this will occur.
This is because, as discussed in [81], there are $n-1$ Phylo2Vec vectors
which, when the root is removed, give the same unrooted tree. Thus, if
multiple rooted vectors giving the same unrooted tree $\mathcal{U}$ are in the
same space of ordered trees, then, if $\mathcal{U}$ is the optimal tree under
this ordering, the algorithm may converge to a non-discrete $W$. In this case,
taking the argmax safely recovers an optimal rooted tree as, by Lemma 0.8, all
possible trees according to $W$ will have the optimal objective value.

The tree-space induced by our continuous objective function will have local
minima whenever changing any single entry of the vector $\bm{v}$ causes the
objective function to increase. As the Hamming distance between vectors
$\bm{u}$ and $\bm{v}$ is comparable to SPR distance [81], we therefore expect
that the discrete subset of our continuous tree-space will be similar in
structure to the space induced by SPR moves. However, by starting from a
uniformly distributed tree, where all possible ordered trees contribute to the
objective, we expect that our algorithm is better able to pick up the “signal”
from the true optimum, and avoid moving towards suboptimal local minima.

A Python-like algorithm to compute $\mathbb{E}(2^{-e})$ is shown in Algorithm
1. To find the $E^{m}_{ij}$ terms, there are $\mathcal{O}(m)$ steps of
$\mathcal{O}(m)$ (finding the $E_{m-1,j}^{m}$ terms) and $\mathcal{O}(m^{2})$
steps of $\mathcal{O}(1)$ (finding the other $E_{ij}^{m}$ terms). There are
$\mathcal{O}(n)$ values of $m$ that need to be considered, and hence this
system can be solved in $\mathcal{O}(n^{3})$ time.

Algorithm 1 Compute $E:=\mathbb{E}(2^{-e})$

1:$n$ $\triangleright$ fixed number of taxa

2:$W$ $\triangleright$ stochastic input matrix[81]

3:$E=\text{{zeros}}((n,n))$

4:for $i\leftarrow 1$ to  $n-1$ do

5: $E^{*}=\text{{zeros}}((n,n))$

6: for $k\leftarrow 0$ to  $i-1$ do

7: for $j\leftarrow 0$ to  $k-1$ do

8: $E^{*}_{[k,j]}=E_{[k,j]}(1-\frac{1}{2}(W_{[i-1,j]}+W_{[i-1,k]}))$

9:
$E^{*}_{[i,k]}=\sum_{l=0}^{i-1}\frac{1}{2}E_{[l,k]}W_{[i-1,l]}+\frac{1}{4}W_{[i-1,k]}$

10: $E=E^{*}+E^{*T}$

11:Return  $E$

#### Orderings

The continuous objective function is only defined for ordered trees, which,
for a given labelling of the nodes, is a subset of the whole tree space. Thus,
we define the concept of an ordering of the nodes, whereby changing the
ordering allows the full space of trees to be explored.

######  Definition 0.1.

Ordering: Suppose that the nodes correspond to taxa with names
$N_{0},N_{1},...N_{n-1}$. We then define an ordering of the nodes to be a
permutation, $\sigma$ of the set $\\{0,...,n-1\\}$ such that the node with
name $N_{i}$ is processed as node $\sigma(i)$ by the Phylo2Vec algorithm. It
is necessary that the associated Phylo2Vec vector $\bm{v}(\sigma)$ is ordered.

Note that we use the phrase “node $x$” to mean “the node with name $N_{x}$”
and will always be explicit if we refer to a node by its label.

Given a tree, it is possible to generate a possible ordering of the nodes,
$\sigma$, as well as the associated vector $\bm{v}(\sigma)$. We will do this
by labelling the tree as follows, again distinguishing between the leaf node
names $N_{i}$ and their labels (which we will call $l(i)$).

Consider labelling the root node as 0. Then, label the children of this node
as $0$ and $1$. One can continue this process inductively, choosing a node,
labelled $x$, with unlabelled children and then labelling its children as $x$
and $y$, where $y$ is the smallest unused label. This process terminates when
every node has been labelled. As discussed above, suppose that the label of
the leaf node with name $N_{i}$ is $l(i)$. Then, Lemmas 0.9, 0.10 and 0.11
show that $l$ is a possible ordering of the tree, and provide a method for
calculated the associated ordered Phylo2Vec vector $\bm{v}(l)$.

Given a tree, one can use the previous labelling algorithm to show that there
are at least $2^{n-1}$ possible orderings for the tree (as the children of
each node can be labelled in either order). The exact number depends on the
choice of which internal node is processed at each step (and so, a “balanced”
tree where the leaves have a low generation has more possible orderings than a
“ladder” tree where each leaf has a distinct generation). However, in general,
it will be substantially larger than $2^{n-1}$ (we conjecture that for flat
trees, it may be factorial in size) and hence there are numerous possibilities
which will allow for a global minimum of the objective to be found.

#### Queue Shuffle: changing orderings to explore all the tree space

The number of ordered trees (different from ranked trees [17]), $(n-1)!$, is
substantially smaller than the possible number of trees, $(2n-3)!!$ (albeit
with a comparable growth pattern in $n$), and hence, while the optimal ordered
tree will be closer to the true tree than a very large proportion of trees, it
is very unlikely to be exactly equal to the true tree.

To fully explore tree space, one must shuffle the labels of the leaf nodes in
the optimal ordered tree. Simply choosing a uniformly random permutation will
lead to extremely inefficient optimisation, as each tree is only possible in
approximately $1/2^{n}$ of the possible orderings. Instead, we use the
topology of the optimal tree to inform our choice of permutation through a
novel approach we call the Queue Shuffle. This ensures that the previous
optimal tree can be written as an ordered tree in the new ordering, while also
ensuring a smooth and efficient path through the space of orderings.

The Queue Shuffle is motivated by the labelling procedure discussed in the
previous section, but ensures that the set of internal nodes with a given
generation (that is, a given distance from the root) are processed
consecutively. That is, we begin by processing all nodes with generation 0
(i.e., the root), then all internal nodes with generation 1, then all internal
nodes with generation 2, and continue in this fashion until all internal nodes
have been processed.

Algorithmically, this can be achieved by a “queue” of internal nodes to be
processed. When an internal node is processed, any of its children that are
also internal nodes are added to the back of this queue. Thus, the queue is
always in ascending order of generation, and it is simple to show that this
ensures that nodes are processed in non-decreasing order of generation.

A crucial feature of this queue is that the child given the same label as its
parent is placed ahead of the other child in the queue. This ensures that one
can, in some way, control the order of processing by choosing the labelling of
the children of each node. Moreover, it is vital for the theoretical result
presented in the following section.

To add randomness into the labelling procedure, every time an internal node is
processed, we randomly choose which child is given the label of their parent,
and which child is given the next available label. This provides $2^{n-1}$
possible orderings for each tree. This stochasticity is helpful in ensuring
that the algorithm does not get stuck – as discussed in the subsequent
section, it ensures that a large class of similar trees will be considered
after a few ordering proposals.

An algorithmic description of the Queue Shuffle is provided in Algorithm 2.

Algorithm 2 The Queue Shuffle

1:$\mathcal{T}$ $\triangleright$ Current Tree

2:$\mathcal{N}=\\{\nu_{0},\nu_{1},...\\}$ $\triangleright$ set of all non-root
nodes

3:$Q=[\nu_{0},\nu_{1}]$ $\triangleright$ "queue" of nodes to process

4:$L=\\{\nu_{0}:0,\nu_{1}:1\\}$ $\triangleright$ node:label mapping

5:$l_{\text{next}}=2$$\triangleright$ next available label

6:$P=[]$ $\triangleright$ processed nodes

7:while $Q\neq[]$ do

8: $\nu=Q[0]$ $\triangleright$ node to process

9: Q = Q[1:] $\triangleright$ $\nu$ will be processed

10: $\text{append}(P,\nu)$ $\triangleright$ $\nu$ will be processed

11: if $\text{isLeaf}(\nu)$ then

12: continue $\triangleright$ move to next node

13: $a,b=\text{randChildren}(\nu)$ $\triangleright$ get randomly ordered
children of $\nu$

14: $L[a]=L[\nu]$ $\triangleright$ label $a$ with $\nu$’s label

15: $L[b]=l_{\text{next}}$ $\triangleright$ give $b$ next available label

16: $l_{\text{next}}=l_{\text{next}}+1$

17: append($Q,a$) $\triangleright$ add $a$ to the queue

18: append$(Q,b)$ $\triangleright$ add $b$ to the queue

19:Return  $L$ $\triangleright$ Ordering determined by values of $L$ for leaf
nodes

#### GradME

The Queue Shuffle completes our optimisation algorithm. We iteratively find
the best ordered tree according to the current ordering and then use Queue
Shuffle to change ordering, changing the space of explorable trees. The
algorithm terminates when the optimal tree has not been improved upon for a
fixed number of iterations (note that, by construction, the previous optimal
tree will always be in the new space of ordered trees). In the examples
presented in this paper, only tens iterations are needed from some random
starting order, and less if a sensible starting ordering (such as from a
neighbour joining tree) is used.

We refer the resulting system, combining the continuous tree representation,
Queue Shuffle reordering, and the gradient-based optimisation framework using
BME, as GradME.

#### Why does Queue Shuffle work?

A given tree is in the space of ordered trees for at least $2^{n-1}$
orderings. This means that we do not need to find a single optimal ordering,
but have exponentially many which will return the true optimal tree. Very
loosely considered, being able to explore $n!$ tree space reliably and
efficiently with continuous optimisation, Queue Shuffle reduces the
inferential task to one that is exponential.

However, while the number of optimal orderings grows exponentially, their
proportion tends quickly to zero as $n$ grows. It is therefore, perhaps,
surprising that we are able to find an optimal ordering so quickly from merely
tens of shuffles. The proportion of optimal orderings (approximately the ratio
of ordered trees to total trees, $\frac{(n-1)!}{(2n-3)!!}$, ranges from
$8\times 10^{-4}$ in our smallest dataset (14 taxa) to $6\times 10^{-29}$ in
the largest (99 taxa; see Table 2).

This efficiency comes from the topology-dependence of the Queue Shuffle
algorithm, which allows us to plot a relatively “smooth” path through the
space of possible orderings. That is, the majority of trees in the new ordered
space will have similar properties to the previous optimal tree and so, unless
the previous tree was a local minimum of the objective, it is likely that one
of these “close” trees will have a lower objective value.

Lemma 0.12 shows that the expected distance from the root grows harmonically
as the label increases. For large trees, the node with label $n-1$ has an
expected distance from the root of approximately twice the expected distance
from the root of the node label $0$. This property is noticeable even for
small trees – if $n=10$, then the ratio of the expected distance to the root
of node 9 and node 0 is approximately $1.65$. Thus, nodes which are close to
the root in the current optimum, will also be closer on average to the root in
the new space of ordered trees. In essence, this means that “fewer slots are
wasted” in the new ordered space – that is, there are fewer trees in the new
space of ordered trees which are topologically far from the previous optimum
(a tree that, after the first few iterations, is likely to be far closer to
the true optimum than a randomly-chosen tree) and hence, more trees which are
reasonable candidates for having lower objective values.

Lemma S1 proves another example of the smoothness in transitions induced by
the Queue Shuffle, based on nearest neighbour interchange (NNI) moves. An NNI
move considers the four subtrees attached to two non-root nodes that share an
edge and swaps two of these subtrees. Lemma S1 shows that, starting from a
tree $\mathcal{T}$, any tree which is one NNI move away from $\mathcal{T}$
will be in the new space of ordered trees with probability at least
$\frac{1}{4}$.

This ensures that this new space contains many sensible proposal trees.
Perhaps the most surprising aspect of this theorem is that this probability is
bounded below, independently of the topology. Thus, with high probability, the
optimal tree will only remain the same for more than a few iterations if large
sets of similar trees yield lower objective values than the current optimum.

That being said, the Queue Shuffle does not guarantee that the global minimum
will be found, even if the gradient-based algorithm for optimising $F(W)$
always converges to the optimal tree. If a tree is “far” from the nearest tree
with a better objective value, then it may take a very large number of
shuffles (or, indeed, it may be impossible) to find a better tree. However,
while the only theoretical guarantee is that Lemma S1 shows it will quickly
find better trees that can be formed by NNI, we expect that stronger
conditions hold on its ability to “escape” from local minima.

Table 2: Evaluation datasets. rRNA/rDNA: ribosomal RNA/DNA, mtDNA: mitochondrial DNA. AA: amino acid. For the Jawed dataset, several subsets of the original dataset [50] were used (from 1,460 to 18,406 sites; cf. Fig. 1(c)). Dataset | # Sites | # Taxa | Type | Taxonomic rank
---|---|---|---|---
(Reference) |  |  |  |
DS1 | 1,949 | 27 | rRNA (18S) | Tetrapods
[43] |  |  |  |
DS2 | 2,520 | 29 | rRNA (18S) | Acanthocephalans
[36] |  |  |  |
DS3 | 1,812 | 36 | mtDNA | Mammals; mainly Lemurs
[115] |  |  |  |
DS4 | 1,137 | 41 | rDNA (18S) | Fungi; mainly Ascomycota
[45] |  |  |  |
DS5 | 378 | 50 | DNA | Lepidoptera
[10, 58] |  |  |  |
DS6 | 1,133 | 50 | rDNA (28S) | Fungi; mainly Diaporthales
[119] |  |  |  |
DS7 | 1,824 | 59 | mtDNA | Mammals; mainly Lemurs
[116] |  |  |  |
DS8 | 1,008 | 64 | rDNA (28S) | Fungi; mainly Hypocreales
[88] |  |  |  |
DS9 | 955 | 67 | DNA | Poaecae (grasses)
[49] |  |  |  |
DS10 | 1,098 | 67 | DNA | Fungi; mainly Ascomycota
[101] |  |  |  |
DS11 | 1,082 | 71 | DNA | Lichen
[55] |  |  |  |
Eutherian | 1,338,678 | 37 | DNA | Eutherian Mammals
[97] |  |  |  |
Jawed | 1,460-18,406 | 99 | AA | Gnathostomata (jawed vertebrates)
[50] |  |  |  |
Primates | 232 | 14 | mtDNA | Mammals; mainly Primates
[42] |  |  |  |
[78] |  |  |  |

#### Computational complexity

The computational complexity for all distance-based algorithms requires an
upfront computational cost of $\mathcal{O}(n^{2})$ to compute the distance
matrix. We will disregard this cost from subsequent comparisons. The standard
neighbour-joining algorithm [91] has an overall computational complexity of
$\mathcal{O}(n^{3})$. FastME has a computational complexity of
$\mathcal{O}(kn^{2}_{\text{Diam}}(T))$ (where Diam$(T)$ is the maximum path
length in a tree, which is generally much smaller than $n$) for $k$ iterations
where $k<n$ when $n$ is large. When fully discrete, our algorithm also has
same complexity but with added mechanisms for escaping optima via Queue
Shuffle. Therefore, a discrete setting is as computationally efficient as
FastME (see Appendix 0.9 for details).

Computing the expectation in Algorithm 1 has complexity $\mathcal{O}(n^{3})$.
A single gradient evaluation (that is, calculating $\frac{\partial F}{\partial
W_{ij}}$ for some $i$ and $j$) is also $\mathcal{O}(n^{3})$ and therefore
computing the full Jacobian is $\mathcal{O}(n^{5})$. Our Queue Shuffle
algorithm runs in $\mathcal{O}(n)$. Therefore, our optimisation for $k$ steps
and $l$ shuffles yields a complexity of $\mathcal{O}(kln^{5})$. The size of
$l$ is dependent on the choice of gradient optimiser, and the size of $k$
varies if a sensible ordering is initialised.

Thus, the computational complexity of GradME is substantially higher than that
of FastME and closer to that of FITCH [32]. This is due to the far greater
mathematical complexity of the continuous objective function, $F(W)$. As it is
an expectation over all possible ordered trees, the explicit formula for
$F(W)$ is a polynomial in $W$ with $(n-1)!$ different terms. Intuitively, the
continuous space always considers a path between any two trees, something that
becomes impossible with discrete settings. Thus, being able to compute it in
polynomial time is a vast improvement on a naive approach, although it is
still considerably less than the innovative FastME greedy approach. More
savings should be possible, we hope to make further efficiency gains in future
work.

#### Evaluation

We evaluate GradME on a diverse corpus of 14 empirical molecular sequence
datasets (Table 2). The first 11 are commonly used to assess phylogenetic
inference performance [109], whereas the last three were used to assess
inference on rooted trees. For each dataset, we start from a random tree and
optimise the $W$ matrix to a tolerance of 1e-10 using gradient descent with
Adafactor [96] optimisation. The distance matrix $D$ is computed using the
GTR+$\Gamma$ substitution model for DNA and an LG model [59] for amino acids.
Substitution model parameters for the GTR+$\Gamma$ are also estimated using
gradient descent with Adafactor using a pairwise maximum likelihood approach
[112]. Jukes-Cantor [51], F81 [30] and TN93 [102] models were also tested for
DNA, while stochastic gradient descent (SGD), RMSprop [103], and AdamW [53,
63] were also considered for optimisation (see Fig. S3). To fairly assess the
performance of GradME, we compare our framework to two well-established
distance-based methods: BioNJ [37], based on the neighbour joining algorithm
[91], and FastME [61], based on balanced minimum evolution.

#### Implementation

Implementation of the BME criterion and the optimisation framework was written
in Python using Jax [8] and Optax [5]. Optimisation was performed on a Xeon
2.30GHz (CPU; Intel Corporation) or on a single GeForce GTX 1080 (GPU; Nvidia
Corporation). Evaluation of the BioNJ [37] and FastME [61] methods was
performed via the R package ape [78] using rpy2 [39]. Tree manipulation and
visualisation scripts were written using ete3 [48] and NetworkX [41]. An
implementation is available at: https://github.com/Neclow/GradME.

## Data availability

All code relevant to reproduce the experiments is available online:
https://github.com/Neclow/GradME.

## References

  * [1] Jozsef Arato and W Tecumseh Fitch  “Phylogenetic signal in the vocalizations of vocal learning and vocal non-learning birds”  In _Philos. Trans. R. Soc. Lond. B Biol. Sci._ 376 The Royal Society, 2021, pp. 20200241
  * [2] K Atteson  “The Performance of Neighbor-Joining Methods of Phylogenetic Reconstruction”  In _Algorithmica_ 25 Springer, 1999, pp. 251–278
  * [3] Stephen W. Attwood et al.  “Phylogenetic and phylodynamic approaches to understanding and combating the early SARS-CoV-pandemic”  In _Nat. Rev. Genet._ 23, 2022, pp. 547–562  DOI: 10.1038/s41576-022-00483-8
  * [4] Daniel L Ayres et al.  “BEAGLE: an application programming interface and high-performance computing library for statistical phylogenetics”  In _Syst. Biol._ 61 Oxford University Press, 2012, pp. 170–173
  * [5] Igor Babuschkin et al.  “The DeepMind JAX Ecosystem”, 2020  URL: http://github.com/deepmind
  * [6] Michael Betancourt  “A Conceptual Introduction to Hamiltonian Monte Carlo”, 2018  arXiv:1701.02434
  * [7] Louis J Billera, Susan P Holmes and Karen Vogtmann  “Geometry of the Space of Phylogenetic Trees”  In _Adv. Appl. Math._ 27, 2001, pp. 733–767
  * [8] James Bradbury et al.  “JAX: composable transformations of Python+NumPy programs”, 2018  URL: http://github.com/google/jax
  * [9] Lindell Bromham and David Penny  “The modern molecular clock”  In _Nat. Rev. Genet._ 4 Nature Publishing Group UK London, 2003, pp. 216–224
  * [10] Andrew VZ Brower  “Phylogenetic relationships among the Nymphalidae (Lepidoptera) inferred from partial sequences of the wingless gene”  In _Proc. R. Soc. Lond. B Biol. Sci._ 267 The Royal Society, 2000, pp. 1201–1211
  * [11] Daniele Catanzaro, Martin Frohn, Olivier Gascuel and Raffaele Pesenti  “A tutorial on the balanced minimum evolution problem”  In _Eur. J. Oper. Res._ 300 Elsevier, 2022, pp. 1–19
  * [12] Daniele Catanzaro, Martine Labbé, Raffaele Pesenti and Juan-José Salazar-González  “The Balanced Minimum Evolution Problem”  In _INFORMS J. Comput._ 24 INFORMS, 2012, pp. 276–294
  * [13] Luigi L Cavalli-Sforza and Anthony WF Edwards  “Phylogenetic analysis. Models and estimation procedures”  In _Am. J. Hum. Genet._ 19 Elsevier, 1967, pp. 233
  * [14] H.-C. Chen and Y.-L. Wang  “An efficient algorithm for generating Prüfer codes from labelled trees”  In _Theory Comput. Syst._ 33 Springer, 2000, pp. 97–105
  * [15] Olga Chernomor, Arndt Haeseler and Bui Quang Minh  “Terrace Aware Data Structure for Phylogenomic Inference from Supermatrices”  In _Syst. Biol._ 65, 2016, pp. 997–1008
  * [16] Mark Collard, Stephen J Shennan and Jamshid J Tehrani  “Branching, blending, and the evolution of cultural similarities and differences among human populations”  In _Evol. Hum. Behav._ 27 Elsevier, 2006, pp. 169–184
  * [17] Lena Collienne and Alex Gavryushkin  “Computing nearest neighbour interchange distances between ranked phylogenetic trees”  In _J. Math. Biol._ 82, 2021, pp. 8
  * [18] William HE Day, David S Johnson and David Sankoff  “The computational complexity of inferring rooted phylogenies by parsimony”  In _Math. Biosci._ 81 Elsevier, 1986, pp. 33–42
  * [19] Nicola De Maio et al.  “Maximum likelihood pandemic-scale phylogenetics”  In _Nat. Genet._ Nature Publishing Group US New York, 2023, pp. 1–7
  * [20] Richard Desper and Olivier Gascuel  “Fast and accurate phylogeny reconstruction algorithms based on the minimum-evolution principle”  In _J. Comput. Biol._ 9, 2002, pp. 687–705
  * [21] Richard Desper and Olivier Gascuel  “Theoretical foundation of the balanced minimum evolution method of phylogenetic inference and its relationship to weighted least-squares tree fitting”  In _Mol. Biol. Evol._ 21, 2004, pp. 587–598
  * [22] Persi W Diaconis and Susan P Holmes  “Matchings and phylogenetic trees”  In _Proc. Natl. Acad. Sci. U.S.A._ 95 National Acad Sciences, 1998, pp. 14600–14602
  * [23] Vu Dinh, Arman Bilge, Cheng Zhang and Frederick A. Matsen  “Probabilistic Path Hamiltonian Monte Carlo”  In _ICML_ 70 PMLR, 2017, pp. 1009–1018
  * [24] Mario Dos Reis, Philip CJ Donoghue and Ziheng Yang  “Bayesian molecular clock dating of species divergences in the genomics era”  In _Nat. Rev. Genet._ 17 Nature Publishing Group UK London, 2016, pp. 71–80
  * [25] Alexei J Drummond and Andrew Rambaut  “BEAST: Bayesian evolutionary analysis by sampling trees”  In _BMC Evol. Biol._ 7, 2007, pp. 214
  * [26] Sebastián Duchêne et al.  “Cross-validation to select Bayesian hierarchical models in phylogenetics”  In _BMC Evol. Biol._ 16, 2016, pp. 115
  * [27] Nuno R Faria et al.  “Genomics and epidemiology of the P. 1 SARS-CoV-2 lineage in Manaus, Brazil”  In _Science_ 372 American Association for the Advancement of Science, 2021, pp. 815–821
  * [28] Joseph Felsenstein  “Cases in which Parsimony or Compatibility Methods will be Positively Misleading”  In _Syst. Biol._ 27 Oxford Academic, 1978, pp. 401–410
  * [29] Joseph Felsenstein  “The Number of Evolutionary Trees”  In _Syst. Biol._ 27 Oxford Academic, 1978, pp. 27–33
  * [30] Joseph Felsenstein  “Evolutionary trees from DNA sequences: a maximum likelihood approach”  In _J. Mol. Evol._ 17 Springer, 1981, pp. 368–376
  * [31] Joseph Felsenstein  “Statistical inference of phylogenies”  In _J R Stat Soc Ser A_ 146 Wiley Online Library, 1983, pp. 246–262
  * [32] Joseph Felsenstein  “An alternating least squares approach to inferring phylogenies from pairwise distances”  In _Syst. Biol._ 46 Society of Systematic Biologists, 1997, pp. 101–111
  * [33] Joseph Felsenstein  “Inferring phylogenies”  Sunderland, MA: Sinauer Associates, 2004
  * [34] W M Fitch and E Margoliash  “Construction of phylogenetic trees”  In _Science_ 155 American Association for the Advancement of Science (AAAS), 1967, pp. 279–284
  * [35] Les R Foulds and Ronald L Graham  “The Steiner problem in phylogeny is NP-complete”  In _Adv. Appl. Math._ 3 Elsevier, 1982, pp. 43–49
  * [36] James R Garey, Thomas J Near, Michael R Nonnemacher and Steven A Nadler  “Molecular evidence for Acanthocephala as a subtaxon of Rotifera”  In _J. Mol. Evol._ 43 Springer, 1996, pp. 287–292
  * [37] Olivier Gascuel  “BIONJ: an improved version of the NJ algorithm based on a simple model of sequence data.”  In _Mol. Biol. Evol._ 14, 1997, pp. 685–695
  * [38] Olivier Gascuel and Mike Steel  “Neighbor-joining revealed”  In _Mol. Biol. Evol._ 23, 2006, pp. 1997–2000
  * [39] Laurent Gautier and Michał Krassowski  “rpy2: Python interface to the R language”, 2021  URL: https://github.com/rpy2/rpy2
  * [40] Bryan T Grenfell et al.  “Unifying the epidemiological and evolutionary dynamics of pathogens”  In _Science_ 303 American Association for the Advancement of Science (AAAS), 2004, pp. 327–332
  * [41] Aric A. Hagberg, Daniel A. Schult and Pieter J. Swart  “Exploring network structure, dynamics, and function using NetworkX”  In _Proc. SciPy_ , 2008, pp. 11–15
  * [42] Masami Hasegawa and Hirohisa Kishino  “Confidence limits of the maximum-likelihood estimate of the hominoid three from mitochondrial-DNA sequences”  In _Evolution_ 43 JSTOR, 1989, pp. 672–677
  * [43] S Blair Hedges, Kirk D Moberg and Linda R Maxson  “Tetrapod phylogeny inferred from 18S and 28S ribosomal RNA sequences and a review of the evidence for amniote relationships.”  In _Mol. Biol. Evol._ 7, 1990, pp. 607–633
  * [44] M.. Hendy and David Penny  “Branch and bound algorithms to determine minimal evolutionary trees”  In _Math. Biosci._ 59, 1982, pp. 277–290
  * [45] Daniel A Henk, Alex Weir and Meredith Blackwell  “Laboulbeniopsis termitarius, an ectoparasite of termites newly recognized as a member of the Laboulbeniomycetes”  In _Mycologia_ 95 Taylor & Francis, 2003, pp. 561–564
  * [46] John P Huelsenbeck, Jonathan P Bollback and Amy M Levine  “Inferring the root of a phylogenetic tree”  In _Syst. Biol._ 51, 2002, pp. 32–43
  * [47] John P Huelsenbeck and Fredrik Ronquist  “MRBAYES: Bayesian inference of phylogenetic trees”  In _Bioinformatics_ 17 Oxford University Press, 2001, pp. 754–755
  * [48] Jaime Huerta-Cepas, François Serra and Peer Bork  “ETE 3: reconstruction, analysis, and visualization of phylogenomic data”  In _Mol. Biol. Evol._ 33 Society for Mol. Biol. Evol., 2016, pp. 1635–1638
  * [49] Amanda L Ingram and Jeff J Doyle  “Is Eragrostis (Poaceae) monophyletic? Insights from nuclear and plastid sequence data”  In _Syst. Bot._ 29 American Society of Plant Taxonomists, 2004, pp. 545–552
  * [50] Iker Irisarri et al.  “Phylotranscriptomic consolidation of the jawed vertebrate timetree”  In _Nat. Ecol. Evol._ 1 Nature Publishing Group UK London, 2017, pp. 1370–1378
  * [51] Thomas H Jukes and Charles R Cantor  “Evolution of protein molecules”  In _Mamm. Protein Metab._ 3, 1969, pp. 21–132
  * [52] K K Kidd and L A Sgaramella-Zonta  “Phylogenetic analysis: concepts and methods”  In _Am. J. Hum. Genet._ 23, 1971, pp. 235–252
  * [53] Diederik P. Kingma and Jimmy Ba  “Adam: A Method for Stochastic Optimization”  In _ICLR_ , 2015
  * [54] Alexey M Kozlov et al.  “RAxML-NG: a fast, scalable and user-friendly tool for maximum likelihood phylogenetic inference”  In _Bioinformatics_ 35 Oxford University Press, 2019, pp. 4453–4455
  * [55] Scott Kroken and John W Taylor  “Phylogenetic species, reproductive mode, and specificity of the green alga Trebouxia forming lichens with the fungal genus Letharia”  In _Bryologist_ JSTOR, 2000, pp. 645–660
  * [56] Mary K Kuhner and Joseph Felsenstein  “A simulation comparison of phylogeny algorithms under equal and unequal evolutionary rates.”  In _Mol. Biol. Evol._ 11, 1994, pp. 459–468
  * [57] Sudhir Kumar and Sudhindra R Gadagkar  “Efficiency of the neighbor-joining method in reconstructing deep and shallow evolutionary relationships in large phylogenies”  In _J. Mol. Evol._ 51 Springer Nature BV, 2000, pp. 544
  * [58] Clemens Lakner et al.  “Efficiency of Markov chain Monte Carlo tree proposals in Bayesian phylogenetics”  In _Syst. Biol._ 57 Taylor & Francis, 2008, pp. 86–103
  * [59] Si Quang Le and Olivier Gascuel  “An improved general amino acid replacement matrix”  In _Mol. Biol. Evol._ 25, 2008, pp. 1307–1320
  * [60] Michael SY Lee and Alessandro Palci  “Morphological phylogenetics in the genomic age”  In _Curr. Biol._ 25 Elsevier, 2015, pp. R922–R929
  * [61] Vincent Lefort, Richard Desper and Olivier Gascuel  “FastME 2.0: a comprehensive, accurate, and fast distance-based phylogeny inference program”  In _Mol. Biol. Evol._ 32 Oxford University Press, 2015, pp. 2798–2800
  * [62] Frédéric Lemoine, Luc Blassel, Jakub Voznica and Olivier Gascuel  “COVID-Align: Accurate online alignment of hCoV-19 genomes using a profile HMM”  In _Bioinformatics_ 37 Oxford University Press, 2021, pp. 1761–1762
  * [63] Ilya Loshchilov and Frank Hutter  “Decoupled Weight Decay Regularization”  In _ICLR_ , 2019
  * [64] Matthew Macaulay, Aaron Darling and Mathieu Fourment  “Fidelity of hyperbolic space for Bayesian phylogenetic inference”  In _PLoS Comput. Biol._ 19 Public Library of Science San Francisco, CA USA, 2023, pp. e1011084
  * [65] Ruth Mace and Clare J Holden  “A phylogenetic approach to cultural evolution”  In _Trends Ecol. Evol._ 20 Elsevier, 2005, pp. 116–121
  * [66] Hirotaka Matsumoto, Takahiro Mimori and Tsukasa Fukunaga  “Novel metric for hyperbolic phylogenetic tree embeddings”  In _Biol. Methods Protoc._ 6, 2021, pp. bpab006
  * [67] Radu Mihaescu, Dan Levy and Lior Pachter  “Why neighbor-joining works”  In _Algorithmica_ 54 Springer ScienceBusiness Media LLC, 2009, pp. 1–24
  * [68] Radu Mihaescu and Lior Pachter  “Combinatorics of least-squares trees”  In _Proc. Natl. Acad. Sci. U.S.A._ 105, 2008, pp. 13206–13211
  * [69] Takahiro Mimori and Michiaki Hamada  “GeoPhy: Differentiable Phylogenetic Inference via Geometric Gradients of Tree Topologies”, 2023  arXiv:2307.03675 [cs.LG]
  * [70] Bui Quang Minh et al.  “IQ-TREE 2: new models and efficient methods for phylogenetic inference in the genomic era”  In _Mol. Biol. Evol._ 37 Oxford University Press, 2020, pp. 1530–1534
  * [71] Hélène Morlon, Todd L Parsons and Joshua B Plotkin  “Reconciling molecular phylogenies with the fossil record”  In _Proc. Natl. Acad. Sci. U.S.A._ 108, 2011, pp. 16327–16332
  * [72] David A Morrison  “Are phylogenetic patterns the same in anthropology and biology?”  In _bioRxiv_ , 2014  BIORXIV: https://www.biorxiv.org/content/early/2014/06/24/006486.full.pdf
  * [73] Suha Naser-Khdour, Bui Quang Minh and Robert Lanfear  “Assessing Confidence in Root Placement on Phylogenies: An Empirical Study Using Nonreversible Models for Mammals”  In _Syst. Biol._ 71, 2022, pp. 959–972
  * [74] Luca Nesterenko, Bastien Boussau and Laurent Jacob  “Phyloformer: towards fast and accurate phylogeny estimation with self-attention networks”, 2022  BIORXIV: https://www.biorxiv.org/content/early/2022/06/28/2022.06.24.496975.full.pdf
  * [75] Lam-Tung Nguyen, Heiko A Schmidt, Arndt Von Haeseler and Bui Quang Minh  “IQ-TREE: a fast and effective stochastic algorithm for estimating maximum-likelihood phylogenies”  In _Mol. Biol. Evol._ 32 Oxford University Press, 2015, pp. 268–274
  * [76] Brian C O’Meara  “Evolutionary inferences from phylogenies: a review of methods”  In _Annu. Rev. Ecol. Evol. Syst._ 43 Annual Reviews, 2012, pp. 267–285
  * [77] Áine O’Toole et al.  “Assignment of epidemiological lineages in an emerging pandemic using the pangolin tool”  In _Virus Evol._ 7 Oxford University Press UK, 2021, pp. veab064
  * [78] Emmanuel Paradis, Julien Claude and Korbinian Strimmer  “APE: analyses of phylogenetics and evolution in R language”  In _Bioinformatics_ 20 Oxford University Press, 2004, pp. 289–290
  * [79] Hyun Jung Park, Seung-Jin Sul and Tiffani L Williams  “Large-scale analysis of phylogenetic search behavior”  In _Adv. Exp. Med. Biol._ 680, 2010, pp. 35–42
  * [80] Y Pauplin  “Direct calculation of a tree length using a distance matrix”  In _J. Mol. Evol._ 51, 2000, pp. 41–47
  * [81] Matthew J Penn et al.  “Phylo2Vec: a vector representation for binary trees”, 2023  arXiv:2304.12693
  * [82] Morgan N Price, Paramvir S Dehal and Adam P Arkin  “FastTree: computing large minimum evolution trees with profiles instead of a distance matrix”  In _Mol. Biol. Evol._ 26 Oxford University Press, 2009, pp. 1641–1650
  * [83] Alan Queiroz and John Gatesy  “The supermatrix approach to systematics”  In _Trends Ecol. Evol._ 22 Elsevier, 2007, pp. 34–41
  * [84] David F Robinson and Leslie R Foulds  “Comparison of phylogenetic trees”  In _Math. Biosci._ 53 Elsevier, 1981, pp. 131–147
  * [85] Sebastien Roch  “A short proof that phylogenetic tree reconstruction by maximum likelihood is hard”  In _IEEE/ACM Trans. Comput. Biol. Bioinform._ 3, 2006, pp. 92–94
  * [86] Antonis Rokas, Barry L Williams, Nicole King and Sean B Carroll  “Genome-scale approaches to resolving incongruence in molecular phylogenies”  In _Nature_ 425 Nature Publishing Group, 2003, pp. 798–804
  * [87] Jonathan Rolland et al.  “Using phylogenies in conservation: New perspectives”  In _Biol. Lett._ 8, 2011, pp. 692–694
  * [88] Amy Y Rossman, John M McKemy, Rebecca A Pardo-Schultheiss and Hans-Josef Schroers  “Molecular studies of the Bionectriaceae using large subunit rDNA sequences”  In _Mycologia_ 93 Taylor & Francis, 2001, pp. 100–110
  * [89] Andrey Rzhetsky and Masatoshi Nei  “A simple method for estimating and testing minimum-evolution trees”  In _Mol. Biol. Evol._ 9, 1992, pp. 945–967
  * [90] N Saitou and M Nei  “The neighbor-joining method: a new method for reconstructing phylogenetic trees”  In _Mol. Biol. Evol._ 4, 1987, pp. 406–425
  * [91] Naruya Saitou and Masatoshi Nei  “The neighbor-joining method: a new method for reconstructing phylogenetic trees.”  In _Mol. Biol. Evol._ 4, 1987, pp. 406–425
  * [92] Michael J Sanderson, Michelle M McMahon and Mike Steel  “Terraces in phylogenetic tree space”  In _Science_ 333 American Association for the Advancement of Science, 2011, pp. 448–450
  * [93] Michael J Sanderson and H Bradley Shaffer  “Troubleshooting molecular phylogenetic analyses”  In _Annu. Rev. Ecol. Syst._ 33 Annual Reviews 4139 El Camino Way, PO Box 10139, Palo Alto, CA 94303-0139, USA, 2002, pp. 49–72
  * [94] Theo Sanderson  “Taxonium, a web-based tool for exploring large phylogenetic trees”  In _eLife_ 11 eLife Sciences Publications, Ltd, 2022, pp. e82392
  * [95] Charles Semple and Mike Steel  “Cyclic permutations and evolutionary trees”  In _Adv. Appl. Math._ 32, 2004, pp. 669–680
  * [96] Noam Shazeer and Mitchell Stern  “Adafactor: Adaptive learning rates with sublinear memory cost”  In _ICML_ , 2018, pp. 4596–4604  PMLR
  * [97] Sen Song, Liang Liu, Scott V Edwards and Shaoyuan Wu  “Resolving conflict in eutherian mammal phylogeny using phylogenomics and the multispecies coalescent model”  In _Proc. Natl. Acad. Sci. U.S.A._ 109, 2012, pp. 14942–14947
  * [98] Katherine St.  “The Shape of Phylogenetic Treespace”  In _Syst. Biol._ 66, 2017, pp. e83–e94
  * [99] Alexandros Stamatakis  “RAxML version 8: a tool for phylogenetic analysis and post-analysis of large phylogenies”  In _Bioinformatics_ 30 Oxford University Press, 2014, pp. 1312–1313
  * [100] Marc A Suchard, Christina M R Kitchen, Janet S Sinsheimer and Robert E Weiss  “Hierarchical phylogenetic models for analyzing multipartite sequence data”  In _Syst. Biol._ 52, 2003, pp. 649–664
  * [101] Sung-Oui Suh and Meredith Blackwell  “Molecular phylogeny of the cleistothecial fungi placed in Cephalothecaceae and Pseudeurotiaceae”  In _Mycologia_ 91 Taylor & Francis, 1999, pp. 836–848
  * [102] Koichiro Tamura and Masatoshi Nei  “Estimation of the number of nucleotide substitutions in the control region of mitochondrial DNA in humans and chimpanzees.”  In _Mol. Biol. Evol._ 10, 1993, pp. 512–526
  * [103] Tijmen Tieleman and Geoffrey Hinton  “Lecture 6e - rmsprop: Divide the gradient by a running average of its recent magnitude” Slides, Coursera: Neural networks for machine learning, 2012  URL: https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
  * [104] Thu-Hien To, Matthieu Jung, Samantha Lycett and Olivier Gascuel  “Fast Dating Using Least-Squares Criteria and Algorithms”  In _Syst. Biol._ 65, 2016, pp. 82–97
  * [105] Fernando Domingues Kümmel Tria, Giddy Landan and Tal Dagan  “Phylogenetic rooting using minimal ancestor deviation”  In _Nat. Ecol. Evol._ 1, 2017, pp. 193
  * [106] Yatish Turakhia et al.  “Pandemic-scale phylogenomics reveals the SARS-CoV-2 recombination landscape”  In _Nature_ 609 Nature Publishing Group UK London, 2022, pp. 994–997
  * [107] Mezzalina Vankan, Simon Y W Ho and David A Duchêne  “Evolutionary Rate Variation among Lineages in Gene Trees has a Negative Impact on Species-Tree Inference”  In _Syst. Biol._ 71, 2022, pp. 490–500
  * [108] Jakub Voznica et al.  “Deep learning from phylogenies to uncover the epidemiological dynamics of outbreaks”  In _Nat. Commun._ 13 Nature Publishing Group, 2022, pp. 1–14
  * [109] Chris Whidden and Frederick A Matsen  “Quantifying MCMC exploration of phylogenetic tree space”  In _Syst. Biol._ 64, 2015, pp. 472–491
  * [110] James C Wilgenbusch and David Swofford  “Inferring evolutionary trees with PAUP”  In _Curr. Protoc. Bioinformatics_ Wiley Online Library, 2003, pp. 6.4.1–6.4.28
  * [111] Benjamin Wilson  “Learning phylogenetic trees as hyperbolic point configurations”, 2021  arXiv:2104.11430
  * [112] Ziheng Yang  “Computational molecular evolution”  Oxford University Press, 2006
  * [113] Ziheng Yang  “PAML 4: phylogenetic analysis by maximum likelihood”  In _Mol. Biol. Evol._ 24, 2007, pp. 1586–1591
  * [114] Ziheng Yang and Bruce Rannala  “Molecular phylogenetics: principles and practice”  In _Nat. Rev. Genet._ 13 Nature Publishing Group UK London, 2012, pp. 303–314
  * [115] Ziheng Yang and Anne D Yoder  “Comparison of likelihood and Bayesian methods for estimating divergence times using multiple gene loci and calibration points, with application to a radiation of cute-looking mouse lemur species”  In _Syst. Biol._ 52 Society of Systematic Zoology, 2003, pp. 705–716
  * [116] Anne D Yoder and Ziheng Yang  “Divergence dates for Malagasy lemurs estimated from multiple gene loci: geological and evolutionary context”  In _Mol. Ecol._ 13 Wiley Online Library, 2004, pp. 757–773
  * [117] Chao Zhang, Maryam Rabiee, Erfan Sayyari and Siavash Mirarab  “ASTRAL-III: polynomial time species tree reconstruction from partially resolved gene trees”  In _BMC Bioinform._ 19 BioMed Central, 2018, pp. 15–30
  * [118] Cheng Zhang and Frederick A Matsen IV  “Variational Bayesian phylogenetic inference”  In _ICLR_ , 2018
  * [119] Ning Zhang and Meredith Blackwell  “Molecular phylogeny of dogwood anthracnose fungus (Discula destructiva) and the Diaporthales”  In _Mycologia_ 93 Taylor & Francis, 2001, pp. 355–365
  * [120] Emile Zuckerkandl and Linus Pauling  “Horizons in Biochemistry”  New York: Academic Press, 1962, pp. 189–225

## Author contributions

S.B and M.J.P conceived of the study. S.B and C.A.D supervised. S.B, N.S,
M.J.P and D.A.D designed the study. S.B and N.S performed optimisation runs.
S.B, M.J.P, N.S and D.A.D performed analysis. All authors contributed to
writing the original draft. M.J.P and J.P drafted the appendix.

## Competing interests

The authors declare no competing interests.

### Funding

S.B. and C.A.D acknowledges support from the MRC Centre for Global Infectious
Disease Analysis (MR/R015600/1), jointly funded by the UK Medical Research
Council (MRC) and the UK Foreign, Commonwealth & Development Office (FCDO),
under the MRC/FCDO Concordat agreement, and also part of the EDCTP2 programme
supported by the European Union. SB is funded by the National Institute for
Health Research (NIHR) Health Protection Research Unit (HPRU) in Modelling and
Health Economics, a partnership between the UK Health Security Agency,
Imperial College London and LSHTM (grant code NIHR200908). Disclaimer: “The
views expressed are those of the author(s) and not necessarily those of the
NIHR, UK Health Security Agency or the Department of Health and Social Care.”
S.B. acknowledges support from the Novo Nordisk Foundation via The Novo
Nordisk Young Investigator Award (NNF20OC0059309). S.B. acknowledges support
from the Danish National Research Foundation via a chair grant which also
supports N.S. S.B. acknowledges support from The Eric and Wendy Schmidt Fund
For Strategic Innovation via the Schmidt Polymath Award (G-22-63345). S.B and
N.S acknowledge the Pioneer Centre for AI, DNRF grant number P1 as affiliate
researchers. C.A.D receives support from the NIHR HPRU in Emerging and
Zoonotic Infections, a partnership between the UK Health Security Agency,
University of Liverpool, University of Oxford and Liverpool School of Tropical
Medicine (grant code NIHR200907). D.A.D. is funded by a European Research
Council Marie Sklodowska-Curie fellowship (H2020-MSCA-IF-2019-883832). M.P.
was funded by a DTP Studentship from the Engineering and Physical Sciences
Research (EPSRC)

## Appendix

###  0.1 BME for rooted trees

#### Comparing the rooted and unrooted objectives

######  Lemma 0.2.

Consider adding an extra taxon, $n$, to the set of taxa such that, for some
$D^{*}$ and $\delta$.

$|D_{ni}-D^{*}|<\delta\quad\forall i\neq n$ (9)

This creates an unrooted tree $\mathcal{T}^{u}$, and we can create a rooted
tree $\mathcal{T}^{r}$ by removing node $n$. If $e^{u}_{ij}$ denotes inter-
taxa distance in $\mathcal{T}^{u}$ and $e_{ij}^{r}$ denotes inter-taxa
distance in $\mathcal{T}^{r}$, then

$\bigg{|}\sum_{i=0}^{n}\sum_{j=0}^{n}D_{ij}2^{-e^{u}_{ij}}-\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}D_{ij}2^{-e^{r}_{ij}}-D^{*}\bigg{|}\leq\delta$
(10)

Proof: Firstly, note that for $i,j<n$,

$e_{ij}^{r}=e_{ij}^{u}$ (11)

as the path between $i$ and $j$ will not contain any leaf nodes other than $i$
and $j$, and therefore will not contain node $n$. Then, using the fact that
$D_{nn}=0$,

$\displaystyle\sum_{i=0}^{n}\sum_{j=0}^{n}D_{ij}2^{-e^{u}_{ij}}$
$\displaystyle=\sum_{j=0}^{n-1}D_{nj}2^{-e^{u}_{nj}}+\sum_{i=0}^{n-1}D_{in}2^{-e^{u}_{in}}+\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}D_{ij}2^{-e^{u}_{ij}}$
(12)
$\displaystyle\leq\sum_{j=0}^{n-1}(D^{*}+\delta)2^{-e_{nj}}+\sum_{i=0}^{n-1}(D^{*}+\delta)2^{-e^{u}_{in}}+\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}D_{ij}2^{-e^{r}_{ij}}$
(13)

We can make progress with this sum by noting the Kraft Equality, found
throughout the literature [12]

$\sum_{j}2^{-e^{u}_{nj}}=\frac{1}{2}$ (14)

This means that

$\sum_{i=0}^{n}\sum_{j=0}^{n}D_{ij}2^{-e^{u}_{ij}}\leq\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}D_{ij}2^{-e^{r}_{ij}}+D^{*}+\delta$
(15)

Similarly, one can show

$\sum_{i=0}^{n}\sum_{j=0}^{n}D_{ij}2^{-e^{u}_{ij}}\geq\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}D_{ij}2^{-e^{r}_{ij}}+D^{*}-\delta$
(16)

and hence the result follows.

#### Understanding the BME rooting

######  Definition 0.3.

Distance to root heuristic: Under the assumption that the tree is ultrametric,
we estimate the distance between the root and the leaf taxa using the
following algorithm:

1) Find two leaf nodes that share a parent that is not the root. If no such
pair exists, then move to step 3.

2) Let $d_{1}$ and $d_{2}$ be the distance between each leaf node and its
parent. Suppose that $d_{3}$ is the distance between the parent and its parent
(i.e. the grandparent of the leaves). Remove the leaves from the tree, update
the distance between the parent and its parent to be
$d_{3}+\frac{d_{1}+d_{2}}{2}$ and return to step 1

3) If $d_{1}$ and $d_{2}$ are the distances between the two remaining children
and the root, then the distance between the root and the taxa is
$\frac{d_{1}+d_{2}}{2}$.

######  Lemma 0.4.

Consider the true unrooted tree $\mathcal{T}$ and suppose that the distance
matrix $D$ gives the true (time) distances between each pair of taxa in
$\mathcal{T}$.

Then, the optimal rooting — that is, the edge in the tree such that placing
the root on this edge minimizes the BME objective — minimizes the distance to
root heuristic defined in Definition 0.3.

Proof: Choose an edge, $b$ on which to place the root. We define an indicator
function $X_{b}$ such that $X_{b}(A,B)$ is 1 if $b$ is on the path between
nodes $A$ and $B$ and 0 otherwise. Adding a root to edge $b$ changes the
objective function by halving the weight assigned to each $D_{AB}$ such that
$X_{b}(A,B)=1$ as the path length between these nodes will increase by 1.
Thus, using $f_{r}$ to denote the rooted objective and $f_{u}$ to denote the
unrooted objective,

$f_{r}=f_{u}-\frac{1}{2}\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}$ (17)

where the $e_{AB}$ are the path lengths in the original, unrooted tree and
both $A$ and $B$ are allowed to vary in the sum. Hence, for a fixed unrooted
tree topology, the optimal rooting will solve

$\max_{b}\bigg{\\{}\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}\bigg{\\}}$ (18)

By Lemma 0.16,

$\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}=2\sum_{A}2^{-g_{A}}D_{Ar}$ (19)

while, by Lemma 0.17, the root-to-tip heuristic, $\mathcal{D}$, satisfies

$\mathcal{D}=\sum_{A}2^{-g_{A}}D_{Ar}$ (20)

and hence, the optimal rooting solves

$\max_{b}\bigg{\\{}\mathcal{D}\bigg{\\}}$ (21)

as required.

###  0.2 Ordered Trees

######  Definition 0.5.

Left-to-right construction algorithm An ordered tree can be constructed as
follows:

1) Begin with nodes $0$ and $1$, each joined to a root. Label the edge joining
node $0$ to the root as edge $0$, and the edge joining node $1$ to the root as
edge $1$

2) Process the nodes in order 2, 3, 4… When processing node $k$, join it to
edge $v_{k}$ (creating a new internal node — this is well-defined as
$v_{k}<k$). Label the edge joining node $k$ to the tree as edge $k$.

######  Lemma 0.6.

The left-to-right algorithm in Definition 0.5 and the standard Phylo2Vec
algorithm defined in [81] produce the same tree, provided $\bm{v}$ is ordered.

Proof: To show that the left-to-right algorithm gives an equivalent tree,
define $\mathcal{T}$ to be the tree resulting from the standard Phylo2Vec
algorithm and $\mathcal{T}^{\prime}$ to be the tree resulting from this new
left-to-right algorithm. We proceed by induction on the number of nodes, $n$,
noting that the case $n=2$ is trivial.

Suppose now that the algorithms are equivalent for $n=m$. Choose some ordered
$\bm{v}$ of length $m+1$ (so that this corresponds to $n=m+1$) and consider
processing the first node using the Phylo2Vec algorithm, so that (using the
fact that $\bm{v}$ is ordered so that no nodes are skipped), node $m$ merges
with node $v_{m}$.

From this step, the Phylo2Vec algorithm proceeds as if there were $n-1$ nodes
and the vector was $\tilde{\bm{v}}=(v_{0},v_{1},...,v_{m-1})$. Define
$\tilde{\mathcal{T}}$ to be the tree given by $\tilde{\bm{v}}$. Hence,
$\tilde{\mathcal{T}}$ can be created from $\mathcal{T}$ by removing nodes
$v_{m}$ and $m$ (and the edges connecting them to the tree) and relabelling
their parent as $v_{m}$. Equivalently (by reversing this process),
$\mathcal{T}$ can be created from $\tilde{\mathcal{T}}$ by adding a node to
the edge joining leaf node $v_{m}$ to the tree, and connecting node $m$ to
this edge.

Moreover, from the inductive hypothesis, $\tilde{\mathcal{T}}$ can be
constructed by using the left-to-right algorithm on $\tilde{\bm{v}}$. As this
algorithm processes $v_{m}$ last, this means that after $m$ nodes have been
added to the tree by the left-to-right algorithm applied to $\bm{v}$, the
current tree is given by $\tilde{\mathcal{T}}$.

The final step of the left-to-right algorithm applied to $\bm{v}$ is to add a
node to the edge joining leaf node $v_{m}$ to the tree, and to connect node
$m$ to this edge. As previously discussed, this creates the tree $\mathcal{T}$
and hence, $\mathcal{T}=\mathcal{T}^{\prime}$ as required.

###  0.3 A continuous objective function

#### Construction

######  Lemma 0.7.

For a randomly chosen tree with distribution $W$, define, for $i,j,<k$,
$e_{ij}^{k}$ to be the path length between taxa $i$ and $j$ when $k$ nodes
have been added to the tree (using the left-to-right construction algorithm).
Define $E_{ik}^{k}:=\mathbb{E}(2^{-e_{ij}^{k}})$. Then, for $i<j$

$E_{ij}^{k+1}=\left\\{\begin{matrix}E^{k}_{ij}\bigg{[}1-\frac{1}{2}(W_{ki}+W_{kj})\bigg{]}&\text{if
$i<j<k$}\\\ \bigg{[}\frac{1}{2}\sum_{x\neq
i}E_{ix}^{k}W_{kx}\bigg{]}+\frac{1}{4}W_{ki}&\text{if
$i<k$}\end{matrix}\right.$ (22)

with the remaining values following by symmetry.

Proof: Adding node $k$ to the tree increases the path length between nodes $i$
and $j$ by 1 if and only if $V_{k}=i$ or $V_{k}=j$. As this condition is
independent of other values of $\bm{V}$, using $e_{ij}^{k}$ to be the path
length after the $k$ nodes $\\{0,...,k-1\\}$ have been added, one can write

$2^{-e^{k+1}_{ij}}=2^{-(e^{k}_{ij}+\mathds{I}\\{V_{k}\in\\{i,j\\}\\})}=2^{-e^{k}_{ij}}\times
2^{-\mathds{I}\\{V_{k}\in\\{i,j\\}\\}}$ (23)

This is a product of independent random variables and so, defining
$E^{k}_{ij}:=\mathbb{E}(2^{-e^{k}_{ij}})$ and noting that
$\mathds{I}\\{V_{k}\in\\{i,j\\}\\}$ is a Bernoulli random variable with
probability $W_{ki}+W_{kj}$, this equation becomes

$E^{k+1}_{ij}=E^{k}_{ij}\bigg{[}(1-(W_{ki}+W_{kj})+\frac{1}{2}(W_{ki}+W_{kj})\bigg{]}=E^{k}_{ij}\bigg{[}1-\frac{1}{2}(W_{ki}+W_{kj})\bigg{]}$
(24)

To close this iterative system, note that when node $j$ is added to the tree,
it will be a path length 2 from the leaf node $V_{j}$ connecting the edge it
is joined to. Moreover, the distance between node $j$ and any other nodes
$x\neq V_{j}$ in the tree will be equal to one plus the distance between node
$x$ and node $V_{j}$. That is,

$2^{-e_{ij}^{j+1}}=\left\\{\begin{matrix}2^{-(1+e_{i,V_{j}}^{j})}&\text{if
$V_{j}\neq i$}\\\ \frac{1}{4}&\text{if $V_{j}=i$}\end{matrix}\right.$ (25)

Thus, conditioning on the value of $V_{j}$,

$E_{ij}^{j+1}=\bigg{[}\frac{1}{2}\sum_{x\neq
i}E_{ix}^{j}W_{jx}\bigg{]}+\frac{1}{4}W_{ji}$ (26)

Finally, by symmetry, $E_{ji}^{j+1}=E_{ij}^{j+1}$. Noting that $E_{ij}^{k}$ is
undefined (and unnecessary) for $k<\max(i,j)+1$, as nodes $i$ and $j$ have not
both been added to the tree, (24) and (26) hence form a closed system. This
can be solved inductively, finding all $E_{ij}^{m}$ terms for $m=2,3,...,n$.

#### Discrete minima

######  Lemma 0.8.

Define $f(\bm{v})$ to be the objective function for the discrete tree given by
the Phylo2Vec vector $\bm{v}$. Then, if $V^{*}$ is the set of vectors $\bm{v}$
which minimize f, any optimal $W$ satisfies

$\mathbb{P}(\bm{V}\in V^{*}|W)=1$ (27)

Moreover if $|V^{*}|=1$ then there is a unique $\bm{v}$ minimizing $f$, then
there is a unique $W$ minimizing $F$, which is the matrix such that

Proof: Define $V$ to be the set of ordered tree vectors. Then,

$\displaystyle F(W)=\sum_{\bm{u}\in
V}\mathbb{P}(\bm{V}=\bm{u}|W)f(\bm{u})=\sum_{\bm{u}\in
V}\prod_{m=1}^{n-1}W_{m,u_{m}}f(\bm{u})$ (28)

As

$\sum_{\bm{u}\in V}\mathbb{P}(\bm{V}=\bm{u}|W)=1$ (29)

we see that $F(W)$ is a weighted average of the values of $f(\bm{u})$ for
$\bm{u}\in V$. Thus, for any $\bm{v}\in V^{*}$

$F(W)\geq f(\bm{v})$ (30)

and

$F(W)=f(\bm{v})\quad\Rightarrow\quad\mathbb{P}(\bm{V}\in V^{*}|W)=1$ (31)

as required. If $V=\\{\bm{v}\\}$, this minimum requires

$\mathbb{P}(\bm{V}=\bm{v}|W)=1$ (32)

and therefore, using (28) $W$ is uniquely defined by

$W_{m,j}=\mathds{I}\\{v_{m}=j\\}\quad\forall m,j$ (33)

as required.

###  0.4 Orderings

This section considers the labelling algorithm introduced in the main text.

######  Lemma 0.9.

Two nodes have the same label only if they share an ancestor with that label.

Proof: Suppose that this is false, and that nodes $x$ and $y$ have the same
label, $L$, but do not share an ancestor with that label. Define $a(x)$ and
$a(y)$ to be the nodes of lowest generation (that is, the nodes closest to the
root) with label $L$ such that they are ancestors of $x$ and $y$ respectively.
Note that, by assumption, $a(x)\neq a(y)$ and also, neither can be the root
(as the root is the ancestor of all nodes). Moreover, they cannot share a
parent, as the children of a parent are labelled differently. Thus, without
loss of generality, one can assume that $a(x)$ was labelled first. By
definition, the parent of $a(x)$ does not have label $L$ and hence, $L$ must
have been the smallest unused label when node $a(x)$ was labelled. A similar
argument for $a(y)$ shows that $L$ must have been the smallest unused label
when node $p(y)$ was labelled. However, when node $a(y)$ was labelled, $L$ had
been used to label $a(x)$, giving the required contradiction.

######  Lemma 0.10.

$l$ is a permutation of the set $\\{0,1,...,n-1\\}$.

Proof: Firstly, note that as there are $n-1$ internal nodes, and a single new
label is introduced every time the children of an internal node are labelled,
the set of labels used (across all nodes) must be $\\{0,1,...,n-1\\}$.

Suppose that leaf nodes $x$ and $y$ have the same label $L$. By Lemma 0.9,
they must share an ancestor $a$ with that label. Define $b$ to be the shared
ancestor with the highest generation (that is, furthest from the root) such
that $b$ has label $L$. Then, either nodes $x$ and $y$ are the children of $b$
(in which case, they have distinct labels) or they are descendants of distinct
children, $c$ and $d$, of $b$. In the second case, one can impose without loss
of generality that $c$ does not have label $L$ and, as label $L$ has already
been used to label $b$, we know that all descendants of $c$ do not have label
$L$. Hence, in both cases, one of $x$ and $y$ does not have label $L$ as
required.

######  Lemma 0.11.

For each $i\in\\{0,1,...,n-1\\}$, define $x(i)$ to be the node of the highest
generation with label $i$. Define (for $i>0$), $y(i)$ to be the label of the
parent of $x(i)$. Then, with ordering $l$, the tree is given by

$v_{0}=0\quad\text{and}\quad v_{i}=y(i)\quad\forall i>0$ (34)

Proof: Firstly, note that $\bm{v}$ is ordered, as the label of a child is
greater than or equal to that of its parent. Hence, as the parent of $x(i)$
does not have label $i$, it must have a label strictly less than $i$ and so
$v_{i}<i$ as required.

Consider constructing the tree according the “standard” right-to-left
Phylo2Vec algorithm. One can then proceed by induction on the number of nodes.
The case $n=2$ is trivial, and so suppose it holds for $n=m$ and consider a
tree with $n=m+1$.

The first node to be processed has label $m$. No other node can have label $m$
cannot also have label $m$ (as, otherwise, one of its children would have
label greater than $m$, contradicting Lemma 0.10) and hence, the first node to
be processed is $x(m)$. The parent of $x(m)$ must have label $y(m)$. This must
also be the label of its other child and hence, $v_{m}=y(m)$ ensures that the
node with label $m$ merges with its sibling from the original tree (which is
correct).

From this point, the tree now has $m$ nodes, and our right-to-left
construction algorithm considers the parent of the node labelled $m$ to now be
a leaf node with label $y(m)$. The values of $y(i)$ for this tree are
unchanged (in particular, $y(y(m))$ depends on the label of an ancestor of the
parent of node $m$) and hence, by induction, the remaining values of $\bm{v}$
correctly generate the rest of the tree. Thus, the correct tree is generated
by $\bm{v}$ as required.

###  0.5 Queue Shuffle: generating principled ordering proposals

#### Asymmetry of ordered tree spaces

######  Lemma 0.12.

Define $g_{m}^{k}$ to be the expected distance from the root of the node
labelled $m$ in a tree with $k$ nodes. Define the harmonic sum function

$H(m)=\left\\{\begin{matrix}\sum_{j=1}^{m}\frac{1}{j}&\text{if $m\geq 1$}\\\
1&\text{if $m=0$}\end{matrix}\right.$ (35)

Then,

$g^{k}_{m}=H(k-1)+H(m)-1$ (36)

Proof: From our left-to-right construction algorithm,

$g^{k}_{k-1}=\sum_{m=0}^{k-2}\frac{1}{k-1}(1+g^{k-1}_{m})$ (37)

as adding node $k-1$ according to $v_{k-1}=m$ means that the path length
between node $k-1$ and the root will be one more than the path length of
between $m$ and the root. If $v_{k-1}=m$, it will also increase the path
length between node $m$ and the root by 1 and so

$g^{k}_{m}=\frac{k-2}{k-1}g^{k-1}_{m}+\frac{1}{k-1}(g^{k-1}_{m}+1)=g^{k-1}_{m}+\frac{1}{k-1}$
(38)

The initial conditions of this system are that

$g_{0}^{2}=g_{1}^{2}=1$ (39)

as in a two-node rooted tree, the leaves are distance 1 from the root.

We claim by induction on $k$ that the solution to this system is (36).

Note that this holds for $k=2$ as $H(1)=H(0)=1$. Moreover, under the inductive
hypothesis that it holds for a tree with $k-1$ nodes, for $m<k-1$,

$g^{k}_{m}=g^{k-1}_{m}+\frac{1}{k-1}=H(k-2)+H(m)-1+\frac{1}{k-1}=H(k-1)+H(m)-1$
(40)

and

$\displaystyle g^{k}_{k-1}$
$\displaystyle=\sum_{m=0}^{k-2}\frac{1}{k-1}(1+g^{k-1}_{m})$ (41)
$\displaystyle=\sum_{m=0}^{k-2}\frac{1}{k-1}(1+H(k-2)+H(m)-1)$ (42)
$\displaystyle=\sum_{m=0}^{k-2}\frac{1}{k-1}(H(k-2)+H(m))$ (43)
$\displaystyle=H(k-2)+\frac{1}{k-1}\sum_{m=0}^{k-2}H(m)$ (44)
$\displaystyle=H(k-1)+\frac{1}{k-1}\sum_{m=1}^{k-2}\sum_{j=1}^{m}\frac{1}{j}$
(45)

Now,

$\sum_{m=1}^{k-2}\sum_{j=1}^{m}\frac{1}{m}=\sum_{j=1}^{k-2}\sum_{m=j}^{k-2}\frac{1}{j}=\sum_{j=1}^{k-2}\frac{k-1-j}{j}=(k-1)H(k-2)-(k-2)$
(46)

and hence

$g^{k}_{k-1}=H(k-1)+H(k-2)-\frac{(k-2)}{(k-1)}=H(k-1)+H(k-1)-1$ (47)

as required.

#### Nearest Neighbour Interchange

######  Definition 0.13.

Define $\tau(\sigma)$ to be the space of possible trees given an ordering
$\sigma$.

######  Definition 0.14.

For a tree $\mathcal{T}$, define $Q(\mathcal{T})$ to be the random ordering
generated by Queue Shuffle.

######  Lemma 0.15.

Consider a tree $\mathcal{T}$ and suppose that another tree,
$\mathcal{T}^{\prime}$ is one NNI move away from $\mathcal{T}^{\prime}$. Then

$\mathbb{P}\bigg{[}\mathcal{T}^{\prime}\in\tau\bigg{(}Q(\mathcal{T})\bigg{)}\bigg{]}\geq\frac{1}{4}$
(48)

To facilitate the proof, we first note that for any permutation $\sigma$, by
Lemma 0.18, that

$\mathbb{P}\bigg{[}Q(\mathcal{T})=\sigma\bigg{]}\in\bigg{\\{}0,\frac{1}{2^{n-1}}\bigg{\\}}$
(49)

Figure S1: A simple representation of the swapping property of Queue Shuffle.
The swapped subtrees have roots labelled $S_{i}$ and $S_{j}$ while the node
label $c$ is the subject of the equivalent subtree-prune and regraft
operation. Note that in step 4, the internal nodes are renamed to illustrate
that the required swap has indeed occurred.

Now, we consider labelling the tree as shown in the left panel of Fig. S1.
Suppose that the edge to which the four subtrees are joined has nodes $i$ and
$j$, with node $i$ being a higher generation than node $j$. Suppose that the
subtrees rooted at the child of $j$ are labelled as $\mathcal{S}_{j}$ and
$\mathcal{S}_{c}$. Suppose that the other child of $i$ (that is, the child not
equal to $j$) is the root of a subtree $\mathcal{S}_{i}$. Finally, suppose
that the fourth subtree is labelled $\mathcal{S}_{r}$ (this is the subtree
containing the root). A single NNI move in this context involves swapping a
pair of subtrees from the set
$\\{\mathcal{S}_{i},\mathcal{S}_{j},\mathcal{S}_{r},\mathcal{S}_{c}\\}$.

By symmetry, swapping the pair $\mathcal{S}_{j}$ and $\mathcal{S}_{c}$ or the
pair $\mathcal{S}_{r}$ and $\mathcal{S}_{i}$ has no impact, as the new tree
will be topologically equivalent to $\mathcal{T}$ (which is always in the new
shuffled space). Thus, one needs only to prove the result for the tree formed
by swapping $\mathcal{S}_{c}$ and $\mathcal{S}_{r}$. This is topologically
equivalent to the tree in the right panel of Fig. S1. We assume that this tree
is $\mathcal{T}^{\prime}$

For a given permutation $\sigma$, define $l_{i}({\sigma})$ to be the label
assigned to node $i$, $l_{j}({\sigma})$ to be the label assigned to node $j$,
$l^{\mathcal{S}}_{i}({\sigma})$ to be the label assigned to the root of
$\mathcal{S}_{i}$, $l^{\mathcal{S}}_{j}({\sigma})$ to be the label assigned to
the root of $\mathcal{S}_{j}$ and $l^{\mathcal{S}}_{c}(\sigma)$ to be the
label assigned to the root of $\mathcal{S}_{c}$

Then, by Lemma 0.19,

$\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})\bigg{]}=\frac{1}{4}$
(50)

Now, suppose for a permutation $\sigma$ that
$l^{\mathcal{S}}_{j}(\sigma)=l_{i}(\sigma)$. Necessarily (as $r_{B}=0$),
$l^{\mathcal{S}}_{c}(\sigma)$ was the smallest available label when processing
node $j$, which must be bigger than the smallest available label when node $i$
was processed (as $j$ is the child of $i$ and hence processed later). Thus,
$l^{\mathcal{S}}_{c}(\sigma)>l^{\mathcal{S}}_{i}(\sigma)$.

A further important result proved in Lemma 0.19 is that node $j$ must have
been placed ahead of the root of $\mathcal{S}_{i}$ in the queue. Thus, all
nodes in $\mathcal{S}_{i}$ are labelled either as
$l^{\mathcal{S}}_{i}(\sigma)$ or with a label that is greater than
$l^{\mathcal{S}}_{c}(\sigma)$, as all internal nodes in $\mathcal{S}_{i}$ are
processed after node $j$. Moreover, all nodes in $\mathcal{S}_{j}$ are
labelled as either $l^{\mathcal{S}}_{j}(\sigma)$ or a label greater than
$l_{c}(\sigma)$ as they were processed after node $j$.

Define $\bm{v}$ to be the vector giving $\mathcal{T}$ under $\sigma$. Then,
define a new vector $\bm{v}^{\prime}$ by

$v^{\prime}_{m}=\left\\{\begin{matrix}l^{\mathcal{S}}_{i}(\sigma)&\text{if
$m=l^{\mathcal{S}}_{c}(\sigma)$}\\\ v_{m}&\text{otherwise}\\\
\end{matrix}\right.$ (51)

Then, $\bm{v}^{\prime}$ is ordered as
$l^{\mathcal{S}}_{c}(\sigma)>l_{j}^{\mathcal{S}}(\sigma)$. We now show that
$\bm{v}^{\prime}$ generates $\mathcal{T}^{\prime}$ by using the left-to-right
construction algorithm. Note that up to the point that the node labelled
$l^{\mathcal{S}}_{c}(\sigma)$ is processed by this algorithm, no nodes have
appended to the edges connecting either the root of $\mathcal{S}_{i}$ or
$\mathcal{S}_{j}$ to the tree. This holds because all other nodes in
$\mathcal{S}_{i}$ and $\mathcal{S}_{j}$ have labels greater than
$l^{\mathcal{S}}_{c}(\sigma)$. Changing from $\bm{v}$ to $\bm{v}^{\prime}$
means that the leaf node labelled $l^{\mathcal{S}}_{c}(\sigma)$ is joined to
the edge connecting the root of $\mathcal{S}_{i}$ rather than the root of
$\mathcal{S}_{j}$. The rest of the tree is then constructed identically as the
vectors $\bm{v}$ and $\bm{v}^{\prime}$ are the same.

Thus, changing from $\bm{v}$ to $\bm{v}^{\prime}$ generates the tree in the
right panel of Fig. S1, and therefore the tree formed by the NNI move swapping
the trees $\mathcal{S}_{c}$ and $\mathcal{S}_{r}$.

This completes the proof as it shows that

$\mathbb{P}\bigg{[}\mathcal{T}^{\prime}\in\tau\bigg{(}Q(\mathcal{T})\bigg{)}\bigg{|}l^{\mathcal{S}}_{j}(\sigma)=l_{i}(\sigma)\bigg{]}=1$
(52)

and hence, by Lemma 0.19

$\mathbb{P}\bigg{[}\mathcal{T}^{\prime}\in\tau\bigg{(}Q(\mathcal{T})\bigg{)}\bigg{]}\geq\mathbb{P}\bigg{[}\mathcal{T}^{\prime}\in\tau\bigg{(}Q(\mathcal{T})\bigg{)}\bigg{|}l^{\mathcal{S}}_{j}(\sigma)=l_{i}(\sigma)\bigg{]}\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}(\sigma)=l_{i}(\sigma)\bigg{]}=\frac{1}{4}$
(53)

###  0.6 Eutherian Mammal phylogeny [97]

Figure S2: Comparison of the best unrooted FastME tree that has been midpoint
rooted to an optimised rooted tree via Queue Shuffle on an Eutherian mammal
dataset [97]. Queue Shuffle correctly places _Gallus gallus_ as the outgroup
of mammals. Branch lengths are ignored and trees are displayed as ultrametric.

###  0.7 Convergence analysis

Figure S3 compares the performance of different optimisers (Adafactor [96],
AdamW [63], RMSprop [103], and SGD) under different DNA substitutions models
(JC69 [51], F81 [30], TN93 [102]) on a single optimisation step (no subsequent
reordering with Queue Shuffle) for a maximum of 5000 steps. Four learning
rates were considered, logarithmically spaced from 0.001 to 1.0. Whereas
convergence speed appears to be independent of the chosen DNA substitution
model, the results varied widely with respect to the optimisation algorithm.
In particular, Adafactor optimisation produced the best performance, with
increasing convergence speed as learning rates were higher.

(a)

(b)

Figure S3: Convergence analysis of different optimisers and DNA substitution
models for the datasets DS1-DS11 (see Table 2). (a) Number of steps needed to
reach convergence (tolerance: 1e-10). (b) Loss reached at convergence. For
each dataset, the log BME losses were scaled using min-max normalisation.
Error bars denote 95% confidence intervals computed with 1000 bootstraps.

###  0.8 Miscallaneous Lemmas

#### Supplementary lemmas for Lemma 0.4

######  Lemma 0.16.

Using the notation of Lemma 0.4, define $D_{Ar}$ to be the distance between
node $A$ and the root, and $g_{A}$ to be the path length (i.e. the generation
of node $A$). Then,

$\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}=2\sum_{A}2^{-g_{A}}D_{Ar}$ (54)

where the $e_{AB}$ terms refer to the path lengths in the original unrooted
tree.

Proof: We proceed by induction on the number of leaf nodes, $n$. For $n=2$,
the root must be placed on the edge connecting nodes $0$ and $1$ and
$e_{01}=1$. Then

$\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}=\frac{1}{2}(D_{01}+D_{10})=\bigg{(}2^{1-1}D_{0r}+2^{1-1}D_{1r}\bigg{)}$
(55)

as required, noting that $D_{0r}+D_{1r}=D_{01}$. Suppose that our claim holds
for $n=k$ and consider a tree with $n=k+1$. Choose a pair of sibling leaf
nodes (i.e. leaf nodes joined to the same internal node) $y$ and $z$ such that
the root is not on an edge joining one of these nodes to the tree (this must
exist as, for $n>2$, there are at least two pairs of sibling leaf nodes).
Suppose that $y$ and $z$ are connected to the internal node $w$. Now, note
that the contribution to the objective of node $y$ is

$\displaystyle 2\sum_{B:X_{b}(y,B)=1}D_{yB}2^{-e_{yB}}$
$\displaystyle=2\sum_{B:X_{b}(w,B)=1}(D_{wB}+D_{wy})2^{-e_{yB}}$ (56)
$\displaystyle=2\sum_{B:X_{b}(w,B)=1}D_{wB}2^{-e_{yB}}+2D_{wy}\sum_{B:X_{b}(w,B)=1}2^{-e_{yB}}$
(57)
$\displaystyle=\sum_{B:X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}+2D_{wy}\sum_{B:X_{b}(w,B)=1}2^{-e_{yB}}$
(58)

where the factor of $2$ comes from the fact that both $(y,B)$ and $(B,y)$ must
be considered. The final term in this equation can be simplified. Consider the
subtree whose leaves are the root and all nodes such that $X_{b}(w,B)=1$.
Then, by the Kraft Equality on this subtree

$\sum_{B:X_{b}(w,B)=1}2^{-g_{B}}=\frac{1}{2}$ (59)

as $g_{B}$ gives the distance between the root and node $B$. Moreover, as
$e_{yB}=g_{y}+g_{B}-1$ (where the $-1$ accounts for the fact that
$\mathcal{T}$ is unrooted)

$\sum_{B:X_{b}(w,B)=1}2^{-e_{yB}}=2^{1-g_{y}}\sum_{B:X_{b}(w,B)=1}2^{-g_{B}}=2^{-g_{y}}$
(60)

Thus, (58) becomes

$2\sum_{X_{b}(y,B)=1}D_{yB}2^{-e_{yB}}=\sum_{B:X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}+2D_{wy}2^{-g_{y}}$
(61)

Thus, the change to $\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}$ caused by adding
nodes $y$ and $z$ from the tree is to subtract

$2\sum_{X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}$ (62)

(as node $w$ is no longer a leaf node) and to add the contributions from nodes
$y$ and $z$

$2\bigg{(}\sum_{B:X_{b}(y,B)=1}D_{yB}2^{-e_{yB}}+\sum_{B:X_{b}(z,B)=1}D_{zB}2^{-e_{zB}}\bigg{)}=\sum_{X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}+2D_{wy}2^{-g_{y}}+\sum_{X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}+2D_{wz}2^{-g_{z}}$
(63)

There is no $(y,z)$ term to consider as $X_{b}(y,z)=0$. Note that the
$\sum_{X_{b}(w,B)=1}D_{wB}2^{-e_{wB}}$ terms cancel. Define
$\mathcal{T}^{\prime}$ to be the tree when node $y$ and $z$ are removed so, as
it now has $k$ leaf nodes, the inductive hypothesis can be used to show

$\displaystyle\sum_{A,B\in\mathcal{T}:X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}$
$\displaystyle=\sum_{A^{\prime},B^{\prime}\in\mathcal{T}^{\prime}:X^{\prime}_{b}(A^{\prime},B^{\prime})=1}D_{A^{\prime}B^{\prime}}2^{-e^{\prime}_{A^{\prime}B^{\prime}}}+2D_{wy}2^{-g_{y}}+2D_{wz}2^{-g_{z}}$
(64)
$\displaystyle=2\sum_{A^{\prime}\in\mathcal{T}^{\prime}}2^{-g^{\prime}_{A^{\prime}}}D_{A^{\prime}r^{\prime}}+2D_{wy}2^{-g_{y}}+2D_{wz}2^{-g_{z}}$
(65)
$\displaystyle=2\sum_{A^{\prime}\in\mathcal{T}^{\prime}/\\{w\\}}2^{-g^{\prime}_{A^{\prime}}}D_{A^{\prime}r^{\prime}}+2D_{wr}2^{-g_{w}}+2D_{wy}2^{-g_{y}}+2D_{wz}2^{-e_{zr}}$
(66)

where dashes are used to denote quantities in $\mathcal{T}^{\prime}$. Note
that

$\sum_{A^{\prime}\in\mathcal{T}^{\prime}/\\{w\\}}2^{-g^{\prime}_{A^{\prime}}}D_{A^{\prime}r^{\prime}}=\sum_{A\in\mathcal{T}/\\{y,z\\}}2^{-g_{A}}D_{Ar}$
(67)

as only the nodes $w$, $y$ and $z$ are affected by changing from $\mathcal{T}$
to $\mathcal{T}^{\prime}$. Moreover, note that

$2^{-g_{w}}=2^{-g_{y}}+2^{-g_{z}}$ (68)

as $g_{w}=g_{y}-1=g_{z}-1$, and, as the distances are additive,

$D_{wy}+D_{wr}=D_{yr}\quad\text{and}\quad D_{wz}+D_{wr}=D_{zr}$ (69)

Hence

$\displaystyle 2D_{wy}2^{-g_{y}}+2D_{wz}2^{-g_{z}}+2D_{wr}2^{-g_{w}}$
$\displaystyle=2D_{wy}2^{-g_{y}}+2D_{wz}2^{-g_{z}}+2D_{wr}(2^{-g_{y}}+2^{-g_{z}})$
(70) $\displaystyle=2(D_{wy}+D_{wr})2^{-g_{y}}+2(D_{wz}+D_{wr})2^{-g_{z}}$
(71) $\displaystyle=2D_{yr}2^{-g_{y}}+2D_{zr}2^{-g_{z}}$ (72)

Thus,

$\sum_{X_{b}(A,B)=1}D_{AB}2^{-e_{AB}}=\sum_{A}2^{-g_{A}}D_{Ar}$ (73)

and hence the claim holds by induction as required.

######  Lemma 0.17.

Using the notation of Lemma 0.4,

$\mathcal{D}=\sum_{A}2^{-g_{A}}D_{Ar}$ (74)

Proof: This can be proved through induction on the number of nodes. It clearly
holds if there are only two nodes in the tree. Assume it holds whenever there
are $k$ nodes, and consider a tree with $k+1$ nodes. We process the first pair
of leaf nodes, $y$ and $z$, and assume that their associated distances to
their parent, $w$ are $D_{wy}$ and $D_{wz}$ and that the new tree created is
$\mathcal{T}^{\prime}$ with distances $D^{\prime}$. Then, we have, by our
inductive hypothesis

$\displaystyle\mathcal{D}$
$\displaystyle=\sum_{A^{\prime}\in\mathcal{T}^{\prime}}2^{-g^{\prime}_{A^{\prime}}}D^{\prime}_{A^{\prime}r^{\prime}}$
(75)
$\displaystyle=\sum_{A^{\prime}\in\mathcal{T}^{\prime}/\\{w\\}}2^{-g^{\prime}_{A^{\prime}}}D^{\prime}_{A^{\prime}r^{\prime}}+2^{-g^{\prime}_{w}}D^{\prime}_{wr^{\prime}}$
(76)
$\displaystyle=\sum_{A\in\mathcal{T}/\\{y,z\\}}2^{-g_{A}}D_{Ar}+2^{-g_{w}}\bigg{(}D_{wr}+\frac{D_{wy}+D_{wz}}{2}\bigg{)}$
(77)
$\displaystyle=\sum_{A\in\mathcal{T}/\\{y,z\\}}2^{-g_{A}}D_{Ar}+2^{-g_{w}}\bigg{(}\frac{D_{yr}+D_{zr}}{2}\bigg{)}$
(78) $\displaystyle=\sum_{A\in\mathcal{T}}2^{-g_{A}}D_{Ar}$ (79)

as required, where we have used the fact that $D_{wr}+D_{wy}=D_{yr}$ and that
$2^{-g_{w}}=2\times 2^{-g_{y}}=2\times 2^{-g_{z}}$.

#### Supplementary Lemmas for Lemma 0.15

######  Lemma 0.18.

Using the notation of Lemma 0.15, for a given permutation $\sigma$,

$\mathbb{P}\bigg{[}Q(\mathcal{T})=\sigma\bigg{]}\in\bigg{\\{}0,\frac{1}{2^{n-1}}\bigg{\\}}$
(80)

Proof Suppose that every (internal and external) node in $\mathcal{T}$ is
assigned a fixed unique position, such that one can consistently define a
“leftmost” and “rightmost” node from a pair of nodes. Then, the randomness in
the Queue Shuffle algorithm can be represented by a uniform random vector
$\bm{r}\in\\{0,1\\}^{n-1}$ such that when the $m^{\text{th}}$ internal node is
processed, the leftmost child is given the same label as its parent if and
only if $r_{m+1}=1$. Each distinct $\bm{r}$ results in a distinct ordering
(as, given the leaf labels, one can generate the unique labelling of the tree
as a parent must have a label equal to the minimum of the labels of its two
children). Hence, each possible labelling is equally likely, giving (80) and
completing the proof.

######  Lemma 0.19.

Using the notation of Lemma 0.15

$\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})\bigg{]}=\frac{1}{4}$
(81)

Moreover, $l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})$ if
and only if node $j$ was placed ahead of the root of $\mathcal{S}_{i}$ in the
queue

Proof Using the $\bm{r}$ defined in the proof of Lemma 0.18, define the random
indices $A$ and $B$ such that $r_{A}$ corresponds to processing node $i$ and
$r_{B}$ corresponds to processing node $j$. Note that $r_{A}$ and $r_{B}$ are
independent (though $B$ will in general depend on $r_{A}$).

Now, suppose without loss of generality that $j$ is the leftmost child of $i$
and that the root of $\mathcal{S}_{j}$ is the leftmost child of $j$. As
children have labels greater than or equal to their parents,

$\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})\bigg{]}=\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})\bigg{]}$
(82)

and hence

$\mathbb{P}\bigg{[}l^{\mathcal{S}}_{j}({Q(\mathcal{T})})=l_{i}({Q(\mathcal{T})})\bigg{]}=\mathbb{P}(r_{A}=r_{B}=1)=\frac{1}{4}$
(83)

which is the first result. The second result follows immediately from the fact
that $r_{A}=1$.

###  0.9 Estimation of GTR+$\Gamma$ distances

To estimate distances under a GTR+$\Gamma$ substitution model we use the
approach outlined in [112]. Assuming a general time reversible rate symmetric
matrix $Q$, the transition probability matrix over time is found via the
matrix exponential $P(t)=e^{Qt}$. This matrix exponential can be readily
computed via eigendecomposition.

Including variable rates among sites for a Gamma distribution $g$, $P(t)=\int
e^{Qu}g(u)du$, which can again be estimated via eigendecomposition. Given
parameters for rates, $S$, frequencies, $\pi$, the time between two sequences
$t_{ij}$, and genetic sequence data $\mathcal{G}$, the log-likelihood for the
transitions between a pair of taxa $i$ and $j$ is

$\mathcal{L}_{ij}(\mathcal{G}|S,\pi,t_{ij})=\sum_{a}\sum_{b}\kappa^{ij}_{ab}\log(P_{ab}(t_{ij};S,\pi))$
(84)

where $\kappa^{ij}_{ab}$ is the number of $a\to b$ transitions from taxon $i$
to taxon $j$. We approximate the optimal parameters by maximizing the total
log-likelihood (that is, the sum over $i$ and $j$ of $\mathcal{L}_{ij}$) using
gradient descent in Jax.

###  0.10 Fast discrete hill-climbing with Phylo2Vec

The computational complexity of GradME is substantially higher than that of
FastME due to the continuous nature of the algorithm. Thus, particularly for
large numbers of taxa, using a similarly fast algorithm, at least to get close
to the optimal tree, may be preferable.

Because of this, we have developed an alternative, discrete algorithm which
has the same computational complexity as FastME. At each step, our algorithm
outputs a matrix $\Delta$ such that $\Delta_{ij}$ is the change in the
objective function if the value of $v_{i}$ were changed to be equal to $j$.
$\Delta$ allows us to perform a hill-climbing optimisation, as the change
corresponding to the minimum value of $\Delta$ is made.

Analogously to FastME, $\Delta$ can be calculated in
$\mathcal{O}(n^{2}\text{diam}(\mathcal{T}))$ time, where
$\text{diam}(\mathcal{T})$ is the maximum inter-taxa path length of the tree
(this is generally substantially smaller than $n$). This algorithm works
exclusively for unrooted trees, though a similar algorithm could be developed
using our rooted objective.

Motivated by the method of [20], our algorithm begins by calculating directed
edge weight vectors $\bm{w}_{e}^{\pm}$ for the edges in $\mathcal{T}$. Each
edge $e$ naturally partitions $\mathcal{T}$ into two disjoint subtrees,
$\mathcal{S}_{1}^{e}$ and $\mathcal{S}_{2}^{e}$, with each one being rooted at
a node connected to $e$. We suppose that $\mathcal{S}_{1}^{e}$ is the tree not
containing some fixed node $X$. Then, we assign $\bm{w}_{e}^{+}$ to be the
balanced distance between the root of $\mathcal{S}_{1}^{e}$ and each of the
individual leaf nodes in $\mathcal{S}_{2}^{e}$, and $\bm{w}_{e}^{-}$ to be the
balanced distance between the root of $\mathcal{S}_{2}^{e}$ and each of the
individual leaf nodes in $\mathcal{S}_{1}^{e}$. These edge weights can be
calculated efficiently by using an iterative scheme, using the fact that the
weight on a given edge from an internal node $x$ to another node $y$ is the
mean of the weights of the two edges from nodes $z$ and $a$ (both distinct
from $y$) to $x$ (recalling that these weights are directional). This
iterative scheme is closed by the fact that the weights on edges from leaf
nodes to the tree are given by the appropriate columns of the distance matrix
$D$.

This weighted tree can be used to calculate the distance between any pair of
subtrees. For each fixed subtree, $\mathcal{S}$, one can iteratively calculate
the distance $d_{\mathcal{S}\mathcal{R}}$ between it and (disjoint) subtrees
$\mathcal{R}$, using the fact that the precalculated $\bm{w}$ terms give the
distance between each subtree and each leaf, and that for any pair of subtrees
$\mathcal{R}_{1}$ and $\mathcal{R}_{2}$ which share a parent node and can
therefore be combined in a subtree $\mathcal{R}_{3}$,

$d_{\mathcal{S}\mathcal{R}_{3}}=\frac{1}{2}\bigg{(}d_{\mathcal{S}\mathcal{R}_{1}}+d_{\mathcal{S}\mathcal{R}_{2}}\bigg{)}$
(85)

From [20], there exists a simple formula for the difference in objective
function caused by pruning a subtree and regrafting it to an adjacent edge. If
an internal node $x$ is joined to nodes which are the roots of disjoint
subtrees $A$, $B$ and $C$, then the difference between attaching a subtree $K$
to the edge joining $x$ and $C$ and attaching the subtree $K$ to the edge
joining $x$ and $B$ is

$\frac{1}{4}\bigg{(}\delta_{AB}+\delta_{KC}-\delta_{AC}-\delta_{KB}\bigg{)}$
(86)

where here, $\delta$ denotes inter-subtree distance. These $\delta$ terms are
not equivalent to the $d$ terms, as removing $K$ from the original tree
creates a different set of possible subtree pairs. However, they can be
calculated in $\mathcal{O}(n^{2}\text{diam}(\mathcal{T}))$ time from the
distances $d$ following the methods of [20].

Explicitly, our algorithm moves outwards from the parent of the subtree which
we are removing. We take $B$ to be the node we are currently processing, $x$
to be the previously processed node on this path, $C$ to the node processed
two iterations ago, and $A$ to be the other node connected to $x$. Then, one
has simply

$\delta_{AB}=d_{AB}\quad\text{and}\quad\delta_{KB}=d_{KB}$ (87)

Calculating $\delta_{AC}$ is slightly more complicated as the subtree $K$ has
been removed from $C$. If the original root of $K$ was a path length $l$ from
$C$, and if that root shares a parent node (in $C$) with a subtree $F$, then

$\delta_{AC}=d_{AC}+2^{-l}(d_{AF}-d_{AK})$ (88)

as the subtree $F$ now has double the weight in $C$. Finally, to calculate
$d_{CK}$, one must calculate the distance between $K$ and all subtrees
connected to, but disjoint from the path taken from the tree root to $C$. If
these subtrees are $\mathcal{S}_{1},...,\mathcal{S}_{m}$ respectively and are
distance $1,...,m$ from $C$, then

$\delta_{KC}=\sum_{q=1}^{m}2^{-q}d_{C\mathcal{S}_{q}}$ (89)

It is this step which pushes the complexity of our algorithm above $n^{2}$

Thus, one can calculate the objective differences caused by each subtree-prune
and regraft move by repeatedly applying this formula for each pruned subtree,
essentially “moving” this subtree around the remaining tree.

The final step of our algorithm is to find the corresponding subtree-prune and
regraft move for each possible change of $\bm{v}$. When $v_{i}$ is changed to
be equal to $j$, the subtree that is moved is the largest subtree of
$\mathcal{T}$ containing $i$ such that the leaves all have labels greater than
or equal to $i$ (in many cases, this will simply be the node $i$). For each
value of $i$, one can find the edges which were connected to each leaf node
when node $i$ was attached in the left-to-right construction algorithm. The
edge connected to node $j$ at this stage gives the location of the regraft
position of the subtree.

###  0.11 Code: continuous BME objective function

⬇

import jax.numpy as np

\parfrom jax import jit, lax

from jax.scipy.special import logsumexp

\par@jit

def get_edges_exp_log(W, rooted):

"""Calculate the log-expectation of the objective value of a tree drawn with
distribution W.

We calculate and update E_ij throughout the left-to-right construction
procedure

\parArgs:

W (jax.numpy.array): Tree distribution

rooted (bool): True is the tree is rooted, otherwise False

\parReturns:

E (jax.numpy.array): Log of the expected objective value of a tree drawn with
W

"""

# Add jnp.finfo(float).eps to W.tmp to avoid floating point errors with
float32

W_tmp = (

jnp.pad(W, (0, 1), constant_values=jnp.finfo(float).eps) +
jnp.finfo(float).eps

)

\parn_leaves = len(W) + 1

\parE = jnp.zeros((n_leaves, n_leaves))

E = E.at[1, 0].set(

0.5 * E[0, 0] * W_tmp[0, 0] + jnp.log(0.25 * (2 - rooted) * W_tmp[0, 0])

)

E = E + E.T

\partrindx_x, trindx_y = jnp.tril_indices(n_leaves - 1, -1)

\pardef body(carry, _):

E, i = carry

\parE_new = jnp.zeros((n_leaves, n_leaves))

\partrindx_x_i = jnp.where(trindx_x < i, trindx_x, 1)

trindx_y_i = jnp.where(trindx_x < i, trindx_y, 0)

\parindx = (trindx_x_i, trindx_y_i)

\parE_new = E_new.at[indx].set(

E[indx] + jnp.log(

1 + jnp.finfo(float).eps - 0.5 * (W_tmp[i - 1, indx[1]] + W_tmp[i - 1,
indx[0]])

)

)

\par# exp array

mask_Ei = jnp.where(jnp.arange(n_leaves) >= i, 0, 1)

exp_array = E * jnp.where(jnp.arange(n_leaves) >= i, 0, mask_Ei.T)

\par# coef array

mask_Wi = jnp.where(jnp.arange(n_leaves) >= i, 0, 0.5 * W_tmp[i - 1])

coef_array = (jnp.zeros_like(W_tmp) + mask_Wi).at[:, i].set(0.25 * W_tmp[i -
1])

coef_array = coef_array * (1 - jnp.eye(W_tmp.shape[0]))

\par# logsumexp

tmp = logsumexp(exp_array, b=coef_array, axis=-1) * mask_Ei

\parE_new = E_new.at[i, :].set(tmp)

\par# Update E

E = E_new + E_new.T

\parreturn (E, i + 1), None

\par(E, _), _ = lax.scan(body, (E, 2), None, length=n_leaves - 2)

\parreturn E

\par@jit

def bme_loss_log(W, D, rooted):

"""Log version of the BME loss function

\parArgs:

W (jax.numpy.array): Tree distribution

D (jax.numpy.array): Distance matrix

rooted (bool): True is the tree is rooted, otherwise False

\parReturns:

loss (float): BME loss

"""

E = get_edges_exp_log(W, rooted)

loss = logsumexp(E, b=D)

return loss