](https://deep-paper.org/en/paper/2302.12563/images/cover.png)
蛋白质表征学习领域目前正在经历一场巨大的范式转变。多年来，理解蛋白质结构和功能的“黄金标准”一直是研究其进化历史。通过将蛋白质序列与其成千上万个进化近亲进行比对——这一过程被称为多序列比对 (Multiple Sequence Alignment, MSA) ——模型可以推断出蛋白质的哪些部分至关重要，以及哪些部分在空间上相互作用。
然而，MSA 付出了昂贵的代价: 它的计算成本高昂、速度缓慢，并且依赖于僵化的预计算数据库。
在这篇深度文章中，我们将探讨一篇挑战这一僵化过程必要性的论文: “Retrieved Sequence Augmentation for Protein Representation Learning” (用于蛋白质表征学习的检索序列增强) 。 研究人员提出了一种名为 RSA 的新框架，该框架借鉴了自然语言处理 (NLP) 中的一个概念——检索增强生成 (Retrieval Augmented Generation) ，旨在取代 MSA。结果如何？该模型不仅在数学上十分优雅，而且比目前最先进的方法快 373 倍 , 同时实现了更优越的性能。
瓶颈: 为什么进化计算如此昂贵
要理解为什么 RSA 是一个突破，我们需要先了解它解决的问题。蛋白质是生物学的核心驱动力。从一维的氨基酸链预测其三维结构和功能是生物学的巨大挑战之一。
AlphaFold 和 MSA Transformer 等最先进的模型严重依赖于 多序列比对 (MSAs) 。 其基本直觉如下: 如果你想了解一个特定的蛋白质 (我们称之为查询序列) ，你会在海量数据库中搜索拥有共同祖先的其他蛋白质。然后将它们在一个网格中进行比对。如果序列第 10 位发生的突变总是伴随着第 50 位发生的突变，模型就会学习到这两个残基可能在三维空间中相互接触 (协同进化) 。
虽然有效，但这种方法有三个主要缺陷:
- 计算复杂度: 构建 MSA 需要将查询序列与数百万个序列进行比较。其复杂度大致为 \(O(LD)\)，其中 \(L\) 是蛋白质长度， \(D\) 是数据库大小。这非常慢。
- 僵化性: 它需要显式的“比对”步骤。如果一个蛋白质是“孤儿蛋白” (没有已知的亲属) 或“从头设计 (de novo) ” (由人类工程设计) ，MSA 方法就会失效。
- 存储: 它需要存储海量的隐马尔可夫模型 (HMM) 配置文件。
RSA背后的研究人员提出了一个发人深省的问题: 我们真的需要比对吗? 或者我们是否可以直接检索相似的序列，让深度神经网络去解决剩下的问题？
理论转变: 将 MSA 视为检索
这篇论文最有深刻见解的贡献之一是对问题进行了理论上的重构。作者认为，基于 MSA 的模型本质上只是一种特定的、僵化的检索增强语言模型 。
在 NLP 中，像 REALM 或 RAG 这样的模型通过在回答问题之前“阅读”来自维基百科的相关文档来改进其预测。这篇论文证明了 MSA Transformer 也在做同样的事情，只是受到生物学条件的限制。
我们可以将给定序列 \(x\) 预测蛋白质属性 \(y\) 的概率视为一个两步过程:
- 检索: 从数据库中找到相关的序列 \(r\)。
- 预测: 基于 \(x\) 和 \(r\) 进行预测。
在数学上，这看起来像这样:
作者拆解了现有的方法，展示了它们如何适应这个框架。传统的 MSA Transformer 根据比对得分选择序列，并使用列式注意力机制聚合它们。
如上表所示，提出的 RSA 方法改变了两个关键的设计支柱:
- 检索器形式: 它使用稠密检索 (Dense Retrieval) (向量相似度) ，而不是离散的 MSA 搜索。
- 比对形式: 它使用无比对 (No Alignment) , 而不是强制序列对齐。
核心方法: 检索序列增强 (RSA)
那么，RSA 实际上是如何工作的呢？其工作流程令人惊讶地优雅，并且模仿了现代搜索引擎。
1. 架构
该过程以“先检索后预测”的方式运行。
步骤 A: 稠密检索器 RSA 没有缓慢地扫描基因数据库寻找匹配项，而是预先对数据库建立了索引。它使用预训练的蛋白质语言模型 (具体是 ESM-1b) 将数据库中的每个蛋白质转换为稠密向量。 为了找到新查询的相关蛋白质，模型只需将查询编码为向量，并执行快速的最近邻搜索 (使用 Faiss) 。
相似度度量非常直接——即嵌入之间负的 L2 距离:
步骤 B: 增强编码器 一旦检索到前 \(K\) 个相关序列，模型不会尝试对齐它们。相反，它将查询序列 \(x\) 和检索到的序列 \(r\) 拼接成一个长的输入。
然后将这个组合序列输入到 Transformer 中。这就是自注意力机制 (Self-Attention) 发挥魔力的地方。注意力机制自然地允许模型“查看”检索到的序列，从而为查询序列收集上下文信息。
通过允许模型在处理查询 (\(H_x\)) 时关注检索到的序列 (\(H_r\))，网络学会了自动执行“软比对”。它无需人类设计的算法就能弄清楚检索到的蛋白质的哪些部分对应于查询序列。
为什么这很重要: 可解释性
你可能会想: “如果不使用进化比对，模型到底发现了什么？”作者分析了检索到的序列，发现稠密检索器捕获了两种截然不同的生物学知识: 同源性 (Homology) 和结构 (Structure) 。
检索同源性
尽管该模型使用向量相似度，但它成功地检索到了同源序列 (具有共同祖先的序列) ，这与传统的 BLAST 或 MSA 工具非常相似。
上图显示，对于大多数任务，稠密检索器发现的序列具有非常低的 E 值 (表明同源性的统计显著性很高) ，这与缓慢、传统的 MSA 方法相当。
检索结构
这正是 RSA 大放异彩的地方。有时，蛋白质在序列上看起来不同 (低同源性) ，但折叠成完全相同的三维形状。传统的 MSA 工具经常会错过这些“结构邻居”。然而，RSA 能找到它们。
可视化搜索结果让这一点变得清晰。在下图中，你可以看到左侧的查询蛋白质和右侧的检索结果。即使序列不同，三维折叠结构也惊人地相似。
实验结果
研究人员在一套标准的蛋白质任务上测试了 RSA，包括二级结构预测 (SSP)、接触预测和同源性预测。他们将 RSA 与普通 Transformer (如 ProtBERT) 和最先进的 MSA Transformer 进行了比较。
1. 性能对比 SOTA
结果令人印象深刻。RSA 不仅达到了基准水平，而且经常超越它们。
表 3 的一个关键亮点是, RSA (ProtBERT backbone) 在所有任务中取得了 0.723 的平均分，明显高于 MSA Transformer 的 0.672 。 它在不需要像 MSA Transformer 和 PMLM 那样昂贵的预训练步骤的情况下实现了这一点。
2. 对 “De Novo” (从头设计) 蛋白质的泛化
进化模型的最大弱点是当蛋白质没有历史时它们会失效。科学家们正越来越多地设计 de novo 蛋白质——自然界中不存在的合成蛋白质。MSA 工具对这些蛋白质返回的结果为空。
由于 RSA 依赖于向量嵌入空间而不是严格的序列匹配，它即使对于合成蛋白质也能找到“结构类似物”。
上面的散点图比较了 RSA 和 MSA Transformer 在 de novo 蛋白质上的表现。对角线以下的点表示 RSA 表现更好的蛋白质。正如你所见，RSA 在大多数情况下获胜。
我们还可以将这种改进可视化。在下图中，请看二级结构预测。与 MSA Transformer (下排) 相比，RSA (上排) 产生的预测更加清晰，且与真实值更加一致。
3. 消融实验: 比对是必须的吗？
为了解决理论争议，作者进行了一项消融研究。他们采用了标准的 MSA 序列，但没有先对其进行比对就输入到模型中 (未比对 MSA 增强) 。
结果 (表 5) 显示，去除比对仅导致性能轻微下降。这证实了论文的假设: 深度学习模型足够聪明，可以自己学习比对。 我们不需要手动将其喂给模型。
结论与未来启示
这篇“Retrieved Sequence Augmentation”论文为改变我们对蛋白质建模的方式提供了令人信服的理由。通过检索增强生成的视角看待蛋白质分析，作者开发了一种具有以下特点的方法:
- 更快: 绕过 \(O(LD)\) 的比对瓶颈，实现了高通量分析。
- 更简单: 不需要复杂的 HMM 配置文件或比对算法。
- 更鲁棒: 它适用于基于进化的方法失效的孤儿蛋白和合成蛋白。
这意味着，蛋白质语言模型的未来可能不在于更大的模型或更深层的进化挖掘，而在于更好的检索 。 就像搜索引擎改变了人类获取信息的方式一样，检索增强模型正在改变人工智能理解生命语言的方式。我们正在从僵化的、预计算的比对转向灵活、动态的生物知识查找。