人工智能时代的超级计算加速蛋白质结构预测

通过与 SEAS 合作，研究人员可以获得安装 Python 包的帮助，学习为他们的项目选择最佳的并行计算引擎，或者——得益于突破性的工作 发布 在 PNAS 杂志上学习成功部署人工智能模型。该论文的标题是“APACE：AlphaFold2 和高级计算作为加速生物物理学发现的服务”。

PNAS 的研究论文由 SEAS 小组高级研究程序员 Roland Haas、美国能源部 (DOE) 阿贡国家实验室转化人工智能负责人兼芝加哥大学 CASE 高级科学家 Eliu Huerta、时任伊利诺伊州生物物理学博士生 Hyun Park 和 NCSA 研究生助理 Parth Patel 共同撰写，描述了一种新颖的计算框架，可简化和加快使用人工智能工具和算法了解三维蛋白质结构的过程。

该框架还预测了蛋白质的构象多样性，这是一个重要的特性，因为蛋白质是可塑性结构，可以在不同构象之间翻转来完成其工作。

该团队开发了 APACE，这是一种计算工具，可以有效处理 AlphaFold2，这是一种用于在高性能计算系统上预测蛋白质结构的 AI 程序。他们在 三角洲 NCSA 的超级计算机来测量其预测四种典型蛋白质结构的表现。

他们使用分布在 300 个 NVIDIA A100 GPU 上的多达 300 个集成，发现 APACE 比现成的 AlphaFold2 实现快两个数量级。

此外，同样的方法可以用于各种科学学科，并可以与机器人实验室相连，实现科学发现的自动化和加速。该团队后来在 北极星 位于阿贡领导计算中心的超级计算机，该中心是美国能源部科学办公室的用户设施。

“如果基础人工智能模型能够被更广泛的科学界找到、访问和使用，那么它们就有可能改变科学实践，”Huerta 说道。“这个项目展示了如何创建和共享所需的科学数据基础设施，以真正实现尖端人工智能的民主化，并利用现代计算环境最大限度地扩大其科学影响力。”

生物医学研究人员通过研究蛋白质来了解各种生物功能。蛋白质是氨基酸链，其三维结构的排列顺序决定了生物功能。

了解蛋白质的形成方式（通常称为蛋白质折叠，其中氨基酸组合成能够执行特定功能的结构化链）对于理解正常的生物功能以及折叠错误如何导致严重疾病至关重要。

预测蛋白质折叠需要大量的计算，因为一个典型的蛋白质可以有数百种氨基酸和数千种可以以不同方式结合的细胞。

研究蛋白质结构的常用方法是 X 射线晶体学（一种用于确定晶体原子和分子结构的工具）和低温电子显微镜（涉及在液氮中快速冷冻分子并用电子轰击它们以用特殊相机捕捉它们的图像）。

AlphaFold 和AlphaFold2表明AI软件可以根据氨基酸序列准确快速地预测蛋白质结构，APACE的开发正是建立在这一突破之上。

APACE 优化了 AlphaFold2，使其能够在高性能计算平台上大规模运行，并有效处理其数 TB 的蛋白质数据库。这项工作表明，大型 AI 模型可以与高性能计算的强大功能相结合，使科学家能够研究多蛋白质复合物并快速、准确、高分辨率地获得结果——所有这些因素都可能导致更全面地了解蛋白质结构并启动可以治疗许多疾病的新药的开发。

哈斯说：“新药研究极其耗时，而且由于需要合成不同的候选化合物以在实验室中测试其医疗效果，因此存在瓶颈。”

APACE 可让药物研究人员大幅缩短筛选潜在候选化合物所需的时间，从而将精力集中在最有前景的物质上。这样一来，研究人员就可以测试更多化合物，并缩短开发新药（例如针对特定病毒株的药物）的时间。

APACE 的一个关键特性是更好的数据管理，这是通过在超级计算机上托管 AlphaFold2 的多 TB 模型和数据库来实现的，框架的神经网络可以从中轻松访问数据。其他改进包括 CPU 优化和 GPU 优化，以并行化 GPU 密集型神经网络蛋白质结构预测步骤。

“使用人工智能模型的第一个难题是数据的存储，”帕克说道，他和帕特尔一样，在 APACE 的工作完成时也在阿贡国家实验室实习。

“我们需要传递 2.6 TB（AlphaFold2 数据库的大小）以及从序列到结构预测的计算。一些大学实验室可能能够做到这一点，但重要的是你要扩大规模，以便世界各地的科学家都可以使用它。”

Patel 补充道：“这就是为什么 HPC 利用率如此重要，尤其是对于 AI 模型而言。任何能够进入 HPC 系统的人都可以访问数据，并拥有进行实际 AI 模型计算的计算能力。更不用说，速度有了巨大的提升。”

韦尔塔表示，团队之所以选择与 AlphaFold2 合作，是因为它广泛应用于不同的研究领域，包括生物物理学、化学以及药物设计和发现。

他说：“APACE 提供了原始 AlphaFold2 模型的所有功能，使研究人员能够利用超级计算机来缩短解决问题的时间，并将该工具与自动驾驶实验室连接起来，实现自动化和加速发现。”

Huerta 表示，该团队将继续建立一个 APACE 用户社区，以最大限度地提高 AI 模型在 HPC 平台上的可用性。Haas 表示，该团队目前正专注于攻克系统中剩余的瓶颈，以进一步提高速度。他还希望让 APACE 在更多的计算集群上可用，以便更多的科学家可以利用它。

“我们还想探索使用我们开发的方法来加速 Alphafold2 以及其他基础机器学习模型，这些模型过于复杂，无法在普通桌面工作站上使用，”哈斯说。“这一切都是为了提供最好的工具，并尽可能易于使用。”