设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计 原创

发布于 2025-8-7 11:57
浏览
0收藏

目前,蛋白质基座模型领域的相关研究仍停留在无法充分适配蛋白质序列生物学特性的「BERT」时代。此前,AlphaFold 和 ESM 等 AI 模型显著推动了包括结构预测、逆向折叠、功能特性预测、突变效应评估以及蛋白质设计等多个领域的发展。但这些模型仍缺乏类似于前沿大语言模型(LLMs)的可扩展、系统化方法论,其能力无法实现随数据量、模型规模和计算资源的增加而持续提升。

此类模型的通用性缺陷,带来了蛋白质设计领域难以解决的挑战:模型无法捕捉蛋白质的构象异质性,对蛋白质设计的预测也无法超越训练数据的范畴;而对 NLP 方法论迁移的过度依赖,则导致了针对蛋白质特性的原创性架构设计的缺失。

在此背景下,清华大学智能产业研究院(AIR)周浩课题组联合上海人工智能实验室,基于贝叶斯流网络提出经系统化训练的蛋白质基座模型 AMix-1,为蛋白质设计提出了一种可扩展的通用路径。该模型首次采用了「Pretraining Scaling Law」、「Emergent Ability」、「In-Context Learning」以及「Test-time Scaling」的系统化方法论,并在这一基础上设计了一种基于多序列比对(MSA)的上下文学习策略,实现了蛋白质设计通用框架的一致性,同时确保了模型的可扩展性。

相关研究成果以「AMix-1: A Pathway to Test-Time Scalable Protein Foundation Model」为题,发布于 arXiv 平台。

研究亮点:

为基于贝叶斯流网络的蛋白质生成模型建立了可预测的 Scaling Law;

AMix-1 模型仅通过序列级训练目标,自发涌现出对蛋白质结构的「感知理解」能力,无需显式结构监督;

基于多序列比对(MSA)的上下文学习框架解决了功能优化中的对齐难题,升级了模型在进化语境中的推理和设计能力,使 AMix-1 能够生成结构和功能保守的新蛋白质; 

* 提出验证成本导向的测试时扩展算法,在验证预算增加时开启基于演化的设计新方式。

论文地址:

https://go.hyper.ai/6Lz0c

关注公众号,后台回复「AMix」获取完整 PDF

更多 AI 前沿论文: 

https://go.hyper.ai/owxf6

UniRef50 数据集:预处理与迭代聚类

研究人员在模型预训练阶段使用了预处理的 UniRef50 数据集。该数据集由 EvoDiff 提供,源自 UniProtKB,通过迭代聚类从 UniParc 序列中筛选得出(UniProtKB+UniParc → UniRef100 → UniRef90 → UniRef50),包含 41,546,293 条训练序列和 82,929 条验证序列。其中,对于长度超过 1,024 个残基的序列,研究人员使用随机裁剪的策略,将其长度裁减至 1,024 个残基以降低计算成本并生成多样化子序列。该迭代流程确保了 UniRef50 代表性序列高质量、非冗余且多样化的特性,为蛋白质语言模型提供了覆盖广泛的蛋白质序列空间。

UniRef50 数据集下载地址:

https://go.hyper.ai/JuYEM

系统性技术方案

AMix-1 为蛋白质基座模型实现 Test-time Scaling 提供了一整套系统性的技术方案:

Pretraining Scaling Law:明确了参数、样本数和计算量如何权衡,才能最大化模型的能力; 

Emergent Ability:显示随着训练的推进,模型会涌现出对蛋白结构的「感知理解」;

In-Context Learning:解决了功能优化中的对齐难题,让模型学会在进化语境中推理与设计; 

Test-time Scaling:AMix-1 在验证预算增加时开启基于演化的设计新方式。

从训练、推理到设计,AMix-1 论证了其作为蛋白质基座模型的通用性和可扩展性,为走向实际落地做好了铺垫。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

模型能力对比

Pretraining Scaling Law:可预测的蛋白质模型能力

该研究为实现 AMix-1 的可预测 Scaling Law,在实验中设计了从 800 万到 17 亿参数的多尺度模型组合,并利用训练浮点操作数(FLOPs)作为统一衡量指标,精确拟合、预测了模型交叉熵损失与计算量的幂律关系。

从结果表征来看,模型损失与计算量之间的幂律曲线吻合度极高,证实了基于贝叶斯流网络的模型训练过程具有高度可预测性。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

不同噪声尺度下 AMix-1 的 Scaling Law

Emergent Ability:实现模型高级能力跃升

在蛋白质序列学习中,结构信息的涌现研究通常基于「序列-结构-功能」的范式。为验证蛋白质建模中优化动力学与功能结果的关联,研究团队基于可预测的 Scaling Law 采用以损失函数为核心的视角分析涌现行为,在将预测性交叉熵损失作为锚点的同时,通过经验方法将训练损失映射至蛋白质的生成性能。该研究对模型 Emergent Ability 的评估聚焦于 3 个层面:

* 从序列一致性观照模型从受损序列分布中恢复序列水平的能力; 

* 从可折叠性看模型从序列理解到结构可行性的转变;

* 从结构一致性判断模型保持结构特征的能力。

AMix-1 在训练中的相关数据完整展现了蛋白质基座模型「序列一致性、可折叠性、结构一致性」能力的涌现过程。数据显示,模型在训练过程中的所有能力指标与交叉熵损失呈高度相关性,验证了通过 Scaling Law 和交叉熵损失预测模型能力的可能性。同时,在仅以序列级自监督目标训练、完全不引入结构信息的情况下,模型依然在交叉熵损失下降至阈值后展现出 Emergent Ability,表现出 pLDDT 与 TM-score 的非线性跃迁。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

AMix-1 模型在不同指标下的涌现现象

In-Context Learning:蛋白质设计的通用范式

通过计算机模拟案例,研究人员验证了 AMix-1 的 In-Context Learning 机制。模拟案例实验显示,AMix-1 能够从输入样本中准确提取并泛化结构或功能约束,而无需依赖显式标签或结构监督。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

LLM 与 AMix-1 在 In-Context Learning 层面的比较

相较于需要根据任务类型进行定制流程、缺乏统一的蛋白设计框架的传统蛋白质设计,AMix-1 引入了大语言模型中的上下文学习(ICL)机制,以完成结构和功能导向的蛋白设计。实验显示,AMix-1 在结构任务上,能够以常规同源蛋白甚至在近乎无同源的蛋白为提示生成在预测结构上高度一致的新蛋白,在功能任务中则能在输入蛋白的酶学功能和化学反应引导的酶设计上,生成出功能高度一致的蛋白酶。

在这一通用机制下,模型能够在给定的一组蛋白质中,自动推理其中的共性信息和规律,并运用此规律来指导生成新的符合共性规律的蛋白。该机制将一组蛋白质 MSA 压缩成一个位置级的概率分布(Profile)输入模型,模型在快速辨析所输入蛋白质的结构及功能规律后,即可生成符合意图的新蛋白。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

验证 AMix-1 In-Context Learning 的案例

Test-time Scaling:可扩展的通用智能

基于 Test-time Scaling 方法,研究人员采用提议者-验证者框架(Proposer-Verifier Framework)构建了 EvoAMix-1,通过持续增加验证预算提升了 AMix-1 的模型性能。研究团队在放大模型的设计效能的同时,实现了模型的可扩展性。同时,为确保兼容性,研究团队减去了对验证者的性质预设。

EvoAMix-1 基于概率模型固有的随机性促进探索,通过整合特定任务的计算机模拟奖励函数或实验检测反馈,在进化约束下迭代生成并筛选候选蛋白质序列,无需模型微调即可实现高效的定向蛋白质进化,在蛋白质设计中实现了稳健且测试时可扩展的性能。在全部 6 项设计任务中,EvoAMix-1 在 In-Context Learning 及各类强基线方法上始终优于 AMix-1 。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

AMix-1 在推理阶段进行进化扩展算法的工作流

和通过重要性采样生成新蛋白质变体的传统方法相比,EvoAMix-1 并不更新模型参数,而是通过上下文范例构建提议分布。在每一轮中,AMix-1 都会以一组多序列比对(MSA)或其谱作为提示,将其视为蛋白质基础模型的输入条件,随后基础模型对邻近序列进行采样,从而有效定义新的条件提议分布。

研究团队在多个代表性蛋白质定向进化任务上系统验证了 EvoAMix-1 的通用性和可扩展性,包括酶的最适 pH 和温度进化、功能保持与增强、孤儿蛋白设计,以及通用的结构引导优化等。实验结果显示,EvoAMix-1 Test-time Scaling 的扩展能力强劲,在跨任务、跨目标的情况下仍表现出强大的通用性。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

EvoAMix-1 在六个定向进化基准中的 Test-time Scaling 性能

湿实验验证:AMix-1 辅助研发蛋白 AmeR 变体活性提升 50 倍

该研究在实际湿实验中检验了「上下文提示设计」策略,进一步验证了 AMix-1 在高效设计高活性 AmeR 变体方面的优势。研究人员选取目标蛋白 AmeR,使用 AMix 模型基于 AmeR 家族的概率分布生成 40 个变体,并通过荧光报告基因实验评估每个变体的抑制能力。其中,每个变体仅包含 ≤10 处氨基酸突变,倍数抑制(Fold Repression)的数值越高功能越强。此外,该研究还提出了进化测试时的缩放算法,以增强 AMix-1 在蛋白质定向进化中的适用性,并通过多种计算机模拟靶区指标验证其性能。

最终结果显示,由 AMix-1 产生的最优变体活性提升高达 50 倍,其效能与当前 SOTA 模型相比提升了约 77% 。此外,AMix-1 不依赖反复筛选或手动设计,完全由模型自动生成,实现了「从模型到实验」的完整闭环,同时达成了 AI 在功能蛋白设计上的首次量级突破。

设计蛋白变体活性提升50倍!清华AIR周浩团队基于贝叶斯流网络提出AMix-1,实现可扩展通用的蛋白质设计-AI.x社区

AMix-1 模型在湿实验中的生成迭代和相应的变异过程

全局拓扑,感知开启蛋白质设计的全新维度

目前,AI 与蛋白质设计的结合研究正方兴未艾。除了 AMix-1,清华大学生命学院龚海鹏课题组联合北京生命科学研究所徐纯福研究团队提出的几何感知扩散模型 TopoDiff,同样在蛋白质设计领域取得了重大突破。

RFDiffusion 等传统扩散模型在生成免疫球蛋白等特定折叠类型时,不仅存在覆盖偏差,而且缺乏对蛋白质全局拓扑的量化评估指标。该研究基于 CATH 、 SCOPe 等结构数据库,提出了无监督系统 TopoDiff 框架,通过学习和利用全局几何感知的潜在表征,实现了基于扩散模型的无条件与可控蛋白质生成。该研究提出「覆盖度」(Coverage)的新评估指标,通过编码器-扩散模型双阶段框架,将蛋白质结构解耦为全局几何蓝图与局部原子坐标生成,突破了蛋白质折叠覆盖的研究难题。

此外,NVIDIA 联合加拿大魁北克人工智能研究所 Mila,还基于 AlphaFold 架构升级的全原子生成模型则攻克了长链预测难题。传统方法不仅难以生成超长链(>500 残基)的全原子结构,而且无法探索膜蛋白特异性口袋等非经典折叠构象。研究团队引入的概率决策机制,以量子场论中的路径积分采样替代了确定性折叠轨迹,从而使膜蛋白设计成功率提升至 68% 。

几何感知解构蛋白质折叠,攻克 500+ 残基长链设计,自然语言指令驱动蛋白设计,靶向「不可成药」IDPs 蛋白……AI 在拓展蛋白质设计能力边界的同时,为蛋白质设计领域的研究路径提供了新的范式。未来,AI 驱动的蛋白质设计有望为开发创新疗法、酶制剂和生物材料开拓更广阔的空间。

参考链接:

1.https://hub.baai.ac.cn/view/23883

2.https://www.webpronews.com/nvidia-ai-model-boosts-protein-prediction-to-800-residues-with-75-accuracy/

3.https://www.nature.com/articles/s42256-025-01059-x

©著作权归作者所有,如需转载,请注明出处,否则将追究法律责任
收藏
回复
举报
回复
相关推荐