华为推出全流程蛋白质结构预测工具 MEGA-Protein（MindSpore for Evolutionary Generation & Assessment Protein），其中业界首创 AI MSA 引擎使得蛋白质结构预测突破「孤儿序列」等高精预测限制
日前，华为与北京昌平实验室、北京大学生物医学前沿创新中心（BIOPIC）和化学与分子工程学院、深圳湾实验室高毅勤教授课题组，在中国算力网智算网络关键节点——西安未来人工智能计算中心的昇腾 AI 强大算力支持下，在全场景 AI 框架昇思 MindSpore 上推出全流程蛋白质结构预测工具 MEGA-Protein。
MEGA-Protein 包含 AI MSA 引擎、蛋白质折叠训练推理流程、蛋白质结构打分、蛋白质结构预测数据集 PSP 等关键技术，该工具提供了高精度高性能蛋白质结构和功能预测，其中 AI MSA 引擎能显著提升单序列的预测速度，并且能够在 MSA 少样本（few shot）甚至零样本（zero-shot，即单序列）的情况下，帮助 AlphaFold 2 等模型维持甚至提高推理精度，突破了在「孤儿序列」、高异变序列和人造蛋白等 MSA 匮乏场景下无法做出准确预测的限制。
这是团队在昇腾 AI 基础软硬件平台上实现 AlphaFold 2 从训练到推理全流程打通且效率同比提升 2 至 3 倍后，取得的又一次成功。相关代码和模型参数后续会在昇思 MindSpore 社区开源，算法和相关的论文也会尽快公布。
昇思 MindSpore 首席架构师金雪锋评论称：「该工作是产学界合作的一次重大成功实践，也是科学计算与人工智能结合潜力的展示。」金雪锋补充说，支持多计算范式的 MindSpore 能够有效提升科研工作的生产力，大幅度提升科学计算场景的性能，希望后续能够看到更多、更好的成果涌现出来。
高性能蛋白质结构预测引擎 AI MSA，解决「孤儿序列」问题
蛋白质是生命活动中最主要的功能体之一，在人体结构和功能中发挥着核心作用。然而，能够解析蛋白质结构的实验手段通常都价格不菲且周期较长，很难直接通过它们去探索整个「蛋白质宇宙」。
蛋白质的 3D 结构完全由组成它的氨基酸序列所决定（1972 年诺贝尔奖得主 Christian Anfinsen 曾通过实验证明），这就从理论上确认了：存在一个「理想函数」，能够将氨基酸序列 x 映射到其 3D 结构 y。因此，蛋白质结构预测问题就转化成典型的数学建模问题。AlphaFold 2 则是目前为止人们所能找到的、最接近「理想函数」的一个模型。
但是，为了达到高精度的预测，AlphaFold 2 必须引入额外的信息，也即蛋白多序列比对（MSA）。有多少 MSA 可用，直接决定了 AlphaFold 等模型的预测精度。
图1. AlphaFold 2 预测精度会随着可用的 MSA个数不足时而显著降低。图片来源: AlphaFold 2 论文
但自然界中仍然存在大量「孤儿序列」，如人类蛋白组中不少蛋白以及大量的病毒等其它抗原相关蛋白，都存在着 MSA 不够深的问题。此外，在蛋白质设计领域中，研究人员会普遍面临人造序列没有 MSA 可用的情况。在这些场景下，AlphaFold 2 等模型对相关的结构预测准确度会大幅下降，变得不再适用。
另一方面，AlphaFold 2 等标准检索 MSA 流程，是在数据库中使用搜索工具对目标蛋白及相关序列进行多重序列比对，该流程需要配置数据库，其中数据库规模大概在 1.3T 左右，配置耗时长且繁琐，并且使用数据检索工具（如 HHblits，JackHMMER，MMseqs2 等）搜索时间长，不利于科研人员展开研究。
昇思 MindSpore 社区与昌平实验室、北京大学生物医学前沿创新中心（BIOPIC）和化学与分子工程学院、深圳湾实验室高毅勤教授课题组新提出的 AI MSA 引擎，对于原始 MSA 质量不高或者数量少的蛋白，将 AI MSA 引擎接入 AlphaFold 2 后可以明显改善结构预测的质量，如下图所示：
图2. 不同方案预测的蛋白质结构效果对比图
图3. 零样本预测精度对比
并且端到端的推理性能相较于 AlphaFold 2 标准检索流程有大约 40+ 倍的提升，相较于 MMseqs2 约 5+ 倍的提升，可以大幅提升结构预测模型的推理通量，且训练完后的 AI MSA 引擎不需要额外配置数据库。
图4. AI MSA 引擎相关性能对比图
团队表示，AI MSA 引擎是一种对下游结构模型普适的预训练方案，可以在免于微调的情况下直接接入下游的结构预测模型，例如 AlphaFold、RoseTTAFold 等，助力科研人员在蛋白质结构预测、蛋白质设计、蛋白质互作等方向进行研究，同时推动生物制药产业发展。
北京大学李兆基讲席教授谢晓亮认为，创新自研的 AI MSA 引擎和基于昇思 MindSpore 的蛋白质结构预测程序的开源，标志着我们国家在全球蛋白质结构预测领域名列前茅，该长足进步源于高毅勤团队在相关底层技术上多年的深耕和积累。
昇腾 AI 基础软硬件平台、昇思 MindSpore AI 框架，支持 AI for Science 研究新范式
AI MSA 引擎训练参数量有 47M，数据量总共达 4.4T，其训练对框架提出了不小的要求，昇思 MindSpore 在昇腾 AI 基础软硬件平台上与昇腾 CANN 深度结合，通过深度协同优化的高性能算子库，充分释放硬件的算力。
昇思 MindSpore 采用了多段并行流水线的方式来构建数据处理 pipeline，更加细粒度地规划 NPU、CPU 等计算资源的使用，天然支持各段使用异构硬件进行流水处理，大大提高了数据处理过程的吞吐量。
同时，昇思 MindSpore 支持大集群高效训练，实现了优质的计算通信比，并且通过三层 AI 分布式编程范式（手动并行＋半自动并行＋全自动并行），大幅提升分布式并行程序开发效率。
相关代码和模型参数后续会在昇思 MindSpore 社区开源，算法和相关的论文也会尽快公布。
团队还发布了首个具有高覆盖度和多样性的百万级蛋白质结构预测数据集 PSP。该数据集由 570k 个真实结构序列（10TB）和 745k 个互补蒸馏序列（15TB）组成。昇思 MindSpore 基于该数据集可进行蛋白质结构预测训练，同时提供了丰富的处理及使用该数据集的接口，欢迎试用。
了解更多可查看论文：https://arxiv.org/pdf/2206.12240.pdf
