🧬 biology

Approximating SPR Distance Between Phylogenetic Trees with Graph Neural Networks

This paper introduces a Graph Neural Network approach that approximates the NP-hard Subtree Prune and Regraft (SPR) distance between phylogenetic trees in near-constant time, supported by a newly released dataset of bacterial trees and a validated heuristic surrogate for training.

Original authors: Renata Martins Castanheira, Miguel Bugalho, Cátia Vaz

Published 2026-07-22

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Renata Martins Castanheira, Miguel Bugalho, Cátia Vaz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to solve a massive, ancient family mystery. You have a pile of old, handwritten letters (the DNA) from thousands of people, and you want to draw a family tree to see who is related to whom. Scientists do this all the time to track how viruses or bacteria spread during an outbreak. But here's the tricky part: there isn't just one way to draw the tree. Depending on which math trick you use to connect the dots, you might get a slightly different picture of the family history. Sometimes, one tree says "Cousin A is close to Cousin B," while another tree says "Actually, they are distant strangers."

To figure out which tree is better, or how different two trees are from each other, scientists need a way to measure the "distance" between them. Think of it like comparing two different maps of the same city. One map might show a road that the other map misses. In the world of biology, the most accurate way to measure this difference is called the "Subtree Prune and Regraft" (SPR) distance. It's like asking, "How many times do I have to cut a branch off one tree and glue it onto a different spot to make it look exactly like the other tree?" This method is incredibly precise because it counts the actual biological moves, like genetic swaps, that happened over time. However, there's a catch: doing this math is so hard and slow that for a tree with thousands of leaves (representing thousands of bacteria), it would take a supercomputer years to finish the calculation. It's like trying to count every single grain of sand on a beach to measure the beach's size; it's possible in theory, but practically impossible.

This is where the story gets exciting. The researchers in this paper asked a simple question: "Can we teach a computer to guess this distance quickly, without doing the impossible math?" They built a special kind of artificial intelligence called a Graph Neural Network (GNN). You can think of this AI as a super-smart student who has studied thousands of tree pairs and learned to spot the patterns of how different they are. Instead of doing the slow, heavy lifting of cutting and pasting branches, the AI looks at the shape of the trees and makes a very fast guess. The paper shows that this AI can learn to predict the distance between trees with surprising accuracy, but it also reveals that the AI has a specific weakness: it struggles when it sees trees that are much bigger than the ones it studied in school.

The Big Experiment: Teaching AI to Compare Trees

The team behind this study, led by Renata Martins Castanheira and her colleagues, set out to build a system that could compare these biological family trees in near-instant time. Their goal was to create a tool that could handle the massive datasets used in modern disease tracking, where scientists might be looking at up to 9,500 different bacterial samples at once.

First, they had to build a training ground. They gathered genetic data from four different types of bacteria: Clostridium, Salmonella, Vibrio, and Streptococcus pneumoniae. They sliced this data into groups of different sizes, ranging from small groups of 90 isolates to massive groups of up to 9,500. Using standard computer methods, they generated 864 different family trees from this data. To make the training interesting, they didn't just use the trees as they were; they shuffled the data slightly to create "twins" with slightly different shapes. This gave them 388 pairs of trees to compare.

Here is the clever part: they needed a "teacher" to tell the AI what the correct answer was. Since the perfect, exact math (the "rooted SPR" distance) is too slow to calculate for big trees, they used a faster, slightly less perfect method called a "heuristic" (a smart shortcut) to generate the answers for the AI to learn from. Before trusting this shortcut, they tested it on small trees where they could do the perfect math. They found that the shortcut was almost perfectly aligned with the real answer, with a correlation of 0.98 to 0.99. It wasn't perfect—it tended to guess numbers that were a bit too low (about 65% to 71% of the true value)—but the ranking was spot on. If Tree A was more different from Tree B than Tree C was from Tree D, the shortcut knew it. This made it a perfect "surrogate" teacher for the AI.

Next, they built the AI. They used a specific type of neural network called a Siamese Graph Isomorphism Network (GIN). Imagine two identical twins (the "Siamese" part) who are both looking at a tree. The AI breaks each tree down into a graph, where every node (branch point or leaf) has a little ID card with information like "I am a leaf," "I am far from the root," or "I belong to this species." The AI then processes these trees, looking for patterns in how the branches are arranged. Finally, it takes the "thoughts" from both trees, combines them, and spits out a single number: the predicted distance between them.

What They Found

The results were a mix of great success and clear limitations, depending on the situation.

When the AI stayed in its comfort zone:
When the researchers tested the AI on trees that were similar in size and type to the ones it had studied (the "in-distribution" scenario), it performed brilliantly. It could explain about 87% to 90% of the differences between the trees. In simple terms, if you asked the AI to guess the distance between two trees it had seen before, it was right most of the time. It was about four times better than just guessing the average distance for every pair. This suggests the AI genuinely learned the rules of how tree shapes relate to each other.

When the AI met new friends:
The team then asked, "Can this AI understand a new type of bacteria it has never seen?" They trained the AI on two species and tested it on two completely different ones. The AI didn't fail completely, but it got much worse. Its accuracy dropped, explaining only about 37% of the variance. This tells us that while the AI learned some general rules about tree shapes, it also got used to the specific "flavor" of the bacteria it studied. It's like a student who is great at math but gets confused when the teacher switches from algebra to geometry.

When the AI faced giants:
The most important finding, however, was what happened when the AI faced trees much larger than the ones in its training set. The researchers trained the AI on small and medium trees and then asked it to guess the distance for the massive trees with up to 9,500 isolates. The result? The AI collapsed. Its accuracy didn't just drop; it became worse than just guessing the average. The model simply could not "extrapolate" or figure out how to handle sizes it had never seen before. It's like teaching a child to count to 100 and then asking them to count to a million; they don't know how to scale up the logic.

The Takeaway

This paper doesn't claim to have solved the problem of comparing phylogenetic trees forever. Instead, it offers a powerful new tool that works incredibly well for trees of a size and type it has already seen. The authors show that we can train an AI to approximate these complex biological distances in a fraction of a second, which is a huge step forward for tracking epidemics.

However, they are very clear about the boundaries. The AI is not a magic wand that works on any tree of any size. It struggles with new species and fails completely when the trees get too big. The authors suggest that to fix this, future work needs to feed the AI even bigger trees during training and perhaps teach it to handle size differences better. They also point out that right now, the AI only guesses the number of moves needed to change one tree into another. A future goal would be to make the AI point out exactly which branches need to be moved, giving scientists a roadmap of the changes rather than just a distance score.

In short, the paper proves that Graph Neural Networks can be excellent "fast guessers" for tree distances, provided the trees aren't too weird or too huge. It's a promising start, but the AI still has a lot of growing up to do before it can handle the biggest challenges in the field.

Technical Summary: Approximating SPR Distance Between Phylogenetic Trees with Graph Neural Networks

Problem Statement

Comparing phylogenetic tree topologies is critical for understanding epidemic dynamics, particularly for bacterial and viral outbreaks. While simple distance metrics like Robinson–Foulds are computationally cheap, they often saturate on divergent isolates or overreact to single taxon displacements, failing to capture biologically meaningful evolutionary events. The Subtree Prune and Regraft (SPR) distance offers a superior biological alternative by counting the specific edit operations (prune-and-regraft) required to transform one tree into another, directly modeling reticulate events like recombination and horizontal gene transfer.

However, computing the exact rooted SPR distance is NP-hard, rendering it intractable for large datasets involving thousands of isolates. Existing heuristics (e.g., phangorn::SPR.dist) scale better but lack formal approximation guarantees and must be recomputed for every tree pair, preventing near-constant-time comparisons at inference. The authors investigate whether a Graph Neural Network (GNN) can learn to approximate SPR distances efficiently, offering a fast, learned estimator that generalizes across tree sizes and species.

Methodology

1. Dataset Construction and Pre-processing
The authors constructed a reproducible pipeline to generate a benchmark dataset:

Data Source: Core-genome MLST (cgMLST) profiles for four bacterial species (Clostridium, Salmonella, Vibrio, and Streptococcus pneumoniae) from EnteroBase and PubMLST.
Stratification: Profiles were sliced into three size tiers: small (90–1,000 isolates), medium (2,000–6,000), and large (7,000–9,500).
Topology Generation: For each sub-dataset, trees were inferred using both UPGMA and Neighbor-Joining (NJ) algorithms. To introduce controlled topological variations, input profiles were shuffled to force tie-breaking differences in distance-based methods.
Rooting: Since exact SPR computation and the model's features require rooted trees, all NJ trees (originally unrooted) underwent midpoint re-rooting. This step also reduced topological depth, lowering the cost of heuristic label computation.
Labeling: 388 tree pairs were formed with a specific distribution (e.g., UPGMA vs. NJ, NJ vs. shuffled NJ). The supervision target was the unrooted SPR distance estimated by phangorn::SPR.dist. For a subset of small trees, the exact rooted SPR distance was computed using rspr to validate the heuristic.

2. Model Architecture
The authors employed a Siamese Graph Isomorphism Network (GIN) regressor:

Graph Representation: Newick trees were parsed into bidirectional graphs. Each node was assigned a 19-dimensional feature vector comprising:
- Node degree.
- Binary leaf indicator.
- Topological distance to the root.
- A categorical taxonomic identifier (embedded via a trainable 16-dimensional layer).
Encoder: Both trees in a pair passed through a shared GIN encoder consisting of two GIN convolution layers (hidden dimension 128). Each layer utilized an MLP (Linear-ReLU-Linear) followed by batch normalization and ReLU.
Pooling and Regression: A global sum-pooling layer aggregated node features into a fixed-size graph embedding (128 dimensions). The embeddings of the two trees were concatenated and fed into a Multi-Layer Perceptron (MLP) head (256→128→64→1) to predict the SPR distance.
Training: The model was trained using the Huber loss (robust to right-skewed SPR distributions) with the Adam optimizer. Predictions were clamped at zero to ensure biological validity.

Key Contributions

Reproducible Pipeline: A complete pre-processing workflow for slicing genomic data, inferring trees, applying midpoint re-rooting, and generating SPR labels.
Public Dataset: Release of a dataset containing 864 phylogenetic trees and 388 labeled pairs across four bacterial species, spanning up to 9,500 isolates.
Heuristic Validation: A rigorous comparison showing that the unrooted phangorn::SPR.dist heuristic correlates almost perfectly (Pearson 0.98–0.99) with the exact rooted SPR distance on small trees, validating its use as a monotonic surrogate for training despite systematic magnitude underestimation.
Learned Approximation: The first application of a Siamese GIN to approximate NP-hard SPR distances on rooted phylogenies, demonstrating high accuracy within the training distribution.

Experimental Results

Supervision Target Validation: On small trees (20–90 leaves), the unrooted heuristic showed a Pearson correlation of 0.983 with the exact rooted distance. While it systematically underestimated the magnitude (returning ~65–71% of the exact value), its high rank correlation confirmed its suitability as a training target.
In-Distribution Performance: On held-out test data from the same species and size range as training, the model achieved an $R^2$ of approximately 0.87–0.90. It reduced the Mean Absolute Error (MAE) by roughly four times compared to a mean-predictor baseline.
Cross-Species Transfer: When trained on two species and tested on two unseen species, performance dropped significantly ( $R^2 \approx 0.37$ ), indicating that while the model captures some species-agnostic structural features, cross-species generalization remains challenging.
Size Extrapolation: The model failed to extrapolate to tree sizes larger than those seen in training. When trained on small and medium trees and tested on large trees, performance fell below the baseline ( $R^2 \approx -0.14$ ).
Calibration: The model tended to overestimate short distances and underestimate long ones, compressing predictions toward the mean.

Significance and Limitations

The paper establishes that GNNs can learn to approximate SPR distances with near-constant inference time, offering a scalable alternative to recomputing heuristics for every pair. The released dataset and validated heuristic relationship provide a foundation for scaling learned SPR approximations.

However, the authors modestly note that the current approach is limited by size extrapolation; the model cannot generalize to tree sizes outside its training distribution. Additionally, the model inherits a calibration bias from the heuristic labels and the sum-pooling strategy. The authors suggest that future work should focus on enriching training distributions with larger trees, exploring size-aware normalization, and potentially predicting the specific SPR trace (the leaves to be moved) rather than just the distance magnitude.