# PhyloGFN: Phylogenetic inference
with generative flow networks

Mingyang Zhou1*, Zichao Yan2*, Elliot Layne1,2, Nikolay Malkin2,3,
Dinghuai Zhang2,3, Moksh Jain2,3, Mathieu Blanchette1,2& Yoshua Bengio2,3,4
1McGill University, 2Mila – Québec AI Institute, 3Université de Montréal,
4CIFAR
*Equal contribution. Correspondence to<EMAIL_ADDRESS>
###### Abstract

Phylogenetics is a branch of computational biology that studies the
evolutionary relationships among biological entities. Its long history and
numerous applications notwithstanding, inference of phylogenetic trees from
sequence data remains challenging: the high complexity of tree space poses a
significant obstacle for the current combinatorial and probabilistic
techniques. In this paper, we adopt the framework of generative flow networks
(GFlowNets) to tackle two core problems in phylogenetics: parsimony-based and
Bayesian phylogenetic inference. Because GFlowNets are well-suited for
sampling complex combinatorial structures, they are a natural choice for
exploring and sampling from the multimodal posterior distribution over tree
topologies and evolutionary distances. We demonstrate that our amortized
posterior sampler, PhyloGFN, produces diverse and high-quality evolutionary
hypotheses on real benchmark datasets. PhyloGFN is competitive with prior
works in marginal likelihood estimation and achieves a closer fit to the
target distribution than state-of-the-art variational inference methods.

##  1 Introduction

Phylogenetic inference has long been a central problem in the field of
computational biology. A broad set of methods has been developed to estimate
the evolutionary history (phylogenetic tree) relating a set of biological
entities. Accurate phylogenetic inference is critical for a number of
important biological analyses, such as understanding the development of
antibiotic resistance (Aminov & Mackie, 2007; Ranjbar et al., 2020; Layne et
al., 2020), assessing the risk of invasive species (Hamelin et al., 2022), and
characterizing tumor progression (Lähnemann et al., 2020). Accurate
phylogenetic trees can also be used to improve downstream computational
analyses, such as multiple genome alignment (Blanchette et al., 2004),
ancestral sequence reconstruction (Ma et al., 2006), protein structure and
function annotation (Celniker et al., 2013).Despite its strong medical
relevance and wide applications in life science, phylogenetic inference has
remained a standing challenge, in part due to the high complexity of tree
space — for $n$ species, $(2n-5)!!$ unique unrooted bifurcating tree
topologies exist. This poses a common obstacle to all branches of phylogenetic
inference; both maximum-likelihood and maximum-parsimony tree reconstruction
are NP-hard problems (Day, 1987; Chor & Tuller, 2005). Under the Bayesian
formulation of phylogenetics, the inference problem is further compounded by
the inclusion of continuous variables that capture the level of sequence
divergence along each branch of the tree.One line of prior work considers
Markov chain Monte Carlo (MCMC)-based approaches, such as MrBayes (Ronquist et
al., 2012). These approaches have been successfully applied to Bayesian
phylogenetic inference. However, a known limitation of MCMC is scalability to
high-dimensional distributions with multiple separated modes (Tjelmeland &
Hegstad, 2001), which arise in larger phylogenetic datasets. Recently,
variational inference (VI)-based approaches have emerged, which offer
computationally efficient and scalable alternatives to MCMC. Among these
methods,some model only a limited portion of the space of tree topologies,
while others are weaker in marginal likelihood estimation due to simplifying
assumptions. In parsimony analysis, state-of-the-art methods such as PAUP*
(Swofford, 1998) have extensively relied on heuristic search algorithms that
are efficient but lack theoretical foundations and guarantees.Coming from the
intersection of variational inference and reinforcement learning is the class
of models known as generative flow networks (GFlowNets) (Bengio et al., 2021).
The flexibility afforded by GFlowNets to learn sequential samplers for
distributions over compositional objects makes them a promising candidate for
performing inference over the posterior space of phylogenetic tree topologies
and evolutionary distances.In this work, we propose PhyloGFN, the first
adaptation of GFlowNets to the task of Bayesian and parsimony-based
phylogenetic inference. Our contributions are as follows:

  1. (1)

We design an acyclic Markov decision process (MDP) with fully customizable
reward functions, by which our PhyloGFN can be trained to construct
phylogenetic trees in a bottom-up fashion.

  2. (2)

PhyloGFN leverages a novel tree representation inspired by Fitch and
Felsenstein’s algorithms to represent rooted trees without introducing
additional learnable parameters. PhyloGFN is also coupled with simple yet
effective training techniques such as using mixed on-policy and dithered-
policy rollouts, replay buffers and cascading temperature-annealing.

  3. (3)

PhyloGFN explores and samples from the entire phylogenetic tree space,
achieving a balance between exploration in this vast space and high-fidelity
modeling of the modes. While PhyloGFN performs on par with the state-of-the-
art MCMC- and VI-based methods in the summary metric of marginal log-
likelihood, it substantially outperforms these approaches in terms of its
ability to estimate the posterior probability of suboptimal trees.

##  2 Related work

Markov chain Monte Carlo (MCMC)-based algorithms are commonly employed for
Bayesian phylogenetics, with notable examples including MrBayes and RevBayes
(Ronquist et al., 2012; Höhna et al., 2016), which are considered state-of-
the-art in the field. Amortized variational inference (VI) is an alternative
approach that parametrically estimates the posterior distribution. VBPI-GNN
(Zhang, 2023) employs subsplit Bayesian networks (SBN) (Zhang & Matsen IV,
2018a) to model tree topology distributions and uses graph neural networks to
learn tree topological embeddings (Zhang, 2023). While VBPI-GNN has obtained
marginal log likelihood competitive with MrBayes in real datasets, it requires
a pre-generated set of high-quality tree topologies to constrain its action
space for tree construction, which ultimately limits its ability to model the
entire tree space.There exist other VI approaches that do not limit the space
of trees. VaiPhy (Koptagel et al., 2022) approximates the posterior
distribution in the augmented space of tree topologies, edge lengths, and
ancestral sequences. Combined with combinatorial sequential Monte Carlo (CSMC;
Moretti et al., 2021), the proposed method enables faster estimation of
marginal likelihood. GeoPhy (Mimori & Hamada, 2023) models the tree topology
distribution in continuous space by mapping continuous-valued coordinates to
tree topologies, using the same technique as VBPI-GNN to model tree
topological embeddings. While both methods model the entire tree topology
space, their performance on marginal likelihood estimation underperforms the
state of the art.For the optimization problem underpinning maximum parsimony
inference, PAUP* is one of the most commonly used programs (Swofford, 1998);
it features several fast, greedy, and heuristic algorithms based on local
branch-swapping operations such as tree bisection and reconnection.GFlowNets
are a family of methods for sampling discrete objects from multimodal
distributions, such as molecules (Bengio et al., 2021) and biological
sequences (Jain et al., 2022), and are used to solve discrete optimization
tasks such as NP-hard scheduling and optimization problems (Zhang et al.,
2023a; b). With their theoretical foundations laid out in Bengio et al.
(2023); Lahlou et al. (2023), and connections to variational inference
established in Malkin et al. (2023); Zimmermann et al. (2023), GFlowNets have
been successfully applied to tackle complex Bayesian inference problems, such
as inferring latent causal structures in gene regulatory networks (Deleu et
al., 2022; 2023), and, most similar to the problems we consider, parse trees
in hierarchical grammars (Hu et al., 2023).

##  3 Background

###  3.1 Phylogenetic inference

Here we introduce the problems of Bayesian and parsimony-based phylogenetic
inference. A weighted phylogenetic tree is denoted by $(z,b)$, where $z$
represents the tree topology with its leaves labeled by observed sequences,
and $b$ represents the branch lengths.The tree topology can be either a rooted
binary tree or a bifurcating unrooted tree. For a tree topology $z$, let
$L(z)$ denote the labeled sequence set and $E(z)$ the set of edges. For an
edge $e\in E(z)$, let $b(e)$ denote its length. Let
${\bm{Y}}=\\{{\bm{y}}_{1},{\bm{y}}_{2}\dots\bm{y}_{n}\\}\in\Sigma^{n\times m}$
be a set of $n$ observed sequences, each having $m$ characters from alphabet
$\Sigma$, e.g., $\\{A,C,G,T\\}$ for DNA sequences.We denote the $i^{\rm th}$
site of all sequences by
${\bm{Y}}^{i}=\\{{\bm{y}}_{1}[i],{\bm{y}}_{2}[i]\dots\bm{y}_{n}[i]\\}$. In
this work, we make two assumptions that are common in the phylogenetic
inference literature: (i) sites evolve independently; (ii) evolution follows a
time-reversible substitution model. The latter implies that an unrooted tree
has the same parsimony score or likelihood as its rooted versions, and thus
the algorithms we introduce below (Fitch and Felsenstein) apply to unrooted
trees as well.

####  3.1.1 Bayesian inference

In Bayesian phylogenetic inference, we are interested in sampling from the
posterior distribution over weighted phylogenetic trees $(z,b)$, formulated
as:

$\displaystyle P(z,b\mid{\bm{Y}})=\frac{P(z,b)P({\bm{Y}}\mid
z,b)}{P({\bm{Y}})}$

where $P({\bm{Y}}\mid z,b)$ is the likelihood, $P({\bm{Y}})$ is the
intractable marginal, and $P(z,b)$ is the prior density over tree topology and
branch lengths. Under the site independence assumption, the likelihood can be
factorized as: $P({\bm{Y}}\mid z,b)=\prod_{i}P({\bm{Y}}^{i}\mid z,b)$, and
each factor is obtained by marginalizing over all internal nodes $a_{j}$ and
their possible character assignment:

$\displaystyle P({\bm{y}}_{1}[i]\dots\bm{y}_{n}[i]\mid
z,b)=\sum_{a_{n+1}^{i},\dots
a_{2n-1}^{i}}P(a_{2n-1})\prod_{j=n+1}^{2n-2}P(a_{j}^{i}|a^{i}_{\alpha(j)},b(e_{j}))\prod_{k=1}^{n}P({\bm{y}}_{k}[i]|a^{i}_{\alpha(k)},b_{z}(e_{k}))$

where $a_{n+1}^{i},\dots a_{2n-2}^{i}$ represent the internal node characters
assigned to site $i$ and $\alpha(i)$ represent the parent of node $i$.
$P(a_{2n-1})$ is a distribution at the root node, which is usually assumed to
be uniform over the vocabulary, while the conditional probability
$P(a_{j}^{i}|a^{i}_{\alpha(j)},b(e_{j}))$ is defined by the substitution model
(where $e_{j}$ is the edge linking $a_{j}$ to $\alpha(a_{j})$).

##### Felsenstein’s algorithm

The likelihood of a given weighted phylogenetic tree can be calculated
efficiently using Felsenstein’s pruning algorithm (Felsenstein, 1973) in a
bottom-up fashion through dynamic programming. Defining $L^{i}_{u}$ as the
leaf sequence characters at site $i$ below the internal node $u$, and given
its two child nodes $v$ and $w$, the conditional probability
$P(L^{i}_{u}|a^{i}_{u})$ can be obtained from $P(L^{i}_{v}|a^{i}_{v})$ and
$P(L^{i}_{w}|a^{i}_{w})$:

$P(L^{i}_{u}\mid a^{i}_{u})=\sum_{a^{i}_{v},a^{i}_{w}\in\Sigma}P(a^{i}_{v}\mid
a^{i}_{u},b(e_{v}))P(L^{i}_{v}\mid a^{i}_{v})P(a^{i}_{w}\mid
a^{i}_{u},b(e_{w}))P(L^{i}_{w}\mid a^{i}_{w}).$ (1)

This equation can be used to recursively compute $P(L^{i}_{u}|a^{i}_{u})$ at
all nodes of the tree and sites $i$. The algorithm performs a post-order
traversal of the tree, thus aggregating the conditional probabilities.
Finally, the likelihood $P({\bm{Y}}^{i}|z,b)$ is calculated using the root-
level conditional probability:
$P({\bm{Y}}^{i}|z,b)=\sum_{a^{i}_{2n-1}\in\Sigma}P(a_{2n-1}^{i})P({\bm{Y}}^{i}|a^{i}_{2n-1})$.The
data structure used by Felsenstein’s algorithm in the evaluation of the
recurrence (1) is a representation of node $u$ at site $i$ by a real vector
${\bm{f}}^{i}_{u}\in[0,1]^{|\Sigma|}$. For a leaf $u$, ${\bm{f}}^{i}_{u}$ is
the one-hot vector encoding the symbol at the $i^{\rm th}$ site, and for non-
leaves, ${\bm{f}}^{i}_{u}$ encodes
${\bm{f}}_{u}^{i}[c]=P(L_{u}^{i}|a_{u}^{i}=c)$ and is recursively computed by
(1). This motivates our modeling choices in PhyloGFN (§4.1).

####  3.1.2 Parsimony analysis

The problem of finding the optimal tree topology under the maximum parsimony
principle is commonly referred as the Large Parsimony problem, which is NP-
hard. For a given tree topology $z$, the parsimony score is the minimum number
of character changes between sequences across branches obtained by optimally
assigning sequences to internal nodes. Let $M(z|{\bm{Y}})$ be the parsimony
score of tree topology $z$ with leaf labels ${\bm{Y}}$. Due to site
independence, $M(z|{\bm{Y}})=\sum_{i}M(z|{\bm{Y}}^{i})$. The trees with the
lowest parsimony score, or most parsimonious trees, are solutions to the Large
Parsimony problem. Note that the Large Parsimony problem is a limiting case of
the maximum likelihood tree inference problem, where branch lengths are
constrained to be equal and infinitesimally short.

##### Fitch algorithm

Given a rooted tree topology $z$, the Fitch algorithm assigns optimal
sequences to internal nodes and computes the parsimony score in linear time.
At each node $u$, the algorithm tracks the set of possible characters labeling
for node $u$ that can yield a most parsimonious solution for the subtree
rooted at $u$. This character set can be represented by a binary vector
${\bm{f}}^{i}_{u}\in\\{0,1\\}^{|\Sigma|}$ for site $i$. As in Felsenstein’s
algorithm, this vector is a one-hot encoding of the sequences at the leaves
and is computed recursively for non-leaves. Specifically, given a rooted tree
with root $u$ and two child trees with roots $v$ and $w$, the Fitch character
set ${\bm{f}}_{u}^{i}$ is calculated as:

$\displaystyle{\bm{f}}_{u}^{i}=\begin{cases}{\bm{f}}_{v}^{i}\land{\bm{f}}_{w}^{i}&\text{if}\
{\bm{f}}^{i}_{v}\cdot{\bm{f}}^{i}_{w}\neq 0\\\
{\bm{f}}_{v}^{i}\lor{\bm{f}}_{w}^{i}&\text{otherwise}\end{cases},$

where $\land$ and $\lor$ are element-wise conjunctions and disjunctions.The
algorithm traverses the tree two times, first in post-order (bottom-up) to
calculate the character set at each node, then in pre-order (top-down) to
assign optimal sequences. The total number of character changes between these
optimal sequences along the tree’s edges is counted as the parsimony score.

###  3.2 GFlowNets

Generative flow networks (GFlowNets) are algorithms for learning generative
models of complex distributions given by unnormalized density functions over
structured spaces. Here, we give a concise summary of the the GFlowNet
framework.

##### Setting

A GFlowNet treats generation of objects $x$ lying in a sample space
${\mathcal{X}}$ as a sequential decision-making problem on an acyclic
deterministic MDP with set of states ${\mathcal{S}}\supset{\mathcal{X}}$ and
set of actions ${\mathcal{A}}\subseteq{\mathcal{S}}\times{\mathcal{S}}$. The
MDP has a designated _initial state_ $s_{0}$, which has no incoming actions,
and a set of _terminal states_ (those with no outgoing actions) that coincides
with ${\mathcal{X}}$. Any $x\in{\mathcal{X}}$ can be reached from $s_{0}$ by a
sequence of actions $s_{0}\rightarrow s_{1}\rightarrow\dots\rightarrow
s_{n}=x$ (with each $(s_{i},s_{i+1})\in{\mathcal{A}}$). Such sequences are
called _complete trajectories_ , and the set of all complete trajectories is
denoted ${\mathcal{T}}$.A _(forward) policy_ $P_{F}$ is a collection of
distributions $P_{F}(s^{\prime}\mid s)$ over the children of each nonterminal
state $s\in{\mathcal{S}}\setminus{\mathcal{X}}$. A policy induces a
distribution over ${\mathcal{T}}$:

$P_{F}(\tau=(s_{0}\rightarrow s_{1}\rightarrow\dots\rightarrow
s_{n}))=\prod_{i=0}^{n-1}P_{F}(s_{i+1}\mid s_{i}).\vspace*{-1mm}$

A policy gives a way to sample objects in ${\mathcal{X}}$, by sampling a
complete trajectory $\tau\sim P_{F}$ and returning its final state, inducing a
marginal distribution $P_{F}^{\top}$ over ${\mathcal{X}}$; $P_{F}^{\top}(x)$
is the sum of $P_{F}(\tau)$ over all complete trajectories $\tau$ that end in
$x$ (a possibly intractable sum). The goal of GFlowNet training is to estimate
a parametric policy $P_{F}(\cdot\mid\cdot;\theta)$ such that the induced
$P_{F}^{\top}$ is proportional to a given _reward function_
$R:{\mathcal{X}}\to\mathbb{R}_{>0}$, i.e.,

$P_{F}^{\top}(x)=\frac{1}{Z}R(x)\quad\forall x\in{\mathcal{X}},$ (2)

where $Z=\sum_{x\in{\mathcal{X}}}R(x)$ is the unknown normalization constant
(partition function).

##### Trajectory balance objective

The direct optimization of $P_{F}$’s parameters $\theta$ is impossible since
it involves an intractable sum over all complete trajectories. Instead, we
leverage the trajectory balance (TB) training objective (Malkin et al., 2022),
which introduces two auxiliary objects: an estimate $Z_{\theta}$ of the
partition function and a _backward policy_. In our experiments, we fix the
backward policy to uniform, which results in a simplified objective:

${\mathcal{L}}_{\rm
TB}(\tau)=\left(\log\frac{Z_{\theta}\prod_{i=0}^{n-1}P_{F}(s_{i+1}\mid
s_{i};\theta)}{R(x)P_{B}(\tau\mid x)}\right)^{2},\quad P_{B}(\tau\mid
x):=\prod_{i=1}^{n}\frac{1}{|{\rm Pa}(s_{i})|},$ (3)

where ${\rm Pa}(s)$ denotes the set of parents of $s$. By the results of
Malkin et al. (2022), there exists a unique policy $P_{F}$ and scalar
$Z_{\theta}$ that simultaneously make ${\mathcal{L}}_{\rm TB}(\tau)=0$ for all
$\tau\in{\mathcal{T}}$, and at this optimum, the policy $P_{F}$ satisfies (2)
and $Z_{\theta}$ equals the true partition function $Z$. In practice, the
policy $P_{F}(\cdot\mid s;\theta)$ is parameterized as a neural network that
outputs the logits of actions $s\rightarrow s^{\prime}$ given a representation
of the state $s$ as input, while $Z_{\theta}$ is parameterized in the log
domain. ${\mathcal{L}}_{\rm TB}(\tau)$ is minimized by gradient descent on
trajectories $\tau$ chosen from a behaviour policy that can encourage
exploration to accelerate mode discovery (see our behaviour policy choices in
§4.3).

##  4 Phylogenetic inference with GFlowNets

###  4.1 GFlowNets for Bayesian phylogenetic inference

This section introduces PhyloGFN-Bayesian, our GFlowNet-based method for
Bayesian phylogenetic inference. Given a set of observed sequence ${\bm{Y}}$,
PhyloGFN-Bayesian learns a sampler over the joint space of tree topologies and
edge lengths ${\mathcal{X}}=\\{(z,b)|L(z)={\bm{Y}}\\}$ such that the sampling
probability of $(z,b)\in{\mathcal{X}}$ approximates its posterior
$P_{F}^{\top}(z,b)=P(z,b|{\bm{Y}})$.We follow the same setup as (Koptagel et
al., 2022; Zhang, 2023; Mimori & Hamada, 2023): (i) uniform prior over tree
topologies; (ii) decomposed prior $P(z,b)=P(z)P(b)$; (iii) exponential
($\lambda=10$) prior over branch lengths; (iv) Jukes-Cantor substitution
model.

##### GFlowNet state and action space

The sequential procedure of constructing phylogenetic trees is illustrated in
Fig. 1. The initial state $s_{0}$ is a set of $n$ rooted trees, each
containing a single leaf node labeled with an observed sequence.Each action
chooses a pair of trees and join them at the root, thus creating a new tree.
The number of rooted trees in the set is reduced by 1 at every step, so after
$n-1$ steps a single rooted tree with $n$ leaves is left. To obtain an
unrooted tree, we simply remove the root node. Thus, a state $s$ consists of a
set of rooted trees $s=((z_{1},b_{1}),\dots,(z_{l},b_{l})),\ \ l\leq n$ and
$\bigcup_{i}L(z_{i})={\bm{Y}}$. Given a nonterminal state with $l>1$ trees, a
transition action consists of two steps: (i) choosing a pair of trees to join
out of the $l\choose 2$ possible pairs; and (ii) generating branch lengths for
the two introduced edges between the new root and its two children. The
distribution over the pair of branch lengths is modeled jointly as a discrete
distribution with fixed bin size.

Figure 1: Left: The state space of PhyloGFN on a four-sequence dataset. The
initial state $s_{0}$ is a forest of rooted trees with only leaf nodes. At
every step, a pair of trees are chosen to join at the root, and the process is
repeated until a single tree is left. To obtain an unrooted tree, the root
node is removed and its two children are connected. Right: Illustration of the
policy model for PhyloGFN-Bayesian. Given states
$s_{i}=((z_{1},b_{1})\dots(z_{l},b_{l}))$, the tree-level features are first
processed by a Transformer encoder. The encoded features are then used to form
pairwise features $\\{e_{ij}\\}$, which are fed to the tree topology MLP to
generate probability logits to sample a pair of trees to join. Then, the
corresponding pairwise feature is fed to the edge length MLP to sample branch
lengths.

##### Reward function

We define the reward function as the product of the likelihood and the edge
length prior: $R(z,b)=P({\bm{Y}}|z,b)P(b)$, implicitly imposing a uniform
prior over tree topologies. By training with this reward, PhyloGFN learns to
approximate the posterior, since
$P(z,b|{\bm{Y}})=R(z,b)\frac{P(z)}{P({\bm{Y}})}$ and $P(z),P({\bm{Y}})$ are
both constant.It is worth emphasizing that in our bottom-up construction of
trees, the set of possible actions at the steps that select two trees to join
by a new common root is never larger than $n^{2}$, even though the size of the
space of all tree topologies – all of which can be reached by our sampler – is
superexponential in $n$.This stands in contrast to the modeling choices of
VBPI-GNN (Zhang, 2023), which constructs trees in a top-down fashion and
limits the action space using a pre-generated set of trees, therefore also
limiting the set of trees it can sample.

##### State representation

To represent a rooted tree in a non-terminal state, we compute features for
each site independently by taking advantage of the Felsenstein features
(§3.1.1). Let $(z,b)$ be a weighted tree with root $u$ which has two children
$v$ and $w$. Let
${\bm{f}}^{i}_{u},{\bm{f}}^{i}_{v},{\bm{f}}^{i}_{w}\in[0,1]^{|\Sigma|}$ be the
Felsenstein feature on nodes $u,v,w$ at site $i$. The feature $\rho^{i}_{u}$
for site $i$ is computed as following:

$\rho^{i}_{u}={\bm{f}}^{i}_{u}\prod_{e\in E(z)}P(b_{z}(e))^{\frac{1}{m}}$ (4)

where $P(b_{z}(e))=\prod_{e\in b_{z}}P(b(e))$ is the edge length prior.The
tree-level feature is the concatenation of site-level features
$\rho=[\rho^{1}\dots\rho^{m}]$. A state $s=(z_{1}\dots z_{l})$, which is a
collection of rooted trees, is represented by the set
$\\{\rho_{1},\dots,\rho_{l}\\}$.

##### Representation power

Although the proposed feature representation $\rho$ does not capture all the
information of tree structure and leaf sequences, we show that $\rho$ indeed
contains sufficient information to express the optimal policy. The proposition
below shows that given an optimal GFlowNet with uniform $P_{B}$, two states
with identical features set share the same transition
probabilities.[][end,restate]Let
$s_{1}=\\{(z_{1},b_{1}),(z_{2},b_{2})\dots(z_{l},b_{l})\\}$ and
$s_{2}=\\{(z^{\prime}_{1},b^{\prime}_{1}),(z^{\prime}_{2},b_{2})\dots(z^{\prime}_{l},b_{l})\\}$
be two non-terminal states such that $s_{1}\neq s_{2}$ but sharing the same
features $\rho_{i}=\rho^{\prime}_{i}$. Let ${\bm{a}}$ be any sequence of
actions, which applied to $s_{1}$ and $s_{2}$, respectively, results in full
weighted trees $x=(z,b_{z}),x^{\prime}=(z^{\prime},b^{\prime})$, with two
partial trajectories $\tau=(s_{1}\rightarrow\dots\rightarrow
x),\tau^{\prime}=(s_{2}\rightarrow\dots\rightarrow x^{\prime})$. If $P_{F}$ is
the policy of an optimal GFlowNet with uniform $P_{B}$, then
$P_{F}(\tau)=P_{F}(\tau^{\prime})$.Let
${\mathcal{G}}_{s_{1}},{\mathcal{G}}_{s_{2}}$ be sub-graphs of the GFlowNet
state graph ${\mathcal{G}}=({\mathcal{S}},{\mathcal{A}})$ defined by reachable
states from $s_{1}$ and $s_{2}$ in ${\mathcal{G}}$. Since $s_{1}$ and $s_{2}$
have the same number of trees, ${\mathcal{G}}_{s_{1}}$ and
${\mathcal{G}}_{s_{2}}$ have the graph structure. Let
${\mathcal{X}}_{1},{\mathcal{X}}_{2}\subseteq{\mathcal{X}}$ be the terminal
states reachable from $s_{1}$ and $s_{2}$. There is thus a bijective
correspondence between ${\mathcal{X}}_{1}$ and ${\mathcal{X}}_{2}$: for every
action set ${\bm{a}}$ applying on $s_{1}$ to obtain $x\in{\mathcal{X}}_{1}$,
we obtain $x^{\prime}\in{\mathcal{X}}_{2}$ by applying ${\bm{a}}$ on $s_{2}$.
Let $\tau$ and $\tau^{\prime}$ be the partial trajectories created by applying
${\bm{a}}$ on $s_{1}$ and $s_{2}$,
$P_{B}(\tau|x)=P_{B}(\tau^{\prime}|x^{\prime})$. Moreover,
$R(x)=R(x^{\prime})$ since $s_{1}$ and $s_{2}$ share the same set of features.
We have the following expressions for $P_{F}(\tau)$ and
$P_{F}(\tau^{\prime})$:

$\displaystyle
P_{F}(\tau)=\frac{R(x)P_{B}(\tau|x)}{\sum_{\tau_{j},x_{j}}R(x_{j})P_{B}(\tau_{j}|x_{j})},\quad
P_{F}(\tau^{\prime})=\frac{R(x^{\prime})P_{B}(\tau^{\prime}|x^{\prime})}{\sum_{\tau^{\prime}_{j},x^{\prime}_{j}}R(x^{\prime}_{j})P_{B}(\tau^{\prime}_{j}|x^{\prime}_{j})}.$

Hence $P_{F}(\tau)=P_{F}(\tau^{\prime})$.All proofs are in Appendix A. The
proposition shows that our proposed features have sufficient representation
power for the PhyloGFN-Bayesian policy. Furthermore, Felsenstein features and
edge length priors are used in calculating reward by Felsenstein’s algorithm.
Therefore, computing these features does not introduce any additional
variables, and computation overhead is minimized.

###  4.2 GFlowNets for parsimony analysis

This section introduces PhyloGFN-Parsimony, our GFlowNet-based method for
parsimony analysis. We treat large parsimony analysis, a discrete optimization
problem, as a sampling problem from the energy distribution
$\exp\left(\frac{-M(z|{\bm{Y}})}{T}\right)$ defined over tree topologies.
Here, $M(z|{\bm{Y}})$ is the parsimony score of $z$ and $T$ is a pre-defined
temperature term to control the smoothness of distribution. With sufficiently
small $T$, the most parsimonious trees dominate the energy distribution.To
state our goals formally, given observed sequences ${\bm{Y}}$, PhyloGFN-
Parsimony learns a sampling policy $P_{F}$ over the space of tree topologies
$\\{z|L(z)={\bm{Y}}\\}$ such that $P_{F}^{\top}(z)\propto
e^{-\frac{M(z|Y)}{T}}$. As $T\to 0$, this target distribution approaches a
uniform distribution over the set of tree topologies with minimum parsimony
scores.PhyloGFN-Parsimony can be seen as a reduced version of PhyloGFN-
Bayesian. The tree shape generation procedure is the same as before, but we no
longer generate branch lengths. The reward is defined as
$R(z)=\exp\left(\frac{C-M(z|Y)}{T}\right)$, where $C$ is an extra
hyperparameter introduced for stability to offset the typically large
$M(z|{\bm{Y}})$ values. Note that $C$ can be absorbed into the partition
function and has no influence on the reward distribution. Similar to PhyloGFN-
Bayesian, a state (collection of rooted trees) is represented by the set of
tree features, with each tree represented by concatenating its site-level
features. With $z$ the rooted tree topology with root $u$, we represent the
tree at site $i$ by its root level Fitch feature ${\bm{f}}_{u}^{i}$. The
proposition below, analogous to Proposition 4.1, shows the representation
power of the proposed feature.[][end,restate]Let $s_{1}=\\{z_{1},z_{2},\dots
z_{l}\\}$ and $s_{2}=\\{z^{\prime}_{1},z^{\prime}_{2},\dots z^{\prime}_{l}\\}$
be two non-terminal states such that $s_{1}\neq s_{2}$ but sharing the same
Fitch features ${\bm{f}}_{z_{i}}={\bm{f}}_{z^{\prime}_{i}}\ \forall i$. Let
${\bm{a}}$ be any sequence of actions, which, applied to $s_{1}$ and $s_{2}$,
respectively, results in tree topologies $x,x^{\prime}\in{\mathcal{Z}}$, with
two partial trajectories $\tau=(s_{1}\rightarrow\dots\rightarrow
x),\tau^{\prime}=(s_{2}\rightarrow\dots\rightarrow x^{\prime})$. If $P_{F}$ is
the policy of an optimal GFlowNet with uniform $P_{B}$, then
$P_{F}(\tau)=P_{F}(\tau^{\prime})$We use the same notation as in the proof of
of Proposition 4.1. Since $s_{1}$ and $s_{2}$ have the same number of trees,
${\mathcal{G}}_{s_{1}}$ and ${\mathcal{G}}_{s_{2}}$ have the same graph
structure. Let ${\mathcal{X}}_{1},{\mathcal{X}}_{2}\subseteq{\mathcal{X}}$ be
the terminal states reachable from $s_{1}$ and $s_{2}$. There is a bijective
correspondence between ${\mathcal{X}}_{1}$ and ${\mathcal{X}}_{2}$: for every
action set ${\bm{a}}$ applying on $s_{1}$ to obtain $x\in{\mathcal{X}}_{1}$,
we obtain $x^{\prime}\in{\mathcal{X}}_{2}$by applying ${\bm{a}}$ on $s_{2}$.
Let $\tau$ and $\tau^{\prime}$ be the partial trajectories created by applying
${\bm{a}}$ on $s_{1}$ and $s_{2}$,
$P_{B}(\tau|x)=P_{B}(\tau^{\prime}|x^{\prime})$.For simplicity, we denote
$M(x)=M(x|L(x))$. Note that $R(x)\neq R(x^{\prime})$ even $s_{1}$ and $s_{2}$
share the same Fitch feature because two trees can have different parsimony
score when their root level Fitch feature equals. However, when applying
${\bm{a}}$ on $s_{1}$ and $s_{2}$, the additional parsimony scores introduced
are the same:

$\displaystyle M(x)-\sum_{i}M(z_{i})=M(x^{\prime})-\sum_{i}M(z^{\prime}_{i})$

We have the following expressions for $P_{F}(\tau)$ and
$P_{F}(\tau^{\prime})$:

$\displaystyle
P_{F}(\tau)=\frac{e^{\frac{-M(x)}{T}}P_{B}(\tau|x)}{\sum_{\tau_{j},x_{j}}e^{\frac{-M(x_{j})}{T}}P_{B}(\tau_{j}|x_{j})},\quad
P_{F}(\tau^{\prime})=\frac{e^{\frac{-M(x^{\prime})}{T}}P_{B}(\tau|x^{\prime})}{\sum_{\tau_{j}^{\prime},x_{j}^{\prime}}e^{\frac{-M(x_{j}^{\prime})}{T}}P_{B}(\tau_{j}^{\prime}|x_{j}^{\prime})}$

We multiply $\frac{\frac{\sum_{i}M(z_{i})}{T}}{\frac{\sum_{i}M(z_{i})}{T}}$ by
$P_{F}(\tau)$ and
$\frac{\frac{\sum_{i}M(z_{i}^{\prime})}{T}}{\frac{\sum_{i}M(z_{i}^{\prime})}{T}}$
by $P_{F}(\tau^{\prime})$ and obtain:

$\displaystyle
P_{F}(\tau)=\frac{e^{\frac{\sum_{i}M(z_{i})-M(x)}{T}}P_{B}(\tau|x)}{\sum_{\tau_{j},x_{j}}e^{\frac{\sum_{i}M(z_{i})-M(x_{j})}{T}}P_{B}(\tau_{j}|x_{j})},\quad
P_{F}(\tau^{\prime})=\frac{e^{\frac{\sum_{i}M(z^{\prime}_{i})-M(x^{\prime})}{T}}P_{B}(\tau|x^{\prime})}{\sum_{\tau_{j}^{\prime},x_{j}^{\prime}}e^{\frac{\sum_{i}M(z^{\prime}_{i})-M(x_{j}^{\prime})}{T}}P_{B}(\tau_{j}^{\prime}|x_{j}^{\prime})}$

Since
$e^{\frac{\sum_{i}M(z^{\prime}_{i})-M(x_{j}^{\prime})}{T}}P_{B}(\tau_{j}^{\prime}|x_{j}^{\prime})=e^{\frac{\sum_{i}M(z_{i})-M(x_{j})}{T}}P_{B}(\tau_{j}|x_{j})$
for all $j$, $P_{F}(\tau)=P_{F}(\tau^{\prime})$.This shows that the Fitch
features contain sufficient information for PhyloGFN-Parsimony. Furthermore,
the Fitch features are used in the computation of the reward by Fitch’s
algorithm, so their use in the policy model does not introduce additional
variables or extra computation.

##### Temperature-conditioned PhyloGFN

The temperature $T$ controls the trade-off between sample diversity and
parsimony scores. Following Zhang et al. (2023a), we extend PhyloGFN-Parsimony
by conditioning the policy on $T$, with reward
$R(z;T)=\exp\left(\frac{C-M(z|{\bm{Y}})}{T}\right)$, and we learn a sampler
such that $P^{\top}(z;T)\propto R(z;T)$. See Appendix E for more details.

###  4.3 Model architecture and training

##### Parameterization of forward transitions

We parameterize the forward transitions of tree topology construction using a
Transformer-based neural network, whose architecture is shown in Fig. 1. We
select Transformer because the input is a set and the model needs to be order-
equivariant. For a state consisting of $n$ trees, after $n$ embeddings are
generated from the Transformer encoder, $n\choose 2$ pairwise features are
created for all possible pairs of trees, and a common MLP generates
probability logits for jointing every tree pair. PhyloGFN-Bayesian
additionally requires generating edge lengths. Once the pair of trees to join
is selected, another MLP is applied to the corresponding pair feature to
generate probability logits for sampling the edge lengths. More details are in
Appendix D.

##### Off-policy training

The action model $P_{F}(\cdot\mid s;\theta)$ is trained with the trajectory
balance objective. Training data are generated from two sources: (i) A set of
trajectories constructed from the currently trained GFlowNet, with actions
sampled from the policy with probability $1-\epsilon$ and uniformly at random
with probability $\epsilon$. The $\epsilon$ rate drops from a pre-defined
$\epsilon_{start}$ to near 0 during the course of training.(ii) Trajectories
corresponding to the best trees seen to date (replay buffer). Trajectories are
sampled backward from these high-reward trees with uniform backward policy.

##### Temperature annealing

For PhyloGFN-Parsimony, it is crucial to choose the appropriate temperature
$T$. Large $T$ defines a flat target distribution, while small $T$ makes the
reward landscape less smooth and leads to training difficulties. We cannot
predetermine the ideal choice of $T$ before training, as we do not know a
priori the parsimony score for the dataset. Therefore, we initialize the
training with large $T$ and reduce $T$ periodically during training. See
Appendix E for details.

###  4.4 Marginal log-likelihood estimation

To assess how well the GFlowNet sampler approximates the true posterior
distribution, we use the following importance-weighted variational lower bound
on the marginal log-likelihood (MLL):

$\log P({\bm{Y}})\geq\mathbb{E}_{\tau_{1},\dots,\tau_{k}\sim
P_{F}}\log\left(P(z)\frac{1}{K}\sum_{\begin{subarray}{c}\tau_{i}\\\
\tau_{i}:s_{0}\rightarrow\dots\rightarrow(z_{i},b_{i})\end{subarray}}^{k}\frac{P_{B}(\tau_{i}|z_{i},b_{i})R(z_{i},b_{i})}{P_{F}(\tau_{i})}\right).$
(5)

Our estimator is computed by sampling a batch of $K$ trajectories and
averaging $\frac{P(z)P_{B}(\tau\mid z,b)R(z,b)}{P_{F}(\tau)}$ over the batch.
This expectation of this estimate is guaranteed to be a lower bound on $\log
P({\bm{Y}})$ and its bias decreases as $K$ grows (Burda et al.,
2016).PhyloGFN-Bayesian models branch lengths with discrete multinomial
distributions, while in reality these are continuous variables. To properly
estimate the MLL and compare with other methods defined over continuous space,
we augment our model to a continuous sampler by performing random
perturbations over edges of trees sampled from PhyloGFN-Bayesian. The
perturbation follows the uniform distribution
${\mathcal{U}}_{[-0.5\omega,0.5\omega]}$, where $\omega$ is the fixed bin size
for edge modeling in PhyloGFN-Bayesian. The resulting model over branch
lengths is then a piecewise-constant continuous distribution. We discuss the
computation details as well as the derivation of (5) in Appendix B.

Table 1: Marginal log-likelihood estimation with different methods on real
datasets DS1-DS8. PhyloGFN outperforms $\phi$-CSMC across all datasets and
GeoPhy on most. *VBPI-GNN uses predefined tree topologies in training and is
not directly comparable.

| MCMC |  |  | ML-based / amortized, full tree space
---|---|---|---|---
Dataset | MrBayes | VBPI-GNN* |  $\phi$-CSMC | GeoPhy | PhyloGFN
DS1 | $-7108.42$ | $\pm 0.18$ | $-7108.41$ | $\pm 0.14$ | $-7290.36$ | $\pm 7.23$ | $-7111.55$ | $\pm 0.07$ | $\bf-7108.95$ | $\pm 0.06$
DS2 | $-26367.57$ | $\pm 0.48$ | $-26367.73$ | $\pm 0.07$ | $-30568.49$ | $\pm 31.34$ | $\bf-26368.44$ | $\pm 0.13$ | $\bf-26368.90$ | $\pm 0.28$
DS3 | $-33735.44$ | $\pm 0.5$ | $-33735.12$ | $\pm 0.09$ | $-33798.06$ | $\pm 6.62$ | $\bf-33735.85$ | $\pm 0.12$ | $\bf-33735.6$ | $\pm 0.35$
DS4 | $-13330.06$ | $\pm 0.54$ | $-13329.94$ | $\pm 0.19$ | $-13582.24$ | $\pm 35.08$ | $-13337.42$ | $\pm 1.32$ | $\bf-13331.83$ | $\pm 0.19$
DS5 | $-8214.51$ | $\pm 0.28$ | $-8214.64$ | $\pm 0.38$ | $-8367.51$ | $\pm 8.87$ | $-8233.89$ | $\pm 6.63$ | $\bf-8215.15$ | $\pm 0.2$
DS6 | $-6724.07$ | $\pm 0.86$ | $-6724.37$ | $\pm 0.4$ | $-7013.83$ | $\pm 16.99$ | $-6733.91$ | $\pm 0.57$ | $\bf-6730.68$ | $\pm 0.54$
DS7 | $-37332.76$ | $\pm 2.42$ | $-37332.04$ | $\pm 0.12$ |  |  | $\bf-37350.77$ | $\pm 11.74$ | $-37359.96$ | $\pm 1.14$
DS8 | $-8649.88$ | $\pm 1.75$ | $-8650.65$ | $\pm 0.45$ | $-9209.18$ | $\pm 18.03$ | $-8660.48$ | $\pm 0.78$ | $\bf-8654.76$ | $\pm 0.19$

##  5 Experiments

We evaluate PhyloGFN on a suite of 8 real-world benchmark datasets (Table S1
in Appendix C) that is standard in the literature. These datasets feature pre-
aligned sequences and vary in difficulty (27 to 64 sequences; 378 to 2520
sites). In the following sections we present our results and analysis on
Bayesian and parsimony-based phylogenetic inference.

Figure 2: Model sampling log-density vs. unnormalized posterior log-density
for high/medium/low-probability trees on DS1. We highlight that PhyloGFN-
Bayesian performs significantly better on medium- and low-probability trees,
highlighting its superiority in modeling the entire data space.

###  5.1 Bayesian phylogenetic inference

PhyloGFN is compared with a variety of baselines in terms of sampling-based
estimates of marginal log-likelihood (MLL; see details in §4.4). The baselines
we compare to are the MCMC-based MrBayes combined with the stepping-stone
sampling technique (Ronquist et al., 2012), and three variational inference
methods: VBPI-GNN (Zhang, 2023), $\phi$-CSMC proposed in VaiPhy (Koptagel et
al., 2022), and GeoPhy (Mimori & Hamada, 2023). The sampling setup for MrBayes
follows Zhang & Matsen IV (2018b) and otherwise show the highest MLL reported
in the respective papers.The results are summarized in Table 1. Our PhyloGFN
is markedly superior to $\phi$-CSMC across all datasets and outperforms GeoPhy
on most, with the exception of DS2 and DS3 where the two perform similarly,
and DS7 where GeoPhy obtains a better result. VBPI-GNN, is the only machine
learning-based method that performs on par against MrBayes, the current gold
standard in Bayesian phylogenetic inference. However, it should be emphasized
that VBPI-GNN requires a set of pre-defined tree topologies that are likely to
achieve high likelihood, and as a consequence, its training and inference are
both constrained to a small space of tree topologies, thus severely
undermining its applicability to, for example, postulating alternative
phylogenetic theories.On the other hand, PhyloGFN operates on the full space
of tree topologies and, in fact, achieves a closer fit to the true posterior
distribution. To show this, for each dataset, we created three sets of
phylogenetic trees with high/medium/low posterior probabilities and obtained
the corresponding sampling probabilities from PhyloGFN and VBPI-GNN. The three
classes of trees are generated from VBPI-GNN by randomly inserting uniformly
sampled actions into its sequential tree construction process with $0\%$,
$30\%$, or $50\%$ probability, respectively, which circumvents VBPI-GNN’s
limitation of being confined to a small tree topology space.Table 2 and Fig. 2
show that PhyloGFN achieves higher Pearson correlation between the sampling
log-probability and the unnormalized ground truth posterior log-density for
the majority of datasets and classes of trees. In particular, while VBPI-GNN
performs better on high-probability trees, its correlation drops significantly
on lower-probability trees. On the other hand, PhyloGFN maintains a high
correlation for all three classes of trees across all datasets, the only
exception being the high-probability trees in DS7. See Appendix F for details
and extended results.

Table 2: Pearson correlation of model sampling log-density and ground truth unnormalized posterior log-density for each dataset on high/medium/low posterior density trees generated by VBPI-GNN. PhyloGFN achieves a good fit on both high and low posterior density regions. | No random | 30% random | 50% random
---|---|---|---
Dataset | PhyloGFN | VBPI-GNN | PhyloGFN | VBPI-GNN | PhyloGFN | VBPI-GNN
DS1 | 0.994 | 0.955 | 0.961 | 0.589 | 0.955 | 0.512
DS2 | 0.930 | 0.952 | 0.948 | 0.580 | 0.919 | 0.538
DS3 | 0.917 | 0.968 | 0.963 | 0.543 | 0.950 | 0.499
DS4 | 0.942 | 0.960 | 0.945 | 0.770 | 0.966 | 0.76
DS5 | 0.969 | 0.965 | 0.937 | 0.786 | 0.939 | 0.773
DS6 | 0.993 | 0.887 | 0.973 | 0.816 | 0.934 | 0.702
DS7 | 0.624 | 0.955 | 0.787 | 0.682 | 0.764 | 0.678
DS8 | 0.978 | 0.955 | 0.913 | 0.604 | 0.901 | 0.463

###  5.2 Parsimony-based phylogenetic inference

Figure 3: A temperature-conditioned PhyloGFN is trained on DS1 using
temperatures sampled between 4.0 and 1.0. (A) The distribution of parsimony
scores can be controlled by varying the statistical temperature – an input
variable to the PhyloGFN policy – from 8.0 to 1.0. 10,000 trees are randomly
sampled at each temperature. (B) PhyloGFN achieves high Pearson correlation
for trees sampled within each temperature range.

As a special case of Bayesian phylogenetic inference, the parsimony problem is
concerned with finding the most-parsimonious trees – a task which is also
amenable to PhyloGFN. Here, we compare to the state-of-the-art parsimony
analysis software PAUP* (Swofford, 1998). For all datasets, our PhyloGFN and
PAUP* are able to identify the same set of optimal solutions to the Large
Parsimony problem, ranging from a single optimal tree for DS1 to 21 optimal
trees for DS8. Although the results are similar between PhyloGFN and PAUP*,
once again we emphasize that PhyloGFN is based on a rigorous mathematical
framework of fitting and sampling from well-defined posterior distributions
over tree topologies. whereas PAUP*’s relies on heuristic algorithms. To put
it more concretely, we show in Fig. S4 that PhyloGFN is able to (i) learn a
smooth echelon of sampling probabilities that distinguish the optimal trees
from suboptimal ones; (ii) learn similar sampling probabilities for trees
within each group of equally-parsimonious trees; and (iii) fit all $2n-3$
rooted trees that belong to the same unrooted tree equally well.Finally, Fig.
3 shows that a single temperature-conditioned PhyloGFN can sample phylogenetic
trees of varying ranges of parsimony scores by providing suitable temperatures
$T$ as input to the model. Also, PhyloGFN is able to sample proportionately
from the Boltzmann distribution defined at different temperatures and achieves
high correlation between sampling log-probability and log-reward. Although the
temperature-conditioned PhyloGFN has only been trained on a small range of
temperatures between 4.0 and 1.0, Fig. 3 shows it can also approximate the
distribution defined at temperature 8.0. Further results are presented in
Appendix G.

##  6 Discussion and future work

In this paper, we propose PhyloGFN, a GFlowNet-based generative modeling
algorithm, to solve parsimony-based and Bayesian phylogenetic inference. We
design an intuitive yet effective tree construction procedure to efficiently
model the entire tree topology space. We propose anovel tree representation
based on Fitch’s and Felsenstein’s algorithms to represent rooted trees
without introducing additional learnable parameters, and we show the
sufficiency of our features for the purpose of phylogenetic tree generation.
We apply our algorithm on eight real datasets, demonstrating that PhyloGFN is
competitive with or superior to prior works in terms of marginal likelihood
estimation, while achieving a closer fit to the target distribution compared
to state-of-the-art variational inference methods. Future work should consider
continuous edge length sampling in place of the current binning procedure,
e.g., using a continuous GFlowNet policy (Lahlou et al., 2023). In addition,
we plan to explore the use of conditional GFlowNets to amortize the dependence
on the sequence dataset itself. This would allow a single trained GFlowNet to
sample phylogenetic trees for sets of sequences that were not seen during
training.

## Acknowledgments

We thank Oskar Kviman and Jens Lagergren for insightful discussions at the
inception of this project.

## References

  * Aminov & Mackie (2007) Rustam I Aminov and Roderick I Mackie.  Evolution and ecology of antibiotic resistance genes.  _FEMS microbiology letters_ , 271(2):147–161, 2007.
  * Bengio et al. (2021) Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio.  Flow network based generative models for non-iterative diverse candidate generation.  _Neural Information Processing Systems (NeurIPS)_ , 2021.
  * Bengio et al. (2023) Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio.  GFlowNet foundations.  _Journal of Machine Learning Research_ , (24):1–76, 2023.
  * Blanchette et al. (2004) Mathieu Blanchette, W James Kent, Cathy Riemer, Laura Elnitski, Arian FA Smit, Krishna M Roskin, Robert Baertsch, Kate Rosenbloom, Hiram Clawson, Eric D Green, et al.  Aligning multiple genomic sequences with the threaded blockset aligner.  _Genome research_ , 14(4):708–715, 2004.
  * Burda et al. (2016) Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov.  Importance weighted autoencoders.  _International Conference on Learning Representations (ICLR)_ , 2016.
  * Celniker et al. (2013) Gershon Celniker, Guy Nimrod, Haim Ashkenazy, Fabian Glaser, Eric Martz, Itay Mayrose, Tal Pupko, and Nir Ben-Tal.  Consurf: using evolutionary data to raise testable hypotheses about protein function.  _Israel Journal of Chemistry_ , 53(3-4):199–206, 2013.
  * Chor & Tuller (2005) Benny Chor and Tamir Tuller.  Maximum likelihood of evolutionary trees is hard.  In _Annual International Conference on Research in Computational Molecular Biology_ , pp. 296–310. Springer, 2005.
  * Day (1987) William HE Day.  Computational complexity of inferring phylogenies from dissimilarity matrices.  _Bulletin of mathematical biology_ , 49(4):461–467, 1987.
  * Deleu et al. (2022) Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio.  Bayesian structure learning with generative flow networks.  _Uncertainty in Artificial Intelligence (UAI)_ , 2022.
  * Deleu et al. (2023) Tristan Deleu, Mizu Nishikawa-Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Charlin, and Yoshua Bengio.  Joint Bayesian inference of graphical structure and parameters with a single generative flow network.  _arXiv preprint arXiv:2305.19366_ , 2023.
  * Felsenstein (1973) Joseph Felsenstein.  Maximum likelihood and minimum-steps methods for estimating evolutionary trees from data on discrete characters.  _Systematic Biology_ , 22(3):240–249, 1973.
  * Garey et al. (1996) J R Garey, T J Near, M R Nonnemacher, and S A Nadler.  Molecular evidence for acanthocephala as a subtaxon of rotifera.  _J Mol Evol_ , 43(3):287–292, September 1996.
  * Hamelin et al. (2022) Richard C Hamelin, Guillaume J Bilodeau, Renate Heinzelmann, Kelly Hrywkiw, Arnaud Capron, Erika Dort, Angela L Dale, Emilie Giroux, Stacey Kus, Nick C Carleson, et al.  Genomic biosurveillance detects a sexual hybrid in the sudden oak death pathogen.  _Communications Biology_ , 5(1):477, 2022.
  * Hedges et al. (1990) S B Hedges, K D Moberg, and L R Maxson.  Tetrapod phylogeny inferred from 18S and 28S ribosomal RNA sequences and a review of the evidence for amniote relationships.  _Molecular Biology and Evolution_ , 7(6):607–633, 11 1990.
  * Henk et al. (2003) Daniel A Henk, Alex Weir, and Meredith Blackwell.  Laboulbeniopsis termitarius, an ectoparasite of termites newly recognized as a member of the laboulbeniomycetes.  _Mycologia_ , 95(4):561–564, July 2003.
  * Hu et al. (2023) Edward J Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio.  GFlowNet-EM for learning compositional latent variable models.  _International Conference on Machine Learning (ICML)_ , 2023.
  * Höhna et al. (2016) Sebastian Höhna, Michael J. Landis, Tracy A. Heath, Bastien Boussau, Nicolas Lartillot, Brian R. Moore, John P. Huelsenbeck, and Fredrik Ronquist.  RevBayes: Bayesian Phylogenetic Inference Using Graphical Models and an Interactive Model-Specification Language.  _Systematic Biology_ , 65(4):726–736, 05 2016.  ISSN 1063-5157.  doi: 10.1093/sysbio/syw021.  URL https://doi.org/10.1093/sysbio/syw021.
  * Jain et al. (2022) Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure F.P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio.  Biological sequence design with GFlowNets.  _International Conference on Machine Learning (ICML)_ , 2022.
  * Koptagel et al. (2022) Hazal Koptagel, Oskar Kviman, Harald Melin, Negar Safinianaini, and Jens Lagergren.  VaiPhy: a variational inference based algorithm for phylogeny.  _Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Lahlou et al. (2023) Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin.  A theory of continuous generative flow networks.  _International Conference on Machine Learning (ICML)_ , 2023.
  * Lähnemann et al. (2020) David Lähnemann, Johannes Köster, Ewa Szczurek, Davis J McCarthy, Stephanie C Hicks, Mark D Robinson, Catalina A Vallejos, Kieran R Campbell, Niko Beerenwinkel, Ahmed Mahfouz, et al.  Eleven grand challenges in single-cell data science.  _Genome biology_ , 21(1):1–35, 2020.
  * Lakner et al. (2008) Clemens Lakner, Paul Van Der Mark, John P Huelsenbeck, Bret Larget, and Fredrik Ronquist.  Efficiency of Markov chain Monte Carlo tree proposals in Bayesian phylogenetics.  _Systematic biology_ , 57(1):86–103, 2008.
  * Layne et al. (2020) Elliot Layne, Erika N Dort, Richard Hamelin, Yue Li, and Mathieu Blanchette.  Supervised learning on phylogenetically distributed data.  _Bioinformatics_ , 36(Supplement_2):i895–i902, 2020.
  * Ma et al. (2006) Jian Ma, Louxin Zhang, Bernard B Suh, Brian J Raney, Richard C Burhans, W James Kent, Mathieu Blanchette, David Haussler, and Webb Miller.  Reconstructing contiguous regions of an ancestral genome.  _Genome research_ , 16(12):1557–1565, 2006.
  * Malkin et al. (2022) Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio.  Trajectory balance: Improved credit assignment in GFlowNets.  _Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Malkin et al. (2023) Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio.  GFlowNets and variational inference.  _International Conference on Learning Representations (ICLR)_ , 2023.
  * Mimori & Hamada (2023) Takahiro Mimori and Michiaki Hamada.  Geophy: Differentiable phylogenetic inference via geometric gradients of tree topologies.  _arXiv preprint arXiv:2307.03675_ , 2023.
  * Moretti et al. (2021) Antonio Khalil Moretti, Liyi Zhang, Christian A. Naesseth, Hadiah Venner, David Blei, and Itsik Pe’er.  Variational combinatorial sequential Monte Carlo methods for Bayesian phylogenetic inference.  _Uncertainty in Artificial Intelligence (UAI)_ , 2021.
  * Ranjbar et al. (2020) Reza Ranjbar, Sedigheh Nazari, and Omid Farahani.  Phylogenetic analysis and antimicrobial resistance profiles of escherichia coli strains isolated from uti-suspected patients.  _Iranian Journal of Public Health_ , 49(9):1743, 2020.
  * Ronquist et al. (2012) Fredrik Ronquist, Maxim Teslenko, Paul Van Der Mark, Daniel L Ayres, Aaron Darling, Sebastian Höhna, Bret Larget, Liang Liu, Marc A Suchard, and John P Huelsenbeck.  MrBayes 3.2: efficient Bayesian phylogenetic inference and model choice across a large model space.  _Systematic biology_ , 61(3):539–542, 2012.
  * Rossman et al. (2001) Amy Y. Rossman, John M. McKemy, Rebecca A. Pardo-Schultheiss, and Hans-Josef Schroers.  Molecular studies of the bionectriaceae using large subunit rdna sequences.  _Mycologia_ , 93(1):100–110, 2001.
  * Swofford (1998) David L Swofford.  Phylogenetic analysis using parsimony.  1998\.
  * Tjelmeland & Hegstad (2001) Hakon Tjelmeland and Bjorn Kare Hegstad.  Mode jumping proposals in mcmc.  _Scandinavian journal of statistics_ , 28(1):205–223, 2001.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Neural Information Processing Systems (NIPS)_ , 2017.
  * Yang & Yoder (2003) 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.  _Systematic Biology_ , 52(5):705–716, 10 2003.
  * Yoder & Yang (2004) Anne D. Yoder and Ziheng Yang.  Divergence dates for malagasy lemurs estimated from multiple gene loci: geological and evolutionary context.  _Molecular Ecology_ , 13(4):757–773, 2004.
  * Zhang (2023) Cheng Zhang.  Learnable topological features for phylogenetic inference via graph neural networks.  _arXiv preprint arXiv:2302.08840_ , 2023.
  * Zhang & Matsen IV (2018a) Cheng Zhang and Frederick A Matsen IV.  Generalizing tree probability estimation via bayesian networks.  _Neural Information Processing Systems (NIPS)_ , 2018a.
  * Zhang & Matsen IV (2018b) Cheng Zhang and Frederick A Matsen IV.  Variational bayesian phylogenetic inference.  _International Conference on Learning Representations (ICLR)_ , 2018b.
  * Zhang et al. (2023a) David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan.  Robust scheduling with GFlowNets.  _International Conference on Learning Representations (ICLR)_ , 2023a.
  * Zhang et al. (2022) Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio.  Generative flow networks for discrete probabilistic modeling.  _International Conference on Machine Learning (ICML)_ , 2022.
  * Zhang et al. (2023b) Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron Courville, Yoshua Bengio, and Ling Pan.  Let the flows tell: Solving graph combinatorial problems with GFlowNets.  _arXiv preprint arXiv:2305.17010_ , 2023b.
  * Zhang & Blackwell (2001) Ning Zhang and Meredith Blackwell.  Molecular phylogeny of dogwood anthracnose fungus (discula destructiva) and the diaporthales.  _Mycologia_ , 93(2):355–365, 2001.
  * Zimmermann et al. (2023) Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, and Christian A. Naesseth.  A variational perspective on generative flow networks.  _Transactions on Machine Learning Research (TMLR)_ , 2023.

##  Appendix A Proofs of propositions 1 and 2

Before proving proposition 4.1, we first prove the following two lemmas.
First, we show that for two states sharing the same tree features, applying
the same action to the states results in two new states still sharing the same
features.

######  Lemma 1.

Let $s_{1}=\\{(z_{1},b_{1}),(z_{2},b_{2})\dots(z_{l},b_{l})\\}$ and
$s_{2}=\\{(z^{\prime}_{1},b^{\prime}_{1}),(z^{\prime}_{2},b_{2})\dots(z^{\prime}_{l},b_{l})\\}$
be two non-terminating states sharing the same features
$\rho_{i}=\rho^{\prime}_{i}$. Leta be the action that joins the trees with
indices $(v,w)$ to form a new tree indexed $u$ with edge lengths
$(b(e_{uv}),b(e_{uv}))$. By applying $a$ on $s_{1}$, we join $(z_{v},b_{v})$
and $(z_{w},b_{w})$ to form new tree $(z_{u},b_{u})$. By applying $a$ on
$s_{2}$, we join $(z^{\prime}_{v},b^{\prime}_{v})$ and
$(z^{\prime}_{w},b^{\prime}_{w})$ to form new tree
$(z^{\prime}_{u},b^{\prime}_{u})$. Then the new trees’ features are equal:
$\rho_{u}=\rho_{u}^{\prime}$.

###### Proof.

We show that $\rho_{u}$ can be calculated from $\rho_{v}$ and $\rho_{w}$:

$\displaystyle\rho_{u}^{i}[j]=P(b(e_{uv}))^{\frac{1}{m}}\times
P(b(e_{uw}))^{\frac{1}{m}}\times\sum_{k\in\Sigma}P(a^{i}_{u}=j|a^{i}_{v}=k,b_{z}(e_{uv}))\rho_{v}^{i}[k]$
$\displaystyle\times\sum_{k}P(a^{i}_{u}=j|a^{i}_{w}=k,b(e_{uw}))\rho^{i}_{w}[k]$

since $\rho_{v}=\rho_{v}^{\prime}$, $\rho_{w}=\rho_{w}^{\prime}$ and
$(b(e_{uv}),b(e_{uw}))$ are new branch lengths for both two trees. Therefore
$\rho_{u}=\rho_{u}^{\prime}$∎

Next, we show that for two states sharing the same tree features, applying the
same action sequences results in two phylogenetic trees with the same reward.

######  Lemma 2.

Let $s_{1}=\\{(z_{1},b_{1}),(z_{2},b_{2})\dots(z_{l},b_{l})\\}$ and
$s_{2}=\\{(z^{\prime}_{1},b^{\prime}_{1}),(z^{\prime}_{2},b_{2})\dots(z^{\prime}_{l},b_{l})\\}$
be two non-terminating states sharing the same features
$\rho_{i}=\rho^{\prime}_{i}$. Let ${\bm{a}}$ be any sequence of actions to
apply on $s_{1}$ and $s_{2}$ to form full trees
$x=(z,b_{z}),x^{\prime}=(z^{\prime},b_{z}^{\prime})$,
$R(z,b)=R(z^{\prime},b^{\prime})$.

###### Proof.

Let $\rho_{u}$ denote the tree feature for $(z,b)$,
$\rho^{i}_{u}={\bm{f}}^{i}_{u}\prod_{e}P(b(e))$. We first show that the reward
can be directly calculated from the root feature $\rho_{u}$:

$\displaystyle\prod_{i}P(a_{2n-1})\cdot\rho_{u}$
$\displaystyle=\prod_{e}P(b(e))\prod_{i}P(a_{2n-1})\cdot f^{i}_{u}$
$\displaystyle=P(b)P({\bm{Y}}|z,b)$ $\displaystyle=R(z,b),$

where $P(a_{2n-1})$ is the constant root character assignment probability. As
${\bm{a}}$ is applied to $s_{1}$ and $s_{2}$ in a sequential manner, at every
step we obtain two state swith the same tree features (by Lemma 1), until, at
the final state, $\rho_{u}=\rho_{u}^{\prime}$. It follows that
$R(z,b)=R(z^{\prime},b^{\prime})$.∎

We are now ready to prove the propositions.

##  Appendix B Marginal log-likelihood estimation

##### Estimating the sampling likelihood of a terminal state

In the PhyloGFN state space (in both the Bayesian and parsimony-based
settings), there exist multiple trajectories leading to the same terminal
state $x$, hence the sampling probability of $x$ is calculated as:
$P^{\top}(x)=\sum_{\tau:s_{0}\rightarrow\dots\rightarrow x}P_{F}(\tau)$. This
sum is intractable for large-scale problems. However, we can estimate the sum
using importance sampling(Zhang et al., 2022):

$P_{F}^{\top}(x)\approx\frac{1}{K}\sum_{\tau_{i}:s_{0}\rightarrow\dots\rightarrow
x}\frac{P_{F}(\tau_{i})}{P_{B}(\tau_{i}|x)},$ (S1)

where the trajectories $\tau_{i}$ are sampled from $P_{B}(\tau_{i}\mid x)$.
The logarithm of the right side of (S1) is, in expectation, a lower bound on
the logarithm of the true sampling likelihood on the left side of (S1).

##### Estimating the MLL

The lower bound on MLL can be evaluated using the importance-weighted bound
$\log P({\bm{Y}})\geq\mathbb{E}_{x_{1}\dots x_{k}\sim
P_{F}}\log\frac{1}{K}\sum\frac{P(x,{\bm{Y}})}{P^{T}(x)}$(Burda et al.,
2016).However, we cannot use it for PhyloGFN since we cannot compute the exact
$P^{T}(x)$, only get a lower bound on it using (S1). Therefore, we propose the
following variational lower bound:

$\displaystyle\log P({\bm{Y}})\geq\mathbb{E}_{\tau_{1},\dots,\tau_{k}\sim
P_{F}}\log\left(P(z)\frac{1}{K}\sum_{\begin{subarray}{c}\tau_{i}\\\
\tau_{i}:s_{0}\rightarrow\dots\rightarrow(z_{i},b_{i})\end{subarray}}^{k}\frac{P_{B}(\tau_{i}|z_{i},b_{i})R(z_{i},b_{i})}{P_{F}(\tau_{i})}\right)$

We show the derivation of the estimator and thus prove its consistency:

$\displaystyle P({\bm{Y}})$
$\displaystyle=\int_{z,b}P({\bm{Y}}|z,b)P(b|z)P(z)$
$\displaystyle=\int_{z,b}R(z,b)P(z)$
$\displaystyle=\int_{z,b}R(z,b)P(z)\sum_{\tau:s_{0}\dots
x=(z,b)}P_{B}(\tau|z,b)$
$\displaystyle=\int_{z,b}R(z,b)P(z)\sum_{\tau:s_{0}\dots
x=(z,b)}P_{B}(\tau|z,b)\frac{P_{F}(\tau)}{P_{F}(\tau)}$
$\displaystyle=\int_{z,b}\sum_{\tau:s_{0}\dots
x=(z,b)}P_{F}(\tau)\frac{P_{B}(\tau|z,b)R(z,b)P(z)}{P_{F}(\tau)}$
$\displaystyle=\int_{\tau:s_{0}\dots
x_{\tau}=(z_{\tau},b_{\tau})}P_{F}(\tau)\frac{P_{B}(\tau|z_{\tau},b_{\tau})R(z_{\tau},b_{\tau})P(z)}{P_{F}(\tau)}$
$\displaystyle=P(z)\mathbb{E}_{\tau\sim
P_{F}}\frac{P_{B}(\tau|z_{\tau},b_{\tau})R(z_{\tau},b_{\tau})}{P_{F}(\tau)}$
$\displaystyle\approx P(z)\frac{1}{K}\sum_{\begin{subarray}{c}\tau_{i}\sim
P_{F}\\\
\tau_{i}:s_{0}\dots(z_{i},b_{i})\end{subarray}}^{k}\frac{P_{B}(\tau_{i}|z_{i},b_{i})R(z_{i},b_{i})}{P_{F}(\tau_{i})}.$

One can show, in a manner identical to the standard importance-weighted bound,
that this estimate is a lower bound on $\log P({\bm{Y}})$.PhyloGFN-Bayesian
models edge lengths using discrete distributions. To estimate our algorithm’s
MLL, we augment the sampler to a continuous sampler by modeling branch lengths
with a piecewise-constant continuous distribution based on the fixed-bin
multinomial distribution of PhyloGFN. We can still use the above formulation
to estimate the lower bound. However, each trajectory now has one extra step:
$\tau^{\prime}=s_{0}\rightarrow\dots\rightarrow(z,b)\rightarrow(z,b_{i})$
where $(z,b_{i})$ is obtained from $z,b$ by randomly perturbing each branch
length by adding noise from $U[-0.5\omega,0.5\omega]$, where $\omega$ is the
bin size used for PhyloGFN. Let $\tau=s_{0}\rightarrow\dots\rightarrow(z,b)$
be the original trajectory in the discrete PhyloGFN, we can compute
$P_{F}(\tau^{\prime}),P_{B}(\tau^{\prime})$ from $P_{F}(\tau),P_{B}(\tau)$:

$\displaystyle P_{F}(\tau^{\prime})=P_{F}(\tau)\frac{1}{\omega^{|E(z)|}},\ \ \
P_{B}(\tau^{\prime})=P_{B}(\tau)$

The term $\frac{1}{\omega^{|E(z)|}}$ is introduced in $P_{F}$ because we
additionally sample over a uniform range $\omega$ for all $|E(z)|$ edges. The
backward probability $P_{B}$ stays unchanged because given $(z,b)$ generated
from the discrete GFN, for any $(z,b^{\prime})$ resulting from perturbing
edges, $(z,b)$ is the the only possible ancestral state.

##  Appendix C Dataset information

Table S1: Statistics of the benchmark datasets from DS1 to DS8. Dataset | # Species | # Sites | Reference
---|---|---|---
DS1 | 27 | 1949 | Hedges et al. (1990)
DS2 | 29 | 2520 | Garey et al. (1996)
DS3 | 36 | 1812 | Yang & Yoder (2003)
DS4 | 41 | 1137 | Henk et al. (2003)
DS5 | 50 | 378 | Lakner et al. (2008)
DS6 | 50 | 1133 | Zhang & Blackwell (2001)
DS7 | 59 | 1824 | Yoder & Yang (2004)
DS8 | 64 | 1008 | Rossman et al. (2001)

##  Appendix D Modeling

Given the character set $\Sigma$, we use one-hot encoding to represent each
site in a sequence. To deal with wild characters in the dataset, for parsimony
analysis we consider them as one special character in $\Sigma$, while in
Bayesian inference, we represent them by a vector of $\mathds{1}$.For both
PhyloGFN-Parsimony and PhyloGFN-Bayesian, given a state with $l$ rooted trees,
its representation feature is the set $\\{\rho_{1}\dots\rho_{l}\\}$, where
$\rho$ is a vector of length $m|\Sigma|$. For example, for DNA sequences of
1000 sites, each $\rho$ would have length 4000. Therefore, before passing
these features to the Transformer encoder, we first use a linear
transformation to obtain lower-dimensional embeddings of the input features.We
use the Transformer architecture (Vaswani et al., 2017) to build the
Transformer encoder. For a state with $l$ trees, the output is a set of $l+1$
features $\\{e_{s},e_{1},\dots,e_{l}\\}$ where $e_{s}$ denotes the summary
feature (i.e., the [CLS] token of the Transformer encoder input).To select
pairs of trees to join, we evaluate tree-pair features for every pair of trees
in the state and pass these tree-pair features as input to the tree MLP to
generate probability logits for all pairs of trees. The tree-pair feature for
a tree pair $(i,j)$ with representations $e_{i}$, $s_{j}$ is the concatenation
of $e_{i}+e_{j}$ with the summary embedding of the state, i.e., the feature is
$[e_{s};e_{i}+e_{j}]$, where $[\cdot;\cdot]$ denotes vector direct sum
(concatenation). For a state with $l$ trees, $\binom{l}{2}=\frac{l(l-1)}{2}$
such pairwise features are generated for all possible pairs.To generate edge
lengths for the joined tree pair $(i,j)$, we pass $[e_{s};e_{i};e_{j}]$ – the
concatenation of the summary feature with the tree-level features of trees $i$
and $j$ – as input to the edge MLP. For unrooted tree topologies we need to
distinguish two scenarios: (i) when only two rooted trees are left in the
state (i.e., at the last step of PhyloGFN state transitions), we only need to
generate a single edge; and (ii) when there are more than two rooted trees in
the state, a pair of edges is required. Therefore, two separate edge MLPs are
employed, as edge length is modeled by $k$ discrete bins, the edge MLP used at
the last step has an output size of $k$ (to model a distribution over a single
edge length) whereas the other edge MLP would have an output size of $k^{2}$
(to model a joint distribution over two edge lengths).For the temperature-
conditioned PhyloGFN, as then temperature $T$ is passed to our PhyloGFN as an
input, two major modifications are required: (i) the estimation of the
partition $Z_{\theta}$ is now a function of $T$: $Z_{\theta}(T)$, which is
modeled by a Z MLP; (ii) the summary token to the Transformer encoder also
captures the temperature information by replacing the usual [CLS] token with a
temp MLP that accepts $T$ as input.

##  Appendix E Training details

Here, we describe the training details for our PhyloGFN. For PhyloGFN-
Bayesian, our models are trained with fixed 500 epochs. For PhyloGFN-
Parsimony, our models are trained until the probability of sampling the
optimal trees, or the most parsimonious trees our PhyloGFN has seen so far, is
above 0.001. Each training epoch consists of 1000 gradient update steps using
a batch of 64 trajectories. For ${\epsilon}$-greedy exploration, the
${\epsilon}$ value is linearly annealed from 0.5 to 0.001 during the first 240
epochs. All common hyperparameters for PhyloGFN-Bayesian and PhyloGFN-
Parsimony are shown in Table S3.

##### Temperature annealing

For PhyloGFN-Bayesian, the initial temperature is set to $16$ for all
experiments. For PhyloGFN-Parsimony, $T$ is initialized at 4. Under the
cascading temperature annealing scheme, $T$ is reduced by half per every 80
epochs of training. For PhyloGFN-Bayesian, $T$ is always reduced to and then
fixed at 1, whereas for PhyloGFN-Parsimony, $T$ is only reduced when the most
parsimonious trees seen by our PhyloGFN so far are sampled with a probability
below 0.001.

##### Hyperparameter $C$ selection

For PhyloGFN-Parsimony, the reward is defined as
$R(x)=\exp\left(\frac{C-M(x|{\bm{Y}})}{T}\right)$, where $C$ is a
hyperparameter introduced for training stability and it controls the magnitude
of the partition function $Z=\sum_{x}R(x)$. Given that we cannot determine the
precise value of $C$ prior to training since we do know the value of $Z$ as a
priori, we use the following heuristic to choose $C$: 1000 random tree
topologies are generated via stepwise-addition, and we set $C$ to the $10\%$
quantile of the lowest parsimony score.Similarly, $C$ is employed for
PhyloGFN-Bayesian under the same goal of stabilizing the training and reducing
the magnitude of the partition function. Recall the reward function $R(z,b)$
defined in section 4.1, it can be rewritten as
$R(z,b)=\exp\left(\frac{C-(-\log P({\bm{Y}}|z,b)P(b))}{T}\right)$ when $T=1$.
Note that $\exp\left(\frac{C}{T}\right)$ can be absorbed into the partition
function. For selecting the $C$, once again we randomly sample 1000 weighted
phylogenetic trees via stepwise-addition and with random edge length, followed
by setting $C$ to the $10\%$ quantile of the lowest $-\log
P({\bm{Y}}|z,b)P(z)$.

##### Temperature-conditioned PhyloGFN-Parsimony

The temperature-conditioned PhyloGFN is introduced so that a single trained
PhyloGFN can be used to sample from a series of reward distributions defined
by different $T$. We modify the fixed-temperature PhyloGFN-Parsimony by
introducing $T$ as input in 3 places: (i) the reward function $R(x;T)$; (ii)
the forward transition policy $P_{F}(x;T)$; and (iii) the learned partition
function estimate $Z_{\theta}(T)$. To train the temperature-conditioned
PhyloGFN, the TB objective also needs to be updated accordingly:

$\displaystyle{\mathcal{L}}_{\rm
TB}(\tau;T)=\left(\log\frac{Z_{\theta}(T)\prod_{i=0}^{n-1}P_{F}(s_{i+1}\mid
s_{i};\theta,T)}{R(x,T)P_{B}(\tau\mid x)}\right)^{2},\quad P_{B}(\tau\mid
x):=\prod_{i=1}^{n}\frac{1}{|{\rm Pa}(s_{i})|},$

Note that $P_{B}$ is unaffected.During training, values for $T$ are randomly
selected from the range $[T_{\rm min},T_{\rm max}]$. When training with a
state from the replay buffer, the temperature used for training is resampled
and may be different than the one originally used when the state was added to
the buffer.We also employ a scheme of gradually reducing the average of
sampled $T$ values through the course of training: $T$’s are sampled from a
truncated normal distribution with a fixed pre-defined variance and a moving
mean $T_{\mu}$ that linearly reduces from $T_{\rm max}$ to $T_{\rm min}$.

##### Modeling branch lengths with discrete multinomial distribution

When a pair of trees are joined at the root, the branch lengths of the two
newly formed edges are modeled jointly by a discrete multinomial distribution.
The reason for using a joint distribution is because under the Jukes-Cantor
evolution model, the likelihood at the root depends on the sum of the two
branch lengths.We use a different set of maximum edge length, bin number and
bin size depending on each dataset, and by testing various combinations we
have selected the set with optimal performance. The maximum branch length is
chosen among $\\{0.05,0.1,0.2\\}$, and bin size $\omega$ is chosen among
$\\{0.001,0.002,0.004\\}$. Table S2 shows the our final selected bin size and
bin number for each dataset.

Table S2: Bin sizes and bin numbers used to model branch lengths for DS1 to DS8. Dataset | Bin Size | # Bins
---|---|---
DS1 | 0.001 | 50
DS2 | 0.004 | 50
DS3 | 0.004 | 50
DS4 | 0.002 | 100
DS5 | 0.002 | 100
DS6 | 0.001 | 100
DS7 | 0.001 | 200
DS8 | 0.001 | 100
Table S3: Common hyperparameters for PhyloGFN-Bayesian and PhyloGFN-Parsimony.
Transformer encoder
---
hidden size | 128
number of layers | 6
number of heads | 4
learning rate (model) | 5e-5
learning rate (Z) | 5e-3
tree MLP
hidden size | 256
number of layers | 3
edge MLP
hidden size | 256
number of layers | 3
Z MLP (in temperature-conditioned PhyloGFN)
hidden size | 128
number of layers | 1
temp MLP (in temperature-conditioned PhyloGFN)
hidden size | 256
number of layers | 3

##  Appendix F Assessing sampling density against unnormalized posterior
density

Here, we assess PhyloGFN-Bayesian’s ability to model the entire phylogenetic
tree space by comparing sampling density against unnormalized posterior
density over high/medium/low-probability regions of the tree space. To
generate trees from medium and low probability regions, we use a trained VBPI-
GNN model, the state-of-the-art variational inference algorithm.We first
provide a brief introduction to how VBPI-GNN generates a phylogenetic tree in
two phases: (i) it first uses SBN to sample tree topology; (ii) followed by a
GNN-based model to sample edge lengths over the tree topology. SBN constructs
tree topologies in a sequential manner. Starting from the root node and a set
of sequences ${\bm{Y}}$ to be labeled at the leaves, the algorithm iteratively
generates the child nodes by splitting and distributing the sequence set among
the child nodes. The action at each step is thus to choose how to split a set
of sequences into two subsets.To introduce random perturbations during tree
construction, at every step, with probability $\epsilon$, the action is
uniformly sampled from the support choices. Given a well-trained VBPI-GNN
model, we sample from high/medium/low-probability regions with $0\%$, $30\%$
and $50\%$ probability. We apply PhyloGFN-Bayesian and VBPI-GNN on these sets
to compute sampling density and compare with the ground truth unnormalized
posterior density computed as $P({\bm{Y}}|z,b)P(z,b)$. Note that VBPI-GNN
models the _log-edge length_ instead of the edge length. Hence we perform a
change of variables when computing both the sampling probability and the
unnormalized prior.Fig. S1 shows scatter plots of sampling density versus
ground-trueth unnormalized posterior density of datasets DS2 to DS8,
complementing the result on DS1 shown in Fig. 2. We can see that while VBPI-
GNN performs well on high-probability regions, our method is better at
modeling the target distribution overall. Fig. 3(b) shows that our model
performs poorly at modeling the high-probability region for DS7. The result
aligns with our MLL estimation, shown in Table 1. Further work is required to
investigate the cause of PhyloGFN’s poor modeling on DS7.

((a)) DS2 Figure S1: Sampling log density vs. ground truth unnormalized
posterior log density for DS2-DS7

((a)) DS3

((b)) DS4

((c)) DS5

Figure S2: (cont.)

((a)) DS6

((b)) DS7

((c)) DS8

##  Appendix G Parsimony analysis results

((a)) A single most-parsimonious tree with a score of 4026 has been identified
for DS1.

((b)) A single most-parsimonious tree with a score of 6223 has been identified
for DS2.

((c)) Two most-parsimonious trees with a score of 6659 have been identified
for DS3.

Figure S4: For each dataset, the sampling probabilities for the most-
parsimonious as well as several suboptimal trees are estimated from our
learned PhyloGFNs. Each boxplot represents a single unique unrooted tree and
shows a distribution of the log-probabilities over all its $2n-3$ rooted
versions where $n$ denotes the number of species. The dashed bounding box
groups the set of equally-parsimonious trees.

((a)) Four most-parsimonious trees with a score of 2424 have been identified
for DS4.

((b)) Two most-parsimonious trees with a score of 1491 have been identified
for DS5.

((c)) 12 most-parsimonious trees with a score of 879 have been identified for
DS6.

Figure S5: (cont.)

((a)) 15 most-parsimonious trees with a score of 7150 have been identified for
DS7.

((b)) 21 most-parsimonious trees with a score of 1461 have been identified for
DS8.

Figure S6: (cont.)