作者:梅菜
编辑:十九,李宝珠
转载请联系本公众号获得授权,并标明来源
加州大学伯克利分校、微软研究院等提出了一种多模态蛋白质生成方法 PLAID ,可以从更丰富的数据模态 (例如序列) 生成较稀缺的模态 (例如晶体结构) 来实现多模态生成。
过去多年间,科学家们持续致力于探索蛋白质的结构和组成,以更好的解谜「生命密码」。蛋白质功能由其结构决定,包括侧链和主链原子的身份、位置及其生物物理性质,这些统称为全原子结构。然而,想要确定侧链原子的放置位置,必须先了解序列。因此,全原子结构生成可以看作是一个多模态问题,需要同时生成序列和结构。
但是,现有蛋白质结构和序列生成方法通常将序列和结构视为独立的模态,结构生成方法通常仅生成主链原子;针对全原子设计的方法通常需要借助外部模型在结构预测和反折叠步骤之间交替进行等等……
为应对这些挑战,加州大学伯克利分校 (UC Berkeley) 、微软研究院以及 Genentech 公司的研究团队提出了一种多模态蛋白质生成方法 PLAID (Protein Latent Induced Diffusion),可以从更丰富的数据模态 (例如序列) 映射到较稀缺的模态 (例如晶体结构) 来实现多模态生成。为验证该方法的有效性,研究人员对来自 Gene Ontology 的 2,219 种功能和遍布生命树的 3,617 种生物体进行了实验。尽管训练过程中未使用结构输入,生成的样本表现出很强的结构质量和一致性。
相关研究以「Generating All-Atom Protein Structure from Sequence-Only Training Data」为题,已提交至顶会 ICLR 2025。「AI 教父」杨立昆也在社交平台上转发了该成果。
PLAID 项目开源地址:
github.com/amyxlu/plaid
研究亮点:
* 研究人员专注于大型蛋白质语言模型 ESMFold 和全原子结构生成,提出了一种可控扩散模型,能够同时生成序列和全原子蛋白质结构,在训练过程中仅需要序列输入。
* 该方法利用了预训练权重中编码的结构信息,而非训练数据,并增加了用于可控生成的序列注释的可用性。
* 尽管论文中使用了 ESMFold 模型,但该方法可以应用于任何预测模型。
论文地址:
https://www.biorxiv.org/content/10.1101/2024.12.02.626353v1
开源项目「awesome-ai4s」汇集了百余篇 AI4S 论文解读,并提供海量数据集与工具:
https://github.com/hyperai/awesome-ai4s
研究精华速览
数据集
研究人员使用了 2023 年 9 月发布的 Pfam 数据库,该版本包含 57,595,205 个序列和 20,795 个家族。PLAID 完全兼容更大的序列数据库,例如 UniRef 或 BFD (约 20 亿个序列),然而该研究依然选择使用 Pfam,因为其序列域包含更多的结构和功能标签,这使得对生成样本的计算机模拟评估更加便捷。此外,研究人员还保留了大约 15% 的数据用于验证。
从 Pfam FTP 服务器提供的 Pfam-A.fasta 文件中可以获得 Pfam 域来源的生物体的 UniRef 代码。研究人员分析了数据集中所有的唯一生物体,总计发现了 3,617 个不同的生物体,然后对这些生物体进行实验以验证 PLAID 方法的有效性。
模型架构
PLAID 是一种通过在预测模型的潜在空间中扩散来实现蛋白质多模态、可控生成的新范式,方法概览如下图所示,简而言之分为 4 个步骤:
图:PLAID 概览
(A) ESMFold 潜在空间 (ESMFold Latent Space):潜在空间 p(x) 表示序列和结构的联合嵌入。
(B) 潜在扩散训练 (Training):目标是从 pθ(x) 中学习并进行采样,遵循扩散公式。为了提高学习效率,研究人员使用 CHEAP 编码器 he(·) 来获得压缩嵌入 x′ = he(x),使得扩散目标变为从 pθ(he(x)) 中进行采样。
(C) 推理 (Inference):为了在推理时同时获得序列和结构,研究人员使用训练好的模型来采样 ˜x′ ∼ pθ(x′),然后使用 CHEAP 解码器解压得到 ˜x = hd(˜x′)。该嵌入通过在 CHEAP 中训练的冻结序列解码器解码为相应的氨基酸序列。残基身份序列和 ˜x 被用作输入,传递给在 ESMFold 中训练的冻结结构解码器,得到全原子结构。
(D) DiT 块架构 (DiT block architecture):研究人员使用扩散变换器 (DiT) 架构,结合 adaLN-zero DiT 块来融合条件信息。使用无分类器指导将功能 (即 GO 术语) 和生物体类标签嵌入 。
研究结果
研究人员进行了不同蛋白质长度的结构质量和多样性分析,结果如下图所示,天然蛋白质和 PLAID 生成的样本在不同长度上保持一致的指标,而 ProteinGenerator 和 Protpardelle 出现了特定长度的模式崩塌 (mode collapse),Multiflow 则在较长序列上表现出多样性下降。
图:不同蛋白质长度的结构质量和多样性分析
* 该图对比天然蛋白质和不同生成方法,展示了蛋白质在不同长度 (64-512 个残基) 下的结构质量 (ccRMSD,青色点) 和多样性 (紫色线,以独特结构簇占总样本比例衡量)。红线位于 2Å,表示设计性阈值)
此外,与基准方法相比,PLAID 生成的二级结构多样性更接近于天然蛋白质的分布。如下图所示:ProteinGenerator、Protpardelle 和 Multiflow 在其二级结构分布上表现出偏差,现有蛋白质结构生成模型通常难以生成具有高 β-折叠含量的样本。
图:不同蛋白生成方法的二级结构组成分析
* 该图展示了天然蛋白质及不同生成方法中蛋白质结构的 α-螺旋和 β-片层含量分布。每个点表示一个结构,其坐标表示 α-螺旋残基比例 「x 轴」和 β-片层残基比例「y 轴」)
研究人员还比较了不同模型在全原子蛋白生成任务中跨多种一致性 (Consistency) 和质量指标 (Quality metrics) 的表现,结果如下表:PLAID 生成的样本在序列与结构之间表现出高度的跨模态一致性。
表:各模型在一致性和质量指标上的性能比较
* 加粗的数值表示在全原子生成模型中的最佳性能
研究人员也进一步评估了不同模型的多样性、新颖性和自然性,结果如下表:在全原子模型中,PLAID 在序列和结构空间中均生成了最多的独特且具有设计性的样本。
表:不同模型的多样性、新颖性和分布一致性指标
* 加粗的数值表示在全原子生成模型中的最佳性能
值得强调的是,PLAID 可以轻松扩展至许多下游功能,不仅局限于 ESMFold,还可以应用于任何预测模型。
AI 为蛋白质研究打开「康庄大道」
Diffusion Transformer 在生物领域应用越发广泛
本论文提及,在模型构建过程中,研究人员采用 Diffusion Transformer (DiT) 来执行去噪任务。
DiT 的基本原理是将 Transformer 架构应用于扩散模型中。扩散模型通常通过逐步添加噪声来破坏原始数据,然后通过模型学习恢复这些数据。DiT 通过将 Transformer 块 (如自适应层归一化、交叉注意等) 引入扩散模型,增强了模型的生成能力。
近年来,DiT 在图像和视频生成领域取得了显著进展,诸如 Sora 等前沿生成模型,其主体架构都是 DiT。而在生物医药领域,Diffusion Transformer 的应用也愈发广泛,它可以帮助科研人员快速筛选潜在药物分子,预测其生物活性,还能辅助基因序列分析、蛋白质结构预测等复杂任务,为生命科学研究提供有力的工具。以蛋白质去噪为例,DiT 能够捕获复杂的序列-结构关系,即通过 Transformer 的全局自注意力机制,有效建模蛋白质序列和结构的复杂交互关系,然后利用扩散模型的逆过程,在每个时间步预测去噪后的潜在向量,从噪声中逐步还原蛋白质的结构和序列。
具体到该篇论文中,DiT 为处理混合输入模态的微调提供了更灵活的选择,特别是在蛋白质结构预测模型开始集成核酸和小分子配体复合体的情况下。此外,这种方法更好地利用了 Transformer 的训练基础设施。
早期实验中,研究人员还发现将可用内存分配给更大的 DiT 模型比使用三角自注意力更为有效。其使用 xFormers 实现的优化算法训练模型,在推理阶段的基准测试中,获得了 55.8% 的速度提升和 15.6% 的 GPU 内存使用量降低。
机器学习使得定制化蛋白质「梦想成真」
上文 UC Berkeley 的相关研究可以说在蛋白质定制方面又迈出了重要一步。我们知道,蛋白质通常由 20 种不同的氨基酸组成,这些氨基酸可以被视为生命的构建模块。由于结构极其复杂,数十年前对科学家而言,预测蛋白质三维结构以及设计全新蛋白质为人类所用尚属「异想天开」,但是近年来机器学习的飞速进展,使得设计定制化蛋白质的梦想逐渐成为可能。
除了大家耳熟能详的 AlphaFold,一些研究进展同样值得关注——
2024 年 11 月,美国能源部阿贡国家实验室团队成功开发出一种名为 MProt-DPO 的创新计算框架,该框架结合了人工智能技术与世界顶尖的超级计算机,标志着蛋白质设计走向了一个全新时代。以实际案例来看,科学家们通过 MProt-DPO 设计出了一种新型的酶,这种酶在特定条件下能够高效催化化学反应。相比于以往的设计方法,新酶反应效率提升近 30%,这不仅加快了实验进度,也为工业应用提供了更多的可能性。此外,MProt-DPO 的成功应用也为设计抗病毒蛋白质打开了新思路。相关研究结果以「MProt-DPO: Breaking the ExaFLOPS Barrier for Multimodal Protein Design Workflows with Direct Preference Optimization」为题,发表于 IEEE Computer Society。
论文地址:
https://www.computer.org/csdl/proceedings-article/sc/2024/529100a074/21HUV88n1F6
蛋白质口袋是蛋白质上适合与特定分子结合的点位,在定制蛋白质的过程中,蛋白质口袋设计是重要的方法之一。2024 年 12 月,中国科学技术大学及合作者设计了深度生成算法 PocketGen,可以基于蛋白质框架和结合小分子生成蛋白质口袋序列和结构。实验显示,PocketGen 模型亲和力和结构合理性等指标超过传统方法,在计算效率方面亦有大幅提高。相关研究成果以 「Efficient generation of protein pockets with PocketGen」为题,发表于 Nature Machine Intelligence。
论文地址:
https://www.nature.com/articles/s42256-024-00920-9
未来,随着人工智能在蛋白质领域的进一步应用,相信人们将对蛋白质空间结构的秘密具备更加深刻的认知。