Using Multiple Vector Channels Improves E(n)-Equivariant Graph Neural Networks

TL;DR

This paper introduces a multi-channel extension to E(n)-equivariant graph neural networks, enhancing their performance on physical system tasks with minimal computational overhead.

Contribution

The paper proposes a multi-channel approach to E(n)-equivariant GNNs, demonstrating improved accuracy across various physical system benchmarks.

Findings

Outperforms standard EGNN on N-body dynamics

Improves molecular property prediction accuracy

Enhances trajectory prediction of solar system bodies

Abstract

We present a natural extension to E(n)-equivariant graph neural networks that uses multiple equivariant vectors per node. We formulate the extension and show that it improves performance across different physical systems benchmark tasks, with minimal differences in runtime or number of parameters. The proposed multichannel EGNN outperforms the standard singlechannel EGNN on N-body charged particle dynamics, molecular property predictions, and predicting the trajectories of solar system bodies. Given the additional benefits and minimal additional cost of multi-channel EGNN, we suggest that this extension may be of practical use to researchers working in machine learning for the physical sciences