Deep Generative Models for Phylogenetic Inference with Complex Evolutionary Processes

AI summary

Overview

Research area: Machine learning for computational biology — specifically simulation-based (likelihood-free) Bayesian inference of phylogenetic trees, combining graph neural networks, discrete diffusion, and generative modeling.

Technical level: Advanced. The paper assumes familiarity with Bayesian phylogenetics, MCMC, discrete diffusion models, graph neural networks, and variational bounds. The conceptual framing is accessible, but the methods section is mathematically dense.

Scope: The paper introduces a family of conditional deep generative models that sample phylogenetic tree topologies from a posterior given sequence data, trained by discrete diffusion on simulated data and evaluated against MCMC, Phyloformer 2, and IQ-Tree.

What This Paper Is About

Biologists infer phylogenetic trees from DNA or protein sequences, and Bayesian methods like MCMC give a posterior over trees that represents uncertainty. But MCMC needs a likelihood, and likelihoods are only computable for simple evolutionary models — realistic models involving insertions, deletions, selection, or context-dependent mutation have intractable likelihoods. This paper's goal is to learn the posterior over tree topologies directly from simulated data using flexible neural generative models, so that inference becomes possible for evolutionary models where conventional Bayesian inference is prohibitively expensive, inaccurate, or impossible.

Key Contributions

1. A flexible family of iterative generative models over trees. The authors define three paradigms — bottom-up, middle-out, and top-down — each of which builds a tree over many forward passes rather than producing a full posterior in a single pass, as prior work did.

2. A discrete diffusion training framework with closed-form Rao-Blackwellized losses. The models are trained by reversing a carefully designed noising process, with the target reverse transition computable in closed form for each of the three paradigms.

3. Architectural innovations for conditioning and pooling. Building on the EvoPF architecture from Phyloformer 2, the authors add graph-topology conditioning (initializing species embeddings with graph information and biasing attention by graph distance) and a new "covariance-pooling" mechanism that embeds average pairwise similarity between nodes rather than averaging nucleotide embeddings.

4. Empirical demonstration that simulation-based inference can beat a well-specified conventional method on realistic data. The models match MCMC posteriors more accurately than Phyloformer 2, match IQ-Tree on simple data, and outperform IQ-Tree on complex and intractable evolutionary models.

Main Findings

• Better ELBOs than the prior state of the art. Across datasets and generative paradigms, the authors' models consistently fit training data better than Phyloformer 2, as measured by negative ELBO. A trivial model that scores all trees uniformly has an ELBO of 157 nats.

• Near-exact fit on the simplest model. On JC data, MCMC-estimated true average entropy of the data is reported at >10 nats, which corresponds to the uncertainty of a uniform distribution over e^10 ≈ 20,000 trees. The authors' models show an average KL of <3 nats (the difference between their ELBO and the entropy), which they describe as an absolutely accurate fit.

• Bottom-up outperforms middle-out, which outperforms top-down. On the JC data, this ordering held consistently. The authors suggest the many noising steps in top-down increase loss variance, and that bottom-up's advantage may reflect the bottom-up nature of likelihood calculation on simple phylogenetic data.

• More accurate recovery of MCMC posteriors than Phyloformer 2 on real data. Across eight classical real-world evolutionary datasets from Treebase (originally selected in Lakner et al.), with N = 27 to N = 64 taxa and L = 378 to L = 2520, the authors inferred 1000 trees each using their JC models and compared to MrBayes posteriors using MMD distance, Spearman correlation of posterior probabilities of tree partitions, and coverage of high-probability partitions. Their models recovered posteriors much more closely than Phyloformer 2, though performance was superior on synthetic data than on real data, indicating room for improvement.

• Competitive with or better than IQ-Tree on complex models. On JC data where IQ-Tree is well specified, the authors' models perform just as well. On CAT (selection in proteins) and InDel (insertions/deletions) data, IQ-Tree ignores complex effects and produces lower-quality predictions, while the authors' models outperform it. This also held for IQ-Tree Flex, which in principle models similar site-variability and non-uniform mutation processes: the mismatch between the exact form of the complex data prior and IQ-Tree Flex's modeled effects allows the authors' models to outperform.

• All three generative paradigms beat Phyloformer 2. The paper states this is the first successful application of the top-down and middle-out paradigms to this problem.

• Rao-Blackwellizing the KL estimate matters. Using closed-form reverse transitions to evaluate the full KL, rather than a single sampled trajectory, achieves substantially better ELBOs.

• The noising schedule choice is minor. A mutual-information-based noise schedule mirroring Austin et al. made little difference in practice, so the authors suggest practitioners can use a uniform schedule.

• Runtime. Run-time comparisons in App. H.5 show bottom-up and middle-out models are faster than IQ-Tree for most N, L settings.

• Posterior under-confidence. The authors report in App. H.4 that their posterior predictions were biased to be under-confident, which they characterize as a hallmark of models trained to maximize likelihood.

Methodology in Plain English

The authors treat tree inference as a simulation-based inference problem. Because simulating sequence data along a tree is easy even when computing the likelihood is not, they generate many synthetic (sequences, tree) pairs from a prior and an evolutionary model, then train a neural network to sample trees given sequences.

Instead of predicting a whole tree in one shot, their model builds a tree step by step across many forward passes. Three construction orders are defined: bottom-up (start with each taxon alone and merge subtrees), middle-out (start with two taxa plus a pseudo-taxon for the root and graft one taxon at a time), and top-down (start with all taxa at the root and percolate them down one layer at a time until a full bifurcating tree forms).

Training uses discrete diffusion. A noising process destroys a true tree over several steps — for bottom-up, removing internal nodes and exposing children; for middle-out, pruning taxa; for top-down, percolating taxa upward to the root. The network is trained to reverse that process, maximizing an ELBO written as a sum of KL divergences between the true reverse transition and the model's prediction. Because the noising process was designed so the true reverse transition has a closed form, the entire KL can be evaluated and integrated over paths rather than estimated from a single sampled trajectory, which the authors call Rao-Blackwellizing and which substantially improves training.

The architecture adapts EvoPF from Phyloformer 2: sequences are embedded into a tensor of size [N, L, 128], axial attention is applied separately over the N and L dimensions, and a mean across the length dimension yields [N, 128] embeddings. Two changes were needed. First, the model must condition on the current partial graph, so species embeddings are initialized with topology information and attention is biased by graph distance. Second, averaging over L > 1000 positions collapses embeddings to base composition, which is nearly identical across taxa; covariance-pooling instead embeds average pairwise similarity, effectively Hamming similarity under this encoding. The pair-attention dimension was reduced from 256 to 128, giving the authors' models 13.6 million parameters versus 20.0 million for Phyloformer 2.

Training used Adam with µP-rescaled per-layer learning rates (base lr = 0.01), CompleteP scaled residual layers, 1,000-step linear warmup, EMA 0.9998, gradient-norm clip 500, bf16-mixed precision, and an effective batch size of 32 trees per step. Each model trained for 48 hours on 2 A100 GPUs — approximately 800,000 datapoints per paradigm, and described elsewhere as roughly 100 GPU hours. Phyloformer 2 baselines were trained for 48 hours on 2 A100 GPUs, approximately 1,100,000 datapoints. Training data spanned N = 25 to N = 65 species and L = 300 to L = 2600, based on the sizes of 8 classical evolutionary datasets.

Why This Matters

Impact on research. The paper shows that a single learned generative model can perform posterior inference under evolutionary models whose likelihoods are intractable, sidestepping the assumptions that Felsenstein's pruning algorithm relies on (notably independence of evolution across sites). It also shows that simulation-based inference can beat a widely used, flexible-but-misspecified maximum likelihood tool on realistic data, not just match it on toy data. The authors frame this as opening the door to a phylogenetic inference foundation model trained across many evolutionary models.

Real-world applications (drawn from the paper's stated motivations and settings):

• Tracking pathogen spread and inferring how fast a pathogen is growing.
• Reconstructing the history of life and understanding ancient evolution.
• Studying the function of the immune system.
• Estimating when a most recent common ancestor occurred and whether there is selective pressure.

Industry relevance. The work targets applications in public health genomics, epidemiology, and immunology, where practitioners currently either wait weeks for MCMC on expensive likelihoods, adopt approximations known to be error-prone (using parameters from previously fit models, alternating fitting of parameters and tree, or relying on simpler models), or ignore phenomena like insertions, deletions, site interdependence, and context-dependent mutation altogether. Because their method only requires the practitioner to provide simulated data, it offers an alternative to bespoke algorithms like IQ-Tree. The authors note the models can be immediately applied to most evolutionary and immunological inference tasks on 1-5 genes.

Future Directions

1. A universal default tree prior / phylogenetic foundation model. Training on simulated data from a variety of evolutionary models could let one model select the appropriate evolutionary process as part of posterior inference, addressing the problem that practitioners often do not know which model fits their data and default to something like GTR+G4+I in MrBayes. The authors suggest this may also address their performance decrease on realistic out-of-distribution data, citing Potapczynski et al., who found that training on a wider variety of synthetic data substantially improved real-data performance.

2. Inferring evolutionary parameters beyond topology. The paper handles only tree topology, but branch lengths and other scalar parameters are also of interest. Since the combinatorially hard part has been addressed, off-the-shelf Bayesian methods for Euclidean data could be applied to the remaining parameters.

3. Species trees and reticulations. Real data involve recombination, incomplete lineage sorting, and time-stamped sampling, requiring reticulated or species trees rather than single-gene trees; extending the models to these settings is left open.

4. Scaling and calibration. The models were trained on up to N = 65 sequences of length up to L = 2600; the authors argue these sizes suffice for most tasks because practitioners often infer sub-trees or individual genes (L ≈ 1000) and combine them with super-tree or species-tree methods, but scaling further via large-scale graph foundation model techniques remains possible. Separately, correcting the under-confidence of posterior predictions may require new algorithms, since borrowing reinforcement learning or sampling strategies from language modeling risks miscalibrating uncertainty.

Target Audience

Machine learning researchers working on generative models, discrete diffusion, and graph neural networks applied to structured combinatorial objects; computational biologists and phylogeneticists interested in likelihood-free inference; and statisticians working on simulation-based inference. The results section is readable by practitioners deciding whether to adopt these methods, while the methodology and appendices require familiarity with diffusion ELBOs and graph architectures.

Authors’ abstract

Inference of phylogenetic trees from sequences plays a central role in understanding evolutionary relationships, with applications ranging from tracking pathogen spread to reconstructing the history of life. Conventionally, practitioners obtain posteriors over trees using likelihood-based methods like MCMC. However, likelihoods are only available for simple evolutionary models with restrictive assumptions. For realistic evolutionary models, inference using conventional methods is prohibitively expensive, inaccurate, or even impossible. In practice, practitioners settle for using as flexible a conventional method as possible, even though they expect it is misspecified. In principle, we could use simulation-based inference: simulate tree and sequence data from an evolutionary model and train a neural network to predict the posterior over trees given sequences. However it is a mathematical challenge to represent and train complex posterior distributions over trees. Here we present a flexible family of conditional generative models over trees. Our models iteratively build trees using three natural paradigms -- top-down, middle-out, or bottom-up -- and are trained using a discrete diffusion framework. For real and synthetic datasets of up to $65$ taxa and $&gt;10^5$ tokens, our models recover MCMC posteriors much more accurately than the previous state-of-the-art simulation-based method Phyloformer 2. Finally, we compare to IQ-Tree, a popular flexible-but-misspecified conventional method. Our models match its tree estimates on simple data, but on more realistic data, our method is reliably more accurate.