PhylaFlow: Hybrid flow matching in phylogenetic tree space

This presentation introduces PhylaFlow, a novel hybrid flow-matching framework designed to accelerate Bayesian phylogenetic inference, a challenging generative modeling problem due to the combination of continuous branch lengths and discrete topological structures in phylogenetic trees. The core problem addressed is the computational bottleneck in traditional MCMC methods for exploring posterior distributions of trees, specifically the lengthy “burn-in” phase where MCMC chains search for high-probability regions in tree space.

The Problem: Slow Bayesian Phylogenetics

Bayesian phylogenetics aims to infer the distribution of evolutionary histories (trees) given a set of sequences and an evolutionary model, rather than a single “true” tree. This is crucial for applications like outbreak response (e.g., COVID-19, flu vaccine development) and cancer research (understanding tumor evolution). State-of-the-art methods rely on MCMC chains, which start from random trees and propose moves (e.g., swapping branches, updating branch lengths). A likelihood function evaluates proposed moves, guiding the chain towards higher probability regions. However, this process is notoriously slow due to the vastness and complexity of tree space (e.g., 10^20 possible trees for just 20 leaves). The majority of time is spent in the “burn-in” phase, searching for the posterior basin.

PhylaFlow’s Approach: Navigating BHV Tree Space

PhylaFlow’s key insight is to leverage the Billera–Holmes–Vogtmann (BHV) tree space, a high-dimensional geometric object that organizes tree space, making navigation more intuitive. In BHV space:

• Each topology corresponds to a “plane” or “orphan” where only branch lengths vary.
• Moving within an orphan changes continuous branch lengths.
• When a branch length collapses to zero, a polytomy forms, representing a boundary where multiple topologies (orphans) intersect. Resolving a polytomy transitions to a new orphan.

The analogy used is “hyperspace” – by navigating BHV space, PhylaFlow aims to start MCMC chains much closer to the posterior basin, significantly reducing burn-in time. Unlike existing initialization methods (e.g., maximum likelihood, parsimony) which provide point estimates, PhylaFlow can generate multiple diverse starts within the posterior basin.

How PhylaFlow Works

PhylaFlow is trained on geodesic paths in BHV space, moving from random trees towards short-run posterior samples. The model learns to:

1. Predict velocity vectors within an orphan, guiding continuous branch length changes.
2. Detect and resolve polytomies at boundaries. Given a degenerate polytomy, the model decides the sequence of merges to resolve it into a valid bifurcating tree in a new orphan.

Loss Functions for Training:

• Velocity Loss (MSE): Optimizes the predicted velocity vector in an orphan.
• First Hit Loss: Explicitly predicts which edges are expected to become zero at the boundary, ensuring accurate boundary hits.
• Anchor Loss: Encourages the model to be invariant to small branch length changes within an orphan, consistently predicting the correct boundary.
• Topology Loss: Guides the model in resolving polytomies by predicting the correct merging sequence.

Model Architecture: A graph transformer forms the base, feeding into three distinct heads for velocity prediction, first hit prediction, and topology resolution. The topology and first hit heads are conditioned on the start tree of the path, providing essential context.

Experimental Validation and Results

PhylaFlow was evaluated on eight standard phylogenetic posterior benchmarks (DS1-DS8), ranging from small to large datasets. Training was performed on short MrBayes runs (approx. 1 hour) to generate target paths, and testing involved initializing MrBayes with PhylaFlow’s output and comparing against long MrBayes runs (up to 30 hours).

Key Metrics: Tree KL divergence, Split KL divergence (both lower is better), and likelihood (less negative is better).

Key Findings:

• Learning Motion in BHV Space: PhylaFlow successfully learns to navigate BHV space. Visualizations show that it often “destroys” the initial random tree structure, rebuilding it towards a more probable topology, a behavior distinct from actual geodesics but effective for the model’s goal.
• PhylaFlow MCMC: A variant where MrBayes’ move acceptance is weighted by the frequency of topological structures observed in PhylaFlow’s output significantly outperforms classical initialization strategies. This is particularly evident on harder, larger datasets (DS5-8), often beating even the “short warm-up” distribution itself, indicating PhylaFlow learns something fundamental beyond just copying training paths.
• Computational Efficiency: While PhylaFlow training can take hours (e.g., 25 hours for DS5), its inference is fast. Crucially, PhylaFlow-initialized MrBayes chains achieve better posterior approximation than running MrBayes for the entire same budget (PhylaFlow training time + short MrBayes run), demonstrating a significant acceleration of Bayesian phylogenetics.
• Towards a Foundation Model: Preliminary work on conditioning PhylaFlow’s movement on sequence embeddings (e.g., from the Phylo paper) shows promise for zero-shot inference on new datasets, paving the way for a “phylogenetics foundation model” that could provide initial trees for any sequence alignment without retraining.

Discussion Points

• Polytomies: PhylaFlow sometimes terminates with polytomies, which are typically considered degenerate in standard evolutionary models. This opens a question about their potential biological significance. Actual geodesics do not end in polytomies.
• Loss Function Complexity: Designing the loss functions was challenging due to the need to balance continuous (velocity) and discrete (boundary hits, topology resolution) predictions, requiring careful weighting to prevent the model from over-optimizing simpler losses.
• Comparison to GFNs: Unlike methods like PhyloGFN which use likelihood calculations at every step, PhylaFlow is likelihood-free during inference, contributing to its speed once trained.

This work represents a significant step towards scalable and differentiable phylogenetics, offering a new paradigm for accelerating the crucial and widespread application of Bayesian inference in biology.

• Tags:
• Computational Biology
• Bayesian Inference
• Flow Matching
• Uncertainty Quantification
• Graph Transformers
• Generative Modeling
• Phylogenetics
• MCMC
• BHV Tree Space