Tree-Conditioned Edit Flows for Ancestral Sequence Reconstruction

This paper introduces a tree-conditioned edit-flow model for ancestral sequence reconstruction that handles variable-length sequences by reconstructing ancestors through paired bidirectional edit trajectories, demonstrating reasonable performance on substitution-only benchmarks and superior localization of evolutionary changes in sequences with abundant insertions and deletions.

The Gist

The Big Picture: Rebuilding a Shredded Family Album

Imagine you have a family photo album, but the photos of your great-grandparents are missing. You only have photos of your cousins (the "descendants"). Your goal is to guess what the great-grandparents looked like based on the photos of their children and grandchildren.

In biology, scientists do this with proteins. They try to guess the sequence of amino acids (the "letters" that make up a protein) for ancient, extinct organisms. This is called Ancestral Sequence Reconstruction (ASR).

The Problem: The Old Way Was Too Rigid

For decades, scientists used "classical" methods to solve this puzzle. Think of these methods like a stiff, grid-based spreadsheet.

• They look at one letter at a time (e.g., "Was this spot an 'A' or a 'G'?").
• They assume every letter changes independently of its neighbors.
• They are terrible at handling insertions and deletions (adding or removing letters).

The Analogy: Imagine trying to fix a torn sentence by only guessing the missing letters, but you aren't allowed to add or remove any words. If the ancient sentence was "The cat sat" and the modern one is "The big cat sat," the old methods struggle because they can't easily account for the new word "big" appearing in the middle. They treat the sentence as a fixed grid where letters just swap places, not a flexible string where words can appear or vanish.

The New Solution: Lærad (The "Flowing" Restorer)

The authors introduce a new AI model called Lærad. Instead of a stiff spreadsheet, think of Lærad as a dynamic, flowing river that can reshape itself.

1. The "Edit Flow" Concept
Lærad treats evolution like a video editing process. It doesn't just guess letters; it guesses actions:

• Substitution: Swapping a letter (like changing "cat" to "bat").
• Insertion: Adding a new letter (like adding "big" to "cat").
• Deletion: Removing a letter (like removing "big" from "big cat").

It learns to "flow" from a modern protein back to an ancient one by simulating these edits step-by-step.

2. The "Tree-Conditioned" Trick
The model knows it's working on a family tree. It uses the "branch lengths" (how much time passed between ancestors) as a budget.

• The Analogy: Imagine you are traveling from City A to City B. The map tells you the distance is 100 miles. You have a "fuel budget" of 100 miles. You can't drive 200 miles, and you can't drive 0 miles. Lærad uses this "distance budget" to know exactly how many edits (swaps, adds, or deletes) are allowed to happen between the ancestor and the descendant.

3. The "Paired" Strategy
This is the model's superpower. Instead of looking at one descendant and guessing the ancestor, Lærad looks at two descendants (like two cousins) at the same time.

• The Analogy: Imagine two cousins, Alice and Bob, are trying to reconstruct what their shared grandmother looked like.
  • Alice tries to "rewind" her DNA back to the grandmother.
  • Bob tries to "rewind" his DNA back to the grandmother.
  • Lærad forces Alice's rewind and Bob's rewind to meet in the middle at the exact same point in time (the grandmother). If Alice's guess and Bob's guess don't match up at that meeting point, the model knows it made a mistake and tries again.

How It Performed: The Results

The authors tested Lærad on two different types of puzzles:

Puzzle 1: The "Messy" Family (Proteins with lots of insertions/deletions)

• The Test: They used a dataset of bacteriophage proteins (viruses that infect bacteria) which are known to be very "messy," with lots of letters being added and removed over time.
• The Result: Lærad was the best at figuring out where changes happened. It was like a detective who could point to the exact spot in the sentence where a word was added or removed, better than any previous method. It didn't necessarily get every single letter perfect, but it understood the structure of the changes best.

Puzzle 2: The "Clean" Family (Proteins with mostly simple swaps)

• The Test: They used fluorescent proteins (glowing proteins) where the changes were mostly just simple letter swaps, with very few additions or deletions.
• The Result: Lærad was slower and less accurate here. The "old" classical methods (the stiff spreadsheets) were still better at this specific task.
• Why? Lærad is a heavy-duty tool designed for complex, messy changes. Using it for simple swaps is like using a sledgehammer to crack a nut. The classical tools are optimized for simple swaps and still win in that specific, clean environment.

The Bottom Line

Lærad is a new way to guess ancient protein sequences that treats evolution as a flexible process of adding, removing, and swapping parts, rather than just swapping letters in a fixed grid.

• When it shines: It is the best tool we have for proteins that have grown, shrunk, and changed shape significantly over time (handling "indels" well).
• When it struggles: It is not yet the best tool for proteins that have stayed very stable and only changed a few letters.

The paper concludes that while Lærad isn't perfect yet, it opens a new door for understanding how proteins evolve when they are constantly gaining and losing pieces, a task that previous methods found very difficult.

Technical Summary

Technical Summary: Tree-Conditioned Edit Flows for Ancestral Sequence Reconstruction

Problem Statement

Ancestral Sequence Reconstruction (ASR) aims to infer the protein sequences of extinct ancestors at internal nodes of a phylogenetic tree. Classical ASR methods, typically based on continuous-time Markov substitution models (e.g., PAML, PhyML), treat sites as conditionally independent and handle insertions and deletions (indels) either by excluding them or ignoring them during likelihood calculations. While these methods excel at global inference over a tree, they struggle with the complex, context-dependent nature of sequence evolution, particularly when indels are abundant. Recent neural approaches (e.g., AutoregressiveASR, BetaReconstruct) offer greater expressivity but often fail to incorporate phylogenetic tree topology, branch lengths, or the constraint that an ancestor must simultaneously explain multiple descendants.

Methodology: Lærad

The authors introduce Lærad, a tree-conditioned paired edit-flow model designed for variable-length ASR. Unlike methods that output a single sequence directly, Lærad models ASR as a branch-conditioned edit process, predicting time-dependent rates for substitutions, insertions, and deletions.

Core Architecture

• Edit-Flow Foundation: Lærad builds on discrete flow matching, lifting the concept from fixed-length token spaces to variable-length sequences. It defines transitions through elementary edit operations: insertion, deletion, and substitution.
• Paired Cross-Attention: The model processes two descendant sequences ($x_a, x_b$) simultaneously. It employs a shared ESM-2 backbone for encoding, followed by paired fusion layers that allow cross-attention between the two descendants. This ensures both children inform the edit field for the ancestor.
• Branch Conditioning: The model is conditioned on the ordered branch distances ($d_a, d_b$) from each descendant to their shared Lowest Common Ancestor (LCA). These distances are converted into "edit budgets" using Fitch parsimony estimates, defining the expected location of the ancestor along the evolutionary bridge ($\tau = d_a / (d_a + d_b)$).

Training Objective

Lærad is trained on stochastic bridge states sampled between two descendants without requiring ground-truth ancestral sequences. The loss function ($L$) combines three terms:

1. Bregman Loss ($L_{Bregman}$): A bidirectional loss that trains the model to predict edit rates that move a sampled bridge state toward the target descendant. This teaches local edit mechanics (where edits happen and what residues are plausible).
2. Ancestor Alignment Loss ($L_{ancestor}$): Near the expected ancestral point ($\tau$), the latent representations of the two opposing edit trajectories (from $a \to b$ and $b \to a$) are aligned using cosine and L2 distances. This enforces that both routes imply a compatible ancestral state.
3. Group Consistency Loss ($L_{group}$): For multiple descendant pairs sharing the exact same LCA node, their mean-pooled latent representations are pulled together. This injects explicit local tree consistency, ensuring different views of the same ancestor converge to a consistent representation.

Inference Procedure

Inference proceeds bottom-up on the phylogenetic tree:

1. Decoding: For a pair of children, the model decodes $N$ candidate parent sequences from each child, conditioned on the other child and the branch budgets.
2. Selection & Consensus: A scoring function $S(s)$ evaluates candidates based on branch-budget agreement, parsimony (edit cost), disagreement between the two directional decodes, and model support.
3. Reconciliation: The best-scoring pair of candidates is merged via a consensus strategy (copying matching residues, resolving disagreements via budget compatibility). The final ancestor is selected from the two directional candidates and their consensus merge.

Key Contributions

1. Variable-Length ASR Framework: Lærad extends ancestral inference to variable-length sequence evolution by explicitly modeling substitutions, insertions, and deletions under phylogenetic constraints, moving beyond fixed-alignment assumptions.
2. Tree-Conditioned Edit Flows: The model uniquely integrates phylogenetic topology and branch lengths directly into the edit-flow generation process, using paired cross-attention to ensure descendants jointly inform the ancestral state.
3. Consistency Mechanisms: The introduction of bidirectional bridge losses and exact-LCA group consistency losses ensures that inferred ancestral states are compatible with multiple descendants and consistent across different pairs mapping to the same node.

Results

The authors evaluated Lærad on two distinct benchmarks:

1. Indel-Rich Benchmark (Bacteriophage J Proteins)

On a benchmark of natural homologous sequences with abundant indels (ID95 dataset), Lærad was compared against classical methods (Fitch, PAML, ARPIP) and neural baselines (AutoRegressiveASR).

• Performance: Lærad achieved the highest observed edit correlation (Pearson correlation between inferred branch-edit density and empirical leaf-level variation), with the Tiny variant reaching 0.778. This surpassed the best classical baseline (PHYLO-Γ at 0.765).
• Localization: The results suggest Lærad is superior at localizing inferred evolutionary changes across empirically variable sites in indel-rich contexts.
• Limitations: While strong in localization, Lærad's operation-specific indel correlation was lower than ARPIP, and its normalized budget error (mismatch between inferred edits and tree-implied budgets) remained higher than some baselines.

2. Substitution-Only Benchmark (Fluorescent Proteins)

On a benchmark of experimentally evolved fluorescent proteins with known internal ancestors (effectively substitution-only), Lærad was compared against substitution-specialized methods.

• Performance: As expected, Lærad trailed behind classical likelihood-based methods (PHYLO-Γ: 97.2% accuracy; ARPIP: 97.1%) and the neural baseline AutoRegressiveASR (87.3%). Lærad-Nano achieved 84.4% accuracy.
• Interpretation: The authors note this is a conservative stress test, as the model is designed for complex edit operations while the task is dominated by substitutions.

Significance and Claims

The paper claims that tree-conditioned edit flows represent a viable direction for variable-length ASR, particularly in settings where evolution is driven by insertions and deletions.

• Primary Strength: Lærad demonstrates that modeling sequence evolution as a paired, tree-conditioned edit process can outperform classical methods in localizing evolutionary changes in indel-rich environments.
• Modest Scope: The authors are explicit that the current formulation is not yet superior to classical methods in substitution-dominated settings. They acknowledge that operation-type calibration (accurately predicting the specific type of edit) and branch-budget calibration (matching the exact number of edits to tree distances) remain open problems.
• Future Potential: The work suggests that scaling the model (e.g., using larger ESM-2 backbones) may improve performance in substitution-dominated settings, but the primary contribution remains the successful integration of phylogenetic constraints into a generative edit-flow framework for variable-length sequences.