近日,深势科技与阿里云机器学习 PAI 团队联手,通过全新的蛋白质结构预测推理加速方案 FoldAcc,结合深势 Uni-Fold 最新模型代码和参数,将单次预测能支持的最大氨基酸序列长度提升至 6.6k,覆盖 99.992% 已知的蛋白序列,同时推理速度显著提升,达到目前已知的最佳推理优化效果,将为 AI 预测蛋白质结构落地应用提供重要助力。
以新冠病毒研究重点之一 —— 具有三聚体结构的刺突蛋白为例,其氨基酸序列典型总长度接近 4k,原版 AlphaFold 会因为超出显存限制(OOM)而无法进行预测计算,使用 Uni-Fold + FoldAcc 则能在 10 分钟左右完成 AI 推理计算。
自 2020 年 DeepMind 推出基于深度学习模型的 AlphaFold2 以来,AI 技术辅助蛋白质结构精确预测的相关研究备受关注,产学研界不断涌现创新成果。然而,在推动 AI 蛋白质结构预测规模化、产业化落地进程中,基础设施及工具完善性、AI 模型开发与部署效率等问题,仍然是横亘在研究者面前的难题。
2022 年 8 月,深势科技升级并开源 Uni-Fold 项目,成功复现了 AlphaFold2、AlphaFold-Multimer 模型的全尺寸从头训练,并通过多项效率优化,并将 AlphaFold 训练速度提升 220%,超越 OpenFold、FastFold 等方案,惠及更多研究者。
其中,针对困扰业界已久的 Evoformer 神经网络结构推理性能瓶颈问题,深势科技与阿里云 PAI 团队,基于在 AI 模型系统优化领域的长期积累,融合多卡并行、混合精度、编译优化等多项推理优化技术,使 Uni-Fold 训练的模型能进行多卡推理加速,并支持计算更长的氨基酸序列。
典型加速效果的测试结果如下(基于 A100-80G GPU,并启用 bf16)
Uni-Fold 升级开源,支持复合物训练
作为生命科学领域重要的基础问题,蛋白质结构研究关系到癌症预警、靶向药物研究、衰老等临床医学和生命科学课题。传统的蛋白质结构研究手段,如:X 射线晶体学、冷冻电镜等,需要消耗大量时间和资源。如何快速高效并且能够规模化地预测蛋白质结构,一直是研究者们探寻求解的重要问题。
2021 年 12 月,深势科技推出 Uni-Fold v1.0.0,国际首次复现了 AlphaFold2 官方代码的全尺寸从头训练,并开源了训练与推理代码;2022 年 8 月,Uni-Fold 升级并开源了最新代码与模型参数,完整支持蛋白质单体、复合物结构预测模型的推理与训练。
此次开源的 Uni-Fold 基于 PyTorch 复现并改进的 AlphaFold (-Multimer) 模型,支持从头训练和推理部署,并且在 Protein Data Bank(PDB)最新公布的、模板相似度小于 40% 的单体与复合物测试集上,都取得了与同类开源项目一致或更优的准确率。
此外 Uni-Fold 也进行了多项效率优化,将训练时间由 11 天缩短至约 4 天,显著优于其他同类开源项目。
机器学习平台 PAI 提供全链路 AI 工程支撑
阿里云机器学习平台 PAI 为 Uni-Fold 项目提供了完善的 AI 工程能力。PAI 是国内唯一连续入选 Gartner 数据科学与机器学习平台报告的机器学习 / 深度学习平台,面向 AI 开发及应用全链路提供全面的工程化服务,并具备丰富的场景化落地实践。
针对模型推理优化场景,PAI 自研的通用推理加速器 PAI-Blade 能在不同业务场景下,通过模型系统联合优化,使模型达到最优推理性能,兼容主流机器学习框架,适配 GPU、CPU、端侧设备等多类型加速设备。
其中,PAI-Blade 核心组件 BladeDISC 具备业界领先的动态尺寸模型优化、大颗粒度算子融合等编译优化技术,在阿里巴巴集团内外客户的实际生产场景广泛应用,助力实现高性价比的 AI 应用部署。2022 年 2 月,BladeDISC 项目正式开源。
未来,为前沿科研提供更好用的 AI 基础设施
以 AI 蛋白质结构预测为典型代表,AI for Science 的科学研究新范式正取得越来越多突破,人工智能与生命科学、物理、化学等领域的技术交织,将为科学研究和产业进步带来极大推动,也对 AI 基础技术与平台应用提出了新挑战。
深势科技是 AI for Science 科学研究范式的先行者,阿里云机器学习平台 PAI 是国内应用最广泛的机器学习平台之一,我们将持续为生物医药、能源、材料等领域的研究与产业应用提供更好用的 AI 基础设施,共同推动 AI for Science 领域的技术生态建设。