Phylogenetic tree inference with tropical axial attention

1 Introduction

A phylogenetic tree represents the evolutionary history among a set of species, genes, or sequences, and it is used in a wide range of applications, from the comparative analysis of genes (Soltis and Soltis, 2003) to the analysis of Coronavirus disease (Chen et al., 2023). Phylogenetic tree reconstruction is a combinatorial optimization problem: given a multiple sequence alignment (MSA), the goal is to infer a tree that explains the observed evolutionary relationships. Standard statistical approaches include Maximum Likelihood (ML) (Felsenstein, 1981) and Bayesian Inference (Huelsenbeck and Ronquist, 2001). Distance-based methods form another widely used class of approaches. For example, Neighbor Joining (NJ) (Saitou and Nei, 1987) and Balanced Minimum Evolution (BME) (Desper and Gascuel, 2002; Lefort et al., 2015) takes an estimated pairwise distance matrix from an MSA (for example, via the stochastic speciation process (Kuhner and Felsenstein, 1994)) and then infer a tree from that matrix. The second step in a distance-based method is difficult because the space of unrooted phylogenetic trees is a union of ( n − 3 )-dimensional polyhedral cones in R n 2 (Billera et al., 2001). In fact, estimating a phylogenetic tree via BME is NP-hard (Day, 1987) and also APX-hard (Fiorini et al., 2012). It is also well-known that NJ is an approximation to BME (Eickmeyer et al., 2008). For more details see (Catanzaro et al., 2022). One can find a tutorial on geometric approaches to phylogenetic tree reconstruction in (Haws et al., 2026).

A key for a successful phylogenetic tree estimation from an input MSA via a distance-based method is to estimate a good distance matrix from the MSA. Recent neural approaches to estimate a distance matrix from a MSA, including Phyloformer (Nesterenko et al., 2025) and Phyloformer 2 (EvoPF) (Blassel et al., 2025), replace explicit search with transformer models that take MSAs as input and output distance matrices or tree distributions. These models operate in Euclidean representations and rely on standard attention mechanisms, which may not fully capture the combinatorial structure inherent in phylogenetic inference (Vaswani et al., 2017). For interested readers, definitions and basic concepts on LLM reasoning models can be found at (Wang, 2025).

In this work, we focus on the problem to estimate a distance matrix of all pairwise distances from an input MSA to as close to a tree metric as possible. Here, we propose a geometrically motivated alternative: we replace softmax dot-product attention with tropical (max-plus) attention (Hashemi et al., 2025), yielding a piecewise-linear architecture that is better aligned with the dynamic-programming and combinatorial-optimization structure of tree reconstruction. Tropical algebra provides a natural framework for representing the polyhedral constraints that arise in distance-based phylogenetic inference.

We introduce a Tropical Axial Transformer, which:

• applies tropical attention over pairwise sequence representations,
• uses axial decomposition over sequence and pair dimensions,
• and enforces structure through losses grounded in tropical geometry.

These structural choices are reflected empirically in competitive overall performance and improved branch-length sensitive reconstruction and tree-metric consistency under distribution shift relative to Phyloformer baselines.

2 Background

2.1 Tropical algebra primer

Tropical algebra provides an alternative arithmetic in which addition and multiplication are replaced by max and addition, respectively. This algebraic structure is appropriate for problems involving piecewise-linear geometry, combinatorial optimization, and max constraints, all of which arise naturally in phylogenetics.

The max-plus tropical semiring is defined as the set R ∪ { − ∞ } equipped with the operationsfor a , b , c ∈ R ∪ { − ∞ }. In this setting, tropical addition corresponds to taking a maximum, while tropical multiplication corresponds to ordinary addition. These operations are associative, commutative, and satisfy distributivity:for a , b , c ∈ R ∪ { − ∞ }. This algebra produces functions that are piecewise linear rather than smooth. Since it operates over polyhedral structures vice points in space, tropical geometry naturally describes combinatorial geometric structures over sets of inputs which generate polyhedral structure, such as shortest paths, dynamic programming recursions, and tree reconstruction (Hashemi et al., 2025; Joswig and Schröter, 2022; Jukna, 2014; Axiotis and Tzamos, 2018).

a ⊕ b = max a , b , a ⊙ b = a + b ,

a ⊙ b ⊕ c = a ⊙ b ⊕ a ⊙ c

2.2 Tree space

Phylogenetic trees are the fundamental mathematical representation of evolutionary relationships between species or taxa, yet the space of all possible trees is both combinatorially vast and geometrically complex. Formally, a phylogenetic tree G ∈ T n is a connected, acyclic graph with n leaves with labels, internal nodes represent ancestral divergence without labels, and edges assigned nonnegative weights corresponding to evolutionary distances. In the rooted binary case, each internal node has degree three (one parent and two children), and the number of distinct rooted phylogenetic tree topologies on n labeled leaves is given bywhich grows super-exponentially with n. Even for relatively small n, this space becomes astronomically large, making exhaustive search computationally infeasible and making inference a fundamentally difficult problem (Bhatt et al., 2025). This combinatorial explosion is compounded by the fact that phylogenetic trees are not naturally embedded in a standard Euclidean space, but instead form a structured, non-Euclidean subspace inside Euclidean space.

2 n − 3 ‼ = 2 n − 3 2 n − 5 ⋯ 3 ⋅ 1 ,

A major step toward understanding this structure was the introduction the Billera-Holmes-Vogtmann (BHV) construction (Billera et al., 2001). In this framework, each tree is represented by its internal edge lengths, with each topology corresponding to a Euclidean region (orthant) of dimension n − 3, reflecting the number of internal edges in an unrooted binary tree (Ren et al., 2021). These regions are glued together along lower-dimensional faces corresponding to degenerate trees obtained by collapsing edges (internal edge length 0). While this provides a rigorous geometric foundation with desirable properties, the resulting space is only locally Euclidean. In particular, the presence of non-smooth boundaries complicates optimization and statistical analysis, and averages of trees may lie on these boundaries, producing unresolved topologies. Although BHV space provides an elegant description, it does not fully resolve the challenges posed by the geometry of phylogenetic trees (Monod et al., 2022).

An alternative perspective arises from tropical geometry, which provides a fundamentally different way to represent and analyze tree space. In particular, Speyer and Sturmfels (Speyer and Sturmfels, 2004) showed that the space of phylogenetic trees with n leaves is isomorphic to the tropical Grassmannian. Tropical geometry replaces classical addition and multiplication with max-plus (or min-plus) operations, transforming algebraic objects into piecewise-linear structures (Maclagan and Sturmfels, 2015). Under this framework, polynomials become maxima or minima of linear functions, and geometric objects become polyhedral complexes. This shift from smooth to piecewise-linear geometry aligns naturally with the structure of the space of phylogenetic trees, which is governed by combinatorial constraints and extremal relationships. In particular, tree metrics are characterized by inequalities such as the four-point condition, which depend on maximum operations rather than linear combinations.

This connection becomes precise through the tropical Grassmannian, which provides a complete characterization of phylogenetic tree space in terms of tropical algebra. In this formulation, tree metrics correspond to points in a tropical variety defined by tropical polynomials, linking phylogenetics directly to algebraic geometry. Unlike Euclidean representations, which struggle to capture the discrete and hierarchical nature of trees, the tropical framework encodes these properties directly through its operations and polyhedral structure. Moreover, methods with tropical operations are closely connected to optimization and dynamic programming, where similar max-plus operations arise naturally, offering both computational and theoretical advantages (Maclagan and Sturmfels, 2015).

Taken together, these observations suggest that the difficulty of phylogenetic inference is not merely computational, but structural. Tree space is inherently non-Euclidean, combinatorial, and governed by extremal relationships. Tropical geometry provides a natural mathematical language for this setting, offering representations that align with both the geometry and combinatorics of phylogenetic trees. This perspective motivates the use of tropical methods, including tropical neural architectures, as a principled approach to modeling and inference in phylogenetics.

2.3 Tree inference

A standard approach to phylogenetic inference is to represent a tree through pairwise distances between its leaves. Let G ∈ T n be a phylogenetic tree with n leaves and nonnegative edge weights. For each pair of leaves ( i , j ), define

d i j = length of the unique path between i and j .

Collecting these values yields a symmetric, zero diagonal matrix d = ( d i j ) ∈ R n 2, whose coordinates correspond to unordered leaf pairs ( i , j ).

This representation embeds phylogenetic trees into an ambient Euclidean space of pairwise distances. However, not every vector d ∈ R n 2 corresponds to a valid tree. The subset that does is known as the space of tree metrics, which is defined by additional combinatorial constraints arising from the branching structure of trees (Monod et al., 2022).

The fundamental characterization of tree metrics is given by the four-point condition.

Theorem 1 (Four-Point Condition). A vector d ∈ R n 2 represents a phylogenetic tree metric if and only if, for all distinct indices i , j , k , ℓ,and the maximum among the three quantitiesis attained at least twice.

Proofs for this theorem are widely available in phylogenetic textbooks. Speyer and Sturmfels (Speyer and Sturmfels, 2004) showed this condition is equivalent to the tropical Plücker relations defining the tropical Grassmannian t r o p ( G r ( 2 , n ) ). In particular, the expressionsare tropical linear forms, and the requirement that the maximum is attained at least twice corresponds to membership in a tropical variety. For more detail, see Maclagan and Sturmfels (Maclagan and Sturmfels, 2015) and Monod et al. (Monod et al., 2022). Thus, the space of phylogenetic tree metrics is naturally identified with a subset of the tropical Grassmannian, linking tree inference directly to tropical algebraic geometry.

max d i j + d k ℓ , d i k + d j ℓ , d i ℓ + d j k

A particularly important subclass of tree metrics is given by ultrametrics, corresponding to rooted trees with equal root-to-leaf distances.

Theorem 2 (Ultrametric Condition). A vector d ∈ R n 2 is an ultrametric ( d ∈ U n ) if and only if, for all distinct i , j , k,and the maximum among { d i j , d i k , d j k } is attained at least twice.

The ultrametric condition is both necessary and sufficient for the existence of a rooted equidistant tree and implies U n ⊂ T n. Geometrically, this enforces a hierarchical structure in which distances are determined by the height of the least common ancestor.

2.3.1 Tropical algebra and tree inference

A central feature of tropical geometry is invariance under additive shifts. For a vector x ∈ R D, adding a constant c to all coordinates does not change its essential structure. This leads to the tropical projective torus (Monod et al., 2022) R D / R 1, where 1 = ( 1 , … , 1 ) and, in the phylogenetic setting, D = ( n 2 ) indexes all pairwise distances between n taxa.

Two vectors x , y ∈ R D / R 1 are equivalent iffor some c ∈ R, that is to say whenever their coordinate-wise differences are constant.

x ∼ y ⇔ x = y + c ⋅ 1

In the context of phylogenetics, pairwise distance representations of trees naturally lie in this quotient space. Adding a constant to all pairwise distances corresponds to a uniform shift in evolutionary time and does not alter the underlying tree topology. Consequently, both tree metrics and ultrametrics are more precisely viewed as elements of R D / R 1 rather than R D.

A natural notion of distance on the quotient space is given by the Hilbert projective metric (Maclagan and Sturmfels, 2015; Hashemi et al., 2025), also referred to as the tropical metric. For x , y ∈ R D / R 1, it is defined as

d t r x , y = max i x i − y i − min i x i − y i .

This metric has several important properties (Hashemi et al., 2025): so it is well-defined on R D / R 1.

• Projective invariance: For any c ∈ R,

d t r x + c 1 , y + c 1 = d t r x , y ,

• Symmetry: d t r ( x , y ) = d t r ( y , x ).
• Non-negativity: d t r ( x , y ) ≥ 0, with equality if and only if x ∼ y.
• Value invariance: The metric depends only on the spread of coordinate-wise differences, rather than their absolute values.

In phylogenetic inference pairwise distances are often treated as elements of a Euclidean space, but the space of valid tree metrics is not closed under linear operations and is instead defined by the polyhedral structure of tree space. This mismatch introduces distortions when Euclidean operations and loss functions are used for learning.

In contrast, tropical geometry provides a representation in which this structure is naturally expressed. Tree metrics form a structured subset of R n 2 / R 1, and their defining relations—such as the four-point condition—are governed by tropical algebra. Within this space, the Hilbert projective metric provides a geometry that is invariant to data shifts and sensitive to relative structure making tropical algebra ideal for phylogenetic inference.

3 Related work

Practitioners have used specialized statistical and combinatorial methods for much of the history of phylogenetic tree inference. ML approaches estimate the tree and branch lengths that maximize the probability of the observed sequence alignment under a specified substitution model (Felsenstein, 1981). These methods operate directly on sequence data but require searching a vast, discrete space of tree topologies, making them computationally expensive despite heuristic optimizations used in tools such as RAxML and IQ-TREE (Stamatakis, 2014; Nguyen, 2015).

Distance-based methods instead first takes an estimated pairwise distance matrix from the MSA and then construct a tree consistent with these distances (Saitou and Nei, 1987; Lefort et al., 2015; Price et al., 2010). Geometrically, distance-based methods can be interpreted as projecting an estimated pairwise distance matrix from R n 2 onto the space of tree metrics. However, this projection is typically approximate, as observed distances derived from data rarely lie exactly in tree space.

While these classical methods often produce high-quality trees, they incur substantial computational cost for each inference due to the combinatorial nature of tree search. More recently, deep learning approaches have been proposed to infer phylogenetic structure directly from MSAs (Blassel et al., 2025). These methods shift the computational burden to a training phase, after which inference can be performed rapidly by a forward pass through a trained model. This approach offers a scalable alternative to traditional per-instance optimization.

One representative example is Phyloformer which adopts a transformer-based architecture operating on multiple sequence alignments (MSAs) (Nesterenko et al., 2025). The model encodes sequences into pairwise representations and predicts evolutionary distances, which are then used to reconstruct trees. By replacing explicit likelihood computation with learned representations, Phyloformer achieves competitive or improved accuracy relative to classical distance-based and likelihood-based methods while significantly reducing inference time per instance (Nesterenko et al., 2025). However, its performance depends on the distribution of training data, and like many supervised models, it can degrade under shifts in sequence length, evolutionary rate, or substitution model.

A more recent line of work extends this paradigm by directly modeling distributions over tree topologies. For example, the Phyloformer 2 uses likelihood-free framework to construct a neural parameterization of the posterior over phylogenies (Blassel et al., 2025). The architecture uses a simulation-based approach called Neural Posterior Estimation (NPE) to estimate evolutionary tree distributions without computing complex likelihood functions. The model uses normalizing flows to map summary statistics from observed and simulated sequence data to infer posterior distributions, from which they directly infer trees. The model demonstrates improved robustness to model misspecification, suggesting stronger generalization compared to classical methods that rely on restrictive assumptions. Phyloformer 2 typically outperforms Phyloformer topologically but underperforms on branch-length aware metrics (Blassel et al., 2025).

Despite these advances, both architectures implicitly learn representations in an ambient Euclidean space. As a result, these models produce outputs that violate tree metric constraints or require post hoc correction. Further, although likelihood-free and neural approaches improve computational efficiency and can outperform classical methods they struggle to generalize to MSAs out-of-distribution (OOD) from the training data which is not an issue for classical methods (Blassel et al., 2025; Nesterenko et al., 2025).

4 Tropical deep learning

Our method is motivated by the mismatch between the geometry of phylogenetic tree space and the Euclidean geometry implicitly assumed by standard neural architectures. Although pairwise distances can be embedded in R n 2, valid tree metrics occupy a structured, non-Euclidean subset shaped by the four point condition. To address this mismatch, we adopt tropical attention introduced by (Hashemi et al., 2025), where comparisons are governed by max-plus structure and the resulting maps are piecewise-linear and scale-invariant in the projective sense. This aligns the model’s inductive bias with the geometry of phylogenetic inference since the target objects are themselves governed by combinatorial and hierarchical structure, not arbitrary linear variation. This geometric reality also motivate a linear combination of loss functions that reflect aspects of the tree inference problem: ℓ 1 for coordinate accuracy, tropical symmetric loss to penalize distortions from target tropical geometry, and ultrametric loss for violations of the rooted tree constraints (and implicitly unrooted tree constraints). Thus, both the network mechanism and the training objective are chosen to reflect the underlying geometry of tree inference.

4.1 Loss functions

Our objective combines ambient accuracy with geometry-aware regularization. Let d ̂ ∈ R n 2 denote the predicted vector of pairwise leaf distances and let d ∈ R n 2 denote the ground-truth distance vector. Although both objects are represented in the Euclidean ambient space R n 2, the subset corresponding to valid tree metrics is highly structured and is better understood through tropical and tree constraints than through ordinary Euclidean geometry alone. For this reason, we do not train solely with a regression loss as in the Phyloformer model (Nesterenko et al., 2025). Instead, we combine an ℓ 1 term, which measures coordinatewise accuracy, with a tropical symmetric term, d t r, which measures distortion in tropical geometry, and an ultrametric term, which biases predictions toward the rooted tree regime of interest.

The first component is the ℓ 1 loss,which directly penalizes absolute prediction error in the ambient vector space. This term stabilizes optimization and ensures that the model learns outputs close to the observed targets. However, by itself, ℓ 1 treats d ̂ and d as arbitrary Euclidean vectors and does not account for the projective, polyhedral structure that distinguishes tree-valued outputs from generic points in R n 2. In our setting, this yields the tropical symmetric loss

L 1 d ̂ , d = ‖ d ̂ − d ‖ 1 ,

d t r d ̂ , d = max i d ̂ i − d i − min i d ̂ i − d i .

The tropical metric has been widely used as a natural metric on tree spaces (Monod et al., 2022). Because of its projective nature, this metric depends on relative discrepancies rather than absolute differences, and it encourages the network to preserve the combinatorial shape of the target distance pattern rather than merely minimizing coordinatewise error.

To further bias the output toward phylogenetic tree structure, we add an ultrametric penalty. For each triple of leaves ( i , j , k ), we extract the three predicted distancesidentify the two largest values, and penalize any gap between them. Writing the ordered top two values as a ≥ b, the triplewise violation is

d ̂ i j , d ̂ i k , d ̂ j k ,

max a − b , 0 .

The ultrametric loss L ultra ( d ̂ ) is then defined as the average of this quantity over all triples. Thus, the penalty is zero exactly when, for every triple, the two largest distances are equal, and increases as triples deviate from ultrametric structure. In this way, the loss provides a differentiable measure of how far the prediction lies from U n, the space of ultrametrics.

Equation 1 reflects a practical training paradigm to leverage the benefits discussed previously. During the first 20 epochs, optimization focuses on recovering the coarse distance geometry of the target through Euclidean accuracy and tropical consistency. Once the model has learned a reasonable representation of pairwise relationships, the ultrametric penalty is activated to encourage convergence toward the rooted-tree subspace without constraining early learning.

L e = 1.0 L 1 d ̂ , d + 1.4 d t r d ̂ , d , e ≤ 20 , 1.0 L 1 d ̂ , d + 1.4 d t r d ̂ , d + 0.6 L ultra d ̂ , e > 20 . (1)

4.2 Tropical attention

The central contribution of this model is to replace vanilla softmax attention with the algorithmically aligned tropical attention mechanism. Vanilla softmax dot-product attention operates in smooth Euclidean space and converts logits into dense probability distributions which tend to spread strictly positive mass across all tokens. In problems whose correct solution depends on decisive selection or comparison, this smoothing can blur the sharp boundaries that define the underlying computation. Hashemi et al. (Hashemi et al., 2025) argue that this is a fundamental source of failure for softmax-based reasoning on combinatorial tasks, especially under out-of-distribution changes in sequence length or value scale, because the top logit gaps can shrink and the resulting attention becomes increasingly diffuse. In contrast, many combinatorial and dynamic-programming computations are naturally governed by max, min, and additive operations, which are piecewise-linear and polyhedral rather than smooth (Jukna, 2014).

Tropical attention is designed to build this inductive bias directly into the reasoning mechanism. Following (Hashemi et al., 2025), the input representation is first mapped into tropical projective space through a learnable valuation map, producing a tropicalized representation. Queries, keys, and values are then formed by tropical linear projections in the max-plus semiring. Specifically, let R ∈ R M × P × d model be the token embedding of the input. Then here we define the tropicalization map by going to an amoeba representation of the input with a learnable valuation map, Φ : R M × P × d model → ( T P d model − 1 ) M × P such thatfor each row i ∈ { 1 , … , M } and for each column j ∈ { 1 , … , P }. The constant shift enforces max i j ϕ λ ( x ) i j = ϵ, so the output of ϕ λ always lies in the tropical simplex, Δ d model − 1 ≔ z ∈ R d model | max i j z i j = ϵ, where every vector is projectively equivalent to exactly one point in the tropical simplex.

Φ λ R i j = X i j − max 1 ≤ r ≤ d model X ijr . 1 d model , where X = log max 0 , R ∈ R M × P × d model (2)

For the tropicalized input X, the projectionsare computed by replacing the linear map with tropical matrix multiplication. Instead of computing similarity by Euclidean dot product and normalizing with softmax, tropical attention scores query-key agreement using the tropical metric,

Q = Φ λ R ⊙ W Q , K = Φ λ R ⊙ W K , V = Φ λ R ⊙ W V

d t r Q , K = S i , j = − max q i − k j − min q i − k j . (3)

The attention scores, S, are then aggregated tropically with V step via tropical matrix multiplication

Attn Q , K , V i = max j S i , j + v j .

Conceptually, this replaces the smooth weighting rule “average all values with exponentially reweighted coefficients” by the sharp rule “select the value that aligns best with the query in the tropical projective torus.” The output is then mapped back into ordinary Euclidean coordinates so that the remainder of the network can can be trained via traditional deep learning gradient descent methods and existing packages (Hashemi et al., 2025).

This construction changes the geometry of the computation. In the tropical-attention framework, the computation is piecewise-linear and uses idempotent max-plus aggregation, while its constituent tropical linear maps are non-expansive in the tropical metric ((Talbut et al., 2024), Proposition 9). This preserves the polyhedral decision structure characteristic of combinatorial value functions, rather than smoothing those boundaries into quadratic or spherical regions as softmax attention does.

This inductive bias is attractive for phylogenetic inference even though the task is not identical to the combinatorial benchmarks studied in the paper. Tree metrics, ultrametrics, and tropical tree spaces are governed by polyhedral structure and constraints (Monod et al., 2022). Smooth softmax attention is therefore not the most natural mechanism for propagating information relevant to hierarchical tree structure. Tropical attention encourages sharp relational comparisons and projectively meaningful aggregation, which is better aligned with distance-based tree prediction. In this sense, our use of tropical attention is not merely a novel attention trick, it is a geometric choice intended to match the architecture’s inductive bias to the geometry of the target objects.

5 Model

Our model is an extension of both the tropical attention model proposed by Hashemi et al. (2025) and the Phyloformer model introduced by Nesterenko et al. (2025), using tropical axial attention on MSAs to predict distance matrices. The overall design closely matches that of Phyloformer, and can be broken into three stages: data embedding, tropical axial attention, and prediction.

MSA input data is a set of n taxa, which will be called leaves for the remainder of the section. Each leaf is a sequence of observed residues of length M. Each residue is from an alphabet of size A. The sequence is then one-hot encoded and embedded in dimension d model. Finally, each unordered pair of leaves is concatenated and re-projected to dimension d model. The full embedding is as follows,

M , n → M , n , A → M , n , d model → M , n 2 , 2 d model → M , n 2 , d model . (4)

After this step, each token is then a representation of an unordered leaf pair at a particular alignment position and T = [ M ] × P where P = { ( i , j ) : 1 ≤ i < j ≤ n }. For the remainder of the section, we will denote the unordered pairs P = n 2 and the embedded tensor R ∈ R M × P × d model. Each token r m , ( i , j ) ∈ R d model contains site specific information about the evolutionary relationship between leaves i and j. Evaluating pairwise vice simply leaf-wise information is motivated by the structure of the model objective. As explained previously, the tree pairwise distance matrix is distinguished from that of a general connected graph by its hierarchical structure and compliance with four point and ultrametric conditions, which are constraints on sets of pairwise distances. This property means the target is constrained by relationships between leaves rather than the properties of individual leaves.

The embedded tensor R ∈ R M × P × d model is then input into the main portion of the model which consists of six tropical axial attention blocks of four heads each. Due to the multi-axial nature of the data - sequence positions M and leaf pairs P - we implement axial attention from Wang (Wang, 2021) to attend to both axes (Ho et al., 2019; Wang et al., 2020), see Figure 1. Each block alternates attention over these axes so that information propagates across both sequence positions and pairwise relationships, preserving the biologically and geometrically relevant information Nesterenko et al. (2025).

• Sequence axis attention: For a fixed leaf pair ( i , j ) ∈ P, attention across M sites lets the model learn mutation evidence across the alignment.
• Pair axis attention: For a fixed site m ∈ [ M ], attention across pairs P computes relationships between pairs which are the elements of geometrically relevant tree-constraints which distinguish tree metrics from ordinary distance matrices, see Theorems Equations 1, 2.

In our model, the vanilla softmax dot-product attention is replaced with tropical attention introduced in the previous section. Thus, queries, keys, and values are computed through tropical linear projections, compared through the projectively invariant tropical metric, and aggregated by tropical matrix multiplication rather than convex Euclidean operations. This gives each block the ability to refine pairwise phylogenetic representations while preserving the polyhedral, projectively meaningful comparisons that define the tree inference problem.

After the attention stack, the resulting M , P , d model tensor is mapped to a scalar prediction for each pair of leaves with a pointwise feed-forward projection via ( 1 × 1 ) convolution. This convolution is a channel-wise projection of each sequence position/pair location, ( m , ( i , j ) ) ∈ T from R d model → R. Because the projection is applied independently at each ( m , ( i , j ) ) location, it collapses only the channel dimension while preserving the learned sequence- and pair-dimension information. Softplus then ensures non-negative outputs appropriate for distance prediction, and the mean is taken across the sequence dimension to yield the pairwise distance prediction y ̂ ∈ R ≥ 0 P.

With a pairwise axial structure and the projectively invariant similarity measure of tropical attention, our model closely aligns with the basic relational nature of the phylogenetic tree inference problem. The full architecture combines the biologically sound embedding pipeline of Phyloformer (Nesterenko et al., 2025) with the more geometrically aligned inductive bias of tropical attention, designed to improve neural phylogenetic inference.

6 Experiments

In this section, we evaluate our proposed tropical axial attention model against both the Phyloformer of Nesterenko et al. (Nesterenko et al., 2025) and Phyloformer 2 of Blassel et al. (Blassel et al., 2025). The experimental objective is to determine if replacing Euclidean softmax dot-product attention with tropical axial attention improves phylogenetic distance estimation and aides tree reconstruction, particularly when evaluating data from outside the training data distribution. In this experiments, we use the FastME implementation (Lefort et al., 2015), a heuristic for the BME, to infer a phylogenetic tree from a distance matrix.

6.1 Experimental design

To isolate architectural differences, we train all models on the same set of 10, 000 simulated MSAs with n = 10 leaves and sequence length M = 250. The training data are simulated utilizing the same birth-death process under the Alisim LG + GC model (Ly-Trong et al., 2022) used in both Phyloformer and Phyloformer 2 (Nesterenko et al., 2025; Blassel et al., 2025). The comparison models are trained using the hyperparameter settings described in their respective papers, except we use a standard embedding dimension of d model = 80 for all models to control their representational capacity. We train the tropical model with the same L = 6 axial attention blocks of H = 4 heads each as the Phyloformer model, and we use a learning rate of λ = 0.001. These hyperparameters product a tropical model with 636,001 trainable parameters and Phyloformer model with 477,697 trainable parameters. We train each model on the same set of 5 random seeds and report performance as an average across seeds.

All models follow the same input-output structure, receiving an input MSA and outputting a vector of predicted pairwise leaf distances. These predicted distances are converted into a distance matrix which is then used to infer a tree for topological comparison using FastME (Lefort et al., 2015). This design allows comparison of models both of distance estimators and as components of a full phylogenetic inference pipeline.

6.2 Evaluation metrics

We evaluate the tree inference performance of all models based on topological accuracy by comparing the predicted tree to the true reference tree. The primary topological metric is the Robinson-Foulds (RF) distance (Robinson and Foulds, 1981). Let G ̂ and G be the predicted and true trees, respectively, and S ( G ) be the set of bipartitions of the internal edges. The RF distancewhere | ⋅ | denotes the cardinality of the set. The RF distance exclusively compares the topological difference between G and G ̂.

d R F G ̂ , G = S ( G ) \ S ( G ̂ ) + S ( G ̂ ) \ S ( G ) ,

We also consider extensions of the RF distance. The first normalizes by the total number of possible splits on fully resolved unrooted binary trees with n leaves, yielding the normalized RF distancewhich controls for the size of the tree (Robinson and Foulds, 1981).

d nRF = d R F 2 n − 6 ,

The second RF extension, weighted RF distancewhere σ is a split in either tree, w G ̂ ( σ ) is the branch length of that split in G ̂, w G ( σ ) is the branch length of the split in G, and w i ( σ ) = 0 if σ is not in tree i. Thus, the weighted RF distance incorporates both the topology and branch length deviation between G and G ̂ (Robinson and Foulds, 1981).

d wRF ( G ̂ , G ) = ∑ σ ∈ S ( G ̂ ) ∪ S ( G ) w 1 ( σ ) − w 2 ( σ ) ,

The final tree error distance is the Kuhner-Felsenstein (KF) distance

d K F ( G ̂ , G ) = ∑ σ ∈ S ( G ̂ ) ∪ S ( G ) ( w 1 ( σ ) − w 2 ( σ ) ) 2 .

Like weighted RF, KF penalizes both topological and branch-length differences, but uses an ℓ 2 rather than ℓ 1 penalty (Kuhner and Felsenstein, 1994).

The distances above each evaluate how close an inferred tree is from the simulated trees associated with the input MSA, but we are also interested in geometric performance and how well each model learns outputs in tree space. To do so, we first infer a valid tree G ̂ ∈ T n from d ̂ with the FastME implementation (Lefort et al., 2015). We then calculate the pairwise path lengths for G ̂ which provides a valid pairwise distance vector π F a s t M E ( d ̂ ) ∈ T n, which is obtained via FastME from d ̂ (in order to simplify, we denote as π FastME thereafter). Finally, we evaluate d t r ( d ̂ , π F a s t M E ) and MAE ( d ̂ , π F a s t M E ) = 1 P ∑ 1 ≤ i < j ≤ n ‖ d ̂ i j − proj ( d ̂ i j ) ‖ 1 to measure how far d ̂ is from T n both projectively and absolutely. This assesses the consistency of each model as an estimator for valid tree metrics.

6.3 Results

We evaluate model performance under a variety of input distributions. The first experiment evaluates performance on simulated data generated in the same manner as the training set with 250 MSAs for each of n ∈ { 10,20,30,40,50 } and 50 MSAs for each of n ∈ { 50,60,70,80,90 }. All MSAs are generated with sequence length M = 200. Results of this experiment are shown in Figure 2 and Table 1. We used the Wilcoxon signed-rank test to perform pairwise comparisons between models for each metric (Wilcoxon, 1945; Hollander et al., 2013), where each test instance constituted an observation. Because multiple pairwise comparisons were conducted within each metric and tree size, adjusted p-values were computed using Holm’s sequential multiple-testing procedure (Holm, 1979).

TABLE 1. Pairwise comparison of Tropical with baseline models, averaged across tree sizes. Each tree constitutes one paired observation in the Wilcoxon signed-rank test. The reported model values are arithmetic means across tree sizes. The difference Δ is Tropical minus the baseline, so negative values show Tropical outperformance.

These experiments are designed to evaluate the input size generalization capability of each model. The n = 10 test samples represent in-distribution performance while the n = 20 through n = 100 samples represent increasing length OOD performance. Since the number of pairwise distances grows with n 2, the experiment tests how well each model can represent coherent structure as the output dimension grows beyond training inputs. Consistent with the results in (Blassel et al., 2025), we see from Table 1 that the Phyloformer 2 model shows consistently superior topological performance while Phyloformer outperforms on metrics that consider branch lengths. The tropical model consistently achieves results between those of Phyloformer and Phyloformer 2 across all metrics. Additionally, the tropical model appears to match the performance degradation trend of both baseline models on the topologically focused RF-based metrics as number of leaves increases, but performance degrades on the branch-length aware distances.

The second experiment evaluates model performance on simulated data with a different generation process. This data set is 30 40-taxon MSAs and 10,100-taxon MSAs both with M = 500 length sequences and corresponding true trees from the PhyML benchmark (Guindon et al., 2010). In contrast to the birth-death process used during training, this benchmark was generated using the simulation design of Guindon and Gascuel (Guindon and Gascuel, 2003), which generated 40-taxon trees using the process of Kuhner and Felsenstein (Kuhner and Felsenstein, 1994). Alignments were simulated under a General Time Reversible (GTR) model with 4 gamma-distributed rates (Lanave et al., 1984; Guindon et al., 2010). This benchmark provides a useful test of the ability of each model to generalize to data from unseen distributions.

Tables 2, 3 show results on the PhyML dataset. In contrast to the birth-death generated data, we see that the tropical model performs comparably with Phyloformer 2 topologically while outperforming both other models when considering branch lengths. Table 3 shows that the tropical model is learning distance matrices closer to tree-metric distances than the other methods. Taken together, we see the tropical model achieves better relative performance on the OOD PhyML dataset than in the test data generated with the same birth-death process as the training data. This indicates distributional robustness and supports the hypothesis that the tropical model learns structure which generalizes beyond the training simulation regime.

TABLE 2. PhyML (Guindon et al., 2010) tree statistics for Phyloformer (PF), Phyloformer 2 (PF2), and Tropical (Trop) models on n = 40,100 leaves.

Values are mean ± standard deviation.

TABLE 3. PhyML (Guindon et al., 2010) distance statistics for Phyloformer (PF), Phyloformer 2 (PF2), and Tropical (Trop) models on n = 40,100 leaves.

Values are mean ± standard deviation. Distances were calculated between d ̂ ∈ R n 2 and π F a s t M E ∈ T n.

The final experiment consists of empirical datasets. We use 11 MSAs, known as DS1-DS11, which serve as a benchmark for empirical evaluation (Whidden and Matsen, 2015). Summary details for these data sets are outlined in Table 4, and full details can be found in (Lakner et al., 2008).

TABLE 4. Empirical DS1-DS11 datasets used for phylogenetic evaluation (Lakner et al., 2008; Whidden and Matsen, 2015). Each data set represents a different type of data of interest to the phylogenetic community following different evolutionary patterns.

Empirical inference is particularly important in the phylogenetic setting because real-world data includes biological structure not captured under any simulation regime, and performance on this data indicates whether the model is simply fitting the training simulation prior or learning the structure of tree metrics. Unlike the simulated experiments, the true trees of empirical MSAs are unknown. Therefore, this experiment tests model output behavior on real sequence data rather than on an ability to reconstruct a synthetic tree.

Table 5 shows matrix-tree space results for the DS1-DS11 datasets. We see that the tropical model is much closer both projectively and absolutely to its FastME reconstruction than either of the other models with lower variance. This suggests, even when presented with empirical data outside of the training distribution, the tropical model provides a more consistent estimator for tree metrics than Phyloformer and Phyloformer 2.

TABLE 5. Empirical matrix-to-tree consistency distance metrics for Phyloformer (PF), Phyloformer 2 (PF2), and Tropical (Trop) models.

Values are mean ± standard deviation across seeds. Distances were calculated between d ̂ ∈ R n 2 and π F a s t M E ∈ T n. Lower values indicate greater consistency between the predicted matrix and its induced tree representation.

6.4 Ablation study

To isolate the contributions of the attention mechanism and individual loss terms, we conducted an ablation study using the final six-block, four-head architecture. Each configuration was trained on the same 1,000 simulated 10-taxon alignments using five random seeds and evaluated on the same 200 held-out 10-taxon alignments. We compared tropical attention trained with ℓ 1 loss only, tropical loss only, ℓ 1 plus tropical loss, ℓ 1 plus tropical loss plus the ultrametric penalty, and a softmax attention model of the same construction with just the tropical loss. Table 6 reports mean ± standard deviation across seeds.

TABLE 6. Ablation results on the held-out 10-taxon test set.

Values are mean ± sample standard deviation across five training seeds, after averaging over test alignments within each seed. Lower values are better on all metrics. Bold indicates the best result in each column. L trop denotes the tropical symmetric loss.

The ablation study confirms that the ℓ 1 term is necessary to minimize pairwise distance pair errors, while the tropical loss acts as a geometry-aware regularizer when combined with ℓ 1. The full objective function with ℓ 1, tropical, and ultrametric penalty provide the best overall balance across distance and branch-length-sensitive metrics, although the incremental effect of the ultrametric penalty is modest.

7 Conclusion

We introduced a tropical axial attention model for phylogenetic tree inference on MSAs. The central motivation is that phylogenetic tree inference is not an unconstrained Euclidean regression problem. Instead, valid tree metrics on n leaves occupy a structured, polyhedral subset of the ambient space R n 2. The tropical operations of the attention mechanism and the tropical metric in the loss function are designed to induce a bias in the model which is both accurate and consistent with tree space.

Across experiments with data generated under the same simulation regime as the training data, we see that the tropical model remains competitive while not outperforming the Phyloformer and Phyloformer 2 baseline models. When presented with data from a different simulation distribution, the tropical model improves, achieving topological results comparable to Phyloformer 2, outperforming both models on branch-length sensitive metrics, and providing a more consistent estimator for the tree metric induced by FastME.

The empirical experiments on DS1-DS11 reinforce this pattern. Because true empirical trees are unknown, these results do not directly measure the reconstruction accuracy. Instead, they show that the tropical model produces distance matrices that are much closer to their FastME reconstructed tree metrics than those produced by the comparison models. This suggests that the tropical axial attention model produces outputs more consistent with tree structure, even outside the simulated training distribution.

Overall, these results support the hypothesis that neural models benefit from aligning the inductive bias to the geometry of the problem. The piecewise linear operations of tropical axial attention provide an algebraic framework to learn the polyhedral structure of tree space. While the proposed model does not outperform across every metric, it shows robustness under distribution shift and improves tree metric consistency over the baseline models, indicating opportunities for further research in tropical attention models for structured inference problems.

A practical limitation of the present implementation is its peak memory requirement during training. The tropical metric computation between the Q and K tensors in Equation 3 and the tropical matrix multiplication of these scores with the V tensor in Equation 4 are currently evaluated with Pytorch broadcast operations. This approach creates large intermediate tensors before applying the maximum and minimum reductions, substantially increasing GPU memory. More specialized implementations, such as the optimized tropical matrix-multiplication rountines in TropicalGEMM (GiggleLiu and contributors, 2026), may reduce this memory overhead if tensor broadcasts can be avoided.

Generative AI statement

The author(s) declared that generative AI was not used in the creation of this manuscript.

Any alternative text (alt text) provided alongside figures in this article has been generated by Frontiers with the support of artificial intelligence and reasonable efforts have been made to ensure accuracy, including review by the authors wherever possible. If you identify any issues, please contact us.

Publisher’s note

All claims expressed in this article are solely those of the authors and do not necessarily represent those of their affiliated organizations, or those of the publisher, the editors and the reviewers. Any product that may be evaluated in this article, or claim that may be made by its manufacturer, is not guaranteed or endorsed by the publisher.

Footnotes

Edited by: Joseph Rusinko, Hobart and William Smith Colleges, United States

Reviewed by: Martin Frohn, Maastricht University, Netherlands

Bhavika Gambhava, Dharamsinh Desai University, India