Geometric and Physical Quantities improve E(3) Equivariant Message Passing
Authors: Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, Max Welling
ICLR 2022 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies. We demonstrate the effectiveness of SEGNNs by setting a new state of the art on n-body toy datasets, in which our method leverages the abundance of geometric and physical quantities available. We further test our model on the molecular datasets QM9 and OC20. Although here only (relative) positional information is available as geometric quantity, SEGNNs achieve state of the art on the IS2RE dataset of OC20, and competitive performance on QM9. For all experiments we provide extensive ablation studies. |
| Researcher Affiliation | Collaboration | Johannes Brandstetter University of Amsterdam Johannes Kepler University Linz brandstetter@ml.jku.at Rob Hesselink University of Amsterdam r.d.hesselink@uva.nl Elise van der Pol Uv A-Bosch Delta Lab University of Amsterdam e.e.vanderpol@uva.nl Erik J Bekkers University of Amsterdam e.j.bekkers@uva.nl Max Welling Uv A-Bosch Delta Lab University of Amsterdam m.welling@uva.nl |
| Pseudocode | Yes | C.1 PSEUDOCODE OF SEGNN AND ABLATED ARCHITECTURES. Algorithm 1 Code sketch of an SElinear message passing layer. Algorithm 2 Code sketch of an SEnon-linear message passing layer. Algorithm 3 Code sketch of an SEGNN message passing layer. |
| Open Source Code | Yes | For the reproducibility of the QM9 experiments, we have uploaded our code and included a command which reproduces results of the α variable. |
| Open Datasets | Yes | The QM9 dataset (Ramakrishnan et al., 2014; Ruddigkeit et al., 2012) consists of small molecules up to 29 atoms... The Open Catalyst Project OC20 dataset (Zitnick et al., 2020; Chanussot et al., 2021), consists of molecular adsorptions onto surfaces... The charged N-body particle system experiment (Kipf et al., 2018) consists of 5 particles... |
| Dataset Splits | Yes | We use the dataset partitions from Anderson et al. (2019). (QM9) ...The IS2RE training set consists of over 450,000 catalyst adsorbate combinations with 70 atoms on average. The four test splits contain in-distribution (ID) catalysts and adsorbates, out-of-domain adsorbates (OOD Ads), out-of-distribution catalysts (OOD Cat), and out-of-distribution adsorbates and catalysts (OOD Both). (OC20) ...For the following experiments, we use 10.000 simulation trajectories in the training set, and 2.000 trajectories for the validation and test set, respectively. (Gravitational N-body dataset, Appendix C.2) |
| Hardware Specification | Yes | Results except for EGNN and SEGNN are taken from (Satorras et al., 2021) and verified. Runtimes are re-measured. ...Ge Force RTX 2080 Ti GPU. (Table 1) ...measured for a batch of 128 samples running on a Ge Force RTX 3090 GPU. (Table 3) ...running on an NVIDIA Ge Force RTX 3090 GPU. (Table C.3) ...measured for a batch of size 8 on a Nvidia Tesla V100 GPU. (Table C.4) |
| Software Dependencies | Yes | The implementation of SEGNN s O(3) steerable MLPs is based on the e3nn library (Geiger et al., 2021a). ...Our models were implemented in Py Torch (Paszke et al., 2019) (BSD license) using CUDA (proprietary license). We used Py Torch extensions such as Py Torch Geometric (Fey & Lenssen, 2019) (MIT license) and e3nn (Geiger et al., 2021b) (MIT license). For tracking runs we used Weights & Biases (Biewald, 2020) (MIT license). |
| Experiment Setup | Yes | We optimise models using the Adam optimiser (Kingma & Ba, 2014) with learning rate 1e-4, weight decay 1e-8 and batch size 100 for 10000 epochs and minimise the mean absolute error. ...We use a cutoff radius of 2 A. ...The network consists of three parts (sequentially applied): 1. Embedding network: Input {CG ai Swish Gate CG ai}, wherein CG ai denotes a steerable linear layer conditioned on node attributes ai and which are applied per node. 2. Four message passing layers as described in Sec. 2.1. 3. A prediction network: {CG0 ai Swish Linear Mean Pool Linear Swish Linear}, in which CG0 ai denotes a steerable linear layer conditioned on node attributes ai and which maps to a vector of 64 scalar (type-0) features. |