Accelerate Drug and Material Discovery with New Math Library NVIDIA cuEquivariance

Originally published at: https://developer.nvidia.com/blog/accelerate-drug-and-material-discovery-with-new-math-library-nvidia-cuequivariance/

AI models for science are often trained to make predictions about the workings of nature, such as predicting the structure of a biomolecule or the properties of a new solid that can become the next battery material. These tasks require high precision and accuracy. What makes AI for science even more challenging is that highly…

Hello Development Team,

Fantastic work on this release! I’m currently developing equivariant graph neural network (GNN) models, and it’s exciting to see advancements that improve GPU efficiency and simplify the software stack. The Efficient compilation and execution of equivariant operations has been a major bottleneck during development and training so far, so the speedups and ease-of-use here are great to see!

I also had a couple of questions:

  1. Hardware Support: While the improvements are optimized for NVIDIA GPUs, would we expect to see speedups on other hardware, such as AMD or other devices? My organization currently trains models across multiple platforms and I’m curious about any potential gains, or, at least, lack of losses.
  2. Equivariant Activations: Does cuEquivariance provide built-in equivariant activations for irreps, or would we still rely on libraries like e3nn for this functionality for now? I might have missed it in the documentation.

Thank you for advancing this important area—I believe it holds great promise, with ease of use being the primary roadblock.

  • Rylie