NVIDIA 助力深原质药实现基于深度生成的功能蛋白质设计

作者 英伟达中国

案例介绍

随着 ChatGPT 的出圈,越来越多的人了解了生成式 AI,以及生成式 AI 这个概念。依托 NVIDIA 计算平台,深原质药(dProtein Biotechnology)把生成式 AI 进一步拓展成为了 AIGP,也就是用 AI 来生成蛋白质。这种由 AI 设计得到的蛋白质未来将在药物、检测试剂、合成生物学,乃至医美等领域都扮演重要的角色。深原质药通过自主开发的一系列 AI 算法,配合 NVIDIA 技术,能够为蛋白质相关的设计问题提供全面的解决方案,并已经服务于检测试剂和工具酶、递送载体和蛋白药物等领域的重要合作伙伴。

蛋白质从预测到设计的挑战

AlphaFold2 的成功点燃了生物学家使用 AI 算法的热情。对于生物学中积累的大量数据,人类的智慧还不足以总结出可计算的定理或定律。这正是 AI 可以大显身手的舞台。

AlphaFold2 解决的是“蛋白质结构预测问题”,即给定蛋白质的氨基酸序列,求解这个蛋白质的三维原子坐标结构。由于蛋白质的空间折叠是由其序列所决定的,同时又有很多已知序列信息以及结构信息作为训练数据,所以这方面的 AI 模型已经取得了极大的成功,能够以相当高的精度预测一个蛋白质的空间结构。

不过结构预测并不等同于蛋白质设计,甚至可以说是完全不同的两个问题。如果说结构预测是一次命题考试,那么蛋白质设计则更像是一次开卷考试,很多时候面临着序列、折叠、主链、侧链、结构全都未知的情况。

面对蛋白质设计的挑战,学术界和工业界给出了多种多样的 AI 解决方案。但是这些模型与当前的大模型相比,仍旧显得太过简单,能力也有限。这主要是由于已有的算法主要都基于蛋白质的氨基酸序列进行学习和生成。这种类似字符串的信息虽然便于理解和处理,但是却高度抽象化,与真实的蛋白质结构相去甚远,丢失了大量结构细节信息。因此,当前算法在一些困难问题上的表现不好,比如药物蛋白的设计以及酶的改造设计。

NVIDIA 计算平台助力高效实现“从结构到结构”蛋白质设计思想

深原质药重新解构了蛋白质设计问题,提出了“从结构到结构”的设计思想,注重真实的三维结构信息,而不是高度抽象化、高度简单化的蛋白质序列信息。在这一框架下,深原质药独立开发了自己的蛋白质设计算法 CUTEDGE。

为了实现在三维空间中解决蛋白质设计问题,深原质药的 AI 团队将标准的二维稳定扩散模型拓展到了三维空间中,构成了 CUTEDGE 算法的基础。由于增加了一个维度,且要维持较高的分辨率以体现结构的细节,于是 CUTEDGE 的模型参数激增至 6.9 亿,比很多主流的蛋白质设计算法提高了一至两个数量级。

为了能够实现这样一个模型的训练,深原质药采用了 NVIDIA 高性能 GPU,配合 NVLink 和超高的显存以及 CUDA 等加速技术。在 NVIDIA 上述技术与产品的加速之下,6.9 亿参数的 CUTEDGE 在单块 NVIDIA 高性能 GPU 上完成单次生成任务(50 步去噪过程)仅需 20 秒的时间。这就使深原质药可以在极短的时间内为客户生成大量的设计蛋白质,以供后续的评估与实验验证。

NVIDIA 全栈技术助力蛋白质设计从“画”到“讲”

如果说 Midjourney 是以生成式 AI 为画笔在作画,那么 CUTEDGE 就是以生成式 AI 为画笔画出了一个设计蛋白质。不过,当下最成功的 AI 模型无疑是 Transformer,以及以它为基础的各种语言模型。那么,有没有可能让语言模型“讲”出来一个蛋白质呢?

深原质药给出的答案是:当然可以。

目前已经有一些算法采用语言模型去学习蛋白质的序列信息,把一个蛋白质的序列当成一句话,来学习其中的“语法规则”,最终解决某些任务场景下的蛋白质设计问题。

深原质药的“讲”法则完全不同,基于公司独有的“从结构到结构”的理念,将蛋白质的三维结构表征为对应的一组编码,或者说词元表。这样一来就可以用一组词元来对应一个蛋白质三维结构,并输入大语言模型进行学习和生成。

得益于 NVIDIA 的开源框架 Megatron-LM,深原质药在 NVIDIA 高性能 GPU 上进行了并行训练,且训练速度较之前的框架(TensorFLow)提升 8 倍,显存占用率减少约 73%。同时还利用了 NVIDIA 的开源框架 Transformer Engine Apex 实现了使用 FP16 快速地训练语言模型,以及通过 NVIDIA 的开源框架 TensorRT-LLM 提升了语言模型的推理速度(约 5 倍)。最终,深原质药得以通过其“蛋白质结构语言模型”,成功“讲”出了全新的设计蛋白。

目前,深原质药已经形成了以 CUTEDGE、AtomSeg 等算法为核心的一个 AI 算法矩阵,紧密围绕蛋白质设计中的各类痛点问题,给出 AIGP 的解决方案。

开放、共赢

蛋白质设计领域方兴未艾,深原质药深信,“开放、共赢”才能进一步推动这个领域向前发展。为了践行这一理念,深原质药于 2023 年底开放了其 AI 算法矩阵当中的 CUTEDGEopt 和 AtomSeg 的公测,用户可以在深原质药官网找到算法入口,计算自己想要解决的蛋白质设计问题。在 2023 年,深原质药也加入了 NVIDIA 初创加速计划,成为会员公司,将与 NVIDIA 开展更加深入和全面的合作。

在这样一个变革的时代,AI 带来的不仅仅是有趣的对话机器人,更是生产力的又一次跃升。深原质药基于 AI 算法的蛋白质设计,将为众多生物技术公司和创新药物研发企业带来更丰富的研发管线,让 AIGP 真正产生生产力,用更绿色、更安全的蛋白质产品服务于人们。

NVIDIA 初创加速计划

深原质药是 NVIDIA 初创加速计划 (NVIDIA Inception) 会员企业。NVIDIA 初创加速计划为免费会员制、旨在培养颠覆行业格局的优秀创业公司。该计划联合国内外知名的风投机构,创业孵化器,创业加速器,行业合作伙伴以及科技创业媒体等,打造创业加速生态系统。能够提供产品折扣,技术支持,市场宣传,融资对接,业务推荐等一系列服务,加速创业公司的发展。

了解更多