用于科学的 AI 模型通常经过训练，可以预测自然界的运作方式，例如预测生物分子的结构或可能成为下一种电池材料的新固态的性质。这些任务需要高精度和准确性。科学领域的 AI 更具挑战性的原因在于，高度准确和精确的科学数据往往稀缺，这与从多个资源中大量获取的文本和图像不同。
鉴于对解决方案的高需求和有限的资源，研究人员转而采用创新方法，例如将自然法则嵌入 AI 模型中，提高其准确性以及减少对数据的依赖。
去年取得成功的一种方法是将科学问题的对称性嵌入 AI 模型。这些神经网络架构在等变神经网络（Equivariant Neural Networks，ENN）下广受欢迎，使用对称相关变换下的等变数学概念构建而成。
简而言之，ENN 旨在了解问题的基本对称性。例如，如果 ENN 的输入进行了旋转，输出也将相应旋转。这意味着，即使以不同的方向呈现，模型也可以识别相同的物体或图案。
为了更好地理解这一概念，请考虑如何使用 ENN 来维护 3D 中对称运算后输入和输出之间的关系。例如，如果 ENN 将分子的 3D 模型作为输入，并将其属性预测为输出，那么无需额外的训练数据或数据增强，就可以预测任何旋转版本的分子的相同属性。ENN“理解”分子的旋转不会改变其基本属性（图 1）。
在网络架构中引入这种基本的对称性，可在更改输入数据时提高模型的稳健性和数据效率。与将自然法则嵌入神经网络的其他策略类似，它还提供了一种提高对未见数据的泛化能力的方法。
所有这些优势都是有代价的：从理论上讲，构建 ENNs 并不简单，并且生成的网络在计算上比非等变版本更昂贵。在本文中，我们将介绍新的数学库 NVIDIA cuEquivariance 如何应对挑战并加速 AI 科学模型的开发，并提供药物发现和材料科学应用的示例。
等变神经网络的挑战
许多 AI 模型（包括 Tensor Field Networks、LieConv、Cormorant、SE(3)-Transformer、NequIP 以及 DiffDock 和 Equiformer 等其他模型）都使用独特的方法来确保一致地处理输入数据中的更改。它们使用对称组的基本元素，称为不可约表示（irreps）或这些元素的变体。这些不可约表示在数学上表示为张量，并且以特定的方式组合，通常涉及张量代数（例如张量积），以确保模型的输出适当反映应用于输入的任何对称变换。
采用使用 irreps 的 ENN 的一个瓶颈是，为给定对称组构建和处理这些 irrep 对象的理论复杂性。由于缺乏现有基元或可扩展 API，再加上理论复杂性，因此使用 irreps 形式进行 ENN 创新极具挑战性。在该领域，更容易选择重复使用现有的实现，即使它们不理想也是如此。
此外，在使用基于 irreps 的 ENNs 时存在计算复杂性。数学基础决定了 irreps 的矩阵表示。对于最常用的对称运算 (例如 3D 旋转)，这些大小对于计算优化来说可能并不常见，例如 5×5 或 7×7 矩阵。这不允许利用现有的优化技术 (例如数学运算中的 tensor cores)，将这些对象开箱即用。
更重要的是，涉及 irreps 的张量积运算遵循源于群论的异常稀疏模式，即使 irreps 本身是密集的。Irreps 具有称为 Clebsch-Gordan 系数的特殊混合系数，用于确定如何在代数运算中将两个不同的 irreps 组合为输出 irrep。
例如，相乘两个 irreps 只能得到特定且有限的输出 irreps 列表，并且此 选择规则 由群理论决定。事实上，由于选择规则，不允许进行多种 irreps 组合，这会导致稀疏的 Clebsch-Gordan 系数，其中大多数为零。从计算的角度来看，忽略群理论所决定的稀疏性会导致内存浪费和算法效率低下。
加速等变神经网络
为应对这些挑战，NVIDIA 开发了新的 cuEquivariance 数学库，该库为等变神经网络引入了 CUDA 加速的构建块。cuEquivariance 现已在 GitHub 和 PyPi 上提供公开测试版。
cuEquivariance Python 前端引入了一个名为分段张量积（Segmented Tensor Product，STP）的统一框架，该框架考虑到前面提到的混合系数的稀疏模式，使用 irrreps 组织代数运算。STP 泛化了等变多线性乘积的计算，使用户能够在 irrreps 之间表示各种此类运算。它还允许用户自由定义不一定等变的运算，这可能有助于研究社区尚未探索的应用。
cuEquivariance 基于 STP 框架构建，利用专门的 CUDA 内核来加速最常用的 STP 实例。在 ENNs 中，大多数瓶颈运算是多个受内存限制的运算，逐个执行，导致不必要的中间产品加载和存储。鉴于 irreps 的规模小且数量多，使用不同的内核调用执行每个操作是另一个开销来源。cuEquivariance 使用内核融合将这些单独的运算替换为一些特殊用途的 GPU 内核。
除了内核融合之外，我们还重新组织了特征的内存布局，使内存访问更好地映射到 NVIDIA GPU 架构的单指令多线程 (SIMT) 范式。这种专门的后端经过优化，可在 NVIDIA GPU 上显著加快等变神经网络中的数学运算。
图 2 显示了 cuEquivariance 加速对两个热门科学模型 AI 的影响：DiffDock 和 MACE，前者是预测蛋白质 – 配体结合姿态的扩散模型，后者是机器学习的原子间相互作用势，广泛用于材料科学和生物学，用于控制分子动力学模拟。
这些等变神经网络模型具有多个带 irreps 的张量运算。出于演示目的，为每个模型选择计算要求最高的运算。对于 DiffDock，这是一个基于 irrep 的张量积运算（TP）。对于 MACE，我们会考虑两种影响性能的运算：对称收缩（SC），即 irreps 张量自身的张量收缩；TP，类似于 DiffDock。对于每个运算，都会显示向前和向后性能。
图 3 展示了具有 cuEquivariance 的 MACE-OFF Large 和 MACE-MP Large 模型的端到端性能。最后，图 4 显示了在不同 NVIDIA GPU 之间的性能变化。
结束语
cuEquivariance 的开发标志着加速 AI 科学发展的重要一步。通过解决等变神经网络的理论和计算挑战，cuEquivariance 使研究人员、科学家和学者能够为各种科学应用构建更准确、更高效和更通用的模型。正如其成功集成到广泛使用的模型如 DiffDock 和 MACE 所证明的那样，cuEquivariance 已准备好在药物发现、材料科学等领域推动创新并加速发现。
通过利用对称性和高效计算的强大功能，cuEquivariance 为 AI 带来了新的可能性，有助于实现科学突破。将 cuEquivariance 等开源加速计算工具与系统生成的大规模数据集相结合，可以提高 AI 模型的准确性，从而促进研究和企业产品的更广泛采用和集成。
