EquiTorch: A Modularized Package for Flexibly Constructing Equivariant GNNs Building upon PyTorch-Geometric

Abstract

Equivariant graph neural networks have recently gained significant attention due to their demonstrated effectiveness in geometric deep learning applications for scientific problems. However, despite the significant progress in developing equivariant graph neural networks, the implementations show slight diversity in their conventions, setting up a little barrier for pure AI researchers to get engaged in this fascinating field. Starting from this point, we would like to introduce our package, *EquiTorch*, which aims to collect the operations related to equivariant neural networks in a standardized style, following the framework and idea of "Message Passing Neural Networks" (MPNN) used in Pytorch-Geometric, which is familiar to classical AI researchers on GNNs. Besides, by aligning to the framework of MPNN, we presents the basic operations in a modularized way, which enables researchers to compose these operations flexibly and explore the larger space of design of equivariant GNNs. In near future, more comprehensive documentation and tutorials will be made available.