随着 AlphaFold2 等模型的出现，对加速推理和训练分子 AI 模型的需求激增。对速度的需求带来了独特的计算挑战，包括算法复杂性、内存效率和严格的准确性要求。为解决这一问题，NVIDIA 与合作伙伴合作提供加速解决方案，例如更快的 equivariant 操作和更快的 MSA 生成。
今天，我们发布了 cuEquivariance 中的新内核和 NVIDIA NIM 微服务，以加速分子 AI 模型的训练和推理，例如由 MIT 和 Recursion 开发的开源基础模型 Boltz-2。这些加速有助于开发更复杂的分子 AI 系统，并更快地大规模了解分子结构。
NVIDIA cuEquivariance 扩展以加速新一代蛋白质结构模型
NVIDIA cuEquivariance 是一个 CUDA-X 库，旨在加速 MACE、Allegro、NequIP 和 DiffDock 等几何感知神经网络的苛刻计算。它提供高度优化的 CUDA 内核和全面的 API，可显著加速核心等方差运算，例如涉及 Segmented Tensor Products 的运算。
从 cuEquivariance v0.5 开始，该库现在包括加速的 Triangle Attention 和 Triangle Multiplication 核函数，这对于获得诺贝尔奖的蛋白质结构预测模型 (如 AlphaFold2) 来说至关重要且独一无二。随着加速三角形运算的加入，cuEquivariance 的影响扩展到蛋白质折叠、RNA/DNA 结合、盲对接、蛋白质复合体预测和亲和力评分等应用。
了解蛋白质的 3D 结构至关重要，因为它揭示了蛋白质的工作原理。然而，细胞和生命本身的真正复杂性源于生物复合体内的动态相互作用。这些复合体不仅仅由蛋白质组成，它们是蛋白质、核酸 (如 DNA 和 RNA) 、脂类、碳水化合物和各种小分子的复杂组合，它们协同工作。
预测这些单个分子和配合物的结构和动态行为是分子 AI 的下一个前沿领域。下一个科学突破可以揭示细胞通路、确定疾病机制，并设计能够精确向特定分子相互作用的药物。
蛋白质、RNA 和 DNA 都是由重复单元构建的长分子：用于 DNA 和 RNA 的核酸，以及用于蛋白质的氨基酸。当这些序列在细胞内生成时，其特定的构建块排列会导致它们折叠成复杂的三维结构。这些 3D 形状至关重要，因为它们决定了分子的功能及其与其他细胞成分的相互作用。
在最先进的几何感知神经网络 (如 AlphaFold3、Proteina、Chai-1、Neo-1 和 Boltz-2) 中，三角形乘法和三角形注意力是两个基本的计算密集型运算。在此类模型中，这些组件通常在最耗时的组件中排名靠前。
配对注意力机制通过引入 Transformers 而流行起来，其工作原理是计算 token (Transformers 词汇表中的构建块) 与所有其他 token 的相关性，例如，允许模型理解一个单词在句子中所有其他单词的上下文中的相关性。
由于分子 AI 模型的任务是根据 2D 表征预测 3D 结构，因此配对关系无法提供所有上下文信息。在这种情况下，“Triangular Relationships”可以充当捕捉空间关系的强大代理。例如，如果构建块 i 接近 k，而 k 接近 j，则 i 和 j 可能在空间上相关，即使没有强直接成对信号也是如此。
对于具有 N 个构建块的分子，这些运算简单地显示出 O (N3) 复杂性。这种计算强度给大分子和复杂的多分子组件带来了重大挑战，导致巨大的计算成本和对 AI 模型扩展程度的硬性限制。
cuEquivariance 加速 Triangle Operations
在这里，我们将讨论 cuEquivariance forward Triangle Attention 模块与原版 PyTorch 实现相比的性能。这仅衡量模块运行时，而非完整的端到端推理或训练加速。我们稍后将在博文中讨论端到端性能基准测试。
同样，具有 BF16 精度的 cuEquivariance 三角形乘法内核可在没有任何精度回归的情况下提供高达 5 倍的模块级加速。
麻省理工学院研究员 Gabriele Corso 表示：“这些 kernel 备受期待，并将成为 Boltz 系列模型不可或缺的一部分，帮助解决速度和内存消耗方面的瓶颈问题。
在下一代版本的Boltz-1x 上，我们比较了基于 PyTorch、Trifast 和 cuEquivariance 实现的各种精度 (TF32、FP32 和 BF16)，并比较了它们的端到端推理运行时间。这些运行使用由Boltz-1x 作者发布的默认测试数据集。在保持精度不变的情况下，从 PyTorch BF16 到 cuEquivariance BF16，我们可将性能提升高达 1.75 倍。如果从 PyTorch FP32 升级到 cuEquivariance BF16，则在使用Boltz-1x 时，性能可提升高达 2.5 倍。
从 PyTorch FP32 到使用 Boltz-1x 的 cuEquivariance BF16 配置，端到端训练速度最高可提升 1.35 倍。端到端加速可能因模型架构而异。
VantAI 首席技术官 Luca Naef 表示：“这种对 cuEquivariance 的扩展非常有价值，我们已经看到训练速度和推理速度分别提高了 2 倍和 3 倍以上，大大缩短了模型迭代周期，并对更大的分子实现了一个数量级的推理。
cuEquivariance 提供的加速得到了MIT、VantAI、Molecular Glue Labs (MGL) 、Dyno、Peptone、Genesis 和 Xaira 等多家合作伙伴的好评，他们能够测试早期版本并提供反馈。我们很高兴社区中的其他人能够从这些加速中受益，并提供关键反馈，帮助我们改进工作，从而推动科学创新的发展。
借助适用于数字生物学的 Boltz-2 NIM 实现企业级协同折叠
麻省理工学院的Boltz团队与 Recursion 合作开发的新一代Boltz-2 模型以Boltz-1 等模型的成功为基础，代表着向前迈出的重要一步。“Boltz-2”旨在打造一个更大、功能更强大的模型，将推理时间优化从Boltz-1x中合并，并融合了独特、先进的亲和力预测功能。NVIDIA 将 Boltz-2 封装为 NVIDIA NIM，以提供这种先进模型的易用版本。
NIM 是易于使用的预构建容器，可为先进的 AI 模型提供经过优化的生产就绪型推理。Boltz-2 NIM 将为研究人员和开发者提供对其强大功能的简化访问，为要求严苛的药物研发工作流程实现实时预测和高效的测试时扩展。这种方法普及了先进的分子 AI，允许更广泛的用户利用 Boltz-2 的预测能力。
面向下一个分子 AI 前沿的加速计算
由 cuEquivariance 驱动的更高计算效率至关重要。对于训练，这些加速的内核使研究人员能够构建更大的 基础模型，这些模型可以进一步利用预训练 scaling laws，其中计算吞吐量的增加通常与模型性能的提高相关。此外，由此带来的计算时间和成本效率为更多的模型开发周期释放了资源，进一步突破了新一代功能的极限。在测试时，加速有助于在 silico 实验中实现更广泛的应用，使 virtual screening 活动能够扩展到数十万甚至数百万个 inferences。
Recursion 首席技术官 Ben Mabey 表示：“NVIDIA 的 cuEquivariance 库提供了显著的加速，这对于像 Boltz-2 这样的结构感知生物分子模型至关重要。通过解决关键的计算瓶颈，这将加快制药行业部署这些强大的模型用于药物发现的研发周期。”
除了这些库级加速之外，NVIDIA 还通过 NVIDIA NIM 微服务等产品增强了对先进模型的访问。例如，将 Boltz-2 等高级模型封装为 NIM，可为研究人员和开发者提供精简的生产就绪型解决方案，以部署这些强大的功能，高效扩展要求严苛的药物研发工作流程。
NVIDIA 与更广泛的科学社区合作，开发和完善 cuEquivariance 库等基础软件，并通过 NIM 提供优化的模型访问。这些产品由 NVIDIA 计算平台提供支持，使科学界能够突破研究界限，加速从计算洞察到现实世界对更广泛的药物研发和生物学的影响的过程。
立即试用 cuEquivariance
目前，这些加速可通过 Apache 2.0 许可下的 PyTorch API 前端提供。单击此处详细了解加速、支持的精度和示例。
