数据科学

借助 GPU 加速的 MMseqs2 提升 AlphaFold2 蛋白质结构预测能力

能够比较多个相关蛋白质的序列是许多生命科学研究人员的一项基础任务。这通常以多序列对齐(MSA)的形式完成,从这些对齐中检索到的进化信息可以深入了解蛋白质结构、功能和进化历史。

现在,借助更新的 GPU 加速进化信息检索库 MMseqs2-GPU,从蛋白质序列中获取见解比以往更快。

简而言之,MSA 是一个大矩阵,包含表示蛋白质序列中残基(或氨基酸)的字母。矩阵的第一行包含“查询”序列——用于分析的感兴趣序列——每个残基从左到右都放在一列中。

后续行表示相似的序列,从最相似到最不相似,包含与行中每列的查询对齐的残基。当没有残基与查询序列匹配时,会在对齐中引入占位符差距,通常由“-”表示 (图 1)。

MSA is visualized with a stacked image of a gapped alignment at each position. It shows various alignment characteristics, including sequence logos with stacked letters whose height is proportionate to their prevalence at that position and bar plots. 图 1. MSAs 捕获有关比较序列的信息,可以使用序列徽标、条形图等方法进行可视化和量化。

通过从包含数千个物种序列的蛋白质数据库中检索到的 MSAs 编码的进化信息,可以深入了解蛋白质域,从而突出 跨物种的保守功能 。通过简单分析 MSA 中残基的保守性(即同一氨基酸在一列中的出现频率),可以快速了解序列中的关键残基,这些关键残基如果改变,可能会导致蛋白质发生故障。

这有助于研究人员深入了解他们所研究的蛋白质。 自 1992 年以来 ,多序列比对(MSAs)一直被用作复杂机器学习算法的输入,用于预测结构和功能等复杂的蛋白质特征。

AlphaFold2 彻底改变了计算和结构生物学,利用 MSAs 获得高度准确的 3D 蛋白质结构预测,这只是 MSAs 在药物研发研究中的众多用途之一(图 2)。

This figure illustrates the process of using protein sequences to derive evolutionary information ‌through homology (e.g., MSA with MMseqs2) and applying these insights in drug discovery use cases, including target identification, structure-function analysis, lead optimization, drug resistance, precision medicine, and vaccine design.
图 2、一种工作流 ,显示蛋白质序列被转换为 MSA,然后通过 MMseqs2 获得进化数据,然后作为机器学习模型(如 AlphaFold2)的基础输入,以预测准确的 3D 蛋白质结构

然而,计算 MSAs 具有挑战性,尤其是因为通用 CPU 并非为高度并行的工作流程而构建,例如筛选庞大的蛋白质序列数据库。由于大规模的元基因组学实验和廉价的新一代测序技术,蛋白质序列数据库每天都在增长,因此这个问题越来越难解决。因此,必须开发可快速筛选大型数据库的新算法,以便为蛋白质分析构建信息丰富的 MSAs。

利用 NVIDIA CUDA 克服计算成本高昂的 MSA

传统的 MSA 工具依赖于基于 CPU 的实现,这虽然在顺序处理方面有效,但无法与 GPU 并行处理能力相匹配。

开发 MMseqs2-GPU 的联合研究团队由首尔国立大学、美因茨的约翰内斯·古腾堡大学和 NVIDIA 的研究人员领导。受之前在 CUDASW++4.0 方面的工作启发,他们针对这一问题开发了一种专为 NVIDIA CUDA 定制的新型无缝预过滤算法,该算法能够以超快的速度实现高效、高灵敏度的序列比较。

此 GPU 加速的预过滤器使用无缝评分方法取代了 MMseqs2 中的 k-mer 预过滤。无缝预过滤直接分析完整序列,而不是使用 k-mer 搜索来简化具有粗略表示的序列之间的比较。它采用经典的 Smith-Waterman-Gotoh 算法的修改版本,该版本仅考虑对角依赖项,避免了对齐中的差距。该流程可在数千个 GPU 核心上高效运行。

在查询和参考数据库中的每个序列之间运行此算法的结果是,从数据库到查询的类似序列的排序列表,这些序列可以过滤到顶级候选项,为此可以执行加速的仿射差距 Smith-Waterman-Gotoh。内置于 MMseqs2 库中的这些算法还降低了内存需求,并且与多 GPU 系统原生兼容,克服了单个 GPU 的潜在内存可用性,并提供了额外的加速。

无缝预过滤步骤非常适合 GPU,因为它能够以最少的数据传输实现序列之间的比较,从而减少延迟并最大限度地提高 GPU 利用率。通过这种方法,单个 NVIDIA L40S 上的 MMseqs2 速度比 128 核 CPU 上的标准 JackHMMER 实现速度快 177 倍(图 3)。使用八个 NVIDIA L40 GPU 可将速度提升至 720 倍(每个序列 0.117 秒)。

This figure shows a boxplot indicating the average time required for protein sequence MSAs for a standard JackHMMER implementation (84.295 seconds), MMseqs2 k-mer (9.59 seconds), and MMseqs2-GPU (0.475 seconds). MMseqs2-GPU is illustrated with an arrow and text to have a 178 speedup compared to JackHMMER.
图 3、与常用的 JackHMMER 相比,MMseqs2-GPU 对比蛋白质序列的需求低 177 倍。

这些数字 是通过与包含 3000 万个序列的参考数据库对齐的 6370 个蛋白质序列的平均运行时间获得的。这些算法在配备 128 个核心的 CPU、1 TB RAM、2 TB NVMe 存储和单个 NVIDIA L40S GPU 的系统上运行。

在上下文中,使用 MMseqs2-GPU 计算序列对齐所需的时间(0.475 秒)与人类形成有意识思维(约 0.3 到 0.5 秒)、眨眼(约 0.3 到 0.4 秒)或闪电击中(约 0.2 到 0.5 秒)所需的时间大致相同。

CUDA 如何为 MMseqs2-GPU 提供优化和加速

这种加速的核心是 CUDA,它使 MMseqs2-GPU 能够执行优化的计算内核,以实现无缝和间隙对齐。这些内核利用多线程和内存共享功能,以更快的速度并行对齐多个参考序列。

MMseqs2-GPU 特别兼容最新的 NVIDIA GPU,例如 NVIDIA L40S GPU。GPU 加速内核可利用 GPU 的高并行性实现无缝预过滤和间隙对齐。无缝预过滤器并行处理每一行矩阵,使用共享 GPU 内存优化访问,并使用 32 位字(使用 half2 或 s16x2 数据类型)打包 16 位数字,以最大限度地提高吞吐量。

它使用跨线程束 shuffle 在线程束级别高效处理动态编程依赖项。另外,必要的内存查找通过使用快速 CUDA 共享内存得到加速。结合使用这些技术可以有效地将问题转变为计算受限的问题,并最大限度地减少内存访问的开销。

该工具还支持多 GPU 设置,以确保可扩展性,使研究人员能够通过在多个 GPU 上分配计算负载来处理更大的数据集。这种架构高度适应基于云的环境,因此 MMseqs2-GPU 对希望在不影响准确性的情况下降低计算成本的学术界和行业研究人员而言是一个极具吸引力的选择。

“我们已经等了很长时间了。众所周知,蛋白质结构预测推理受 MSA 计算步骤的限制。这是一项了不起的成就;将 MSA 步骤缩短到执行时间的 20%以下,彻底改变了我们未来处理结构预测工作流程的方式。”VantAI 首席技术官 Luca Naef 说。

MMseqs2-GPU 加速蛋白质结构预测

MMseqs2-GPU 的成功源于对无缝预过滤和间隙对齐算法的重新设计,利用 CUDA 提供快速、经济实惠且可扩展的序列对齐,从而满足当今的生物信息学研究需求。

由于 MMseqs2 已集成到使用 GPU 的许多计算工作流中,包括 使用 Colabfold 进行结构预测 ,因此用户有望获得易于交换的性能提升。

速度提升

使用 MMseqs2-GPU 的 Colabfold 比使用 JackHMMER 和 HHblits 进行蛋白质折叠的 AlphaFold2 快 22 倍(图 4)。在实践中,这意味着您无需等待 40 分钟来使用 HHblits、JackHMMER 和 AlphaFold2 预测蛋白质结构,而可以使用 Colabfold 和 MMseqs2-GPU 在 1.5 分钟内得到相同的预测。

This figure shows a boxplot indicating the average time required to predict 20 CASP14 protein sequence structures for the standard AlphaFold2 pipeline (2,277 seconds), ColabFold MMseqs2 k-mer (325 seconds), and ColabFold MMseqs2-GPU (105 seconds), with each of these methods having the same accuracy (LDDT at ~0.76).
图 4、与标准 AlphaFold2 实现相比,ColabFold-GPU 预测每个序列蛋白质结构所需的时间减少了 22 倍

该图基于 对 20 项 CASP14 查询的预测,且每种方法的准确性 (LDDT) 相同 (~0.76)。预测方法在配备 128 个核心的 CPU、1 TB RAM、2 TB NVMe 存储和单个 NVIDIA L40S GPU 的系统上运行。

显存需求

借助无差距 GPU 预过滤器,MMseqs2-GPU 可以避免使用 CPU 实现所需的大型 k-mer 哈希表索引。这使得此解决方案更适合 GPU,并将总体内存需求减少了一个数量级(参见图 3 和图 4 描述)

成本效益

使用 MMseqs2-GPU 的 Colabfold 在云成本估算方面比使用 JackHMMER 和 HHblits 的 AlphaFold2 便宜 70 倍。这使得实验室,尤其是预算有限的实验室,可以在不损失资金的情况下使用强大的生物信息学工具。更低的计算成本还可以实现持续的大规模分析,而这些分析原本会在财务上令人望而却步。

高吞吐量和可扩展性

新开发的无缝预过滤器可在 8 个 GPU 上实现每秒 102 Tera Cell Updates Per Second (TCUPS),从而快速预过滤大型数据集。该工具支持多 GPU 执行,使用户能够进一步扩展,处理更大的数据集,同时提高总执行速度,这对于大型基因组或蛋白质组研究至关重要。

准确率

MMseqs2-GPU 可在不影响准确性的情况下实现这些速度和成本优势。它保持了与原版相似的灵敏度和蛋白质折叠准确性,确保研究人员在不损失可靠性的情况下快速获得见解。

“我在哥伦比亚大学的实验室开发了 OpenFold,这是 AlphaFold2 的忠实再现,使社区能够训练其蛋白质结构预测模型。我们的应用程序特别关注的是执行迭代配置文件搜索的能力,事实证明,这种能力可以为结构预测提供信息量更大的 MSA。我们很高兴看到 MMseqs2-GPU 支持配置文件搜索,速度比以前的方法更快,”哥伦比亚大学教授 Mohammed AlQuraishi 说。

加速 MMseqs2 意味着更快的发现速度 

展望未来,联合研究团队专注于进一步完善算法和 MMseqs2 集成,将其应用扩展到蛋白质集群和级联数据库搜索。MMseqs2 的推出意味着更快地输入蛋白质结构预测,从而加速药物研发,正如我们在此处所展示的那样,以及许多其他应用(图 2)。

例如,这意味着 更快地输入 GEMME 等蛋白质变异预测变量 ,这些变量可用于加深我们对疾病变异的理解,以及对 PoET 等蛋白质语言模型的实时检索。这意味着 抗生素耐药性分析速度更快 。它甚至意味着 疫苗设计速度更快

对于那些有兴趣深入研究或为这项转变领域的工作做出贡献的人来说,MMseqs2-GPU 是开源的,可在线获取,为全球研究人员提供了宝贵的资源。

有关更多信息,请访问 MMseqs2 GitHub 或阅读 其详细分析和基准测试 。您还可以测试使用 MMseqs2 作为 MSA 步骤的 AlphaFold2 NVIDIA NIM

 

标签