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

Abstract: Phylogenetic trees are hybrid objects: branch lengths vary continuously, while topologies change discretely through edge contractions and expansions. Billera-Holmes-Vogtmann (BHV) tree space provides a canonical geometry for this structure, representing each resolved topology as a Euclidean orthant and topological changes as motion across shared lower-dimensional boundaries. We introduce PhylaFlow, a hybrid flow-matching model that learns posterior-basin transport in BHV tree space. PhylaFlow is trained on BHV geodesic paths from random starting trees to short-run posterior samples, coupling continuous branch-length motion within orthants with learned boundary events and discrete topology transitions. We evaluate the learned geometry operationally: if the flow reaches posterior-relevant regions, finite-budget Bayesian refinement initialized from, or guided by, its terminal trees should recover posterior-supported topologies more efficiently. Across DS1-DS8 phylogenetic posterior benchmarks, PhylaFlow substantially reduces initial Tree-KL relative to classical initializers. After finite-budget MrBayes refinement, direct PhylaFlow improves early and intermediate topology-recovery trajectories on most datasets, while split-guided PhylaFlow-MCMC obtains the strongest hard-case results. The best PhylaFlow variant outperforms short-warmup on seven of eight datasets and PhyloGFN on five of eight under the same refinement budget. In a joint sequence-conditioned experiment, sequence embeddings steer posterior split recovery, although exact posterior topology recovery remains preliminary. These results show that hybrid flow matching can learn actionable transport in BHV tree space and provide a geometry-aware proposal mechanism for Bayesian phylogenetic inference.

• The paper introduces a neural hybrid flow matching model that transports random starting trees to high-posterior regions in BHV tree space.
• It leverages a transformer-based graph encoder with specialized heads to predict continuous branch velocities and discrete topology transitions.
• Empirical evaluations show that PhylaFlow dramatically reduces initial divergence and accelerates convergence compared to traditional MCMC methods.

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

Introduction

Phylogenetic inference is a foundational problem in evolutionary biology, molecular epidemiology, and computational genomics. The statistical and computational challenge lies in modeling distributions over phylogenetic trees, which are hybrid combinatorial objects: they combine a discrete topology space (the tree structure) with a continuous parameter space (branch lengths). The Billera-Holmes-Vogtmann (BHV) metric space provides a principled geometric framework for reasoning about such objects, representing each tree topology as a Euclidean orthant and capturing topological transitions as movements across shared lower-dimensional boundaries. Conventional Bayesian inference in this setting is based on stochastic Markov chain Monte Carlo (MCMC) methods, which face serious limitations in burn-in and poor exploration of the multi-modal posterior in tree space, especially with increasing taxa and model complexity.

"PhylaFlow: Hybrid Flow Matching in Billera-Holmes-Vogtmann Tree Space for Phylogenetic Inference" (2605.21859) introduces a highly novel approach for initializing and guiding posterior exploration in BHV tree space via a learned neural hybrid flow-matching model. Rather than directly approximating the intractable high-dimensional posterior, PhylaFlow is trained to transport diverse random starting trees efficiently towards regions of high posterior density, thereby improving the operational efficiency of downstream Bayesian refinement.

Theoretical Framework and Problem Formulation

Classical MCMC approaches treat phylogenetic inference as a hard combinatorial optimization and sampling problem, navigating the exponentially large space of tree topologies in conjunction with continuous branch-length optimization. The BHV space formalism allows geometric relaxation: each resolved topology is an orthant (a non-negative Euclidean space of dimensionality $N-3$ for $N$ taxa), boundaries between orthants correspond to topological transitions (edge contractions and expansions). Yet, no prior method realized neural generative modeling that can learn actionable global transport in this space.

PhylaFlow reframes phylogenetic inference as an optimal transport problem in BHV space: given a random starting tree, the model learns to transport this tree’s representation, via a geometry-respecting hybrid vector field, towards the empirical distribution of posterior-supported trees. This transport is decomposed into continuous branch-length flows within orthants and discrete boundary crossing decisions (topological transitions).

Figure 1: BHV tree space structure, showing continuous intra-orthant motion and legal geometry-compliant orthant transitions corresponding to edge contraction (boundary crossing) and topology change.

Model Architecture

PhylaFlow implements a graph-based neural transport model leveraging a shared transformer-based graph encoder. Inputs to the model include:

• The current tree state (topology, branch lengths, and node/edge features),
• Sequence-derived conditionings (optional, for downstream amortized inference),
• Time/phase coordinates reflecting continuous intra-orthant evolution and discrete phase (number of topology changes).

On top of a shared encoder, the architecture comprises three specialized heads:

• Velocity Head: Predicts per-branch continuous velocities for branch length modification inside a fixed topology.
• First-Hit Head: Predicts which internal branches are expected to reach zero next, thus identifying the "first-hit" boundary and the ensuing topology change event.
• Autoregressive Topology Head: Given a polytomy at a collapsed boundary, predicts a merge sequence (subset merging) that resolves the polytomy into a new valid topology.

At training time, supervision is obtained via geodesics in BHV space: the model is trained on paths from random start trees towards empirical posterior samples, with explicit branch-length velocity and structural transition targets.

Figure 2: Detailed architecture showing tokenization of tree state and heads for velocity, boundary first-hit, and autoregressive topology transitions via a graph transformer backbone.

Training Regime and Objectives

Training is formulated as a multi-objective supervised learning problem:

• A mean-squared error is used to match predicted branch velocities to the geodesic-derived targets within each orthant.
• A binary cross-entropy loss is used for the first-hit collapse set, predicting which edges should contract to zero at the next boundary event.
• A structured autoregressive loss supervises the discrete sequence of merge actions to resolve polytomies at boundaries.
• Auxiliary losses aid numerical and geometric stability.

The model is applied in a deterministic event-based rollout procedure: beginning with an unseen start tree, at each phase the model predicts motion and transitions until a stopping condition is reached. Terminal trees are then optionally branch-length-relaxed.

Figure 3: Overview of the PhylaFlow initialization and comparison with traditional MCMC/MrBayes methods, showcasing rapid entry into diverse posterior basins using BHV transport.

Empirical Evaluation

The performance of PhylaFlow is rigorously benchmarked against classical initializers (random, maximum-parsimony, IQ-TREE maximum likelihood, UShER-parsimony), and state-of-the-art neural methods (PhyloGFN), across eight diverse phylogenetic posterior datasets (DS1–DS8). Key evaluation metrics include:

• Tree-KL: Kullback-Leibler divergence between the empirical post-burn-in topology distribution and a long-run reference posterior,
• Split-KL: Divergence in split (bipartition) distributions,
• Optimized Likelihood Gap: Distance in maximized log-likelihood between sampled trees and the posterior core.

Core findings:

• PhylaFlow initialization yields dramatically lower initial Tree-KL across all benchmarks, i.e., start trees are much closer, in topology space, to posterior-supported regions before any likelihood-based refinement.
• After fixed-budget downstream MrBayes MCMC, direct PhylaFlow starts improve or match the final Tree-KL on most datasets.
• In hard multi-modal posteriors, Split-Guided PhylaFlow-MCMC refinement achieves the strongest results, with statistically significant reductions in Tree-KL even relative to long MCMC trajectories from strong baselines.
• Likelihood analysis demonstrates that PhylaFlow does not propose topologies in low-likelihood basins; initializations remain in or near the plausible likelihood region.

  Figure 4: Representative learned tree trajectories in BHV space produced by PhylaFlow, demonstrating front-loaded large topology changes followed by local refinements; RF denotes Robinson-Foulds distance to reference.

Sequence-Conditioned and Amortized Proposals

PhylaFlow is also capable of sequence-conditional transport by leveraging per-taxon sequence embeddings (derived from models such as Phyla) as part of the tokenization. Cross-conditioning experiments demonstrate that these sequence-derived signals successfully steer the flow towards the appropriate dataset’s posterior basin, evidenced by strong diagonal dominance in split-KL across datasets.

Learned Transport Dynamics

Inspection of test-time PhylaFlow BHV trajectories reveals important structural insights:

• Motion is typically front-loaded, with the initial phase consisting of a coordinated multi-branch collapse leading to rapid entry into a basin spanned by high-posterior topologies.
• Subsequent phases involve finer local topological refinements and minor branch-length adjustments.
• These rollout patterns differ from naive BHV geodesics and indicate that the learned vector field is amortized to prioritize posterior basin occupancy under fixed compute budgets.

Practical and Theoretical Implications

PhylaFlow demonstrates that geometry-aware, differentiable initialization can substantially accelerate posterior exploration in Bayesian phylogenetics. Rather than replacing MCMC/Bayesian inference, PhylaFlow acts as an intelligent, amortized proposal mechanism, addressing burn-in and local trapping problems endemic to standard samplers.

Theoretically, this is the first demonstration of hybrid flow matching implemented over the combinatorially structured BHV metric, combining Riemannian-inspired continuous transport and discrete action prediction in a non-Euclidean space. Practically, the method is computationally competitive in setup cost relative to prior approaches (notably, PhyloGFN) and obtains superlinear gains in challenging multi-modal regimes. Sequence-conditioned PhylaFlow extends the paradigm towards amortized inference across diverse datasets.

Future Prospects

PhylaFlow opens up several promising research avenues:

• Joint end-to-end training of the sequence and tree modules, facilitating fully end-to-end amortized phylogenetic inference.
• More efficient, non-autoregressive or batched rollout strategies to improve proposal throughput.
• Integration with or guidance of alternative samplers and refiners (e.g., variational Bayes, score-based diffusion samplers) in tree space.
• Generalization to larger trees, incomplete taxa, and integration with unsupervised embeddings for full-scale evolutionary modeling.

Conclusion

PhylaFlow establishes that hybrid neural flow-matching in the BHV metric enables actionable, geometry-consistent transport from arbitrary starts to posterior support in phylogenetic inference. As both a technical and conceptual contribution, this work advances the intersection of deep generative modeling, geometric statistics, and Bayesian phylogenetics, and provides a novel foundation for amortized, sequence-informed inference in combinatorial biological spaces.