美国国家工程院外籍院士沈向洋曾强调:「如果说今天有什么事是我们一定要做的,那就是 AI for Science。难以想象今天还有什么事情比它更重要,今年诺贝尔奖的颁布便是最好证明。」
过去,科学家依赖手动整理资料、基于学科理论提出假设;如今,由 AI 辅助、直接基于海量数据进行研究,AI for Science 不仅提升了科研效率,更改变了整个科研范式,这一点在蛋白质研究领域尤为明显。
在第五期 Meet AI4S 直播中,HyperAI超神经有幸邀请到了浙江大学知识引擎实验室的博士研究生王泽元, 他以「借助扩散去噪过程助力大模型对蛋白质的优化」为题,向大家详细介绍了团队入选 NeurIPS 2024 的一篇论文「DePLM: Denoising Protein Language Models for Property Optimization」。
NeurIPS 作为 AI 领域的顶级会议,被称为难度最大、水平最高、影响力最强的 AI 学术会议之一,今年大会共收到 15,671 篇有效论文投稿,比去年增长 27%,但最终接收率仅有 25.8%,入选论文均具有极大的学习价值。本次分享,王泽元博士详细介绍了去噪蛋白质语言模型 DePLM 的设计理念、实验结论、demo 运行方式及未来展望,此外,他还分享了顶会投稿经验,希望对大家有所帮助。
具体而言,王博士表示,当我们投稿时,可以从选题、创新点、论文写作、应对跨学科评审等方面入手。
首先,在选题方面, 可以广泛阅读顶会文章,了解目前社区关心的、比较有意义的研究方向。例如,准备 DePLM 这篇论文时,王博士就发现,在去年的 ICLR 和 NeurIPS 会议上,蛋白质工程特别是蛋白质预测任务就是一个热门方向。
其次,在创新点方面, 他认为培养发现问题的能力很重要,在 AI for Science 领域,我们应当先深入了解 Science 领域的知识,并将其与 AI 领域的内容进行对照,找出尚未被 AI 探索的空白地带。
在论文写作方面, 他表示,写作必须逻辑清晰、详略得当,确保文章易于理解,还要多与导师和同学交流,避免陷入自我的思维定式。
最后,考虑到 AI for Science 类论文可能会受到来自两种不同背景评审人的审阅,即一方更注重 AI 技术,另一方则关注 Science 应用,因此在写作时需要明确论文的核心定位, 即这篇论文是面向 AI 社区还是 Science 社区,并据此构建逻辑框架,确保内容紧扣主题。
在他看来,目前的大模型研究趋势已经发生了变化,我们从简单的拿来主义,转向了对大模型的深入理解。 以前我们让大模型去适配各种下游任务,而现在更多的是考虑如何让下游任务更好地配合大模型的预训练阶段,两者之间的契合度越高,模型的表现也就越好。
比如,对于预测适应性景观来说,传统的简单微调方法在泛化能力上表现不佳,那我们需要更深入地理解大模型以及无监督学习范式,识别其中的不足并加以改进。此外,我们还必须关注大模型自身的缺陷,比如探索消除模型偏见的方法,以优化模型的性能。
模型开源可测试
今天我想分享我们在 NeurIPS 2024 发表的一篇论文,即探讨如何借助扩散去噪模型来助力大语言模型对蛋白质的优化。在这篇论文中,我们提出了一个全新的去噪蛋白质语言模型 (DePLM), 其核心在于将蛋白质语言模型捕捉到的进化信息看作是与目标特性相关和无关的混合体,其中无关信息被视为「噪音」并消除。我们发现,所提出的基于排序的去噪过程能显著提高蛋白质的优化性能,还能保持强大的泛化能力。
目前,DePLM 已开源,由于模型的配置环境较为复杂,我们在 HyperAI超神经官网的教程板块上线了「DePLM:用去噪的语言模型优化蛋白质(小样本)」。 为了帮助大家更好地理解和复现我们的工作,我将从 DePLM 模型如何运行,其相关配置文件分别是什么;如何微调模型的扩散步骤;如何用自己的数据集运行 DePLM 模型等几个方面来向大家讲解模型的运行方式。
DePLM 开源地址:
https://github.com/HICAI-ZJU/DePLM
DePLM 教程地址:
背景介绍:最大化利用进化信息,最小化引入数据偏差信号
本文的研究对象是蛋白质,这是由 20 种氨基酸串联组成的生物大分子,在体内承担着催化、代谢及 DNA 复制等功能,也是生命活动的主要执行者。生物学家通常将其结构分为 4 个层级。第一级描述蛋白质是如何组成的;第二级描述蛋白质的局部结构,如常见的 α 螺旋和 β 折叠;第三级描述蛋白质的整体三维结构;第四级则考虑蛋白质间的相互作用。
目前,大多数 AI+蛋白质的研究都可以追溯到自然语言处理的研究,这是因为两者之间存在相似性,比如,我们可以将蛋白质的四级结构与自然语言中的字母、单词、句子、段落相对。当一句话中出现字母错误,这句话就会失去意义,同样地,蛋白质的氨基酸发生突变也可能导致蛋白质无法形成稳定结构,进而丧失功能。
如下图所示,在「Controllable protein design with language models」这篇论文中,研究人员将自然语言与蛋白质进行对应,这种方式受到了研究人员们的广泛认可,自 2020 年以来,蛋白质研究呈现出爆发式增长。
论文原文:
https://www.nature.com/articles/s42256-022-00499-z
我们这次讨论的问题是 AI+蛋白质优化,即如果我们拥有一个功能不及预期的蛋白质,如何调整其氨基酸序列以满足预期的功能。
在自然界中,蛋白质通过随机改变(包括点的插入、删除或点突变等)不断自我优化。生物学家通过模仿此过程,提出了定向进化 (Directed evolution) 和深度突变扫描 (Deep Mutational Scanning),让其对蛋白质进行优化,这两种方法的问题是,实验资源消耗过大。因此,我们使用计算的方法来模拟蛋白质与其属性适应度之间的关系,即预测适应性景观 (fitness landscape),这对蛋白质优化至关重要。
为了对这个问题进行建模,我们通常会用到数据集、评估指标 (Metrics)、计算方法。 如下图所示,蛋白质优化数据集通常包含一条野生型序列 xwt,多个突变对 μi、及突变后预测的 fitness 值 yi。评估模型等主要依赖斯皮尔曼相关系数 (Spearman’s correlation coefficient),这个指标不会关注预测的具体值,而是关注因突变导致 fitness 值变化的排序,真实突变 R(Y) 与预测的 fitness score 排序值越接近,则模型训练的越好。
计算方法大致可分为监督学习 (Supervised modeling) 和无监督学习 (Unsupervised modeling)。监督学习依赖于有标签的数据,通过优化损失函数来训练模型,以提高对 fitness 的预测能力。无监督学习则不需要标签数据,而是在与 fitness 无关的大规模蛋白质数据集上进行自监督学习,模型只需训练一次,就可以泛化到各种蛋白质的预测任务中。
例如,掩码语言建模是一种无监督学习方法,训练模型时,我们要为模型提供一个被污染的序列,可以把某个单词 mask(如下图红色框里的单词),或者把它随机变化为别的单词,让语言模型对其进行恢复,即恢复原有的序列。在 NeurIPS 2021 的一篇论文中,研究人员发现,此类模型预测蛋白质突变的概率与适应性景观 (fitness landscape) 存在一定的相关性,为此,他们设计了 4 种突变评分公式,如下图右侧所示。
论文原文:
https://proceedings.neurips.cc/paper/2021/file/f51338d736f95dd42427296047067694-Paper.pdf
总结来说,监督方法表现出色但泛化能力有限,无监督方法虽然表现稍逊,但泛化能力强。 为了结合两者的优点,如下图所示,我们借鉴 NLP 领域的 Pre-train + Fine-tune 策略,做了一些尝试后发现,这种方法虽然表现良好,但泛化能力差,和监督学习类似。于是我们分析了无监督方法为什么具有优异的泛化能力,并假设这种泛化能力来自进化信息(Evolutionary Information, EI),这是因为,生物体可以通过自然进化来优化蛋白质,这种进化的突变也会被保留。于是,我们考虑,突变概率与 fitness landscape 的相关性是正相关。
然而,当我们尝试微调模型时,实际上使用的是 embedding 信息,并没有充分利用进化信息,此外,湿实验数据中还存在无关信息的偏差。我们认为,进化信息包含各个方向的综合信息,比如 stability、activity、expression、binding 等方向,当我们对蛋白质的 stability 进行优化时,activity、expression、binding 的进化就属于不相关信息,若能去除这种不感兴趣信息的概率值,则可提高模型的表现,又因为整个过程在 likelihood 空间中进行,也不会影响模型的泛化能力。因此,我们需要在微调过程中,最大化利用进化信息,同时最小化引入数据集中的偏差信号。
DePLM 算法框架:基于排序空间的去噪模型
基于此,我们提出了 DePLM 模型,其核心观点在于,将蛋白质语言模型捕获到的进化信息视为感兴趣与不感兴趣信号的融合体,后者被视为目标属性优化任务中的「噪声」,需要消除。DePLM 通过在属性值的排序空间中进行扩散过程来去噪进化信息,从而增强模型的泛化能力,预测突变效应。
给定蛋白质的一个氨基酸序列,模型会预测每个位置突变成各种氨基酸的概率,进化似然再通过 Denosing Module 生成感兴趣的性质的 likelihood。具体而言,DePLM 主要包括两个部分:正向扩散过程、学习的反向去噪过程。 在正向过程中,逐步向真实情况添加少量噪声,在逆向去噪过程中,学习逐步消除累积的噪声,恢复真实情况。
如下图所示,DePLM 基于 ESM 系列,采用 Transformer 架构。其 Denosing Module 模块基于扩散过程训练,网络架构包括 Feature Encoder 和 Denosing Block,其中,Feature Encoder 从蛋白质语言模型 (Protein Language Model) 中提取序列特征,并通过 ESM 1v 模型抽取结构特征,将这两种特征作为锚点,通过多轮 Denoising Block 迭代逐步去噪获得 Denoised likelihood。
过去,去噪方法多应用于图像生成领域,尤其是在扩散模型 (Diffusion Model) 中。如下图所示,主要是通过定义好的加噪过程,将原始图片 x0 转换至接近高斯分布的噪声空间中 (xT),然后模型学习反向去噪过程。
然而,将图像去噪模型直接应用在蛋白质领域会存在一些问题。如上图所示,图像去噪模型可以加入随机噪声,进而形成不可分的噪声空间(从 x0 到 xT),但是,蛋白质存在 Fitness score 和进化似然 (Evolutionary likelihood),始末状态已被固定,因此,加噪过程需精心设计。其次,模型会向 fitness score 对齐,造成表现虽好、但泛化能力差的结果。
于是我们提出基于排序空间的去噪模型,专注于最大化排序相关性。 这是因为,我们希望让 Evolutionary likelihood 去噪至感兴趣的 property 空间,虽然不知道该空间的具体情况,但我们知道,其排序与 fitness 排序是一致的。
我们在这个空间里加噪,让模型学习大量数据集,逐渐学会 Denosied likelihood 应有的样子,而不是直接对齐 fitness score。在这个前向加噪过程中,我们利用了一种排序算法,使排序的每一步都更接近最终状态,且包含随机性,模型也会学习反向的逐步排序思想。具体而言,如下图所示,如果我们有了 xt-1,就可以给排序算法喂一个 xt-1 和 xT 让它多次排序,当得到第 t 步的排序空间后,可以从中随机采样第 t 步的排序变量,让模型从 t+1 步预测到第 t 步的 likelihood,并计算斯皮尔曼损失。由于我们不需要像图像去噪那样添加许多步骤,排序过程通常 5-6 步即可完成,也提高了效率。
实验结论:DePLM 性能优越且具备强大的泛化能力
为了评估 DePLM 在蛋白质工程任务中的性能,我们在 ProteinGym、β-lactamase(β-乳酸酶)、GB1、Fluorescence(荧光)数据集上,将其与从零开始训练的蛋白质序列编码器、自监督模型等进行比较。结果如下图所示,DePLM 的表现优于基线模型,通过对比,我们发现,高质量的进化信息能显著改善微调后的结果,这说明我们所提去噪训练过程的有效性,也证实了将进化信息与实验数据集成在蛋白质工程任务中的优势。
DePLM 与基线模型在蛋白质工程任务中的表现 最佳结果和次优结果在表中分别用粗体和下划线标出
接着,为了进一步评估 DePLM 的泛化能力,ProteinGym 根据它们测量的蛋白质属性将 DMS 数据集分为 5 个类别,分别是稳定性、适应度、表达、结合和活性。我们将其与其他自监督模型、基于结构的模型以及监督基线模型进行比较。结果如下图所示,DePLM 优于所有基线模型,这说明仅依赖未过滤进化信息的模型存在不足,其往往由于同时优化多个目标而稀释了目标属性,通过消除无关因素的影响,DePLM 显著提高了性能。
泛化能力评估 最佳结果和次优结果在表中分别用粗体和下划线标出
为了进一步分析泛化性能,确定过滤掉与属性无关的信息的重要性,我们进行了性质间训练与测试的交叉验证,如下图所示,在大多数情况下,当模型在性质 A 上训练而在性质 B 上测试时,其表现相较于在同一性质(即 A)上进行训练和测试有所下降。这表明不同性质的优化方向并不一致,存在相互干扰,证实了我们最初的假设。
此外,我们发现,在其他性质数据集上训练,在 Binding 数据集进行测试,模型性能有所提升。这可能归因于 Binding 数据集的数据量有限以及数据质量不高,导致其自身泛化能力不足。这启发我们,对于新性质的蛋白质优化时,如果该性质相关的数据集较少,可以考虑利用相关性质的数据进行去噪和训练,从而获得更好的泛化能力。
持续深耕蛋白质领域
本次直播分享的嘉宾是浙江大学知识引擎实验室的博士研究生王泽元,他所在的团队在陈华钧教授、张强研究员等人的带领下,致力于知识图谱、大语言模型、AI for Science 等领域的学术研究,在 NeurIPS、ICML、ICLR、AAAI、IJCAI 等 AI 顶会上发表了多篇论文。
张强个人主页:
https://person.zju.edu.cn/H124023
在蛋白质领域,团队不仅提出 DePLM 这样的先进模型来优化蛋白质,还致力于弥补生物序列与人类语言之间的鸿沟。为此,他们提出 InstructProtein 模型, 利用知识指令对齐蛋白质语言与人类语言,探索蛋白质语言和人类语言间的双向生成能力,将生物序列整合到大语言模型中,有效弥补两种语言之间的差距。大量双向蛋白质-文本生成任务的实验表明,InstructProtein 在性能上优于现有的最先进 LLMs。
点击查看更多详情:入选ACL2024主会 | InstructProtein:利用知识指令对齐蛋白质语言与人类语言
此外,团队还提出了一种基于 「预训练和提示」框架的多用途蛋白质序列设计方法 PROPEND。 通过对骨架、蓝图、功能标签及其组合的提示,可以直接控制多种属性,该方法具备广泛的实用性和准确性。在体外实验测试的 5 个序列中,PROPEND 的最大功能恢复率达到了 105.2%,显著超过了经典设计管道的 50.8%。
论文原文:
https://www.biorxiv.org/content/10.1101/2024.11.17.624051v1
目前,团队发布的多项成果均已开源,他们还长期招聘优秀博后、百人、研发工程师等各类专职研究人员,欢迎大家加入~
实验室 Github 主页:
http://github.com/zjunlp