Communication-Efficient Graph Neural Networks with Probabilistic
  Neighborhood Expansion Analysis and Caching
Training and inference with graph neural networks (GNNs) on massive graphs
has been actively studied since the inception of GNNs, owing to the widespread
use and success of GNNs in applications such as recommendation systems and
financial forensics. This paper is concerned with minibatch training and
inference with GNNs that employ node-wise sampling in distributed settings,
where the necessary partitioning of vertex features across distributed storage
causes feature communication to become a major bottleneck that hampers
scalability. To significantly reduce the communication volume without
compromising prediction accuracy, we propose a policy for caching data
associated with frequently accessed vertices in remote partitions. The proposed
policy is based on an analysis of vertex-wise inclusion probabilities (VIP)
during multi-hop neighborhood sampling, which may expand the neighborhood far
beyond the partition boundaries of the graph. VIP analysis not only enables the
elimination of the communication bottleneck, but it also offers a means to
organize in-memory data by prioritizing GPU storage for the most frequently
accessed vertex features. We present SALIENT++, which extends the prior
state-of-the-art SALIENT system to work with partitioned feature data and
leverages the VIP-driven caching policy. SALIENT++ retains the local training
efficiency and scalability of SALIENT by using a deep pipeline and drastically
reducing communication volume while consuming only a fraction of the storage
required by SALIENT. We provide experimental results with the Open Graph
Benchmark data sets and demonstrate that training a 3-layer GraphSAGE model
with SALIENT++ on 8 single-GPU machines is 7.1 faster than with SALIENT on 1
single-GPU machine, and 12.7 faster than with DistDGL on 8 single-GPU machines.
Tim Kaler, Alexandros-Stavros Iliopoulos, Philip Murzynowski, Tao B. Schardl, Charles E. Leiserson, Jie Chen
Communication-Efficient Graph Neural Networks with Probabilistic Neighborhood Expansion Analysis and Caching

###### Abstract

Training and inference with graph neural networks (GNNs) on massive graphs has been actively studied since the inception of GNNs, owing to the widespread use and success of GNNs in applications such as recommendation systems and financial forensics. This paper is concerned with minibatch training and inference with GNNs that employ node-wise sampling in distributed settings, where the necessary partitioning of vertex features across distributed storage causes feature communication to become a major bottleneck that hampers scalability. To significantly reduce the communication volume without compromising prediction accuracy, we propose a policy for caching data associated with frequently accessed vertices in remote partitions. The proposed policy is based on an analysis of vertex-wise inclusion probabilities (VIP) during multi-hop neighborhood sampling, which may expand the neighborhood far beyond the partition boundaries of the graph. VIP analysis not only enables the elimination of the communication bottleneck, but it also offers a means to organize in-memory data by prioritizing GPU storage for the most frequently accessed vertex features. We present SALIENT++, which extends the prior state-of-the-art SALIENT system to work with partitioned feature data and leverages the VIP-driven caching policy. SALIENT++ retains the local training efficiency and scalability of SALIENT by using a deep pipeline and drastically reducing communication volume while consuming only a fraction of the storage required by SALIENT. We provide experimental results with the Open Graph Benchmark data sets and demonstrate that training a 3-layer GraphSAGE model with SALIENT++ on 8 single-GPU machines is \(7.1\times\) faster than with SALIENT on 1 single-GPU machine, and \(12.7\times\) faster than with DistDGL on 8 single-GPU machines.

## 1 Introduction

Graph neural networks (GNNs) are an important class of machine learning models that incorporate relational inductive bias for effective representation learning on graph structured data (Li et al., 2016; Kipf and Welling, 2017; Hamilton et al., 2017; Velickovic et al., 2018; Xu et al., 2019). These neural networks have been successfully applied in a number of use cases, including product recommendation, traffic forecasting, and financial forensics (Ying et al., 2018; Li et al., 2018; Weber et al., 2019). In many applications, data are continuously collected and the resulting graph grows rapidly, calling for efficient GNN training and inference systems that can scale with the explosive increase of graph data.

This work considers minibatch training and inference with neighborhood sampling in the distributed setting, where vertex data (features) are partitioned across machines. As opposed to full-batch optimization, minibatch optimization is typical for training neural networks, but it poses a unique challenge for GNNs because of the exponential increase of neighborhood size across network layers (Chen et al., 2018). Neighborhood sampling is a popular and effective approach to mitigating this issue (Hamilton et al., 2017; Chen et al., 2018; Ying et al., 2018; Zou et al., 2019; Zeng et al., 2020; Ramezani et al., 2020; Dong et al., 2021). In distributed GNN training, each machine computes the training loss for a minibatch of vertices that are local to the machine's partition. At every step of the training optimization, a set of partition-wise minibatches forms a "distributed minibatch" which is used to update the GNN model parameters. Distributed GNN inference is organized similarly.

Even with neighborhood sampling, minibatch neighborhood expansion creates a communication bottleneck as each machine needs to access remote data for sampled vertices. The communication pattern is stochastic because of the random nature of minibatch and neighborhood sampling. Communication time often dominates the time for training computations. An example of this bottleneck effect is shown in Table 1, which we will walk through shortly.
To overcome the communication bottleneck, we propose an analysis of neighborhood access patterns via vertex inclusion probabilities, as well as a caching policy based on the analysis. We refer to this analysis as "VIP analysis." In contrast to commonly used heuristics for estimating vertex access probabilities -- e.g., based on vertex degree and expanded boundary frontiers (Lin et al., 2020), random walks (Dong et al., 2021; Min et al., 2021), or simulated GNN computations (Yang et al., 2022) -- VIP analysis estimates access probabilities for all graph vertices based on an analytical model of the actual stochastic neighborhood expansion process. We derive the specific VIP model for the important class of node-wise sampling schemes (Hamilton et al., 2017; Ying et al., 2018; Velickovic et al., 2018; Xu et al., 2019; Liu et al., 2020) and find that the resulting caching policy drastically reduces the communication volume in distributed GNNs that employ node-wise sampling with only a modest memory overhead. Moreover, VIP analysis offers a means to organize in-memory data by prioritizing GPU storage for the most frequently accessed vertex features, reducing host-to-device data transfers.

Our system, SALIENT++, is developed over SALIENT (Kaler et al., 2022), an efficient GNN system that achieves state-of-the-art performance through performance-engineered neighborhood sampling, shared-memory parallel batch preparation, and data-transfer pipelining. SALIENT achieves per-epoch time that is nearly equal to the GPU time for model computation alone (effectively hiding the cost of minibatch construction, data transfer, and gradient communication), and it scales well with additional machines. But SALIENT has the drawback of replicating the entire data set on each machine, imposing a hard limit on the data set size. Extending SALIENT to distribute vertex feature data requires addressing the feature communication bottleneck.

SALIENT++ uses the VIP-analysis-driven caching policy and a deep pipeline to achieve scalability and efficiency. It nearly matches the performance of SALIENT with only a fraction of SALIENT's memory requirements. To highlight the efficacy of SALIENT++, Table 1 lists the resulting performance of progressive modifications on top of SALIENT. This example uses the ogbn-papers100M data set (Hu et al., 2020) and partitions the graph using METIS (Karypis & Kumar, 1997) with an edge-cut minimization objective and balancing constraints for the number of training, validation, and overall vertices, as well as the total number of edges, in each partition. Accompanying Table 1, Figure 1 illustrates the key differences between SALIENT and the successive optimizations in SALIENT++ using simplified computation profiles and storage requirement depictions.

We first explore the distributed performance of SALIENT with a disjoint partitioning of feature data across machines. We modify SALIENT to only replicate the graph structure and communicate vertex features on demand and in bulk for each minibatch and its sampled neighborhood. Although distributing the vertex features reduces SALIENT's memory requirement substantially, it also leads to an immediate performance degradation. Specifically, we observe slowdown of roughly \(2\)-\(3\times\) on two or four machines and \(3.5\times\) on eight machines relative to SALIENT with full-replication. This slowdown presents itself in the form of high latencies between GNN model computations (row 2 in Figure 1). Overlapping feature communication with other GNN operations (sampling, local feature slicing, host-to-device transfers, training computations, and model updates) through pipelining offers some improvement, but the system remains bottlenecked by communication (row 3 in Figure 1).

If we allow a small memory overhead to cache frequently accessed remote features locally, however, the combined effects of reduced communication volume and pipelining boost performance substantially. The sheer reduction of communication volume due to caching makes it possible for pipelining to (nearly) fully hide communication (4th row in Figure 1). On \(K\) machines, SALIENT++ reduces the total feature memory footprint of SALIENT by almost a factor of \(K\) while matching the performance achieved by SALIENT's full-replication strategy.

In summary, this work makes the following contributions:

1. An analysis of vertex inclusion probabilities (VIP) during GNN neighborhood expansion with node-wise sampling. We show how VIP analysis can be used to minimize the expected communication among machines as well as the host-to-device data transfers within each machine.
2. The design and implementation of SALIENT++, a system for distributed GNN training and inference that is both efficient and scalable. The system is based on two key components: a maximum-likelihood static caching policy for remote and local features based on VIP analysis; and a deep pipeline that overlaps feature communication with other GNN operations, ensuring high utilization

\begin{table}
\begin{tabular}{l c c c c} \hline \hline  & \multicolumn{4}{c}{_No. of machines_} \\ \cline{2-5} _System_ & 1 & 2 & 4 & 8 \\ \hline _SALIENT (Full replication)_ & 20.7s & 10.76s & 6.02s & 3.08s \\ _+ Partitioned features_ & — & 33.04s & 15.98s & 10.85s \\ _+ Pipeline communication_ & — & 16.12s & 8.73s & 5.43s \\ _+ Feature caching_ & — & 10.51s & 5.45s & 2.91s \\ \hline \hline \end{tabular}
\end{table}
Table 1: Per-epoch runtime of a progressively more sophisticated distributed GNN training system on the undirected ogbn-papers100M data set, using a 3-layer GraphSAGE architecture with sampling fanouts (15,10,5) and a hidden layer dimension of 256. For the system with remote feature caching, the size of the cache relative to the size of local-vertex features for each machine was 8% (2 machines), 16% (4 machines), or 32% (8 machines).
of available network bandwidth.
3. A systematic evaluation of the end-to-end training performance and scalability of SALIENT++ on three large benchmark graphs. On the largest of these data sets, which does not fit in-memory on one machine, SALIENT++ achieves a speedup of \(1.75\times\) when scaling from 4 to 8 single-GPU machines and an additional \(1.45\times\) when scaling from 8 to 16 single-GPU machines.

## 2 Background and Related Work

This section briefly recalls background and introduces notation for GNNs and their training/inference. It also discusses related work on systems and distributed training.

### Graph neural networks

This work focuses on the class of _message passing neural networks_ (MPNNs) (Gilmer et al., 2017), which encompasses many GNNs of major interest for massive graphs. Let \(G=(\mathcal{V},\mathcal{E})\) be a graph with \(\mathcal{V}\) being the vertex set and \(\mathcal{E}\) being the edge set. Let \(X\in\mathbb{R}^{N\times D}\) be the vertex feature matrix, where \(N\) is the number of vertices and \(D\) is the feature dimension. A row of \(X\), denoted by \(x_{v}\in\mathbb{R}^{D}\), is the feature vector of a vertex \(v\). Let \(\ell=0,\ldots,L\) denote the layer index and \(\mathcal{N}_{1}(v)\) be the one-hop neighborhood of \(v\). MPNNs use the following update rule to define a layer:

\[h_{v}^{\ell}=\textsc{upd}^{\ell}\Big{(}h_{v}^{\ell-1},\textsc{agg}^{\ell} \big{(}\{h_{u}^{\ell-1}\mid u\in\mathcal{N}_{1}(v)\}\big{)}\Big{)}, \tag{1}\]

where \(h_{v}^{\ell}\) is the layer-\(\ell\) representation of \(v\), \(\textsc{agg}^{\ell}\) is a set aggregation function, \(\textsc{upd}^{\ell}\) is a two-argument update function, and \(h_{v}^{0}=x_{v}\). One sees that the representation of a vertex \(v\) depends on the past-layer representations of \(v\) and its neighbors in \(\mathcal{N}_{1}(v)\); and this dependence is recursive.

GNNs differ in their design of the two functions in (1). For example, in GraphSAGE (Hamilton et al., 2017), \(\textsc{agg}^{\ell}\) is a mean, LSTM, or pooling operator, and \(\textsc{upd}^{\ell}\) concatenates the two arguments and applies a linear layer. In GIN (Xu et al., 2019), \(\textsc{agg}^{\ell}\) is the sum of \(\{h_{u}^{\ell-1}\}\) and \(\textsc{upd}^{\ell}\) is the sum of its arguments followed by an MLP. In GAT (Velickovic et al., 2018), \(\textsc{agg}^{\ell}\) is the identity and \(\textsc{upd}^{\ell}\) computes \(h_{v}^{\ell}\) as a weighted combination of \(W^{\ell-1}h_{u}^{\ell-1}\) for all \(u\in\{v\}\cup\mathcal{N}_{1}(v)\), where the weights are attention coefficients and \(W^{\ell-1}\) is the parameter matrix of the layer.

### Minibatch training

Neural networks can be trained with gradient descent (full-batch training) or stochastic gradient descent methods (minibatch training). A majority of the deep learning literature advocates minibatch training, with theoretical and empirical evidence suggesting that it converges faster, generalizes better, and is more suitable for GPU-oriented computing infrastructures (Bottou et al., 2018).

Minibatch training incurs a unique challenge for GNNs, often termed the "neighborhood explosion" problem: to compute the training loss for a vertex \(v\), equation (1) indicates that it requests information from the neighborhood recursively, growing its size exponentially with the number of layers. For a reasonably sized minibatch, the multi-hop neighborhood quickly covers a large portion of the graph, incurring a prohibitive time cost that eclipses the saving in convergence speed. Moreover, the storage requirement easily exceeds the memory capacity of a GPU.

Figure 1: Illustration of distributed GNN computation profiles and feature data storage for the precursor system SALIENT and successive optimizations leading to our system SALIENT++. **Left:** Utilization of system resources (CPU, GPU, and network) in one machine. The processing of each minibatch is broken into three parts: "MFG": message-flow graph construction. “Feat.”: local feature tensor slicing, feature communication and joining, and host-to-device transfers. “Model”: forward and backward passes, plus communication among machines for model updates. In addition to the highlighted minibatch, the profiles also depict parts of a preceding and a succeeding minibatch in faded colors to show overlap between pipelined operations. **Right:** The color-stacked bars indicate the number of locally stored feature vectors in each of \(K\) total machines.
### Neighborhood sampling

Restricting the neighborhood size by sampling is an effective method for tackling the neighborhood explosion problem. Three major types of sampling strategies are node-wise sampling (Hamilton et al., 2017; Ying et al., 2018), layer-wise sampling (Chen et al., 2018; Zou et al., 2019), and subgraph sampling (Chiang et al., 2019; Zeng et al., 2020).

Node-wise sampling approaches modify the neighborhood for each vertex \(v\) separately by taking a random subset containing at most \(f\) neighbors (called the _fanout_). Layer-wise sampling approaches collect the neighbors of all vertices in a minibatch and sample the neighborhood from their union. Such sampling proceeds recursively layer by layer and can use a nontrivial sampling distribution to control pre-activation variance while preserving unbiasedness. Nonlinear activation functions destroy unbiasedness anyway, but training convergence results can still be established based on asymptotic consistency (Chen and Luss, 2018). Subgraph sampling approaches sample a connected subgraph and compute the minibatch loss restricted to this subgraph. This work focuses on node-wise sampling for its widespread use.

### Minibatch inference

Similar to training, inference can be performed in either full batches or minibatches. Whereas the choice between full batch or minibatch training centers on convergence and generalization, the choice made for inference depends on the computational pattern and implementation effort. Full-batch inference avoids neighborhood sampling but requires a separate implementation of the forward pass for different GNN architectures. On the other hand, minibatch inference can reuse the forward function for training, improving productivity for application developers and architecture designers, but the stochastic nature of neighborhood sampling may return different predictions and does not guarantee the same accuracy as the case of not performing sampling. Nevertheless, the authors of SALIENT demonstrate strong empirical results that show that prediction accuracy is not compromised with a reasonable choice of the fanouts (Kaler et al., 2022). This work follows the practice of minibatch inference.

### Related work

Due to the computational pattern of neighborhood aggregation, GNN training systems on massive graphs differ substantially from those for usual neural networks in design and implementation. Many GNN systems are developed based on full-batch training for simplicity, which avoids the complication of neighborhood sampling. Examples include NeuGraph (Ma et al., 2019), Roc (Jia et al., 2020), DeepGalois (Hoang et al., 2021), FlexGraph (Wang et al., 2021), Seastar (Wu et al., 2021), GNAdvisor (Wang et al., 2021), DistGNN (Md et al., 2021), and BNS-GCN (Wan et al., 2022). Some of these systems are built on common deep learning frameworks, such as TensorFlow (Abadi et al., 2015) and PyTorch (Paszke et al., 2019), while others are built on self-developed programming models.

Examples of systems that perform minibatch training include DistDGL (Zheng et al., 2020), Zero-Copy (Min et al., 2021), GNS (Dong et al., 2021), and \(P^{3}\)(Gandhi and Iyer, 2021). However, these publications report results on either multiple machines with only CPUs or a single machine with one or multiple GPUs. A newer version of DistDGL (DistDGLv2 (Zheng et al., 2021)), SALIENT (Kaler et al., 2022), and our system SALIENT++ demonstrate results in the distributed multi-GPU setting.

SALIENT++ employs edge-cut graph partitioning to distribute data across machines. An alternative approach, adopted by DistGNN (Md et al., 2021), is vertex-cut partitioning, which ensures each edge is local to some machine. A drawback of this approach is that a vertex may be assigned to multiple machines, leading to memory overhead due to feature replication that is reported to be as high as 5\(\times\). SALIENT++, on the other hand, achieves high distributed performance with less than 50% memory overhead due to caching in our experiments.

SALIENT++ employs a caching policy to reduce communication and data transfer. We introduce a policy based on an analysis of vertex access probabilities during neighborhood sampling. Although caching has been explored by several other systems in this context, previous caching policies are heuristic, using proxy measures such as vertex degrees (Lin et al., 2020), random walks (Dong et al., 2021; Min et al., 2021), or simulated access frequencies (Yang et al., 2022). A complementary approach to reducing communication time is taken by DGCL (Cai et al., 2021), which uses a communication-planning algorithm to optimize the communication pattern for a specific network topology.

Marius and MariusGNN (Mohoney et al., 2021; Waleffe et al., 2022) are out-of-core training systems that work on a single machine by exploiting external memory. They operate in a different setting from ours, but our vertex inclusion probability analysis may be used in complement to optimize on-disk data representation and disk I/O.

## 3 Vertex inclusion probabilities for GNNs with node-wise sampling

This section describes a principled approach to estimating and reducing data access costs in distributed GNNs with node-wise sampling. Node-wise sampling was notably introduced in GraphSAGE (Hamilton et al., 2017) and is used with a variety of GNN architectures (Ying et al., 2018; Velickovic et al., 2018; Xu et al., 2019; Liu et al., 2020). We
first introduce _vertex-inclusion probability_ (_VIP_) analysis to calculate the probability that any vertex will be present in the sampled \(L\)-hop expanded neighborhood of a minibatch in some partition. VIP analysis takes into account the distinctive structure of neighbor accesses in GNNs with minibatches and node-wise sampling. We then demonstrate that VIP analysis enables a simple but highly effective strategy for reducing data movement during GNN computations.

The VIP model for node-wise sampling derived in this section does not apply to other sampling schemes, such as layer-wise or subgraph sampling. It may be possible to derive VIP models for these other sampling schemes in a similar way, but this is outside our scope.

### Analysis of vertex inclusion probabilities in \(L\)-hop neighborhoods with node-wise sampling

We analyze the following random process, which models neighborhood expansion in GNN architectures with node-wise sampling. (i) Start from a set of vertices rather than a single vertex. (ii) For each vertex in the starting set, sample a set of direct neighbors without replacement, thereby expanding the walk frontier. (iii) Using the union of sampled neighborhoods as the new starting set, repeat step (ii). The process terminates after \(L\) repetitions (i.e., hops). We assume that sampling is independent across hops and across vertices in the current-hop set. That is, although each vertex samples among its direct neighbors without replacement, different vertices may sample the same direct neighbor.

We seek to estimate the vertex inclusion probabilities, or VIP values, in expanded neighborhoods that are obtained per the above random process. The VIP values form a vector of probabilities \(p(u)\) over all graph vertices. We can express \(p(u)\) via hop-wise VIP vectors \(p^{[h]}(u)\) that contain the probabilities that any vertex \(u\) is sampled exactly \(h\) hops away from the starting set. (The hop-\(h\) neighborhood contains the input vertices for the \(\ell\)-th GNN layer, \(\ell=L-h\).) A vertex \(u\) is _not_ sampled at hop \(h\) only if, for each of its neighbors \(v\in\mathcal{N}_{1}(u)\), either \(v\) was not sampled at hop \(h-1\) or \(v\) was sampled but did not pick \(u\) among its neighbors. The following Proposition shows how to propagate the initial-set probabilities through the graph to compute the VIP values. We give a proof in Appendix C.

**Proposition 1** (Vertex inclusion probabilities in minibatch neighborhood expansion with node-wise sampling).: _The probability that vertex \(u\in\mathcal{V}\) appears in the node-wise sampled \(L\)-hop expanded neighborhood of any minibatch is_

\[p(u)=1-\prod_{h=1}^{L}\big{(}1-p^{[h]}(u)\big{)}, \tag{2}\] \[p^{[h]}(u)=1-\prod_{v\in\mathcal{N}_{1}(u)}\big{(}1-t_{h}(u,v)\, p^{[h-1]}(v)\big{)}, \tag{3}\]

_where \(p^{[0]}(v)\) is the probability that \(v\) is in the minibatch and \(t_{h}(u,v)\) is the transition probability that \(v\) samples \(u\) as a neighbor at hop \(h=1,\ldots,L\)._

The VIP model of Proposition 1 applies to any initial sampling and hop-wise transition probability function for node-wise sampling. For example, if minibatches and vertex-wise neighbors are sampled uniformly at random without replacement as in GraphSAGE, then we have: \(p^{[0]}(u)=B/|\mathcal{T}|\) if \(u\in\mathcal{T}\) and \(0\) otherwise, where \(B\) is the minibatch size and \(\mathcal{T}\) is the set of training vertices; and \(t_{h}(u,v)=\min\{1,f_{h}/d(v)\}\) if \(u\in\mathcal{N}_{1}(v)\) and \(0\) otherwise, where \(d(v)\) is the (outgoing) degree of \(v\). Non-uniform neighbor sampling models are accommodated via the corresponding transition probability matrix or matrices.

The neighborhood expansion random process parametrizes a continuum between a random walk and full neighborhood expansion. If the initial set size and layer-wise fanouts are equal to 1, it becomes a random walk. If the fanouts are greater than the max in-degree of graph vertices, it becomes a full neighborhood expansion. The VIP model for the above two special cases is linear, while the generalized model in Proposition 1 is nonlinear; nonetheless, these models all have the same computational complexity: \(O(L(M+N))\).

### Vertex feature caching

We now show how to embody VIP analysis into a caching policy for minimizing communication among distributed machines and host-to-device data transfers in each machine.

Communication reductionWe consider the optimal static caching policy, in the maximum-likelihood sense, for reducing inter-partition communication volume in distributed GNNs. The policy is straightforward: each machine extends its local vertex feature storage with copies of the remote features that are most likely to be accessed by the machine, thereby minimizing the total expected communication volume. For some given initial partitioning -- and making no assumptions about the order in which minibatch and neighboring vertices are sampled -- this policy is directly related to the VIP analysis. Specifically, we calculate partition-wise VIP values \(p_{k}(u)\) from the corresponding initial probabilities \(p_{k}^{[0]}(u)\). The cache contents for the \(k\)-th partition are then determined by ranking the remote vertices in order of decreasing \(p_{k}(u)\).

We measure the size of a remote-feature cache by its corresponding _replication factor_\(\alpha\). The replication factor is defined such that the number of cached feature vectors stored in each machine is \(\alpha N/K\), where \(K\) is the number of partitions or machines. That is, the fraction of cached to local vertices for each partition is \(\alpha\), and each feature vector is stored in \((1+\alpha)\) machines on average. We may say that a
partition or cache replicates a remote vertex to mean that it copies the corresponding vertex feature data.

Figure 2 compares a number of static feature-caching policies and provides evidence that the analytical VIP model of Proposition 1 effectively yields the optimal policy. The specific policies being compared are as follows: "deg." ranks vertices by degree, after filtering out remote vertices that are not reachable from a partition's training set (Lin et al., 2020); "1-hop" replicates the entire 1-hop halo of each partition; "wPR" ranks vertices by their score after 5 iterations of a weighted reverse-PageRank model (Min et al., 2021) with damping factor 0.85; "#paths" ranks vertices by the number of paths with length \(\leq L\) that reach them from any local training vertex; "sim." ranks vertices by empirical VIP estimates, measured by counting vertex-wise frequencies of access over 2 simulated training epochs (Yang et al., 2022); "VIP" ranks vertices by VIP values as per Proposition 1; and "oracle" retroactively ranks vertices by their actual access frequencies after training, providing a lower bound on communication volume.1 We make the following observations.

Footnote 1: With all caching policies, we calculate vertex rankings with respect to each partition. This is different from the original setting in which some of these policies were introduced, where vertices are cached based on a single, global ranking score.

* _Optimality._ The analytical VIP policy yields near-optimal communication volume (within 5% of "oracle") across fanouts and replication factors. The bigger gap (\(\sim\)30%) for \(f=(5,5,5)\) and \(\alpha=1.0\) is due to variance in the accesses of lower-ranked vertices; the gap narrows with more samples (e.g., bigger fanouts or more epochs).
* _Communication reduction efficacy._ VIP-driven caching is highly effective for reducing communication volume. Compared to no caching, the VIP policy achieves a geometric-mean reduction of \(2.2\times\)-\(5.3\times\) with small replication factors \(\alpha\in[0.05,0.20]\) and more than \(10\times\) with replication factors over \(0.50\). Compared to the other caching policies, it achieves consistently lower communication volume and its relative improvement increases with the replication factor.
* _Local information is not sufficient._ The high-degree and 1-hop halo policies, which do not take neighborhood expansion into account, offer only marginal improvements over no caching. Caching multi-hop neighbors will reduce communication further (Lin et al., 2020) but it will also blow up the replication factor.
* _Impact of tailoring the model to the expansion process._ The VIP-driven policy is up to \(4\times\) and \(2\times\) more effective than the "wPR" and "#paths" policies, respectively. Both of these policies take graph-structural expansion into account, but "wPR" is agnostic to the GNN fanout and number of layers, while "#paths" does not directly account for sampling.
* _Benefits and limitations of empirical estimation._ The empirical and analytical VIP estimates ("sim." and "VIP" policy, respectively) yield comparable results for low replication factors and high fanouts, but the empirical estimates become less effective as the replication factor increases or the fanouts decrease. For replication factors \(0.50\) and \(1.0\), the relative aggregate improvement of the analytical over the empirical policy is \(1.6\times\) and \(3.2\times\), respectively. The empirical approach has the benefit that it can be used for any sampling scheme, but it also requires increasingly many samples to accurately estimate VIP values for infrequently accessed vertices.

Host-to-device data transfer reductionVIP analysis can also be used to reduce the volume of host-to-device data transfers for local features on each machine. For example, assume that each GPU can retain a fraction of the local feature data in memory throughout the GNN computations (i.e., across all minibatches). We may rank the local vertices in the \(k\)-th partition by \(p_{k}(u)\) and keep the highest-ranking-vertex features on the GPU. As with remote-vertex caching, this policy minimizes the expected data transfer volume due to local vertices in each partition.

Figure 2: Comparison of caching policies with respect to remote feature communication volume. GNN: 3-layer GraphSAGE, varying fanouts \(f\), minibatch size 1024. Data set: ogbn-papers100M (undirected). Partition: 8-way METIS with edge-cut objective plus edge and training/validation/test vertex balancing. **(a)**–**(c)** Average per-epoch communication in number of vertices over 100 epochs (log-scale). Lower is better. The feasible region between the communication upper bound (no caching) and lower bound (oracle caching) is shaded gray. **(d)** Geometric-mean improvement across fanouts, relative to no caching (log-scale). Higher is better.
## 4 SalIENT++: Fast and Scalable Distributed GNN training via Caching and Pipelining

This section describes the design of SALIENT++, a fast and scalable system for performing distributed minibatch training of GNNs on large partitioned datasets. The design of SALIENT++ is comprised of three majors components:

**Vertex reordering**: A vertex ordering based on the graph partitions and VIP values that facilitates efficient sub-partitioning of local feature data between CPU and GPU to reduce host-to-device data transfers.
**Data replication**: An economical data replication strategy that uses VIP analysis to reduce communication volume for fetching remote features while using only a small amount of additional memory.
**Pipelined communication**: A pipelined distributed minibatch preparation system that hides latencies due to sampling, data transfers, and inter-machine communication of remote feature data.

These three techniques together enable SALIENT++ to perform GNN training on partitioned data-graphs while matching the efficiency of highly optimized GNN training systems that perform distributed computations with full replication of vertex feature data across all machines. This enables SALIENT++ to perform GNN training efficiently on large datasets where full replication is impractical.

As we shall discuss in our empirical evaluation in Section 5, SALIENT++ achieves high performance for GNN training on large datasets even when operating in clusters with only modest hardware configurations.

### Partitioning and reordering of vertex features

Let us describe how SALIENT++ partitions the vertex features of a graph to reduce inter-machine and host-to-device communication. SALIENT++ employs a two-level strategy for distributing the vertex features across machines and devices, as illustrated in Figure 3. The top level involves distributing vertex features based on a vertex partitioning of the graph, and the bottom level involves dividing each partition's local vertex features between GPU and CPU memory.

SALIENT++ operates on graph data sets that are partitioned using a graph partitioning algorithm such as METIS (Karypis and Kumar, 1997). The graph partitions are generated such that the training, validation, and test vertices are balanced. As a heuristic to balance work-per-partition, an additional criterion is used to balance the amount of edges per partition. Each machine is assigned a subset of training/validation/test vertices and stores vertex feature data locally according to the machine's partition.

The graph is reordered such that: (a) indices for vertices in the same partition are contiguous; and (b) vertices within a partition are ordered based on how beneficial it is for them to be stored on the GPU. Each machine stores a prefix of its ordered list of vertex features on the GPU, which when using VIP ordering corresponds to those vertices that are accessed most frequently by the local machine. This partitioning structure facilitates efficient determination of whether a vertex is remote or local, where a local vertex is stored on the machine, and index calculations using only a constant amount of additional memory. Figure 3 provides a graphical illustration of the VIP-reordered partitioned dataset.

### Replication of remote vertex features

Each machine maintains a replica or static cache of features for vertices in remote partitions, where the size of the cache is determined by a given replication factor \(\alpha\). The vertices whose features are cached in each machine are determined by their VIP rank to reduce network communication, as described in Section 3.2.

Once the neighborhood sampling code prepares a batch, SALIENT++ partitions the vertices in the expanded neighborhood into local and remote vertices. Among the subset of remote vertices, an efficient lookup is performed using a hash table to determine whether the features for a remote vertex can be found in the local machine's cache. Only remote vertices that are not resident in the local cache proceed to be requested from remote machines.

### Minibatch preparation pipeline

SALIENT++ implements a deep multi-stage pipeline which enables overlap between neighborhood sampling, host-to-device data transfers, network communication between machines, and training computation. SALIENT++'s pipeline is more complex than the one used in SALIENT (Kaler et al., 2022), but it maintains the latter's usability and efficiency. SALIENT++'s pipeline can be used in place of standard data loaders for GNNs by changing just a few lines of code.

**Sampling and slicing** SALIENT++ uses a modified version of SALIENT's optimized neighborhood sampling and slicing code to prepare minibatches using shared-memory

Figure 3: Illustration of the partitioned data set with VIP-driven local ordering and caching of remote features across machines.
parallelism. The original SALIENT batch-preparation code fused the neighborhood sampling and feature tensor slicing operations. In SALIENT++, these tensor-slicing operations are only performed on the subset of the vertices in the sampled neighborhood that are stored locally in the CPU memory of each machine. For the remaining vertices in the sampled neighborhood, SALIENT++ distinguishes between vertices in the following categories: vertices belonging to the local partition whose features are stored on GPU; remote vertices that are replicated on the local machine; and remote vertices that belong to other partitions.

**Feature collection pipeline** We briefly summarize the pipeline operations. The expanded neighborhood of each minibatch may include remote vertices, remote vertices whose features were cached on the local machine following VIP analysis, and local vertices which are split across CPU and GPU storage. Immediately following sampling, vertices are categorized by their storage location so that vertex features are properly retrieved. Communication rounds interspersed with host-to-device and device-to-host transfers are enqueued with remote machines to coordinate retrieval of remote features. Replicated and local features partitioned across CPU and GPU storage are sliced locally. Finally, all features are concatenated into a single tensor and reordered for compatability with the message-flow graph generated during neighborhood sampling.

A total of \(10\) mini-batches can be "in-flight" in the SALIENT++ pipeline at any time, each mini-batch being processed by a separate stage of the pipeline. Pipelining enables overlap between host-to-device data transfers and network communication, and it also allows for hiding latencies related to CPU-GPU data transfers and inter-machine communication. A more detailed stage-by-stage description of SALIENT++'s pipeline is provided in Appendix D.

## 5 Evaluation

In this section, we evaluate the performance of SALIENT++ and investigate the impact of its optimizations relating to pipelining and use of VIP analysis (as developed in Section 3). Additionally, we report end-to-end performance results, model accuracies, and contextualize our performance in relation to existing distributed GNN training systems.

### Experimental setup

**Computing environment** Our experiments were run on a cluster consisting of 16 machines, each equipped with a 16-core (2-way hyperthreaded) AMD CPU with 128GB DRAM and 1 NVIDIA A10G GPU with 24GB memory. The cluster was configured using the Amazon Web Services (AWS) ParallelCluster software and employed fleets of AWS _g5.8xlarge_ instances. The network SLA specified for this instance type is 25Gbps. Distributed communications were implemented via PyTorch's DistributedDataParallel module with the NCCL backend. Since our cluster comprises machines with 1 GPU each, experiments with \(K\) GPUs involve \(K\) separate machines communicating over the network.

**Datasets** We evaluate SALIENT++ on three standard benchmark data sets: ogbn-products (_products_), ogbn-papers100M (_papers_) (Hu et al., 2020), and lsc-mag240 (_mag240c_) (Hu et al., 2021). The graph and training set in these data sets vary in size, with _papers_ and _mag240c_ being two of the largest open benchmarks at the time of this work. See Table 2 for detailed information. All graphs were made undirected (if originally not), as is common practice. For _mag240c_, we train on the homogeneous papers-to-papers component of the graph.

**GNN architectures** We evaluated SALIENT++'s performance using standard GraphSAGE (Hamilton et al., 2017) architectures with each data set's most commonly used hyperparameters. Table 3 lists the GNN architectures and hyperparameters that were used for each data set.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline _Data Set_ & _\#Vertices_ & _\#Edges_ & _\#Feat._ & _Train. / Val. / Test_ \\ \hline products & 2.4M & 123M & 100 & 197K / 39K / 2.2M \\ papers & 111M & 3.2B & 128 & 1.2M / 125K / 214K \\ mag240c & 121M & 2.6B & 768 & 1.1M / 134K / 88K \\ \hline \hline \end{tabular}
\end{table}
Table 2: Summary of data sets. Edge counts reflect the number of edges in the graph after standard preprocessing (e.g., making the graph undirected). The _mag240c_ data set is the papers-to-papers citation subgraph of the _lsc-mag240_ heterogeneous graph data set.

Figure 4: Impact of pipelining and VIP optimizations in SALIENT++. Illustrates the performance improvements obtained by successively more optimized versions of SALIENT++ for _products_ (4 partitions), _papers_ (8 partitions), and _mag240c_ (16 partitions). The GNN architectures for each benchmark are listed in Table 3. The replication factors for VIP-based caching were \(0.16\), \(0.32\), and \(0.32\) for _products_, _papers_, and _mag240c_, respectively.

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline _Data Set_ & _GNN_ & _\#Layers_ & _Hidden Dim._ & _Fanout_ & _Batch_ \\ \hline products & SAGE & 3 & 256 & (15, 10, 5) & 1024 \\ papers & SAGE & 3 & 256 & (15, 10, 5) & 1024 \\ mag240c & SAGE & 2 & 1024 & (25, 15) & 1024 \\ \hline \hline \end{tabular}
\end{table}
Table 3: GNN architecture hyperparameters for all experiments in Section 5. Fanout is for training. Batch size is per GPU.
### Performance analysis of SALIENT++

We analyze the performance of SALIENT++ by investigating the impact of its component optimizations, and measuring end-to-end performance and scalability across three large graph data sets.

**Summary of performance improvements** Figure 4 summarizes how SALIENT++'s optimizations result in increasingly better end-to-end training performance on _products_, _papers_, and _mag240c_ in the distributed setting. Significant improvements are observed for both the _papers_ and _mag240c_ benchmarks. In relative terms, the _papers_ benchmark benefits equally from pipelining and VIP-based caching, while the _mag240c_ benchmark benefits slightly more from caching on top of pipelining. This is mainly due to the larger feature dimensionality for _mag240c_, which causes the communication of remote feature data to be relatively more throughput-bound than for _papers_.

**Scalability** Figure 5 (left plot) illustrates the reduction in per-epoch runtime when using SALIENT++ to execute the _products_, _papers_, and _mag240c_ benchmarks on \(2\)-\(16\) GPUs. For all three benchmarks, the reported runtimes include warm-up time to fill up the pipeline at the start of each training epoch, which results in diminished scalability once per-epoch runtimes become less than a second. The _products_ benchmark achieves \(1.7\times\) speedup when going from \(2\) to \(4\) GPUs, but does not significantly benefit from more scaling. The _papers_ benchmark achieves \(1.9\times\) speedup from \(2\) to \(4\) GPUs and another \(1.9\times\) from \(4\) to \(8\) GPUs. The _mag240c_ benchmark achieves \(1.75\times\) speedup going from \(4\) to \(8\) GPUs, and \(1.45\times\) speedup from \(8\) to \(16\) GPUs. SALIENT++'s performance matches that of prior fast systems which perform distributed GNN training with full replication (Kaler et al., 2022), while SALIENT++ uses substantially less memory. Figure 5 (right plot) shows the total memory for storing vertex feature data (including replicated vertices) across all machines. Note that full replication with \(K\) machines corresponds to replication factor \(\alpha=K-1\).

**Local partition storage on CPU versus GPU** Figure 6 illustrates the performance impact of increasing the percentage of the local partition stored on GPU while running the _papers_ benchmark. Given an ordering of the vertices in the local partition and a percentage \(\beta\%\) of local data to store on GPU, SALIENT++ stores the first \(\beta\%\) of vertex features on the GPU. The results labeled "no reorder" show an essentially linear reduction in per-epoch runtime as a function of \(\beta\%\) when ranking local vertices by their initial ordering. The results labeled "VIP reorder," on the other hand, show that when ranking local vertices by their VIP values, data transfer bottlenecks are effectively eliminated with as little as \(10\%\) of the local partition data on the GPU.

**Impact of replication factor** Figure 7 illustrates the impact on per-epoch runtime of increasing the number of replicated vertex features for _papers_ on 4 and 8 partitions (left) and on _mag240c_ for 8 and 16 partitions (right). Figure 7 shows that modest replication factors of \(0.08\)-\(0.16\) and \(0.16\)-\(0.32\) are sufficient for minimizing per-epoch runtime when using \(4\) and \(8\) partitions respectively on the _papers_ dataset.

**Performance breakdown for SALIENT++** Figure 8 illustrates the performance bottlenecks in SALIENT++ before and after incorporating pipelining and caching via VIP analysis.2 A comparison of the "pipelining on" and "pipelining

Figure 5: Scalability and total memory used by all machines for _products_, _papers_, and _mag240c_. Total memory is plotted as a multiple of the unreplicated data set size \(1+\alpha\) where \(\alpha\) is the replication factor used for the \(K\)-GPU execution.

Figure 6: Impact of VIP-based local vertex ordering on per-epoch runtime. Experiment run on \(4\) GPUs with the _papers_ dataset and replication factor \(\alpha=0.15\). Reported numbers are the mean over 10 epochs and vertical lines show the standard deviation.

Figure 7: Replication factor impact on per-epoch runtime. Results are shown for _papers_ (left) for \(4\) and \(8\) partitions, and for _mag240c_ for \(8\) and \(16\) partitions for replication factors \(\alpha\) varied between \(0\) and \(0.32\). The percentage of the local partition stored on GPU for this experiment was \(90\%\) for _papers_ and \(10\%\) for _mag240c_. Reported numbers are the mean value over 10 epochs and vertical bars indicate the standard deviation across epoch runtimes.
off" breakdowns for \(\alpha=0\) illustrates that network communication is the primary bottleneck for distributed GNN training, and it remains the bottleneck even when communication is pipelined to overlap with computation. When using caching with \(\alpha=0.32\), the time spent performing communication is sufficiently small so that it can be overlapped nearly perfectly with other computation in the program.

Performance on a slow networkFigure 9 illustrates the performance of SALIENT++ on a slow network where higher replication factors are needed to alleviate communication bottlenecks. The slower network speeds were enforced using the Linux traffic control subsystem with the token-bucket filter (TBF) queuing discipline (Hubert et al., 2002). In this setting, we compare the two best caching policies from Section 3, VIP (analytic) and VIP (simulation). Recall from Section 3 that these two policies tend to perform very similarly for low replication factors, but diverge as the replication factor grows. On the _papers_ dataset, the relative gap between these two policies grows up to \(30\%\) until \(\alpha=0.64\) and then starts to shrink once communication ceases to bottleneck training time. For the _mag240c_ dataset, which uses node features that are \(6\times\) larger than _papers_, the relative gap between VIP (analytic) and VIP (simulation) is larger (up to \(45\%\) with \(\alpha=0.48\)) and persistent. Whether or not VIP (analytic) is preferred to VIP (simulation) depends on the workload characteristics (e.g., fanout and size of node features) as well as the network bandwidth budget.

### Context for SALIENT++'s performance

To conclude our empirical evaluation, we contextualize SALIENT++'s performance by reporting model accuracy results, discussing preprocessing overheads, and comparing our performance to existing systems.

Model accuracyThe optimizations in SALIENT++ do not impact model accuracy. SALIENT++ is integrates with PyTorch Geometric (PyG) and can use existing GNN models implemented with PyG. Regardless, we report accuracies for _products_, _papers_ and _mag240c_ so that one can contextualize the performance numbers. For computing accuracies, we executed SALIENT++ on \(8\) machines for \(30\) epochs with a fixed learning rate of \(0.001\) and a batch size of \(1024\) per machine.3 Sampling was used during inference with fanouts \((20,20,20)\) for _papers_ and _products_, and \((25,15)\) for _mag240c_. On _products_, we observed a test accuracy of \(0.785\) for the \(3\)-layer SAGE architecture. On _papers_, we observed a test accuracy of \(0.646\) for the \(3\)-layer SAGE architecture. On _mag240c_, the 2-layer GraphSAGE architecture achieves an accuracy of \(0.651\).4

Footnote 3: These hyperparameters were not tuned for accuracy, but are known to be reasonable values.

Footnote 4: Validation accuracy is reported instead of “test-dev” or “test-challenge” accuracy, which we did not measure at the time of different replication factors.

Figure 8: Performance breakdown for SALIENT++ (with added synchronization) for an \(8\)-GPU execution of _papers_ with all local node features on GPU. The “pipelining off” breakdown is best able to isolate the time spent performing each operation. The “pipelining on” breakdown adds minimal synchronization to isolate distributed batch preparation from training computation. _Batch Prep (Comm)_ and _Batch Prep (Comp)_ are the time for communication and computation, respectively, in distributed batch preparation. These categories are combined into _Batch Prep_ for the “pipelining on” breakdown. _Startup_ is the time to prepare the first batch and/or fill the pipeline. _Train (sync)_ is the time for parallel workers to synchronize at the first collective operation of the training backward pass. _Train_ measures GPU computation for training.

Figure 9: Comparison of VIP-analytic versus VIP-simulation on end-to-end runtime on slow networks. Illustrates the per-epoch runtime of \(16\)-node executions on _papers_ and _mag240c_ when using the VIP-analytic and VIP-simulation caching policies with different replication factors.
Preprocessing overheadsSALIENT++'s preprocessing overheads fall into two categories: (a) runtime overheads incurred before each execution on a dataset; and, (b) dataset preparation overheads that are incurred only once and can be amortized over multiple experiments. The runtime overheads for an \(8\)-node execution of SALIENT++ on the _papers_ dataset with replication factor \(\alpha=0.32\) are as follows. Loading the dataset from disk takes approximately \(10\) seconds. Computing the VIP weights for fanouts \((15,10,5)\) takes \(11.8\) seconds when implemented in PyTorch with sparse transition weight matrices and dense hop-wise VIP-vectors. For large graphs, the VIP computation code streams batches of matrix rows to the GPU, overlapping communication and data transfer. The communication of remote feature vectors and the related tensor slicing operations take about \(22\) seconds. The dataset preprocessing overheads are highly dependent on the workflow used for partitioning graphs. Our (unoptimized) workflow for graph partitioning uses serial METIS to partition graphs and, additionally, uses machines with limited memory necessitating the use of swap files. In this setting, partitioning the _papers_ dataset takes \(\sim\!2\) hours and creating a reordered dataset from that partition takes \(30\) minutes. Optimized workflows for partitioning can generate partitions substantially faster. For example, DistDGLv2 (Zheng et al., 2022) reports a \(12\) minute runtime for partitioning when using ParMETIS (Karypis et al., 1998) and \(4\) large multicores. SALIENT++ is agnostic to the source of the graph partitioning, and can be used in conjunction with more scalable graph partitioning codes.

Comparison to DistDGLWe compared the performance of SALIENT++ to DistDGL on _papers_ using the GraphSAGE architecture. A summary of this comparison is provided in Table 4. The DistDGL code was obtained from DGL's public GitHub repository and executed in the same computing environment used to evaluate SALIENT++. DistDGL's public code for distributed multi-GPU training was approximately 12.7\(\times\) slower than SALIENT++ on \(8\) GPUs (1 GPU per machine). Of course, DGL is under active development and its support for efficient GNN training (especially in distributed environments) is evolving rapidly. As such, we also report performance results from the recent work "DistDGLv2" (Zheng et al., 2022) that describes innovations in the DistDGL system for improving distributed GNN training performance. Although not all of these innovations are publicly available, their reported performance can be used to contextualize the performance we achieve with SALIENT++. A full comparison of the hardware used by DistDGLv2 and SALIENT++ shows that SALIENT++ consistently achieves similar (often better) performance than what was reported for DistDGLv2, while using substantially less resources. SALIENT++'s \(8\)-GPU runtime is approximately \(1.7\times\) faster than DistDGLv2's reported numbers for executing a training epoch on _papers_ using \(64\) GPUs. This faster per-epoch time is achieved with \(8\times\) fewer GPUs, \(4\times\) smaller network bandwidth, and \(3.2\times\) cheaper hardware.

## 6 Conclusions and Future Work

We have presented a distributed multi-GPU system, SALIENT++, for GNN training and inference on massive graphs. This system is built on top of SALIENT, a prior state-of-the-art system that attains efficiency and scalability through fast sampling and pipelining, at the cost of full data replication on all machines. In SALIENT++, we distribute the vertex feature data and address the resulting communication bottleneck through analyzing the access pattern of the out-of-machine (i.e., remote) vertices and proposing a caching strategy that replicates a small amount of the most frequently accessed vertex features. Together with a deep pipelining of all operations from communication to computation, SALIENT++ retains the efficiency and scalability of SALIENT while consuming only a fraction of the storage required by SALIENT.

An avenue of future work is to further apply the access pattern analysis to improve the initial graph partitioning, with an aim of reducing the communication volume orthogonally to the use of caching. This would require incorporating the vertex inclusion probabilities in the graph partitioning objective, on top of edge cuts and load balancing. Additionally, a hierarchical graph partitioning may better leverage the higher intra-machine bandwidth among GPUs than inter-machine communication. Another line of future work is to explore distributing the graph structure across machines. Distributing the graph incurs nontrivial challenges in the multi-round communication of node features, but resolving this challenge will further reduce memory consumption and bears the potential of handling graphs that are orders of magnitudes larger than the current largest benchmarks.

\begin{table}
\begin{tabular}{l l l} \hline \hline _System_ & _Time (s)_ & _Notes_ \\ \hline SALIENT++ & 2.9 & 8 NVIDIA AI0G GPUs, 25Gbps network throughput, 519.5 / hr. \\ DistDGL & 37.0 & Same hardware as is used by SALIENT++. Example distributed code from Github3  \\ DistDGLv2 & \(\approx 5\) & 64 NVIDIA T4 GPUs, 100Gbps network throughput, 862.6 / hr. \\ \hline \hline \end{tabular}
\end{table}
Table 4: Comparison of SALIENT++, DistDGL (public code), and DistDGLv2 (Zheng et al., 2022). All numbers are reported for a \(3\)-layer GraphSAGE architecture with fanouts 15,10,5 and \(256\) hidden nodes run on the _papers_ dataset.
## Acknowledgements

This work was conducted on the SuperCloud computing cluster [https://supercloud.mit.edu](https://supercloud.mit.edu) and the Satori computing cluster [https://mit-satori.github.io](https://mit-satori.github.io). This research was sponsored by MIT-IBM Watson AI Lab and in part by the United States Air Force Research Laboratory and the United States Air Force Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750-19-2-1000. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
