编辑 | ScienceAI
作者 | 计算所张海仓团队
近期，中国科学院计算所张海仓带领的研究团队提出了 CarbonNovo，以端到端的方式联合设计蛋白质主链结构和序列。
该研究以「CarbonNovo: Joint Design of Protein Structure and Sequence Using a Unified Energy-based Model」为题发表在机器学习会议 ICML 2024 上。
背景介绍
蛋白质是生物执行功能的重要大分子。蛋白质从头设计旨在创造全新的蛋白质，在药物开发和酶工程中有着广泛的应用。
近几年，基于 AI 的蛋白质从头设计快速发展，已被成功应用于抗体设计、小蛋白药物设计等领域，和传统设计方法相比，其显著提高了设计成功率和效率。
AI 蛋白质设计得益于近几年的两大技术突破：
一是蛋白质结构预测领域的 AlphaFold2 模型，它为蛋白质计算领域包括蛋白质设计领域，提供了基础的神经网络模型架构、蛋白质序列表示和结构表示方法、以及先进的训练策略（蒸馏训练、端到端训练）等技术；
二是 AIGC 在文字、图像、视频生成领域的快速发展，为蛋白质设计提供了成熟的生成模型，例如 DDPM, SDE, Flow Matching，Bayesian Flow Network 等。代表性的蛋白质设计模型，例如 RFDiffusion 和 Chroma 等，主要思路都是将这两大技术融合，将蛋白质的序列和结构表示网络嵌入到基于 AI 的生成模型框架中。
图 1：蛋白质从头设计的「两阶段」框架。(来源，作者)
蛋白质从头设计主要包含两个步骤，蛋白质主链结构设计和序列设计（图1）。相应地，当前的主流模型通常采用「两阶段」框架进行：在训练过程中，结构设计模块和序列设计模块分别训练；在推断过程中，首先生成主链结构，然后为主链结构生成最优序列。代表性地，领域内使用 RFDiffusion 和 ProteinMPNN 这两个软件，依次生成主链结构和序列。
蛋白质从头设计的「两阶段」框架存在固有的局限性：
(1) 序列设计模块面临过拟合风险。序列设计模块是在准确的结晶结构上训练的，而在推断阶段，结构模块生成的结构是包含噪音的，其准确性与晶体结构不一致。
(2) 序列设计模块缺乏与结构设计模块的交互，设计序列无法为结构生成模块提供反馈以进一步优化生成的结构。
CarbonNovo进行端到端的结构和序列联合设计
针对蛋白质设计「两阶段」框架的局限性，中科院计算所张海仓带领的研究团队提出 CarbonNovo，以端到端的方式联合设计蛋白质主链结构和序列。论文近期已于近期发表在机器学习会议 ICML 2024 上。
图 2：CarbonNovo 端到端生成蛋白质结构和序列。（来源：论文）
CarbonNovo 的主要贡献总结如下：
1）设计了基于能量的生成模型，以端到端的方式设计主链结构和序列（图 2）。之前的结构-序列联合设计模型主要针对抗体等特殊的蛋白质家族，CarbonNovo 是第一个针对所有蛋白质家族的结构-序列联合设计模型。
2）基于 network recycling 技术，CarbonNovo 首次将蛋白质语言模型引入蛋白质结构设计任务，以利用海量天然蛋白质序列数据包含的先验信息。
3）CarbonNovo 采用多种技术，以提高结构-序列联合生成模型的训练和推断效率，例如多阶段训练策略和用于序列采样的离散 M-H Langevin 算法。
蛋白质结构-序列的联合能量模型
在经典物理模型下，天然蛋白质构象具有比较低的自由能，这也是蛋白质结构预测和设计的一般假设。基于此，CarbonNovo 建立了蛋白质结构和序列的联合能量模型:
1) 采用SE(3)空间的扩散模型刻画主链结构能量。最近的一些研究，无论是在训练目标还是推理策略方面，都已将扩散模型统一在能量模型（Energy-based model）下。
CarbonNovo也是基于此，将AI模型下的能量和经典物理模型下的能量联系在一起。需要指出的是，主流扩散模型（例如 FrameDiff 和 Genie）只是用IPA (Invariant Point Attention network)作为生成网络。而 CarbonNovo 引入了 AlphaFold2 模型里的三角更新网络，也是该工作的创新之一。
2) 采用 amortized Potts 模型刻画给定主链结构下的序列能量。Potts 模型是经典的能量模型，在蛋白质结构预测领域被用来学习同源序列中的共进化信号。CarbonNovo 利用学习到的结构表示来参数化 Potts 模型，建立给定结构条件的序列能量模型。
图 2 展示了 CarbonNovo 的具体生成过程:
1) 基于扩散模型设计出当前时间步的主链结构。
2) 给定当前的主链结构，基于 Potts 模型设计出可能的序列。
3) 通过 network recycling 机制，将中间序列的语言模型表示回传到结构模块来帮助生成更一致的序列-结构，实现了结构模块、序列模块和预训练语言模型的统一。
CarbonNovo 生成蛋白质结构-序列的性能评测
图 3：CarbonNovo 和「两阶段」设计模型的性能比较。（来源：论文）
论文采用多种指标充分评价了 CarbonNovo 在蛋白质从头设计方面的性能 （图 3），例如，可折叠性、多样性、新颖性是领域内常用的评价指标。此外，论文还采用了 Rosetta 能量和语言模型下的似然概率（Sequence plausibility）作为评价指标。
CarbonNovo 和当前主流的「两阶段」设计模型做了比较，例如 RFdiffusion, Chroma, Genie, FrameDiff 和 FrameFlow。在最关键的可折叠性指标上 CarbonNovo 显著超过了所有基线方法，在其他指标上也显著超过基线方法或者和基线方法相当。
为了展示 CarbonNovo 在联合设计序列与结构方面的优势，作者还对比了使用 ProteinMPNN 生成序列的结果（图 3 a-c）。可以观察到，联合设计模型可以设计出更加匹配的蛋白质主链结构和序列。
图 4：在不同蛋白质长度下的性能比较。（来源：论文）
作者进一步评估了 CarbonNovo 在不同长度的蛋白质设计上的性能（图 4）。在设计比较短的蛋白质时（例如长度 100），各个模型的表现相当。而随着蛋白质长度增长，CarbonNovo 的设计性能显著优于「两阶段」设计模型。
消融实验
图 5：消融实验结果。（来源：论文）
作者训练了多个消融模型，以评估关键组件对 CarbonNovo 性能的相对贡献（图 5）。语言模型、序列设计模块和辅助训练损失对 CarbonNovo 的性能都有贡献。其中，语言模型的引入表现出最显著的贡献。此外，使用基于能量的序列设计模块相比自回归模型，也能显著提升序列设计的性能。
Case study：蛋白质结构「插值」
图 5：从 all beta-sheets 结构到 all alpha-helices 结构的渐变。（来源：论文）
在图像生成领域，人脸图片插值/渐变是生成模型比较经典的应用。作者也尝试利用 CarbonNovo 进行蛋白质结构的插值。
图 5 展示了一个代表例子，随着在隐空间逐渐增加 all alpha-helices 结构向量的权重，生成的 all beta-sheets 结构会逐渐过渡到 all alpha-helices 结构。
这是领域内第一个关于蛋白质结构的插值实验，也体现了 CarbonNovo 学习到的蛋白质隐空间比较紧致。
结语
最后作者指出，虽然 CarbonNovo 主要侧重于蛋白质单体设计，但它也可以很容易被扩展到蛋白质复合物设计和条件设计，如多肽设计、抗体设计等。
作者团队目前在跟生物实验团队合作，通过湿实验的方式验证 CarbonNovo 设计的蛋白质。
作者所在 CarbonMatrix 团队，长期致力于 AI 蛋白质设计和 AI 药物设计，正在建立生物大分子结构设计和预测的统一生成模型。
其研究成果多次发表在 ICML、NeurIPS 等顶级机器学习会议和 Nature Machine Intelligence、Nature Communications 等顶级学术期刊上，目前也在和生物实验室合作，积极推动 AI 模型在药物设计领域的产业化落地。
论文链接：https://openreview.net/pdf?id=FSxTEvuFa7
代码链接：https://github.com/zhanghaicang/carbonmatrix_public
