NVIDIA cuEquivariance: A High-Performance Geometric Deep Learning Library for Equivariant Neural Networks

cuEquivariance is a Python library developed by NVIDIA. It builds high-performance geometric neural networks using piecewise polynomials and trigonometric operations, providing underlying acceleration for mainstream models such as DiffDock, MACE, and Allegro.

In 3D space, physical laws exhibit translation and rotation invariance. Equivariance is the mathematical formalization of this concept—when input data undergoes geometric transformation, the model's output transforms accordingly in a predictable way. This property is crucial for tasks like molecular modeling, protein structure prediction, and materials science.

Traditional neural networks require large amounts of data to learn these symmetries when processing 3D geometric data. In contrast, equivariant neural networks directly embed physical symmetries into their architecture, making them more data-efficient and having stronger generalization capabilities. cuEquivariance is precisely a specialized tool library that provides underlying acceleration for such models.

cuEquivariance provides a complete set of APIs for describing and executing piecewise polynomial operations. Its core components include:

This is the basic operation unit of the library, supporting the decomposition of complex equivariant operations into efficiently executable CUDA kernel functions. The segmented design allows parallel computation on different data segments, fully leveraging the parallel capabilities of GPUs.

The library has built-in CUDA implementations specifically optimized for equivariant operations. Compared to the default implementations of general deep learning frameworks, it can achieve significant performance improvements on specific workloads. These kernel functions are carefully tuned to adapt to the characteristics of NVIDIA GPU architectures.

cuEquivariance provides bindings for both PyTorch and JAX. Developers can choose the familiar framework according to their project needs. This design ensures wide compatibility and easy integration:

• cuequivariance-torch: PyTorch frontend binding
• cuequivariance-jax: JAX frontend binding
• cuequivariance: Core non-ML components only
• cuequivariance-ops-*: CUDA operation kernel package

One of the design goals of cuEquivariance is to accelerate widely used equivariant neural network models in the industry:

A deep learning model for molecular docking prediction, widely used in the field of drug discovery. cuEquivariance can accelerate its geometric feature calculation and structure prediction processes.