超神经HyperAI · 6 天前

西湖大学团队开源SaProt等多款蛋白质语言模型,覆盖结构功能预测/跨模态信息搜索/氨基酸序列设计等

微信图片_20250407174509.png
2025 年 3 月 22—23 日,上海交通大学「AI 蛋白质设计峰会」正式举行。 本次峰会汇聚了来自清华大学、北京大学、复旦大学、浙江大学、厦门大学等知名高校的 300 多位专家学者,以及 200 余位行业领军企业代表和技术研发人员,深入探讨了 AI 在蛋白质设计领域的最新研究成果、技术突破及产业化应用前景等。

在这里插入图片描述

西湖大学原发杰博士在峰会上带来深度分享

在峰会期间,西湖大学原发杰博士以「蛋白质大语言模型研究与应用」为主题,向大家分享了蛋白质语言模型的最新研究进展,并详细介绍了团队的重要成果, 包括蛋白质语言模型 SaProt、ProTrek、Pinal、Evolla 等。HyperAI超神经在不违原意的前提下,对其深度分享进行了整理汇总,以下为演讲精华实录。

值得关注的蛋白质语言模型

蛋白质是由 20 种氨基酸串联而成的生物大分子,它们在体内承担催化、代谢等关键功能,是生命活动的主要执行者。生物学家通常将蛋白质的结构划分为 4 个层级:一级结构描述蛋白质的氨基酸序列、二级结构关注蛋白质局部构象、三级结构代表蛋白质整体的三维构型、四级结构则涉及多个蛋白质分子间的相互作用。在 AI 蛋白质领域,主要是基于这些结构来展开研究。

在这里插入图片描述

例如,从蛋白质序列预测其三维结构,这正是 AlphaFold 2 解决的核心问题——攻克了困扰科学界长达 50 年的蛋白质折叠问题,并因此荣获诺贝尔奖。反之,根据结构和功能设计新的蛋白质序列,即 Protein Design,这一领域的重要贡献者 David Baker 教授同样也获得了诺贝尔奖。

传统上,蛋白质结构通常以 PDB 坐标的形式来表示,近年来,研究人员们探索了将连续空间结构信息转化为离散 token 的方法,比如 Foldseek、ProTokens、FoldToken、ProtSSN、ESM-3 等。
*Foldseek 可以将蛋白质的三维结构编码成一维的离散 token。

我们团队的蛋白质语言模型就是基于这些离散结果来展开研究。

大多数 AI + 蛋白质的研究都可以追溯到自然语言处理的研究,所以我们先来回顾一下自然语言处理 (NLP) 领域的两个经典语言模型:一个是以 GPT 系列为代表的单向语言模型, 其机制是基于从左到右的信息流动方式,根据左侧的数据(上文)预测下一个 token。一个是以 BERT 为代表的双向语言模型, 其通过 Masked Language Model 做预训练,可以看到某个 cooked 左右两侧的信息(上下文),从而预测该 cooked。

在这里插入图片描述

在蛋白质领域,这两类模型都有其对应的蛋白质语言模型。 比如,对应 GPT,有 ProtGPT2、ProGen 等。 对应 BERT,有 ESM 系列模型:ESM-1b、ESM-2、ESM-3,它们主要是通过遮盖一些氨基酸,并预测其「真实身份」,在自然语言任务中,也就是遮住一些单词,然后预测这些单词。如下图左侧所示,在蛋白质社区影响力比较大的语言模型还有 MSA Transformer、GearNet 、ProTrans 等。

在这里插入图片描述

入选 ICLR 2024,蛋白质语言模型 SaProt 融合结构知识

我要向大家介绍的第一个成果是具有结构感知词汇的蛋白质语言模型 SaProt, 这篇论文以「SaProt: Protein Language Modeling with Structure-aware Vocabulary」为题曾入选 ICLR 2024。

这篇论文中,我们提出了结构感知词汇 (structure-aware vocabulary) 的概念,并将氨基酸残基 token 与结构 token 结合,在约 4,000 万条蛋白质序列与结构数据集上训练了一个大规模通用蛋白质语言模型 SaProt,该模型在 10 项重要的下游任务中全面超越已有的成熟基线模型。
SaProt 开源地址:
https://github.com/westlake-repl/SaProt
SaProt 论文地址:
https://openreview.net/forum?id=6MRm3G4NiU

我们为什么要做这个模型?

其实,大多数蛋白质语言模型的输入信息主要基于氨基酸序列。在 AlphaFold 取得突破后,DeepMind 团队与欧洲生物信息研究所 (EMBL-EBI) 合作,发布了 AlphaFold 蛋白质结构数据库 (AlphaFold Protein Structure Database),其中存储了 2 亿的蛋白质结构。于是我们开始考虑:能否将蛋白质的结构信息整合到语言模型中,以提升其性能表现?

在这里插入图片描述

我们的做法非常简单:采用 Foldseek,将蛋白质的结构信息从坐标形式转换为离散 token,从而构建氨基酸词表和结构词表,然后对这两个词表进行交叉组合,生成新的词表,也就是结构感知词汇 (Structure-aware Vocabulary, SA token)。这样一来,原本的氨基酸序列就可以转化为一种新的氨基酸序列——在这个序列中,大写字母代表氨基酸 token,小写字母代表结构 token,然后,就可以继续做 Masked Language Model 了。基于此,我们训练了一个 6.5 亿参数的 SaProt 模型,该模型使用了 64 张 A100 GPU,总训练时间约 3 个月。

在这里插入图片描述

我们为什么选择 Foldseek 转化蛋白质结构 token?

在最终确定使用 Foldseek 3Di token sequence 之前,我们经历了长达半年的探索。直觉上,将结构信息融入蛋白质语言模型应该能带来性能提升,但实际尝试时我们用了多种方法都未成功。比如,我们用了 GNN 方法来建模蛋白质结构。因为蛋白质结构其实是一个图神经网络,我们自然而然就想把蛋白质结构当做图来建模,于是采用了 MIF 方法,但发现训练的模型泛化能力较差,无法推广至真实的 PDB 结构上。我们深入分析后认为,这可能是因为采用 Masked Language Model 的建模方式会出现信息泄露问题。

简单来讲,AlphaFold 预测的蛋白质结构本身存在某种偏差 (bias)、模式、AI 预测的痕迹。当用这些数据去训练语言模型时,模型很容易就能捕捉到这些痕迹,从而导致模型在训练数据上表现很好,但泛化能力极差。

在这里插入图片描述

我们尝试了各种改进,包括使用 Evoformer 方法,但信息泄露问题依然存在,直到尝试了 Foldseek,我们发现得到的 SaProt 模型在 AlphaFold 预测的结构数据上 loss 能够下降,在真实 PDB 结构数据上 loss 同样显著下降,满足了我们的预期。

此外,SaProt 在多个基准测试中表现良好, 去年,它也曾在权威榜单 ProteinGym 上排名第一。与此同时,我们也搜集到社区对 SaProt /ColabSaProt 进行了超过 10 种蛋白质的湿实验验证(如各种酶突变改造,荧光蛋白质改造与荧光性预测等)结果,均表现优异。

在这里插入图片描述

虽然我们认为 SaProt 模型还不错,但考虑到许多生物学家并未接受深度学习相关的训练, 让他们独立微调一个规模达 10 亿参数左右的蛋白质语言模型是非常困难的,于是我们构建了一个交互界面式平台 ColabSaprot + SaprotHub。

在这里插入图片描述

在传统的模型训练流程 (Normal pipeline) 中,用户需要经历数据准备、环境配置、代码编写、数据处理、模型训练、模型评估、模型推理等多个环节。而有了 ColabSaprot 之后,整个过程被大大简化——用户只需点击几下按钮,即可完成环境安装、模型训练和预测等操作,大大降低了使用门槛。

如下图所示,ColabSaprot 主要由训练模块、预测模块和共享模块 3 部分组成。

在这里插入图片描述

  • 在训练模块中,用户只需在左侧描述任务并上传数据,然后点击训练即可。系统会自动选择最优的超参数(如 Batch Size 等)。
  • 在预测模块中,用户可以直接加载自己之前训练的模型,并进行预测。也可以直接录入其他研究人员分享的模型去做预测。
  • 共享模块提供了一种既保护数据隐私、又能进行成果共享的方式。许多实验室的数据极为珍贵,一些研究人员可能还需要用这些数据做后续研究,但他们仍希望分享现有模型。在 ColabSaprot 中,用户可以仅共享模型本身,由于模型本质上是一个黑盒,其他人也无法获取原始数据。

在共享模型时,考虑到语言模型通常规模庞大,直接在线分享一个 10 亿参数的模型几乎是不可行的。因此,我们采用了一种成熟的 Adapter 机制, 用户只需共享极少量的参数——通常仅占原始模型参数的 1% 或 1/1,000 即可。大家可以互相 share Adapter,并加载其他人的 Adapter,在此基础上进行微调或预测,如果改进效果良好,还可以再次分享新的 Adapter,从而形成一种高效的社区合作机制,大大提升了研究效率。

在这里插入图片描述

此外,我们也进行了用户研究 (User Study), 邀请了 12 位没有机器学习背景、编程基础的同学试用 ColabSaprot 平台。我们提供数据,并告诉他们要做的任务,要求他们使用 ColabSaprot 进行模型训练和预测。最终,将他们的结果与 AI 专家的表现进行对比,我们发现这些非专业用户使用 ColabSaprot 后,能够达到与专家相近的水平。

另外,为了推动蛋白质语言模型共享,我们还建立了一个名为 OPMC 的社区, 该领域的国内外知名学者参与其中,鼓励大家共同分享模型,促进合作与交流。
OPMC 地址:

https://theopmc.github.io/

在这里插入图片描述

ProTrek 模型:可查找蛋白质序列、结构、功能之间的对应关系

我要介绍的第二个工作是蛋白质语言模型 ProTrek。

在生物研究中,许多科学家面临这样的需求:手头有一个基因组的很多蛋白质,但并不清楚它们的具体功能。

ProTrek 是一个三模态语言模型,可用于对序列、结构和功能进行对比学习。 借助其自然语言搜索界面,用户能够在几秒内探索庞大的蛋白质空间,并搜索序列、结构、功能所有两两组合的 9 种不同任务之间的关系。也就是说,使用 ProTrek,用户只需输入蛋白质序列并点击按钮,即可快速查找蛋白质功能、结构相关信息。同样,也可以根据功能找序列、结构信息,根据结构查找序列、功能信息等。此外,它还支持 sequence—sequence、structure—structure 类的查找。
ProTrek 使用地址:

http://search-protrek.com/

在这里插入图片描述

我们合作者对 ProTrek 模型进行了干、湿实验的评估, 与现有相关方法相比,ProTrek 在性能上取得了显著提升。此外,我们还利用 ProTrek 生成了大量数据,用于训练我们的生成模型,效果同样表现优异。

在这里插入图片描述

在这里插入图片描述

我们在 Twitter 上注意到,许多用户已经开始使用 ProTrek 参与竞赛, 也收到了不少积极反馈,进一步证明了该模型的实用性。

在这里插入图片描述

Pinal 模型:只需输入文本,即可设计全新蛋白质序列

我们的另一个工作是基于文本描述设计蛋白质的模型 Pinal。

传统的蛋白质设计通常需要考虑复杂的因素,如生物物理学的能量函数模板信息等。我们想探讨的是,既然大语言模型在许多任务中表现出色,是否可以设计一个基于文本的蛋白质语言模型?在这个模型中,我们只需要简单地描述某个蛋白质的信息,就能设计其氨基酸序列?
Pinal 使用地址:
http://www.denovo-pinal.com/
论文地址:
https://www.biorxiv.org/content/10.1101/2024.08.01.606258v1

在这里插入图片描述

我简单介绍一下 Pinal(160 亿参数)的基本原理。 最初,我们的思路是采用编码器-解码器架构,输入文本,然后输出氨基酸序列。但尝试了很长时间,效果一直不理想,主要原因是氨基酸序列的空间过于庞大,导致预测困难。

因此,我们调整了策略,首先设计蛋白质的结构,然后根据结构和文本提示来设计氨基酸序列,这里的蛋白质结构也采用离散化编码进行表示。结果发现,结合结构的设计方法,在 Next Token Prediction Accuracy 上的表现明显优于直接预测氨基酸序列的方法,如下图所示。

在这里插入图片描述

我们近期收到合作者对 Pinal 的湿实验验证, Pinal 设计了 6 条蛋白质序列,其中 3 条表达,2 条序列经过验证具有相应的酶催化活性。值得一提的是,在这个工作中,我们并没有强调一定要设计出比野生型更优秀的蛋白质,我们的主要目标是验证根据文本设计出来的蛋白质是否具有相应的蛋白质功能。

Evolla 模型:解码蛋白质分子语言

最后介绍的成果是 Evolla 模型, 这是一个 800 亿参数的蛋白质语言生成模型,也是最大的开源生物大模型(之一),旨在解码蛋白质的分子语言。

通过整合蛋白质的序列、结构和用户查询信息,Evolla 能够生成准确的蛋白质功能见解。 用户只需输入蛋白质的序列和结构,然后提出问题,比如介绍该蛋白质基本功能或催化活性,直接点击按钮,Evolla 就会生成约 200-500 字的详细描述。
Evolla 使用地址:

http://www.chat-protein.com/
Evolla 论文地址:

https://www.biorxiv.org/content/10.1101/2025.01.05.630192v2

在这里插入图片描述

值得一提的是,Evolla 这个工作训练数据和算力非常巨大,仅仅训练数据收集和处理我们就有两位博士生花了近 1 年时间,最终我们通过合成数据手段,产生了超过 5 亿条高质量蛋白质-文本对数据,覆盖千亿级 word token。模型在酶的功能预测上表现相当准确,但也不可避免地存在一些幻觉问题。

在这里插入图片描述

团队介绍

西湖大学原发杰博士主要从事传统机器学习及交叉学科相关的应用科学研究,并专注于 AI 大模型与计算生物学的探索,在机器学习与人工智能领域的顶会和期刊上(如 NeurIPS、ICLR、SIGIR、WWW、TPAMI、Molecular Cell等)发表了 40 余篇学术论文,团队成员及项目贡献者信息详见论文。

课题组长期开展机器学习和 AI + 生物信息方向的研究,欢迎申请课题组博士生、科研助理、博士后、研究员系列职位,在校生欢迎来实验室访问实习。有意者可发简历到邮箱 yuanfajie@westlake.edu.cn

推荐阅读
关注数
704
文章数
571
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息