🤖 AI Summary
This work addresses the inconsistency between distance matrices and tree metrics in phylogenetic inference, particularly under distributional shifts. The authors propose the Tropical Axial Attention Network, which uniquely integrates tropical geometry with attention mechanisms by replacing the conventional softmax operator with a max-plus formulation. Leveraging the isomorphism between the tropical Grassmannian and the space of phylogenetic trees, the method incorporates geometric inductive bias and jointly optimizes an ℓ1 loss, a tropical symmetric distance loss, and a penalty for ultrametric violation. Evaluated on datasets DS1–DS11, the learned distance matrices exhibit significantly closer alignment with tree metrics induced by the Balanced Minimum Evolution (BME) criterion, outperforming existing baselines and demonstrating enhanced tree consistency and robustness to distributional shifts.
📝 Abstract
In this work, we introduce a Tropical Axial Attention neural reasoning architecture that replaces vanilla softmax dot-product attention with max-plus operators, inducing a piecewise-linear structure aligned with dynamic programming formulations. From multi-species sequence alignments, our model learns all possible pairwise distances and is trained using a combination of $\ell_1$ and tropical symmetric distance metric losses with an ultrametric violation penalty. We leverage the well known isomorphic relationship between the space of all phylogenetic trees with $n$ species and tropical Grassmannian to show that tropical attention provides a natural geometric framework for phylogenetic inference.
On empirical $DS1-DS11$ alignments, where true trees are unknown, the tropical model produces distance matrices that are substantially closer to their BME-induced tree metrics than the baseline models. These results suggest that tropical attention is a useful geometric inductive bias for neural phylogenetic inference, especially under distribution shift and when tree-metric consistency is important.