技术摘要：基于热带轴向注意力的系统发育树推断

问题陈述

系统发育树重建是一个组合优化问题，旨在从多序列比对（MSA）中推断进化关系。尽管存在标准的统计方法（如最大似然法、贝叶斯推断）和基于距离的方法（如邻接法、平衡最小进化法），但它们仍面临重大挑战。基于距离的方法需要估计成对距离矩阵，然后将其投影到有效系统发育树的空间中。该空间并非平滑的欧几里得流形，而是一个由多面体约束（具体为四点条件和超度量不等式）定义的、结构化的非欧几里得子空间。

最近的深度学习方法，如 Phyloformer 和 Phyloformer 2，利用带有标准 softmax 点积注意力的 Transformer 架构，直接从 MSA 预测距离矩阵。然而，这些模型在欧几里得空间中运行，并依赖于平滑的稠密概率分布。作者认为，这种几何上的不匹配未能捕捉树度量固有的组合和分段线性结构，导致输出可能违反树的约束，并在分布偏移（如序列长度或进化模型的变化）下难以泛化。

方法论

本文提出了一种热带轴向注意力（Tropical Axial Attention）神经网络架构，旨在使模型的归纳偏置与系统发育树空间的几何特性相一致。

1. 理论基础：热带几何

核心理论见解在于，包含 $n$ 个物种的系统发育树空间与热带格拉斯曼流形（tropical Grassmannian）之间存在同构关系。热带代数将标准的加法和乘法分别替换为 $\max$ 和加法（即 max-plus 半环）。

• 树度量：有效的树度量由四点条件刻画，该条件涉及最大值运算而非线性组合。
• 超度量：有根等距树满足超度量条件，同样由最大值运算支配。
• 度量空间：作者利用商空间 $\mathbb{R}^D / \mathbb{R}\mathbf{1}$ 上的希尔伯特投影度量（热带度量），该度量对加法平移具有不变性，且对相对结构敏感，这与标准欧几里得度量不同。

2. 模型架构

该模型扩展了 Phyloformer 架构，用热带轴向注意力替换了普通的 softmax 注意力：

• 输入处理：MSA 经过独热编码并嵌入。无序的叶子对被拼接起来，形成表示跨序列位置成对关系的张量。
• 热带注意力机制：
  • 查询（$Q$）、键（$K$）和值（$V$）通过 max-plus 半环中的热带线性投影（使用 $\max$ 和 $+$）计算得出。
  • 注意力分数使用热带度量计算：$S_{i,j} = -(\max(k_j - q_i) - \min(k_j - q_i))$。
  • 聚合通过热带矩阵乘法执行：$\text{Attn}(Q, K, V)_i = \max_j (S_{i,j} + v_j)$。
  • 这取代了 softmax 的“软”平均，采用一种“锐利”的选择机制，从而保留分段线性的决策边界。
• 轴向分解：模型沿两个轴应用注意力：序列轴（用于学习比对位点间的突变证据）和成对轴（用于学习叶子对之间的关系，强制执行几何树约束）。
• 输出：逐点前馈投影将通道维度压缩以预测成对距离，随后通过 Softplus 激活函数确保非负性。

3. 损失函数

训练目标结合了三个组成部分，以同时确保准确性和几何一致性：

1. $\ell_1$ 损失：衡量环境欧几里得空间中的绝对坐标误差，以确保接近真实值。
2. 热带对称损失：使用希尔伯特投影度量（$d_{tr}$）来惩罚热带几何中的失真，鼓励模型保留距离的相对结构，而不仅仅是绝对值。
3. 超度量惩罚：惩罚违反超度量条件的情况（即任意三元组中两个最大距离必须相等）。

• 训练策略：超度量惩罚仅在前 20 个 epoch 之后激活，以允许模型首先学习粗略的距离几何，然后再强制执行严格的树约束。

关键结果

该模型在三种类型的数据集上与 Phyloformer 和 Phyloformer 2 进行了评估：

1. 分布内模拟：在与训练集相同的出生 - 死亡过程生成的数据上（叶子数量 $n$ 从 10 变化到 100），热带模型表现具有竞争力，通常在拓扑结构（Robinson-Foulds）和分支长度指标上取得介于 Phyloformer 和 Phyloformer 2 之间的结果。
2. 分布外（OOD）模拟：在使用不同模拟过程生成的数据上（带有 GTR 模型的 PhyML 基准），热带模型表现出卓越的鲁棒性。其拓扑准确性与 Phyloformer 2 相当，而在对分支长度敏感的指标（加权 RF 和 KF 距离）上显著优于两个基线模型。
3. 实证数据集（DS1–DS11）：在真实生物数据（真实树未知）上，评估侧重于矩阵到树的一致性。与基线模型相比，热带模型生成的距离矩阵与其诱导的 BME（平衡最小进化）树度量在 MAE 和热带距离上都更为接近。这表明热带模型学习到了与树空间潜在多面体结构更为一致的表示。

意义与主张

本文主张，系统发育推断的困难不仅仅是计算上的，更是结构上的，源于树空间的非欧几里得和组合性质。

• 几何归纳偏置：主要贡献在于证明，用热带注意力替换欧几里得注意力为系统发育推断提供了一个自然的几何框架。通过将网络的操作（max-plus）与树度量的约束（四点条件）对齐，模型学习到的表示在本质上与树拓扑更为一致。
• 对分布偏移的鲁棒性：结果表明，热带注意力作为一种鲁棒的归纳偏置，特别是在模型在不同于训练集的数据分布上进行测试时表现突出。
• 一致性优于纯准确性：虽然该模型在分布内设置下并未在所有指标上普遍优于基线，但它始终生成的输出更符合树度量的数学定义。这对于距离矩阵的有效性至关重要的下游任务尤为宝贵。

作者得出结论，热带轴向注意力为神经系统发育推断提供了一种原则性方法，弥合了深度学习架构与进化树代数几何之间的差距。