技术摘要：通过因子图实现面向演化的子采样 MSA 推理

问题陈述

多序列比对（MSA）为蛋白质语言模型（PLMs）提供了显式的演化上下文，这对于接触预测和构象分析等任务至关重要。然而，从基因组数据库中检索到的 MSA 深度往往达到数万条序列，远超基于注意力机制的 PLM（如 MSA Transformer、AlphaFold2）的 Token 预算（通常仅处理数百条序列）。因此，子采样成为一个必经步骤。

现有的子采样策略依赖于启发式方法，例如随机选择、基于身份的过滤（如 HHfilter）或多样性驱动的贪婪采样（如 divMax、divMin）。虽然这些方法在去除近重复序列或强制执行特定偏差方面行之有效，但它们作用于局部相似性关系，缺乏对两个互补演化度量指标的显式控制：查询身份（Query Identity，即序列与目标序列的相似度）和序列多样性（Sequence Diversity，即演化变异的广度）。这种局限性导致无法显式平衡这两个因素，可能导致保留的演化信号对于下游任务而言并非最优。

方法论：AP-REASONER

作者提出了 AP-REASONER，该方法将 MSA 子采样从一种启发式过滤过程重新构建为一个通过因子图上的**亲和力传播（Affinity Propagation, AP）**求解的显式优化问题。

因子图公式化

该问题被建模为因子图上的最大后验概率（MAP）估计，其中：

• 变量： 对于 MSA 中的每一行 $i$，一个离散变量 $s_i$ 表示它由哪一行 $k$ 来代表。当且仅当 $s_k = k$ 时，行 $k$ 是一个“代表元”（Exemplar，即被选入子集的序列）。
• 一元因子 ($L_i$)： 这些因子衡量行 $i$ 被代表元 $k$ 代表的可能性。相似性矩阵 $S$ 使用归一化的、考虑空位的 Hamming 距离 ($d_{H}^{gap}$) 构建，以量化序列冗余度和差异性。
  • 非对角线项 ($S_{ik}, i \neq k$)： 定义为 $-\left(\frac{d_{H}^{gap}(i, k)}{L}\right)^\alpha$。指数 $\alpha$ 作为控制全局多样性的控制旋钮。较大的 $\alpha$ 会平滑小距离带来的惩罚，从而鼓励选择在全局分布更广的代表元。
  • 对角线项 ($S_{kk}$)： 定义为 $p + \beta \frac{d_{H}^{gap}(k, 0)}{L}$，其中 $0$是查询序列。系数$\beta$作为控制**查询身份**的控制旋钮。正值的$\beta$激励选择远离查询序列的代表元，而负值的$\beta$ 则偏向于选择与查询序列相似的序列。查询行被赋予一个较大的固定偏好，以确保其始终被保留。
• 代表元一致性因子 ($f_k$)： 这些因子强制执行约束：如果任何行选择 $k$ 作为代表元，则 $k$ 也必须选择自身（即 $s_k = k$）。

优化与控制

系统执行消息传递推理（对数域最大乘积算法）以推断最优代表元集合。为了严格遵守固定的 Token 预算 $N_B(L)$，作者采用了固定基数封装器（Fixed-Cardinality Wrapper）。这涉及对全局偏好偏移量 $p$ 进行二分搜索，以调节推断出的代表元数量至目标预算，随后根据对角线信念进行前 $N$ 个截断或填充。

核心贡献

1. 优化重构： 本工作首次将 MSA 子采样从启发式采样程序重新构建为显式优化问题，将生物学度量（多样性和查询身份）提升为显式目标。
2. AP-REASONER： 引入了一种基于亲和力传播的因子图方法，通过因子编码子采样目标，并通过消息传递进行求解。它具有两个可解释的控制旋钮（$\alpha$ 和 $\beta$），分别用于显式调节序列多样性和查询身份。
3. 实证验证： 通过全面的评估证明，AP-REASONER 在结构敏感的下游任务上优于基准采样器（随机、divMin、divMax、HHfilter、K-medians），并能够可控地恢复替代蛋白质构象。

实验结果

作者在蛋白质语言模型预训练、长程接触预测（LRCP）和蛋白质多构象预测上评估了 AP-REASONER。

• 预训练 (MLM)： 在掩码语言模型任务上，AP-REASONER 的表现与随机采样和 K-medians 相当，优于 divMin 和 divMax。然而，作者指出 MLM 的性能并不总是与下游结构性能相关联。
• 长程接触预测 (LRCP)： 在 CASP15 目标集上，使用 AP-REASONER 训练的模型在 Top-L 接触精度上始终高于使用基准采样器的模型。在“信息丰富的子集”目标上，这种性能差距进一步扩大。此外，在推理过程中改变多样性旋钮 $\alpha$ 显示出 LRCP 性能随 $\alpha$ 增加而单调提升，证实了该旋钮在促进多样性子集方面的效用。
• 构象系综预测： 在折叠开关蛋白（KaiB, RfaH, MAD2）上，AP-REASONER 比暴力随机采样和 AF-Cluster 展示了更高的效率和精度。
  • 它成功恢复了随机采样会遗漏的罕见功能状态（例如 KaiB 的折叠开关状态）。
  • 它以显著减少的样本量（例如在 KaiB 上为 56 个样本 vs AF-Cluster 的 1008 个样本）实现了相当的结构准确度 (RMSD)。
  • $\beta$ 旋钮允许控制预测的系综在不同构象盆地（例如基态与折叠开关状态）之间进行转向，从而实现对特定功能状态的针对性恢复。

重要性与主张

论文声称，将 MSA 子采样建模为通过因子图推理求解的可控优化问题，为启发式选择提供了一种原则性的替代方案。通过显式平衡查询身份和多样性，AP-REASONER 不仅提高了结构预测的下游性能，还提供了一种高效的机制来探索蛋白质构象景观。通过 $\alpha$ 和 $\beta$ 进行转向采样，可以实现对替代功能状态的可预测恢复，这对于理解蛋白质动力学和变构机制至关重要，且无需承受暴力采样的高昂计算成本。

作者承认了局限性，指出由于消息传递和二分封装的存在，AP-REASONER 比轻量级启发式采样器慢约 2–3 倍。此外，预训练实验是在 260K 规模的数据集而非完整的 26M 规模语料库上进行的，这使得该方法在更大规模场景下的可扩展性成为未来探索的一个开放性问题。