Can SevenNet also benefit from the recent NVIDIA cuEquivariance library?

[yjw0510]

I recently came across NVIDIA's cuEquivariance library. Their benchmarks show significant speedups for MACE when using the cuEquivariance math lib., both in training and inference compared to the base MACE implementation.

The performance was demonstrated on MACE models, but considering most modern GNN-based MLFF architectures share similar constructions, I believe SevenNet could potentially gain computational benefits as well. I'm interested in the SevenNet developers' perspectives on potential performance improvements.

The article references several related models:

Many AI models—including Tensor Field Networks, LieConv, Cormorant, SE(3)-Transformer, NequIP, and others like DiffDock and Equiformer—use a unique approach to ensure that they handle changes in input data consistently.

1. Given this context, I'm curious about the developers' thoughts on whether SevenNet could achieve similar computational speedups using the cuEquivariance library?

2. Are there any plans to integrate cuEquivariance support into SevenNet?

[YutackPark]

Thanks for noticing me. The short answer is yes.

1. I think not, especially for SevenNet-0. The number of channels is different for each 'L' value in SevenNet-0 (128x0e+64+1e+32x2e), which is different from that of MACE-MP-0. It restricts possible memory layout and it is may critical to optimization performance of cuEquivaraince. However, if we use the same number of channels, similar speed-up I expect.

2. Yes. But seeing the library, it replaces most of e3nn parts. As sevenn.nn and sevenn.model_build is on top of e3nn, it take me some times to refactor SevenNet.

Thanks again for this great news!

[yjw0510]

Thanks for the clarification. look forward to seeing how much more powerful SevenNet will be!

[YutackPark]

Initial update in cu_equi branch.

#127