NVIDIA cuEquivariance：为等变神经网络打造的高性能几何深度学习库

cuEquivariance是NVIDIA推出的Python库，通过分段多项式和三角运算构建高性能几何神经网络，为DiffDock、MACE、Allegro等主流模型提供底层加速。

在三维空间中，物理规律具有平移和旋转不变性。等变性（Equivariance）是这一概念的数学形式化表达——当输入数据发生几何变换时，模型的输出会以可预测的方式相应变换。这种特性对于分子建模、蛋白质结构预测、材料科学等任务至关重要。

传统神经网络在处理三维几何数据时，需要大量数据来学习这些对称性。而等变神经网络将物理对称性直接嵌入模型架构，使其更加数据高效、泛化能力更强。cuEquivariance正是为这类模型提供底层加速的专业工具库。

cuEquivariance提供了一套完整的API，用于描述和执行分段多项式运算。其核心组件包括：

这是库的基础运算单元，支持将复杂的等变操作分解为可高效执行的CUDA核函数。分段设计允许在不同数据片段上并行计算，充分利用GPU的并行能力。

库内置了针对等变运算专门优化的CUDA实现，相比通用深度学习框架的默认实现，在特定工作负载上可获得显著性能提升。这些核函数经过精心调优，适配NVIDIA GPU架构特性。

cuEquivariance同时提供PyTorch和JAX的绑定，开发者可以根据项目需求选择熟悉的框架。这种设计确保了广泛的兼容性和易于集成：

• cuequivariance-torch：PyTorch前端绑定
• cuequivariance-jax：JAX前端绑定
• cuequivariance：仅核心非ML组件
• cuequivariance-ops-*：CUDA运算核函数包

cuEquivariance的设计目标之一是加速业界广泛使用的等变神经网络模型：

用于分子对接预测的深度学习模型，在药物发现领域应用广泛。cuEquivariance可加速其几何特征计算和结构预测环节。