PhylaFlow: Hybrid Flow Matching in Billera-Holmes-Vogtmann Tree Space for Phylogenetic Inference

Technical Summary: PhylaFlow – Hybrid Flow Matching in Billera–Holmes–Vogtmann Tree Space

Problem Statement
Phylogenetic inference is fundamentally a hybrid optimization problem: it involves continuous optimization over branch lengths and discrete search over tree topologies. Traditional Bayesian methods, such as Markov Chain Monte Carlo (MCMC) implemented in MrBayes, often struggle with "burn-in" periods where chains spend excessive computational resources exploring low-probability regions of tree space before discovering posterior-relevant basins. This is particularly challenging because the space of phylogenetic trees is not purely Euclidean; topological changes occur through discrete edge contractions and expansions, while branch lengths vary continuously. The Billera–Holmes–Vogtmann (BHV) tree space provides a canonical geometry for this structure, representing each topology as a Euclidean orthant and topological transitions as motion across shared lower-dimensional boundaries. However, existing neural approaches often either attempt to replace MCMC entirely with variational approximations or generate single high-scoring trees without learning the continuous-discrete transport dynamics required to navigate the BHV geometry effectively.

Methodology
The authors introduce PhylaFlow, a hybrid flow-matching model designed to learn "posterior-basin transport" within BHV tree space. Rather than replacing Bayesian inference, PhylaFlow acts as an amortized initializer and proposal mechanism that transports arbitrary starting trees toward regions supported by the posterior distribution.

• Geometric Formulation: The model operates within the BHV metric space. It treats the inference task as an optimal transport problem, learning a map that pushes forward samples from a base distribution (random starting trees) to an approximation of the target posterior distribution.
• Hybrid Architecture: PhylaFlow employs a graph transformer to process tree states (topology and branch lengths) conditioned on sequence embeddings and time. The architecture utilizes three distinct heads to handle the hybrid nature of the space:
  1. Velocity Head: Predicts continuous branch-length velocities within a fixed topology (orthant), modeling motion via a time-dependent vector field.
  2. First-Hit Head: Predicts which internal edges will contract to zero first, identifying the next boundary event where a topology change is imminent.
  3. Topology Head: An autoregressive module that resolves the resulting polytomy (unresolved tree) after a boundary collapse into a new, valid topology by predicting a sequence of structured-subset actions.
• Training Objective: The model is trained on supervised BHV geodesic paths connecting random starting trees to target trees sampled from a short-run posterior. The loss function combines continuous flow matching (for branch lengths), binary classification (for edge collapse events), and autoregressive cross-entropy (for topology transitions).
• Inference and Refinement: At inference, PhylaFlow generates candidate trees from random starts. These trees are then used to initialize standard MrBayes chains for likelihood-based refinement. The authors also propose a PhylaFlow-MCMC variant, where the learned split structure from PhylaFlow is used to bias local NNI (Nearest Neighbor Interchange) proposals during the MCMC run, guiding the chain toward posterior-supported splits.

Key Contributions

1. Hybrid Flow Matching in BHV Space: The paper demonstrates that neural transport models can learn meaningful motion through the hybrid geometry of tree space, coupling continuous branch-length evolution with discrete topological transitions.
2. Amortized Initialization: PhylaFlow provides a geometry-aware proposal mechanism that significantly reduces the distance (measured by Tree-KL) between random starting trees and posterior-relevant basins before any likelihood-based refinement occurs.
3. Split-Guided Refinement: The work introduces a method to use the learned posterior geometry to guide MCMC proposals, improving performance in "hard" cases where standard local moves struggle to escape suboptimal basins.
4. Sequence Conditioning: The authors show that the transport process can be conditioned on sequence embeddings (Phyla embeddings), allowing the model to steer transport toward dataset-specific split structures, although exact topology recovery in a joint setting remains preliminary.

Results
Evaluated across eight benchmark datasets (DS1–DS8) spanning diverse evolutionary regimes, PhylaFlow demonstrates substantial improvements over standard baselines:

• Initial Quality: Before MCMC refinement, PhylaFlow achieves the lowest initial Tree-KL (divergence from the long-run posterior) on all eight datasets compared to random, maximum-parsimony, maximum-likelihood (IQ-TREE), and UShER-based initializers.
• Refinement Efficiency: After finite-budget (100k generations) MrBayes refinement, direct PhylaFlow initialization outperforms other non-guided methods on five of eight datasets.
• Hard-Case Performance: The PhylaFlow-MCMC variant, which uses learned splits to guide proposals, achieves the best final Tree-KL on the most difficult datasets (DS3, DS5, DS6, DS7), significantly outperforming both standard baselines and strong neural controls like PhyloGFN under the same computational budget.
• Comparison to Controls: PhylaFlow outperforms short-warmup baselines on seven of eight datasets and PhyloGFN on five of eight datasets. While PhyloGFN is a strong likelihood-trained sampler, PhylaFlow offers a more computationally efficient path to useful posterior geometry, with setup times generally lower than PhyloGFN on larger datasets.
• Likelihood Plausibility: The generated trees maintain plausible likelihood profiles, indicating that the model learns reusable posterior structure rather than merely producing isolated high-scoring trees.

Significance and Claims
The paper claims that PhylaFlow provides an early example of differentiable, geometry-aware modeling for phylogenetic inference. Its primary significance lies in demonstrating that actionable transport in BHV tree space is learnable. The authors position PhylaFlow not as a replacement for Bayesian inference, but as a complementary mechanism that improves the efficiency of posterior-basin discovery. By learning the "motion" through tree space, the model reduces the burn-in time required for MCMC chains to reach high-posterior regions. The results suggest that hybrid flow matching can capture actionable posterior geometry, offering a new paradigm for initializing and guiding phylogenetic inference that respects the intrinsic combinatorial-continuous structure of tree space. The authors remain modest regarding the sequence-conditioned experiments, noting that while embeddings can steer split recovery, exact posterior topology recovery in a joint, cross-dataset setting remains preliminary.