Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations
Authors: Anuroop Sriram
Abhishek Das
Brandon M Wood
Siddharth Goyal
C. Lawrence Zitnick
Venue: ICLR 2022 Poster (2022)
Keywords: Graph Neural Networks
Atomic Simulations
Computational Chemistry

Abstract: Recent progress in Graph Neural Networks (GNNs) for modeling atomic simulations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are memory intensive as they model higher-order interactions in the graphs such as those between triplets or quadruplets of atoms, making it challenging to scale these models. In this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling us to train very large GNNs with hundreds of millions or billions of parameters. We empirically evaluate our method by scaling up the recently proposed DimeNet++ and GemNet models by over an order of magnitude in the number of parameters. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parallelized models lead to relative improvements of 1) 15% on the force MAE metric on the S2EF task and 2) 21% on the AFbT metric on the IS2RS task, establishing new state-of-the-art results.

Published as a conference paper at ICLR 2022
TOWARDS TRAINING BILLION PARAMETER GRAPH NEURAL NETWORKS FOR ATOMIC SIMULATIONS
Anuroop S Abhishek Dast, Brandon M. Wood -, Siddharth Goyal', C. Lawrence Zitnick
Meet FAIR
i National Energy Research Scientific Computing Center (NERSC)
(anuroops, abhshkdz, bmwood, s sidooyal, zi zitnickkefbbccom
ABSTRACT
Recent progress in Graph Neural Networks (GNNs) for modeling atomic simu- lations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are mem- ory intensive as they model higher-order interactions in the graphs such as those between triplets quadruplets of atoms, making it challenging to scale these mod- els. this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling dreds of millions or billions of parameters. We empirically evaluate our method by scaling up the number of parameters of the recently proposed DimeNet++ and GemNet models by over an order of magnitude. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parailelized models lead to relative improve- ments of 1) 15% on the force MAE metric for the S2EF task and 2) 21% on the AFbT metric for the IS2RS task, establishing new state-of-the-ar results.
train very large GNNs with hun-
INTRODUCTION
Graph Neural Networks (GNNs) GGor et al. 2005 Zhou et al. 2020) have emerged as the standard
architecture of choice for modeling atomic systems, with wide range of applications from protein structure prediction to catalyst discovery and drug design (Schütt et al. 2017b Gilmer et al. 2017 Jørgensen al.] 2018 Zitnick et al. 2020, Schüit et al. 2017a Xie & Grossman 2018). These models operate on graph-structured inputs, where nodes of the graph represent atoms, and edges represent bonds or atomic neighbors. Despite their widespread success and the availability of large molecular datasets, training massive GNNs (with up to billions of parameters) is important but under-explored area. Success of similarly large models in computer vision, natural language pro- cessing, and speech recognition (Shoeybi et al.) 2020 Huang et al.) 2019: Brown al. 2020 Zhai
et al. 2021 suggests that scaling up GNNs could yield significant performance gains.
Most previous approaches scaling GNNs have focused scaling small models (with up few million parameters) to massive graphs. These methods generally assume that we are working with a large, fixed graph, leading the development of methods like neighborhood sampling (Jangda al.) 2021 Zheng et al. 2021) (Jia et al. 20200: Ma et al. 2019 Tripathy et al.) (020)). These methods do not apply atomic simulation datasets that contain millions of smaller graphs where it is necessary to consider the entire graph for prediction. Our focus is on the complementary problem of scaling to very large models for a dataset of many moderately-sine sized graphs (~ 1A nodes. 500k edges).
Another limitation of existing methods is that they focus on scaling simple GNN architectures such as Graph Convolutional Networks (GCNs) that only represent lower-order interactions i.e., repre- sentations for nodes and edges, that are then updated by passing messages between neighboring nodes. In practice, the most successful GNNs used for atomic systems also model higher-order in- teractions between atoms, such as the interactions between triplets or quadruplets of atoms (Klicpera) et al. 2020a 2021 Liu et al. 2022). These interactions are necessary to capture the geometry of the underlying system, critical in making accurate predictions. Scaling such GNN architectures is challenging because even moderately-sizer graphs can contain a large number of higher-order inter- actions. For example, single graph with 1 nodes could contain several million triplets of atoms.

---

Published as a conference paper at ICLR 2022
In this paper, introduce Graph Parallelism, approach to scale up such GNNs with higher-order interactions to billions of parameters, by splitting up the input graph across multiple GPUs.
We benchmark our approach by scaling up two recent GNN architectures DimeNet++ (Klicpera ict al., 2020a) and GemNet-T (Klicpera et al. 202T) on the Open Catalyst (OC20) dataset (Chanus- sot et al. 202T). The OC20 dataset, aimed at discovering new catalyst materials for renewable energy storage, consists of 134M training examples spanning a wide range of adsorbates and cat- alyst materials. A GNN that can accurately predict per-atom forces and system energies on OC20 has the potential to replace the computationally expensive quantum chemistry calculations based on Density Functional Theory (DFT) that are currently the bottleneck computational catalysis. Our large-scale, graph-parale suph-parallelized models lead to relative improvements of 1) 15% for predicting forces on the force MAE metric (S2EF task), and 2) 21 % on the AFbT metric for predicting relaxed structures (IS2RS task), establishing new state-of-ther of--the-rttt-ttt- results on this dataset.
2
GRAPH PARALLELISM
2.1 EXTENDED GRAPH NETS
Battaglia et al. 2018) introduced a framework called Graph Network (GN) that provides a general abstraction for many popular Graph Neural Networks (GNNs) operating on edge and node represen- tations of graphs. We build on their work and define the Extended Graph Network (EGN) framework to include GNNs that also operate on higher order terms like triplets or quadruplets of nodes.
In the Graph Network (GN) framework, a graph is defined as 3-tuple G = (u, V, F),, where represents global attributes about the entire graph: V [ux)= the set of all nodes, with os representing the attributes of node i; and E { Tkk s=))== k=N:N is the set of all edges where e, represents the attributes for the edge from node 8% to rk- A GNN then contains a series of GN blocks that iteratively operate on the input graph, updating the various representations.
In our Extended Graph Network (EGN) framework, a graph is defined as a 4-tule G = (u, E,E,T) T), where V, and E are defined in the Graph Network, and T = {(tm: Cmy Cm.:...)] is the set of higher-order interaction terms involving edges indexed by m1, m2,.
As concrete example, consider atomic system represented as graph in this framework with the nodes representing the atoms and the edges representing atomic neighbors. The node attributes 2 and edge attributes G₁ could represent the atom's atomic numbers and distances between atoms respectively. The higher order interactions could represent triplets of atoms, i.e., pairs of neighboring edges with tm representing the bond angle, which is the angle between edges that share common node. Finally, the global attribute u can represent the energy of the system. For clarity of exposition, we will limit our discussion to triplets in the rest of the paper, but higher order interactions can be handled in similar manner. We denote these triplets as (t.., Smy Em2).
In the EGN framework, GNNs then contain a series of EGN blocks that iteratively update the graph. Each EGN block consists of several update and aggregation functions that are applied to transform an input graph (u, V, T) into an output graph (u', V', E", T²). Each EGN block starts by updating the highest order interactions, which are then aggregated before updating the next highest order interaction. For instance, the triplet representations are first updated (using Triple(Update function) and then aggregated (TripletAggr) at each edge. These aggregated representations are then used to update the edge representation (EdgeeUpdate). Next, the edges going into a node are aggregated and used to update the node representations, and so on. This is illustrated in figure Ta
Many GNNs can be cast in the EGN framework using appropriate update and aggregation functions.
2.2 GRAPH PARALLELISM FOR EXTENDED GRAPH NETS
Training large EGNs can be challenging even on moderately sized graphs because of the large mem- ory footprint required in storing and updating the representation for each triplet, edge, and node. In many applications, the number of edges is one or two orders of magnitude larger than the number of nodes, while the number of triplets is one two orders of magnitude larger than the number of edges. Thus, storing and updating the triplet representations is often the bottleneck in terms of GPU memory and compute. Many recent methods such (Klicpera et al. 2020a 2021) overcome this
2

---

Published as a conference paper at ICLR 2022
Triplet Aggr.
Edge
Node
Aggr.
Aggr.
Triplet Update
Edge Update
Node Update
Global Update
Forward computation of EGN block. First, each triplet representation is up- dated (TripleUUpdate function), followed by aggregation of these updated representations (rrpletAggr)) Next, these aggregated values are used to update edge representations (edge, followed by edge aggregation (EdgeAggr), and finally node and global updates.
Processor
Processor
Processor
Processor
All Reduce
Global Triplet Aggregation
Triplet Update and Local Triplet Aggregation
Edge Update and Local Edge Aggregation
All Reduce
All Reduce
Global Edge
Global Node
Aggregation
Aggregation
Node Update and Local Node Aggregation
Global Update
(b) Distrihuted computation of EGN block. The graph split up among the different GPUs such that processing unit p contains its subset of triplets 7 (1()) in memory, along with all edges E' and nodes V. The triplet attributes updated in parallel, followed by triplet aggregation to locally update the edge attributes. Next, an allreduce operation is performed update edge attributes globally. We continue this process to update the node and global attributes.
problem by using very diwensiniional representation for the triplets. However, we found this to significantly reduce model capacity leading to underfitting for some applications with large training datasets. is necessary to overcome these memory limitations for better generalization.
One way to avoid the memory limits is to distribute the computation across multiple GPUs. For graph neural network, natural choice to distribute the computation is by splitting the graph. The update functions in an EGN are casy to apply in parallel since they are applied independently for each triplet, edge, or node. is substantially more challenging to apply the aggregations in parallel. To simplify parallel aggregation, and associative. This is not limiting in practice since most popular GNN architectures use sum or mean aggregations which can be implemented in this manner.
restrict ourselves to aggregation functions that are commutative
We will now describe the implementation to P processing units that we wish to split the graph computation over. Each unit p is responsible for computing the updates to subset of triplets. edges and nodes, that we denote by T(p), E(P), and V(P), respectively. At the beginning of computation, we split the graph that the processing unit contains its subset of triplets T(P) in memory, along with the entire set of edges E and nodes V. During the forward pass, we first update each set of triplets ( parallel to obtain TT6)), followed by a local triplet aggregation. Next, an all-reduce operation is performed on these locally aggregated triplet represetatiio globally aggregated representations are then used to update edges in parallel, and the process is repeated for nodes and ultimately for the global graph level representation. Figure Th shows this.
the distributed EGN block. Suppose we have access
obtain globally aggregated triplet representations. These
3

---

Published as a conference paper at ICLR 2022
In this framework, the highest order interaction attributes Therefore, the communication time is bound by the number of lower order node interactions. In our example, the triplet representations are communicated once per EGN block, making the total communication cost equal to O(N, , Dr + D2), where Dr and D. are the dimensions of node and edge representations. This makes it possible to work with large number of triplets and large triplet representations. In section 3 we show concrete instantiations of this framework for two contemporary GNNs.
communicated across processors.
not communicated, while the edge and node representations
3 GRAPH PARALLELISM FOR ATOMIC SIMULATIONS
In this section, we present two concrete examples of using GNNs for the problem of predicting the energy and the forces for atomic system, modeled as graph whose nodes represent the atoms and whose edges represent the atoms' neighbors. The GNN takes such graph as input and predicts the energy of the entire system well as 3D force vector each atom. Two paradigms have been proposed in the literature, that we call energy-centric and force-centric approaches. These approaches differ in how they estimate forces and whether they energy conserving, which is an important physical property. We begin by describing the components shared by both approaches, followed by one recent model in each paradigm.
3.1 INPUTS AND OUTPUTS
The inputs to the network are 3D positions x; E IR's and atomic number for each atom i E {I...n). The outputs are the per-atom forces f; R and the energy E E R the entire structure. The distance between atoms and j is denoted as di- ||x; Xj If edges (i, and (j, k) exist, then (i k)) defines a triplet in the graph and we denote the angle between the edges as okj.ji-
The input graph is constructed with each atom & (target) atom's neighbors s N, where N, contains all atoms s (source) within a distance 8, which is treated hyperparameter. Each edge has corresponding message m,t that passes information from source atom s to target atom t.
node and the edges representing the
3.2 ESTIMATING FORCES AND ENERGY
As previously stated, there are two paradigms for estimating the energy and forces for an atomic system: energy-centric and force-centric. In energy-centric models, the model first computes the energy by applying forward pass of the GNN: F. = GNN(x, z), where x, and z represent the atomic positions and atomic numbers respectively. The forces then computed as the negative gradient of the energy with respect to atomic positions by using backpropagation: f Vx
In force-centric models, the energy and forces are both computed directly during the forward prop- agation: E,F GNN(x, z) where F represents the matrix of all atomic forces. Force-centric models tend to be more efficient in terms of computation time and memory usage compared to the energy-centric models. However, energy-centric models guarantee that the forces are energy conserving which is an important physical property satisfied by atomic systems. In this work, demonstrate the benefits of scaling GNNs in both paradigms.
3.3 ENERGY-CENTRIC MODEL: DIMENET+
DimeNet++ (Klicpera al. 2020a) is a recently proposed energy-centric model for atomic sys- tems. In this model, the edges are represented by a feature representation mjj. which are iteratively updated using both directional information (via bond angles) as well as the interatomic distances.
(jz) The edges are initially represented using a radial basis function (RBF) representation CRBP of their (kjj)))
lengths dyr. The triplets are represented using a spherical basis function (SBF) expansion aSBF of the distances dej well as bond angles ((jjjj))) In each block, the messages updated as:
m²: : = fupdlate (mj..
fit(may
RBF ASBF
(j)) (kj.jt)
(1)
kEA
4

---

Published as a conference paper at ICLR 2022
where the interaction function fint corresponds to the TripletUpdate function, the summation cor- responds to the TripletAggr function, and the update function fupdtte to the EdgeUpdate function respectively. The interaction function consists of a Hadamard product of the embeddings, followed by a multi-layer perceptron. To speed up these computations, the embeddings are projected down to smaller dimension before computing interactions, and later projected back up.
The updated messages are then fed as input to an output block that sums them up per atom to
obtain a per-atom representation: h, = -, m/e (EdgeAggr function). These are then transformed by another MLP obtain the node representation of UNodeUpdate function).
Thus, DimeNet++ is case of EGN, and we closely follow the recipe from Sec. 2.2 to parallelize it.
3.4 CEICE-CENTRIC MODEL: GEMNET
GemNet (Klicpera al. 2021) extends DimeNet++ in number of ways, including the addition of quadruplets as well as a force-centric version. Here, we focus only on the force-centric GemNet-T model since was recently shown to obtain state-of-the-art results on the OC20 datase GemNet-T largely follows the structure of the DimeNet++ model, but includes some modifications.
First, GemNet-T uses a bilinear layer instead of the Hadamard product for the interaction function. This is made efficient by optimally choosing the order of operations in the bilinear function to minimize computation. Second, GemNet-T maintains an explicit embedding for each atom that is first updated by aggregating the directional embeddings involving that atom, similar to DimeNet++. Next, the updated atom embedding is used to update each of the edge embeddings. This creates second edge update function, Edge Update', that is after the node embeddings are updated. Third, GemNet makes use of symmetric message passing, that is the messages mjj and m,j that are on the same edge, but in different directions, are coupled. In a parallel implementation, this step requires an additional all-reduce step since messages my and maj could be on different processors.
Thus, GemNet-T largely follows the EGN framework, with few minor deviations from the standard formulation. The distributed EGN implementation described in section 2.2 can used for the GemNet-T well, but with additional communication steps to account for the second edge update function and symmetric message passing.
4 EXPERIMENTS
In this section, we present the results of our scaling experiments on the Open Catalyst 2020 (OC20) dataset (Chanussot* et al. 2021). The OC20 dataset contains over 130 million atomic structures used to train models for predicting forces and energies during structure relaxations. We report results for three tasks: 1) Structure Energy and Forces (S2EF) that involves predicting energy and forces for a given structure; 2) Initial Structure to Relaxed Energy (IS2RF) that involves predicting the relaxed energy for given initial structure; and 3) Initial Structure Relaxed Structure (IS2RS) which involves performing a structure relaxation using the predicted forces. DimeNet++ and GemNet-T the current state-of-the-arr energy-centric and force-centric models respectively.
4.1 EXPERIMENTAL SETUP
DimeNet++-XI. Our DimeNet++ model consists of B = interaction blocks, with hidden dimen- sion of H 2048, output block dimension of D 1536, and intermediate triplet dimension of T = 256. This model has about 240M parameters, which is over 20x larger than the DimeNet++- large model used in *Chanussot et al. (02TT. We call this model DimeNei+ +XX The model was trained with the AdamW optimizer (Kingma & Ba. 2015 Loshchilov & Hutter 2019) starting with an initial learning rate of 10- that was multiplied by 0.8 whenever the validation error plateaus. The model was trained with an effective batch size of 128 on 256 Volta 32GB GPUs with a com- bination of data parallel and graph parallel training: each graph was split over GPUs with data parallel training across groups of 4 GPUs.
spenaatalystprrootttoor///eeeeee
5

---

Published as a conference paper at ICLR 2022
Model
AParams
Training GPU-days
Energy MAE (eV))
Force MAE (eVAA)]
Force Cos?
EPw/T
ForceNer-Large He et al. 2021
34&M
194
2.2628
0,03115
0.5195
0.01%
DimeNatee Killooers ara.. 2020a
10.8M
1600
31,5409
0,03132
0,540
0.00%
SpizCony Shuaihi 2021
8.9M
0.3363
0.02966
GomNet-1 [Klkpera al. 2021
31M
0.2921
0.02/22
06162
1.20%
GemNettXL
100M
1962
0,2701
0,02040
0,6003
181%
Table 1: Experimental results on the S2EF task comparing our GemNet-XL. to the top entries on the Open Catalyst leaderboard, showing metrics averaged across the 4 test datasets.
Training Dataset
IS2RS Test ADwTj
Model
#Params
AFbT
FbTT
SpinConv (Shuaibi et al.
2021
8.9M
S2EF-AIl
16.67%
53.62%
0.05%
DimeNet++ KKicpera al. 2020a
1.8M
S2EF 20M MD
17.15%
47.72%
0.15%
DimeNet++ large (Klcppra al.
2020a
10.8M
S2EF-AII
21.82%
51.68%
0.40%
GemNet-T T Klicpera et al. 2021
31M
S2EF-AII
27.60%
58.68%
0.70%
DimeNet44-XI.
240M
S2EF 20M + MD
33.44%
59.21%
1.25%
GemNet-XI.
300M
S2EF-AII
30.82%
62.65%
0.90%
Table 2: Results the IS2RS task comparing models to the top entries the Open Catalyst leaderboard, showing metrics averaged across the 4 test datasets. The DimeNet++ and DimeNet++- XL models were trained the S2EF 20M MD dataset, that contains additional molecular dynam- ies data and has been shown to be helpful for the IS2RS task (Chanussot* et al. 2021).
GemNet-XL. Our GemNet model consists of B = interaction blocks, with an edge embedding size of E 1536, triplet embedding size of T 384 and embedding dimension of the bilinear layer of B = 192. We found that it was beneficial to use a small atom embedding size of A = 128, much smaller than the previous SOTA GemNet mode This model has roughly 300M parameters, which is about 10x larger than the previous SOTA model. However, since we reduced the atom embedding dimension and increased the edge and triplet dimensions, the total amount of compute and memory usage is significantly larger. We call this model GemNer-XL. We followed the same training procedure as with DimeNet++, except for starting learning rate of 10-4
4.2 STRUCTURE TO ENERGY AND FORCES (S2EF)
The Structure to Energy and Forces task takes an atomic structure as input and predicts the energy of the entire structure and per-atom forces. The S2EF task has four metrics: the energy and force Mean Absolute Error (MAE), the Foree Cosine Similarity, and the Energy and Forces within a Threshold (EFwT). EFWT indicates the percentage of energy and force predictions below a preset threshold.
Table compares the top models on the Open Catalyst Project leaderboard with the GemNet- XL model. GemNet-XL obtains roughly 16% lower force MAE and an 8% lower energy MAE relative to the previous state-of-the-art Further, GemNet-XI. improves the EFWT metric more than 50% relative to the previous best, although the value is still very small. These results indicate that model scaling is beneficial for the S2EF task.
4.3 INITIAL STRUCTURE TO RELAXED STRUCTURE (IS2RS)
The Initial Structure to Relaxed Structure (IS2RS) task involves taking an initial atomic structure and predicting the atomic structure that minimizes the energy. This is performed by iteratively predicting the forces on each atom and then using these forces update the atomic positions. This process is repeated until convergence or 200 iterations. There are three metrics for this task: Average Distance within Threshold (ADwT), that measures the fraction of final atomic positions within a distance threshold of the ground truth; Forces below Threshold (FbT), which measures whether true energy minimum was found (i.e., forces are smaller than preset threshold); and, the Average Forces below Threshold (AFbt), which averages the FbT over several thresholds.
Table shows the results on the IS2RS task, comparing our models with the top few models on the Open Catalyst leaderboard. Both the DimLNett+-LL and GemNet-XL models outperform all
discuss. http:sss...ppaaaaapppppppteeee/e
moddlll mwwee--wiwghts//ee.

---

Published as a conference paper at ICLR 2022
IS2RE Test
Model
Approach
Energy MAE (EV)
EwTh
SpinConv Shuaibi
2021
Relaxation
0.4343
7.90%
GemNet-T Klicpera cl al.
2021
Relaxation
0.3997
9.86%
Noisy Nodes (Godwin
2022
Direct
0,4728
6.50%
3D- AD-Grrphormer (Yigg
2021
Direct
0.4722
6.10%
GemNet-XL GemNet-XL-FT
Relaxation
0.3712
11.13%
Direct
0.4623
5.60%
Table 3: Results on the IS2RE task comparing our GemNet-XL, to the top entries on the Open Catalyst leaderboard, showing metrics averaged across the 4 test datasets.
#Blocks Node Dim
Edge Dim
Trip Dim
Bil Dim
Params
#GP GPUs
#GP+DP GPUs
1280
768
128
64
125M
32
1536
1024
192
96
245M
64
1792
1184
288
160
480M
128
2320
1302
512
288
960M
256
Table 4: Model hyperparameters for the scaling analysis. "#GP GPUs" denotes the number of GPUs which the graph is distributed over for pure graph parallel training single node. "#GP+DP GPUs" denotes the total number GPUs used to train with graph parallel training together with 32-way data parallel training.
existing models. The DimeNet++-XL model obtains a relative improvement of 53% on the AFbT metric, and more than triples the FbT metric compared to the DemeNet-large model. The GemNet- XL model obtains similar improvements compared to the smaller GemNet-T model. These results underscore the importance of model scaling for this task.
4.4 INITIAL STRUCTURE TO RELAXED ENERGY (IS2RE)
The Initial Structure to Relaxed Energy (IS2RE) task takes an initial atomic structure and attempts to predict the energy the structure after has been relaxed. Two approaches can be taken to address this problem, the direct and the relaxation approaches (Chanussot* et al. (02T). In the direct approach, we treat this task as regression problem, and train model to directly estimate the relaxed energy for given atomic structure. The relaxation approach first estimates the relaxed structure (IS2RS task) after which the energy is estimated using the relaxed structure as input. Relaxation approaches typically outperform the direct approaches, though they are generally magnitude slower during inference time due to the need to estimate the relaxed structure.
orders of
Table 3 compares our GemNet-XL model the top three models from the Open Catalyst Project leaderboard. There are two metrics for the IS2RE task: energy Mean Absolute Error (MAE) and the Energy within Threshold (EwT) which measures the percentage of time the predicted energy is that the GemNet-XL. model obtains roughly within threshold of the true energy. EwI compared to the previous best, which is the smaller 8% lower energy MAE and 12% higher GemNet-T model, demonstrating the benefits of scaling up IS2RE models.
Since direct IS2RE models are very fast at inference time, there are applications where they are more useful than relaxation based approaches. It is possible to convert a trained S2EF model into direct IS2RE model by fine-tuning it on the IS2RE training data. We finetune our GemNet-XI. model in this manner for 5 epochs, starting with an initial learning rate of 3 10-5 that is exponentially decayed by multiplying with 0.95 at the end of each epoch. The resuting model GemNet-XL-FFT obtains relative improvement of 2% energy MAE compared to 3D-Graphormer (Ying et al. 2021), the current state-of-te oft------tt-rr--- direct approach (Table B
4.5 SCALING ANALYSIS
Weak scaling studies the effect on the throughput when computation is scaled proportional to the number of processors. We study weak scaling efficiency in terms of scaling the model size propor-

---

Published as a conference paper at ICLR 2022
200
Parsa Gapt da
Grazt Faraist Rorieeeraid
ter
Tera LIPs
Number GPUs
Number of GPUs
(a) Weak scaling efficiency.
(b) Comparison graph and pipeline parallel training.
Figure 2: Left: Weak scaling efficiency measured by scaling the model size proportional to the number of GPUs. Right: Comparing graph parallelism with pipeline parallelism. Solid blue, green and red curves show the scaling performance of graph, pipeline, and graph+pipeline parallel training. Dashed lines show the same with 32-way data parallel training. We do not show the results of training the largest models with pipeline parallelism those runs out of GPU memory.
tional to the number of GPUs. For these experiments, we use 4 different GemNet-T models, ranging from 120M parameters to nearly IB parameters (see Table 4
We train increasingly larger models on multiple GPUs that do not fit on a single GPU. Figure 2a shows that we are able to obtain scaling efficiency roughly 79% with GPUs for our largest model with nearly billion parameters. This shows that our graph parallel training can be scaled up to billion parameter GemNet-T models while obtaining reasonably good scaling performance.
Figure 2a further shows the scaling efficiency for each model combining graph parallelism with 32-way data parallelism. Pure data parallel training with 32 GPUs obtains a scaling efficiency of 75% for the smallest model, showing the effect of network communication and load imbalance between the GPUs. Combining graph and data parallel training with 256 GPUs only reduces the scaling efficiency to 47% for the largest model compared to the 1 GPU case, suggesting the graph parallelism promising for training extremely large models on hundreds of GPUs.
Finally, figure 2b shows the raw performance of running these models V100 GPUs in terms of TeraFL.OPs per second as a function of the number of GPUs. On single GPU, the 120M parameter GemNet-T sustains 32 TeraFLOPs or roughly 25% of the theoretical peak FLOPS for single GPU, and is thus a strong baseline. With 256 GPUs, the largest model sustains 3.5 PetaFLOPs.
Figure 2b compares graph and pipeline parallelism (Huang al. 2019) showing that graph paral- lelism outperforms pipeline parallelism for the models that we consider. Since each graph in the training data contains different numbers of nodes, edges and triplets, load balancing across GPUs is difficult for pipeline parallelism. Graph parallelism is able to overcome this problem since the nodes, edges and triplets of given batch are always distributed evenly across the GPUs, helping outperform pipeline parallelism. It possible, however, that pipeline parallelism might outperform graph parallelism for very deep GNNs since inter-GPU communication overhead for pipeline paral- lelism is independent of the number of blocks. Figure 2blalso shows results with graph and pipeline parallelism combined, indicating that these methods are complementary to each other.
5 RELATED WORK
GNNs for simluating atomic systems Many GNN based approaches have been proposed for the task of estimating atomic properties such as (Schütt et al. 2017b Gilmer et all 2017 Jørgensen
et al. 2018 [Schütt et al., 2017a 2018: Xie & Grossman 2018 Qiao et al. 2020 Klicpera et al., 2020b), where the atoms are represented by nodes and neighboring atoms are connected by edges. An early approach for force estimation was the SchNet model Schütt et al. (2017a). which computed forces using only the distance between atoms without the use of angular information. SchNet pro- posed the use of differentiable edge filters which enabled constructing energy conserrvin models by estimating forces as the gradient the energy. Subsequent work (Klicpera et al. 2020a 2021 Liu et al. 2022) has extended on the SchNet model by adding bond angles and dihedral angles,
8

---

Published as a conference paper at ICLR 2022
which has resulted in improved performance. These models make use of higher order interactions among nodes which make them highly compute and memory intensive. An alternate approach for estimating forces is to directly regress the forces does not enforce energy conservation or rotational equivariance, recent work (Hu et al. 2021) has
output of the network. While this approach
shown that such models can still accurately predict forces.
Distributed GNN Training. Research on distributed GNN training has focused on the regime of training small models on single, large graph, for instance by sampling local neighborhoods around nodes to create mini-batches (Jangda et al., [2021] Zhang et al., 2020 Zhu et al. 2019). While these approaches can scale very large graphs, they are not suitable for the task of modeling atomic
systems where it is important to consider the entire graph for predicting the energy and forces.
An alternate line of work, that is more similar to ours, keeps the entire graph in memory by efficiently partitioning the graph among multiple nodes (Jia et al. 2020 Ma et al. 2019 Tripathy et al.) 2020). These methods still operate a single graph regime that is partitioned ahead of time. Thus the focus of that work is on finding efficient partitions of the graph, which is not applicable to our problem since we operate on millions of graphs. Further, these works do not train very large GNNs, or GNNs that operate on higher-order interactions (e.g. triplets), that are important for atomic systems.
Model Parallelism. methods focus on training large models that do not fit entirely on one GPU (even with a batch size of 1). GPipe (Huang al. 2019) splits different sequences of layers into different processors, and splits each training mini-batch into micro-batches to avoid idle time per GPU. Megatron-LM (Shoeybi et al., (020) splits the model breadth-wise, where Transformer layer weights are partitioned across multiple GPUs to distribute computation. We see model and graph parallelism as complementary approaches that can be combined train even larger models.
6 DISCUSSION
We presented Graph Parallelism, a new method for training large GNNs for modeling atomic simu- lations, where modeling higher-order interactions between atoms (triplets / quadruplets) is critical. We showed that training larger DimeNet++ and GemNet-T models can yield significant improve- ments on the OC20 dataset. Although we demonstrated graph parallelism for just two GNNs, it is broadly applicable a host of message- message-passing GNNs, including equivariant networks, that can be cast in the GN/ / EGN framework (Sec. 2.1 by appropriately picking update and aggregate functions.
Further, it is possible to combine graph parallelism with model parallel methods such as GPipe (Huang et al.] 2019) to train even larger models, which could yield further improvements, as briefly explored in Sec 4.5 For force-centric GNNs, it should be possible to use graph parallel for "breadth-wise" scaling i.e., split higher-order computations (e.g. triplets) across GPUs, and GPipe for 'dipth-wise" scaling iee., to scale to larger number of message- paaaagggg blocks, sequentially split across GPUs For energy-centre GNNs e.g., DimeNet++, this combination is less obvious since these models require an additional backward pass through the network to compute forces as gra- dients of energy with respect to atomic positions. Energy-centric GNNs common for atomic systems because they enforce the physical constraint of energy conservation. As we demonstrate with our DimeNet++-XL experiments, graph parallelism is applicable energy-centric GNNs.
We see scaling to large model sizes as a necessary (but not sufficient) step to effectively model atomic simulations from large, diverse datasets of adsorbates and catalysts. Further progress may require marrying large scale with ways to better capture 3D geometry and physical priors.
The carbon emissions from training large deep learning models are non- and the work we have presented here is no exception (Strubell et al. | Schwartz et al. 2019). We estimate that training our GemNet-Xi model with Tesla v00 32 GB GPUs on cloud resources in the US ranges from 3490 8052 kg of CO2 eq. (Lacoste et al. 2019). In the worst case, the emissions are roughly equivalent to 16 round trip flights from Los Angeles to New York. Assuming that the training time is fixed, emissions largely depend on the carbon intensity of the energy generation in given region and the percentage of emissions offset with other investments by the resource provider. When choosing compute resources we recommend evaluating the stated carbon offset commitments and if possible, consider running experiments in regions that have more sustainable energy generation. The compute resources utilized for this paper were committed to be 100% offset.

---

Published as a conference paper at ICLR 2022
REFERENCES
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzale, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre. Francis Song. Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks, 2018.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are sew-shot learners. In NeurlPS, 2020.
Lowik Chanussot* Abhishek Das*, Siddharth Goyal*, Thibaut Lavril*, Muhammed Shuaibi*, Mor- gane Riviere, Kevin Tran, Javier Heras- cerras-Dmmingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram. Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open Catalyst 2020 (OC20) Dataset and Community Challenges. ACS Catalywis, 2021.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In ICMI.,2111
Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzaler, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. In ICLR, 2022.
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In IJCNN, 2005.
Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram, Jure Leskovec, Devi Parikh, and C Lawrence Zitnick. ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations. arXiv preprint arXiv:2103.01436, 2021.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V. Le. Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In NeurlPS, 2019.
Abhinav Jangda, Sandeep Polisetty, Arjun Guha, and Marco Serafini. Accelerating graph sampling for graph machine learning using gpus. arXiv preprint arXiv:2009.0661
2021.
Zhihao Jia, Sina Lin, Mingyu Gao, Matei A. Zaharia, and Alexander Aiken. Improving the accuracy, scalability, and performance of graph neural networks with roc. In MLSys, 2020.
Peter Bjørn Jørgensen, Karsten Wedel Jacobsen, and Mikkel N Schmidt. Neural message pass- ing with edge updates for predicting properties molecules and materials. arXiv preprint arXiv:1806.02 nat.1880003144 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Johannes Klicpera, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In NeurIPS-W, 2020a.
Johannes Klicpera, Janck Groß, and Stephan Günnemann. Directional message passing for molec- ular graphs. In ICLR, 2020b.
Johannes Klicpera, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules. In NeurIPS, 2021.
Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and Thomas Dandres. Quantifying the carbon emissions of machine learning. arXiv preprint 1910v 2019.
10

---

Published as a conference paper at ICLR 2022
Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In ICLR, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:/111.05101, 2019.
Lingxiao Ma, Zhi Yang, Youshan Miao, Jilong Xue, Ming Wu, Lidong Zhou, and Yafei Dai. Neu- graph: Parallel deep neural network computation on large graphs. In USENIX Annual Technical Conference, 2019.
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R Manby, and Thomas Miller III. Orbnet: Deep learning for quantum chemistry using symmetry-adapter atomic-orbital features. arXiv preprint arXiv:2007.08026 080006 2020.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filte convolutional neural network for modeling quantum interactions. In NewrlPS, 2017a.
Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko. Quantum-chemii insights from deep tensor neural networks. Nature communications, 8:13890, 2017b.
Kristof T Schlitt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
Roy Schwartz, Jesse Dodge, Noah A. Smith, and Oren Etzioni. Green AL arXiv preprint arXiv:1907 2019.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint (etdv::999.00000 2020.
Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick. Rotation Invariant Graph Neural Networks using Spin Convolu- tions. arXiv preprint (t..ix:200000575: 2021.
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in NLP. arXiv preprint /906. 02244, 2019.
Alok Tripathy, Katherine Yelick, and Aydin Bulue. Reducing communication in graph neural net- work training. arXiv preprint 2xX::20005003000 2020.
Tian Xie and Jeffrey C Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical review letters, 120(14): 2018.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do Transformers Really Perform Badly for Graph Representation? In NeurIPS, 2021.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lueas Beyer. Scaling vision transformers. arXiv preprint arXiv:2106.04500 2021.
Dalong Zhang, Xin Huang, Ziqi Liu, Zhiyang Hu, Xianzheng Song. Zhibang Ge, Zhiqiang Zhang, Lin Wang. Jun Zhou, Yang Shuang, and Yuan Qi. AGL: a Scalable System for industrial-purpott Graph Machine Learning, arXiv preprint arXi::2003.02454 2020.
Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, and George Karypis. DistGLI: Distributed Graph Neural Network Training for Billion-Seale Graphs. arXiv preprint arXiv:2010.05337, 2021.
Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changeheng Li, and Maosong Sun. Graph neural networks: A review of methods and applica- tions. AI Open, 1:57-81, 2020.
11

---

Published as a conference paper at ICLR 2022
Rong Zhu, Kun Zhao, Hongxia Yang. Wei Lin, Chang Zhou, Baole Ai, Yong Li, and Jingren Zhou. AliGraph: A Comprehensive Graph Neural Network Platform. arXiv preprint arXiv:1902.087330 2019.
C. L. Zitnick, L. Chanussot, A. Das, S. Goyal, J. Heras-Domingo, C. Ho, W. Hu, T. Lavril, A. Pal- izhati, M. Riviere, M. Shuaibi, A. Sriram, K. Tran, B. Wood, J. Yoon, D. Parikh, and Z. Ulissi. An introduction electrocatalyst design using machine learning for renewable energy storage. arXiv preprint arXiv:2010.09435,, 2020.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1
ADDITIONAL RESULTS
Tables 6 and 7 show results from each of the four test sets for the S2EF, IS2RE and IS2RS tasks respectively. DimeNett+-XL and GemNet-XL. achieve the best results for each test set along each metric.
S2EF Test
Model
Energy MAE (EV)
Force Cos
Force MAE (EVIA)
EFWTf
ID
GemNet-T
0.2257
0.637
0.02099
2.4%
SpinConv
0.2612
0.5479
0.02689
0.82%
LoreeNet-Large
2.0674
0.533
0.02782
0.02%
lameNet++-large
29.3333
0.5633
0,02807
0%
GemNet-XI.
0.2120
0.6759
0.0181
3.30%
OOD Ads
GemNer-T
0.2099
0.6242
0.02186
1.15%
SpinConv
0,2752
0.5345
0,02769
0.38%
ForoeNet-Large
2.4188
0.5212
0.02834
0.01%
laneNet++-large
30.0338
0.5503
0.02896
0%
GeMNet-XL
0.1980
0.6642
0.0186
1.62 %
OOD Cat
GemNet-T
0.3403
0.5813
0.02445
0.93%
SpinConv
0.3501
0.5187
0.02849
0.46%
ForaeNet-LLarge
2.0203
0.4936
0.03089
0.01%
DimeNet++-Lagge
30.0437
0.5109
0,0312
0%
GemNet-XL
0.3083
0.6306
0.0206
1.72%
OOD Both
GemNet-T
0.3937
0.6221
0.02956
0.3%
SpinConv
0,4585
0.5554
0,03556
0.14%
Lorcenet-Large
2.5447
0.5302
0.03754
0%
DimeNet++-Large 36.7529
0.5517
0.03705
0%%
GemNet-XI.
0.362
0.6704
0.0245
0.61 %
Table 5: Full set of results on the S2EF task comparing our GemNet-XL to the top three entries on the Open Catalyst leaderboard, showing metrics from each test set.
13

---

Published as a conference paper at ICLR 2022
IS2RE Test
Model
Approach
Energy MAE (EV)
EwTf
ID
GemNet-T SpinConv
Relaxation
0.3901
12.37%
Relaxation
0.4207
9.4%
Noisy Nodes
Direct
0.4776
5.71%
GemNet-XL
Relaxation
0.3764
13.25%
GemNet-XL-FT
Direct
0.4194
7.52%
OOD Ads
GemNet-T SpinConv
Relaxation
0.3907
9.11%
Relaxation
0.4381
7.47%
Noisy Nodes
Direct
0.5646
3,49%
GemNet-XL
Relaxation
0.3677
10.00
GeLFFt-XLL-FT
Direct
0.5258
3.95%
OOD Cat
GemNet-T SpinConv
Relaxation
0.4339
10.09%
Relaxation
0,4579
8.16%
Noisy Nodes
Direct
0.4932
5.02%
GemNet-XL.
Relaxation
0,4022
11.61 %
GemNet-XL-FT
Direct
0.4373
6.76%
OOD Both
GemNet-T
Relaxation
0.3843
7.87%
SpinConv
Relaxation
0.4203
6.56%
Noisy Nodes
Direct
0.5042
3,82%
GemNet-XL
Relaxation
0.3383
9.65
GemNet-XLL-FT
Direct
0.4665
4.19%
Table Experimental results on the IS2RE task comparing our GemNet-XL. to the top three entries on the Open Catalyst leaderboard, showing metrics from each test set.
14

---

Published as a conference paper at ICLR 2022
Training Dataset
IS2RS Test
Model
AFPТ
ADwTf
FbT
ID 33.75%
GemNet-T
S2EF-ALL
59.18%
2.0%
lareNet++laarge
S2EF-ALL
25.65%
52.45%
1.0%
SpinConv
S2EF-ALL 21.10%
53.68%
0.2%
DimeNet+++
S2EF 20M MD
21.08%
48.6%
0.2%
+X
S2EF 20M +
40.00%
59.90%
2.4%
GemNet-XL.
S2EF-ALL 34.61
62.73%
2.4%
OOD Ads
GemNet-T
S2EF-ALL
26.84%
54.59%
0.2%
DimeNet++llarge
S2EF-ALL
20.73%
48.47%
0.4%
SpinConv
S2EF-ALL 15.70%
48.87%
0.0%
DimeNet++
S2EF 20M MD
17.05%
42.98%
0.0%
DimeNet++-XI
S2EF 20M MD
36.01%
55.68%
1.6
GemNet-XL
S2EF-ALL
30.32%
58.57%
0.6%
OOD Cat
GemNet-T
S2EF-ALL
24.69%
58.71% %
0.4%
large
S2EF-ALL
20.24%
50.99%
0%
SpinConv
S2EF-ALL 15.86%
53.92%
0%
DimeNett+-
S2EF 20M MD
16.43%
48.19%
0%
DimeNet++-XX
S2EF 20M MD
29.62%
58.43%
0.6%
GemNet-XL
S2EF-ALL 29.33%
62.60 %
0.2%
OOD Both
GemNet-T
S2EF-ALL
25.11%
62.23%
0.2%
laree
S2EF-ALL
20.67%
54.82%
0.2%
SpinConv
S2EF-ALL 14.01%
58.03%
0%
DimeNet++
S2EF 20M MD
14.02%
51.09%
0.4%
DiM
S2EF 20M MD
28.14%
62.85%
0.4%
GemNet-XL
S2EF-ALL
29.02%
66.72%
0.4% %
Table 7: Experimental results on the IS2RS task comparing our models the top four entries on the Open Catalyst leaderboard, showing metrics for each test dataset. The DimeNet++ and DimeNet++- XL models were trained on the S2EF 20M MD dataset, that contains additional molecular dynam- ics data, which has been shown to be helpful for the IS2RS task.
15


Reviews:
[{"review_id": "lIV7Mz01KKp", "reviewer": "Reviewer_ZcRk", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is based on the idea that most of the computation costs to extended graphs come from triplets. The authors proposed a way to parallelize the computation of triplets in a distributed way. The author also discussed two models that fit into this framework and showed their increase in performance due to the larger parameter count enabled by their parallel framework.\n\nThe paper is well written and easy to follow. A minor part I am a bit confused about is that in the sentence “In many applications, the number of edges is one or two orders of magnitude larger than the number of nodes, while the number of triplets is one or two orders of magnitude larger still”, do you mean the triplets are one or two orders larger than the node or the edges?\n\nIn the experiment section, network structures and hyperparameters are given and the experiment should be reproducible. The results support the authors' claim that larger graph networks have better performance.\n\nI do not know much about this field.", "summary_of_the_paper": "The paper is based on the idea that most of the computation costs to extended graphs come from triplets. The authors proposed a way to parallelize the computation of triplets in a distributed way. The author also discussed two models that fit into this framework and showed their increase in performance due to the larger parameter count enabled by their parallel framework.", "summary_of_the_review": "The paper is based on the idea that most of the computation costs to extended graphs come from triplets. The authors proposed a way to parallelize the computation of triplets in a distributed way. The author also discussed two models that fit into this framework and showed their increase in performance due to the larger parameter count enabled by their parallel framework.\n\nThe paper is well written and easy to follow. A minor part I am a bit confused about is that in the sentence “In many applications, the number of edges is one or two orders of magnitude larger than the number of nodes, while the number of triplets is one or two orders of magnitude larger still”, do you mean the triplets are one or two orders larger than the node or the edges?\n\nIn the experiment section, network structures and hyperparameters are given and the experiment should be reproducible. The results support the authors' claim that larger graph networks have better performance.\n\nI do not know much about this field.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_3WVdW3Hugt", "reviewer": "Reviewer_4Syd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:  \n\nOverall the paper is well structured  \n\nMethods that enable effective training of large scale GNNs on datasets containing many graphs could be very impactful  \n\nThe large GNN models are benchmarked on a set of tasks that are important in the real world (catalyst design), and show quite large improvements in performance compared to the smaller baselines  \n\nWeaknesses:\n\nNot sure how reproducible this work is without code  \n\nOther comments/questions:  \n\nSome additional information on the computational cost of Gemnet-XL and Dimenet++-XL models in the 3 tasks would be useful. Eg wall time, gpu hours, etc. Although methods to train very large GNN models can be very impactful, if the computational cost is prohibitive for a lot of people then that could limit its impactfulness. \n", "summary_of_the_paper": "This paper proposes a method to train large scale graph neural networks containing up to billions of parameters, called Graph Parallelism. The method is used to train large scale versions of the DimeNet++ and GemNet models containing 10-20x more parameters that the vanilla versions. These large GNN models are evaluated on a set of tasks from the Open Catalyst 2020 (OC20) benchmark and show improved performance compared to the smaller baselines. ", "summary_of_the_review": "Overall, I vote for acceptance. The paper proposes a method to train large scale graph neural networks and shows that very large GNN models can have quite large improvements in the Open Catalyst 2020 (OC20) benchmark", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "WVkKS2gU1-f", "reviewer": "Reviewer_BwW9", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The approach is well-motivated and well-explained. I am familiar with the DimeNet architecture but not very familiar with the recent literature on distributed GNN training, so it is hard for me to judge the novelty of the proposed parallel approach. The results on the OpenCatalyst benchmark are compelling.\n\nAmong graph neural networks for molecular structures, the DimeNet++ architecture is specifically memory and compute intensive (and as such naturally stands to specifically benefit from the proposed approach). It would be great if the authors could include a discussion of the applicability of their approach in the context of other graph neural networks for molecular structure, such as equivariant message passing in Jing et al. (ICLR, 2021) or Schuett et al. (ICML, 2021). \n\nThe paper does not put the scaling results in Figure 2 in context with other approaches for parallel training. ", "summary_of_the_paper": "The authors present an approach to train graph neural networks with many parameters across multiple GPUs. The approach is specifically demonstrated for graphs representing molecular structure and two graph neural networks of similar flavor that have previously been presented to learn from such structures. Higher-order interaction terms make these networks very compute intensive. Leveraging their parallel approach to scale up the existing GNNs, the authors demonstrate compelling results on the OpenCatalyst benchmark. ", "summary_of_the_review": "I did not identify technical concerns with respect to the proposed method. I will defer to other reviewers regarding the technical novelty of the approach. The paper could be made stronger by explicitly discussing how this approach is applicable to other GNN architectures that have been designed for 3D molecular structure.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GAfjPs49I4C", "reviewer": "Reviewer_N2CL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n * Good empirical results: OC20 is an open, comprehensive benchmark and the method achieved decent improvement over strong baselines in the leaderboard.\n\nWeakness and questions:\n * Since this paper is about distributed training, it should be compared with other distributed training methods, including data and model parallelism. Does the new method run faster than data/model parallelism? For data parallelism, one could consider splitting the batch to multiple GPUs. The paper states that model parallelism can be combined with the proposed approach, but I think this paper should compare with standard model parallelism, such as pytorch DDP. If memory is the issue, combining data and model parallelism should be a reasonable baseline.\n * I am not so sure about the novelty of the distributed training approach here. It simply splits the graph to different GPUs and apply \"All Reduce\" after each layer of computation. This strategy doesn't seem to be efficient. Using 8 GPUs, the method can only reduce run time by 50%.", "summary_of_the_paper": "This paper proposes a distributed training method for large graph neural networks (GNN) up to billion parameters. The method first distributes the triplet update operations to multiple GPUs and aggregate the updated vectors by global synchronization. It then distributes the edge update operations to GPUs and aggregate the edge vectors (another global synchronization). Finally, it applies node update in parallel and global node aggregation at the end. The method is applied to GemNet and DimeNet up to 1.12 billion parameters and achieved state-of-the-art results on OC20 benchmark.", "summary_of_the_review": "Despite the good empirical results, it is not clear what is the benefit of the proposed approach over standard distributed training techniques (data/model parallelism). Since this paper proposes a new distributed training algorithm for graphs, it is necessary to compare the speedup of the proposed method against standard distributed training techniques.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers were split about this paper: on one hand they appreciated the clarity and the experimental improvments in the paper, on the other they were concerned about the novelty of the work. After going through it and the discussion I have decided to vote to accept this paper for the following reasons: (a) the potential impact of the work, (b) the simplicity of the idea, and (c) promise of release of open source code. I think these things make the paper a strong contribution to ICLR. The only thing I would like to see added, apart from the suggestions detailed by the reviewers, is a small discussion on the carbon footprint of training such largescale graph networks. The authors motivated the work by saying it could have a beneficial impact for modelling energy which is needed to combat climate change. However, we know from recent results that such large scale models also have a non-trivial emission footprint. So I'd like to see the authors specifically calculate the carbon footprints of the models they trained. There are tools to help with this such as: https://mlco2.github.io/impact/  With this addition I think this paper will not only make a large impact on graph network training but also start a discussion of how to responsibly decide training, taking environmental impact into account.
Decision: Accept (Poster)