TreeSBM: Tree-Valued Schrödinger Bridges for Protein Evolution Forecasting

Protein evolution is a branching stochastic process in which ancestral sequences diversify into lineages that accumulate mutations across complex evolutionary trees. Existing protein language models, diffusion models, and flow-based generators produce individual sequences, while phylogenetic methods reconstruct historical trees from observed descendants. However, these approaches do not forecast distributions of future lineages from a single ancestral sequence. To address this gap, we introduce Tree-Valued Schrödinger Bridge Matching (TreeSBM), a generative model for protein evolution forecasting. We construct a reference evolutionary process from protein language model mutation priors, fitness-aware sequence evolution, and sequence-dependent branching, then use bridge matching to reproduce the distribution of observed evolutionary trees. At inference, TreeSBM transports a single ancestral protein sequence into a distribution over future mutational histories, jointly modeling sequence evolution, lineage diversification, branching topology, branch lengths, and evolutionary uncertainty. Across temporally and geographically held-out viral datasets, TreeSBM improves tree-topology agreement, future-lineage coverage, and mutation recovery relative to sequence-level and tree-generative baselines. Generated SARS-CoV-2 trees recover eleven of twelve Gamma-associated mutations and produce additional substitutions with high predicted immune-escape scores. Overall, our results demonstrate that tree-valued generative modeling can forecast both the structure and sequence content of future viral diversification.