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

Published: 17 Jun 2024, Last Modified: 25 Jul 2024ICML2024-AI4Science PosterEveryoneRevisionsBibTeXCC BY 4.0
Keywords: Geometric Deep Learning; Equivariant Neural Networks, Tools
TL;DR: We presents a modularized package for constructing a wide range of equivariant GNNs that follows the standard of Pytorch-Geometric with examples.
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.
Submission Number: 245
Loading