语言模型革新同源蛋白检测，Nature Biotechnology一作聊聊论文里没有的细节！

笔者：brighthong

编辑：Yimin

新一代测序技术推动的生物序列数据库指数性增长，需要创新方法进行大规模分析。传统方法例如AlphaFold使用的HMMER在应对亿级数据库时常常力不从心，数以小时计的运行时间极大限制了其应用场景。为应对这一挑战，香港中文大学（CUHK）李煜教授领导的研究团队开创性地研发了DHR，这是一种旨在快速、灵敏地检测远缘蛋白质同源物的前沿方法。研究已发表在著名期刊Nature Biotechnology上。

相信大家在各个公众号上已经刷到了关于这篇论文的一些报道，对我们的论文有了一定的了解。那么，作为这篇论文的一作， 我将先为大家简单介绍这篇论文的模型架构，下游任务和主要优势，之后重点和大家介绍一下论文里没有的那些细节和经验！另外，预告一下，明天AIMShare将会有专稿文章介绍这文章背后的故事，包括项目的来源，投稿过程，心路历程等等，敬请期待。

1）模型框架

就像基于人类语言训练的ChatGPT能够理解语义一样，在大规模序列数据上预训练的蛋白质语言模型具有捕获进化信息的潜力。我们基于蛋白质语言模型和密集检索（Dense Retrieval）开发了Dense Homolog Retriever (DHR)方法。该系统利用先进的蛋白质语言模型将查询序列编码为嵌入式向量表示(Embedding)，并通过向量之间的相似度度量对数据库进行搜索和同源性比较。DHR整体采用了双编码器架构，训练时使用正反同源序列对进行基于对比学习的训练，拉近同源序列对之间的距离。在推理时，给定一个蛋白质序列，DHR可以快速检索出具有较多同源信息的序列。首先，DHR检索出的蛋白质同源序列具有极高的速度和灵敏度，此外，这些序列可以被用来构建MSA，用于各种下游任务，如加速Alphafold2的预测，提升蛋白质结构预测的性能。

2）主要优势: 速度与精度

DHR的相比于其他方法的主要优势在于其速度和精度。

在速度方面，蛋白质语言模型的一个优势在于其无需进行序列比对，在处理序列时速度极快。在工具开发时，我们也进行了一些工程优化，在训练时使用正反同源序列对进行训练，而推理前用数据库编码器离线编码，推理时仅使用查询编码器编码并直接与数据库编码进行比较。这使得 DHR的速度远远超过传统方法。在检索同源序列时，DHR的速度是传统方法（如PSI-BLAST和DIAMOND2）的22倍，而当将DHR与JackHMMER串联以加速迭代MSA构建过程时，DHR比传统方法快93倍，并且构建的MSA与AlphaFold2生成的MSA高度一致。

在精度方面，得益于蛋白质语言模型的强大表示能力，与传统方法相比，DHR的灵敏度提高了10%以上，在传统方法难以识别的难样例上，DHR的灵敏度更是提高了超过56%。此外，DHR与JackHMMER生成了更多样化和全面的MSA，在与AlphaFold2生成的MSA合并时，平均能提高0.4 Å的蛋白质结构预测精度。

DHR的初版叫fastMSA，是三年前一篇投稿RECOMB的工作。fastMSA的思路非常直接，因为AlphaFold使用HMMER构建MSA太慢了，所以我们想加速一下。当时2021年有了pLM和一些信息检索方面的进展，于是我们尝试将他们用在HMMER构建之前，先使用Dense Retrieval的方法筛选出最相似的200k个蛋白质序列，再使用HMMER构建MSA，通过缩小需要搜索的序列范围来加速搜索。

到这里为止，fastMSA的主要贡献是在蛋白数据库中用语言模型做信息检索（初步筛选），从而加速传统MSA构建方法。但是在优化fastMSA的过程中，我们发现该模型的能力其实远远超出我们的预期。经过性能调优之后之后的fastMSA可以直接得到高质量的同源序列，且能扩展到更大的数据库，于是我们决定增加更多训练数据，并将我们关注的重点从加速MSA构建转变到更为通用的同源序列检测， 以此为基础，我们开发了fastMSA的升级版本-DHR。

DHR的关注重心是同源序列的检索上，并在后续将它们构建为MSA进一步在下游任务测试其质量。在大量实验验证中，我们有两个意料外的发现：

在接下来几个小节中，我会详细阐述这两个意外产生的原因。首先我会先简单介绍同源的定义以及传统方法是如何检索同源的。然后我会回顾DHR的做法并详细解释。

1）什么是同源，怎么找同源

同源本身定义很简单，就是演化上来自同一祖先，或者说进化上相关。但是我们的序列数据库里并不提供这样的信息，尤其是我们拿到一条不知道是哪里来的序列，该如何找到他的同源序列？

https://en.wikipedia.org/wiki/Sequence_homology

早期的方法直接借鉴了NLP的手段，通过各种序列相似度的计算来获取。但是只关心相似度显然搞错了问题，我们希望找到的是同源序列，而不是相似序列。后续的方法将相似度定义更改为统计定义，即我们找到的同源序列在打分上是统计显著的。但这依旧不够形式化也不完全合理，因为这里仅仅说明我们搜索到的序列相比别的序列在打分上跟搜索序列更接近。当今作为金标准的HMMER更近一步定义了一个H与R的假说检验（H：序列相对查询是同源的，R：序列只是随机序列）。

这样在计算上的同源定义会直接导致的问题则为，同源与否是跟数据库相关的，即数据库里的其他序列也会影响同源序列检索的结果。

2）传统方法的努力和其极限

在同源序列检测方面，早期最有名的方法是BLAST，BLAST借鉴了NLP领域的n-gram方法，定义了生物序列的kmer。这样对于kmer的完全匹配显然是十分粗暴的，但同时也不得不承认，kmer能将检索变为倒表，极大提升了搜索效率。

后续PSI-BLAST引入了PSSM(位置相关打分矩阵)，把每个位点都独立出来，这样考虑到了全序列不同位置的保守性不同，突变方向不同。HMMER包将这一切更进一步定义为了HMM，能够更合理地处理插入删除，也将打分形式化为概率分布。

Durbin, R., Eddy, S. R., Krogh, A., & Mitchison, G. (1998). Biological Sequence Analysis: Probabilistic Models of Proteins and Nucleic Acids. Cambridge University Press. doi: 10.1017/CBO9780511790492

结合前一节同源性的定义，不难看出打分、HMM的构建是从序列数据库统计来的，同源性也是从数据库统计来的，加之HMM的表征能力有限，传统方法被其使用的数据库限制住了。即便他们能进行迭代，如果在最开始的一轮里没能找到合适的同源序列，后续无论如何都无法补救。从另一方面讲，如果最开始得到的同源序列存在偏见，那么在后续也无法修正。而最开始的一轮完全靠的是HMMER内置的打分。

3）DHR的尝试

DHR能够替代所有传统方法需要的第一轮打分，摆脱了对于HMMER内置打分的依赖。我们认为在高维流型（由蛋白质语言模型映射到的）上的相似度比较能够有效避免传统方法基于字母表频次的筛选，并且与数据库统计无关。虽然DHR本身并非完美，但是我们对其筛选能力有信心。在理论上，基于transformer的模型能够有全局的感知，能够发现远程的关联性，这是HMM模型所不具备的。再者，DHR的筛选与HMMER并不冲突，我们可以认为DHR的筛选结果给了HMMER一个较好的起始点，HMMER能够更好的进行第一步统计，并帮助DHR抑制假阳性。

回到两个意料之外的发现，原因已十分明晰：数据库本身就会影响同源序列，DHR的筛选并不仅仅是加速作用，它能够切实地辅助HMMER更好地迭代。传统的同源序列检测方法依赖于HMMER第一轮统计的结果，因此下游的MSA构建和蛋白质结构预测上的性能上限就是HMMER的结果。而DHR的结果不依赖于HMMER的第一轮统计，可以获得更多样的同源序列。基于DHR得到的同源序列构建的MSA和基于HMMER构建的MSA结合之后，带来了更多的信息量，自然能提升下游的蛋白质预测性能。

文章中我们使用了ESM系列模型，早期的fastMSA使用了旧版的ESM1，而最新的DHR已经升级到新版的ESM2。我们在早期实验阶段(fastMSA v1, 2021.8)就观察到ESM的能力大大超出了我们的预计。最开始我们只使用了两万数据进行微调，而那时的模型就已经在多数简单样例上表现不俗，这也促使我们去构建了更大的数据集。ESM2的能力相比ESM1又有一次提升，主要体现在编码信息密度上，用更低维度编码更多信息。现在ESM3发布了，相信ESM3能更加有效。

但是ESM模型本身也存在问题。首先，在长度超过800aa时，我们观测到召回性能存在显著的下降。其次，Dense模型在训练和推理阶段仍十分昂贵。再者，DHR的流程相比通常NLP领域做RAG其召回序列数量要多得多，这对Faiss一类相似度搜索引擎带来了挑战。DHR希望召回的是相对相似但同时多样的序列，如果只是找到了最相似的同质信息其实对下游任务帮助有限，当前DHR并没有显式处理这个问题，但是其编码空间的相似度比对已跳脱于序列相似度的绑定，在排序中增添了多样性，但仍旧不能完全保证远同源序列能不被大量相似序列掩盖，故我们希望在未来能够实现类似hhsuite中限定召回序列间相似度的功能。