EXACT: SCALABLE GRAPH NEURAL NETWORKS TRAINING VIA EXTREME ACTIVATION COMPRESSION
Year: 2021

Abstract:
Training Graph Neural Networks (GNNs) on large graphs is a fundamental challenge due to the high memory usage, which is mainly occupied by activations (e.g., node embeddings). Previous works usually focus on reducing the number of nodes retained in memory. In parallel, unlike what has been developed for other types of neural networks, training with compressed activation maps is less explored for GNNs. This extension is notoriously difficult to implement due to the miss of necessary tools in common graph learning packages. To unleash the potential of this direction, we provide an optimized GPU implementation which supports training GNNs with compressed activations. Based on the implementation, we propose a memory-efficient framework called “EXACT”, which for the first time demonstrate the potential and evaluate the feasibility of training GNNs with compressed activations. We systematically analyze the trade-off among the memory saving, time overhead, and accuracy drop. In practice, EXACT can reduce the memory footprint of activations by up to 32× with 0.2-0.5% accuracy drop and 10-25% time overhead across different models and datasets. In addition, EXACT can trim down the hardware requirement of training a three-layer full-batch GraphSAGE on ogbn-products from a 48GB GPU to a 12GB GPU. The code is available at https://anonymous.4open.science/r/exact-revd.

heading: 1 INTRODUCTION
text: Despite Graph Neural Networks (GNNs) have achieved great success across different graph-related tasks, training GNNs on large graphs is a long-standing challenge due to its extensive memory requirement (Kipf & Welling, 2017; Zhang & Chen, 2018; Cai et al., 2021b). The extensive memory consumption of GNN stems from its recursive neighborhood aggregation scheme, where each node aggregates embeddings of its neighbors to update its new embedding at each layer. Thus, training an L-layer GNN requires storing all L layers’ intermediate node embeddings in GPU memory for computing the gradients, and this typically adds several times more memory than holding the node feature matrix (see Algorithm 1 for a detailed analysis). Hence, storing these node embeddings is the major memory bottleneck for training GNNs on large graphs.
Most of the existing works towards this problem can be roughly divided into two categories. First, some works propose to train GNNs with sampled subgraphs instead of the whole graph at each step. In this way, only node embeddings that are present in the current subgraph will be retained in memory (Chiang et al., 2019; Hamilton et al., 2017; Zeng et al., 2020; Zou et al., 2019; Chen et al., 2018; Huang et al., 2018). Second, another line of work tries to decouple the neighborhood aggregation from prediction, either as a preprocessing step (Wu et al., 2019; Klicpera et al., 2018; Yu et al., 2020) or post-processing step (Huang et al., 2020), where the model is simplified as the Multi-Layer Perceptron (MLP) that can be trained with mini-batch data.
In parallel, another orthogonal direction is to store only the compressed node embeddings (i.e., activations) in memory for computing the gradients. Recent works propose to quantize the activations in lower numerical precision (e.g., using 8-bit integer) during the forward pass (Chakrabarti & Moseley, 2019; Fu et al., 2020; Chen et al., 2021a; Evans & Aamodt, 2021). This framework successfully trims down the memory requirement for training Convolutional Neural Networks (CNNs) by a large margin, at the cost of additional time overhead and loss of accuracy. Ideally, the real-world usage requires the training method should achieve a balanced trade-off among the following three aspects:
1. Space. It should enable to train GNNs on large graphs using off-the-shelf hardwares, such as GPUs and CPUs; 2. Speed. The time overhead should be acceptable, ideally as small as possible; 3. Model Performance. The loss of accuracy should be acceptable, ideally as small as possible.
Although storing the compressed activations successfully saves the memory for CNNs, up to our knowledge, there is no existing work extends this direction to GNNs and evaluates the mentioned trade-off for analyzing its feasibility. Despite the extension is conceptually straightforward, this direction is less-explored since it can be notoriously difficult to implement to fully leverage hardware potentials. This dilemma stems from the fact that the necessary tools for supporting this direction is usually missed in common graph learning packages. For example, operations in popular graph learning packages only support casting tensors down to 8-bit integer on GPUs, significantly limiting the memory saving (Paszke et al., 2019; Fey & Lenssen, 2019; Wang et al., 2019). As a result, previous GNN quantization works either emulates inference-time quantization via “simulated quantization” (Tailor et al., 2021; Zhao et al., 2020), or is impractical to use GPUs for accelerating (Feng et al., 2020). To unleash the potential of this direction, we provide a space-efficient GPU implementation for supporting common operations in GNNs with compressed activations. Equipped with our implementation, this paper asks the following open question:
To what extent can we compress the activations with both acceptable loss in accuracy and time overhead, for scalable GNN training?
To answer the open question, we first explore two different types of compression methods. Namely, one is “quantization” that compresses the activations into lower numerical precision. The other one is called “random projection” (Achlioptas, 2001) that projects the activations into low-dimensional space. Both these two simple strategies can achieve near-lossless accuracy at a non-trivial compression ratio. For example, the loss in accuracy is negligible (0.2%) even under the vanilla 2-bit quantization. However, we cannot further push forward the memory saving by these two methods, e.g., we cannot use a numerical precision below 1-bit. Considering that the real-world graphs often contain hundreds of millions of nodes, our main goal is to trim down the memory consumption to the maximum extent among the three aspects, as long as the other two are acceptable. We then naturally explore the direction of combining random projection and quantization, dubbed “EXACT”, to aggressively maximize the memory saving. EXACT is essentially applies random projection and quantization sequentially for compressing activations. Despite the superior memory saving, another following question is that, whether the combination bring significantly worse model performance and larger time overhead? Following the questions, we make three major contributions as follows:
• We provide a space-efficient GPU implementation for training GNNs with compressed activations as an extension for Pytorch. Based on our implementation, we are the first one training GNNs with compressed activations and demonstrating its potential for real-world usage. EXACT can complement the existing studies, as it can be integrated with most of existing solutions.
• We propose EXACT, a simple-yet-effective framework which applies random projection and quantization sequentially on activation for scalable GNN training. We theoretically and experimentally show that applying random projection and quantization sequentially have an “interaction” effect. Namely, from the model performance aspect, after random projection, applying quantization only has a limited impact on the model performance. From the time aspect, EXACT runs comparable or even faster than quantization only.
• Despite the simplicity, EXACT achieves non-trivial memory saving with both acceptable time overhead and loss in accuracy: EXACT can reduce the memory footprint of activations by up to 32× with roughly 0.5% loss in accuracy and 10− 25% time overhead across models and datasets. Noteworthily, EXACT trims down the hardware requirement of training a full-batch GraphSAGE (Hamilton et al., 2017) on ogbn-products (Hu et al., 2020) from a 48GB GPU to a 12GB GPU, which is affordable to most research labs.

heading: 2 THE MEMORY CONSUMPTION OF GNNS
text: Background. Let G = (V, E) be an undirected graph with V = (v1, · · · , v|V|) and E = (e1, · · · , e|E|) being the set of nodes and edges, respectively. Let X ∈ R|V|×d be the node feature matrix of the whole graph. The graph structure can be represented by an adjacency matrix A ∈ R|V|×|V|, where Ai,j = 1 if (vi, vj) ∈ E else Ai,j = 0. In this work, we are mostly interested
in the task of node classification, where the goal is to learn the representation hv for all v ∈ V such that the label yv can be easily predicted. To obtain such a representation, GNNs follow the neighborhood aggregation scheme. Specifically, GNNs recursively update the representation of a node by aggregating representations of its neighbors. Formally, the lth layer of a GNN can be written
as h(l+1)v = UPDATE ( h (l) v , ⊕ u∈N (v) MSG ( h (l) u ,h (l) v )) , where h(l)v is the representation of node
v at the lth layer. N (v) denotes the neighboring nodes of node v, not including v itself. The full table of notations can be found in Appendix A Table 5. For node v, messages from its neighbors are calculated by MSG(·). Then these messages are aggregated using a permutation-invariant aggregation function ⊕ . The aggregated features at v are updated by UPDATE(·). For example, the Graph Convolutional Network (GCN) (Kipf & Welling, 2017) layer can be defined as
H(l+1) = ReLU(ÂH(l)Θ(l)), (1)
where H(l) is the node embedding matrix consisting of all nodes’ embeddings at the lth layer and H(0) = X . Θ(l) is the weight matrix of the lth GCN layer. Â = D̃− 1 2AD̃− 1 2 is
the normalized adjacency matrix, where D̃ is the degree matrix of A + I . We note that Â is usually stored using the sparse matrix format. For GCN layers, the message and the aggregation function are fused into a Sparse-Dense Matrix Multiplication (SPMM 1) operation. Namely, ÂH(l)Θ(l) = SPMM(Â,H(l)Θ(l)). Thus, the computation graph of Equation 1 can be written as
H(l+1) = ReLU ( SPMM ( Â,MM(H(l),Θ(l)) )) , (2)
where MM(·, ·) is the normal Dense-Dense Matrix Multiplication. Equation 2 resembles how GCNs are implemented in popular packages (Fey & Lenssen, 2019; Wang et al., 2019).
Here we analyze the memory consumption for the forward pass of GCN since most of the memory is occupied during the forward pass. For the memory usage of the backward pass, a detailed analysis is given in Appendix C. Specifically, for an L layer GCN, suppose the hidden dimensions of layer 0, · · · , L − 1 are the same, which are denoted as D. The forward pass of GCN layers is shown in Appendix C Algorithm 1. For layer l, it saves the following four variables in memory. (1) the weight matrix Θ(l) ∈ RD×D whose shape is independent to the graph size and is generally negligible. (2) the normalized adjacency matrix Â (in CSR format) whose space complexity is O(|V|+ |E|). Note that there only needs to store one Â in memory which can be accessed by different layers. Thus, the memory consumption of Â is independent of the number of layers and is not the main memory bottleneck. (3) the intermediate result J (l) = MM(H(l),Θ(l)) ∈ R|V|×D. For an L layer GCN, storing {J (0), · · ·J (L−1)} has a O(L|V|D) space complexity, which is the main memory bottleneck. (4) the node embedding matrix H(l) ∈ R|V|×D. For a L layer GCN, storing {H(0), · · ·H(L−1)} has a O(L|V|D) space complexity, which is the also main memory bottleneck. In this paper, we use the term “activation maps” to encompass H , J , and activation maps of other commonly used layers/operations such as BatchNorm, ReLU and Dropout.

heading: 3 METHODOLOGY
text: From the above analysis, the memory consumption of GCNs can be significantly reduced by storing only the compressed activations. As shown in Figure 1, due to the simplicity and small time overhead of quantization (Section 3.1) and random projection (Section 3.2), we first explore these two methods for compressing the activations of GNNs. We show that these two simple methods can bring moderate compression ratio with negligible loss in accuracy. However, the highest compression rates of these two methods are either limited by hardwares or limited by the accuracy drop. Motivated by GNN’s endless demand on memory, we then explore the direction of combining quantization with random projection into one framework termed “EXACT”, to aggressively maximize the memory saving (Section 3.3). We show that random projection and quantization have an “interaction” effect. Namely, from the model performance aspect, after random projection, applying quantization only has a limited impact on the model performance. From the time aspect, EXACT runs comparable or even faster than quantization only.
1We note that the output of SPMM(·, ·) is a dense matrix.

heading: 3.1 STORING QUANTIZED ACTIVATIONS
text: Inspired by ActNN (Chen et al., 2021a), as shown in Figure 1a, we first explore the direction of compressing activation maps in lower numerical precision. Specifically, during the forward pass, each layer’s accurate activations (e.g., J (l) and H(l)) is used to compute those for the subsequent layer. Then the accurate activations will be quantized into the compressed activations (e.g., H(l)INT and J (l)INT) with lower numerical precision, which overwrites the the accurate activations and is retained in the memory. During the backward pass, we dequantize the compressed activations back to full-precision (e.g., Ĥ(l) and Ĵ (l)). Then the gradients are computed based on the dequantized activations. We note that all operations (e.g., MM and SPMM ) are done in full-precision since most of the GPUs do not support operands with bit-width other than full-precision and half-precision. For convenience, we use the term “precision” to present “numerical precision” in the following of this paper. Below we introduce the technical details of (de)quantization.
Specifically, each node embedding vector h(l)v will be quantized and stored using b-bit integers. Let B = 2b − 1 be the number of quantization bins. The integer quantization can be expressed as
h(l)vINT = Quant(h (l) v ) = b h (l) v − Z(l)v r (l) v Be, (3)
where Z(l)v = min{h(l)v } is the zero-point, r(l)v = max{h(l)v }−min{h(l)v } is the range for h(l)v , and b·e is the stochastic rounding operation (Courbariaux et al., 2015). H(l)INT in Figure 1a is the matrix containing all quantized h(l)vINT . During the backward pass, each h (l) vINT will be dequantized as
ĥ(l)v = Dequant(h (l) vINT ) = r (l) v h (l) vINT/B + Z (l) v . (4)
Ĥ(l) in Figure 1a is the matrix containing all dequantized embeddings ĥ(l)v . The following proposition characterizes the effect of quantization, which is adopted from Chen et al. (2021a).
Proposition 1 (Details in Appendix E) The above quantization and dequantization are unbiased operations, i.e., E[ĥ(l)v ] = E[Dequant(Quant(h(l)v ))] = h(l)v and Var(ĥ(l)v ) = D[r (l) v ] 2
6B2 .
From the unbiased nature of quantization illustrated in Proposition 1, the calculated gradient is also unbiased. The approach imposes extra noise (i.e., the variance term in Proposition 1) to the gradient during the backward pass. From Proposition 1, the noise effect is inversely-correlated with the number of quantization bins B. To evaluate how the noise affects the model performance, we train three popular GNN models, namely, GCN, GraphSAGE, and GAT, on the ogbn-arxiv dataset (Hu
et al., 2020) with different precision (Table 1). The detailed experiment setting is elaborated in Appendix G.2. Here we emphasize that all these three models are trained with full-batch data.
Observations. For all three models, the loss in accuracy is negligible (≈ 0.2%) even using the vanilla 2-bit quantization. In contrast, for CNNs, adopting the vanilla INT1 or INT2 quantization, will cause a significant accuracy drop (usually > 5%) (Chen et al., 2021a; 2020a). This observation can be explained by the following mechanism. We show in Appendix E.4 that the approximation error will compound layer-by-layer during the backward pass. The layer-depth of GNNs is much smaller than CNNs due to the over-smoothing problem. And hence GNNs is much more noisetolerant than CNNs. Our observation suggests that in practical scenarios, there is no need for an expensive sophisticated quantization approach (e.g., mixed precision quantization) as considered in previous works (Chen et al., 2021a; Fu et al., 2020).

heading: 3.2 STORING RANDOMLY PROJECTED ACTIVATIONS
text: Here we explore the direction of random projection. The key idea is to project the activation maps into a low-dimensional space that keeps the original information as much as possible. In this way, similar to the framework in Section 3.1, we only need to store the dimension reduced activation maps in memory. Figure 1b illustrates the workflow of GCN layers. During the forward pass, each layer’s accurate activations (e.g., h(l)v ) are used to compute those for the subsequent layer. Then the accurate activations will be projected into low-dimensional space, which can be expressed as:
h(l)vproj = RP(h (l) v ) = h (l) v R, (5)
where R ∈ RD×R is a random matrix (R < D) which satisfies E[RR>] = I . H(l)proj ∈ R|V|×R is the matrix containing all projected node embeddings h(l)vproj . In this paper, R is set as the normalized Rademacher random matrix (Achlioptas, 2001) due to its low sampling cost (detailed introduction in Appendix D). After projection, we store only H(l)proj in memory and hence the compression ratio is DR . During the backward pass, the projected node embeddings are inversely transformed by
ĥ(l)v = IRP(h (l) vproj) = h (l) vprojR >, (6)
where IRP(·) is the inverse projection operation. Ĥ(l) = H(l)projRR> ∈ R|V|×D is the recovered activation map containing all ĥ(l)v . The following proposition shows the effect of random projection.
Proposition 2 (Proof in Appendix E) The above RP and IRP are unbiased operations, i.e., E[Ĥ(l)] = E[IRP(RP(H(l)))] = H(l). For each ĥ(l)v in Ĥ(l), we have Var(ĥ(l)v ) = D−1R ||h (l) v ||22.
From the unbiased nature of random projection illustrated in Proposition 2, the calculated gradient is also unbiased. The approach here also only imposes extra variance to the calculated gradient, where the variance linearly scales with the DR ratio. To quantitatively study the effect of the extra variance in scenarios of practical interest, following the same setting in Section 3.1, we show the performance of GCN, GraphSAGE, and GAT trained on the ogbn-arxiv dataset with different DR ratio in Table 2.
Observations. We make two main observations: (1) For all three models, the loss in accuracy tends to range from negligible (≈ 0.2%) to moderate (≈ 0.5%) when DR ≤ 8. (2) Under the same compression ratio, the quantization is better than random projection in terms of the loss in accuracy (e.g., compare the DR = 16 result in Table 2 with the INT2 result in Table 1).

heading: 3.3 EXACT: COMBINING RANDOM PROJECTION AND QUANTIZATION
text: We experimentally show that GNNs are noise-tolerant to inaccurate activations compressed by quantization and random projection. However, the highest compression ratio of these two methods is limited. Namely, the precision cannot be less than 1-bit for quantization and DR cannot surpass eight for random projection due to a large accuracy drop. However, GNN’s demand on memory is endless since real-world graphs can contain hundreds of millions of nodes and the graph size is ever growing. Motivated by this fact, we try to further push forward the memory saving as long as the loss of accuracy and the time overhead are both acceptable. We explore the direction of combining these two methods into one framework termed “EXACT” to maximize the memory saving. Specifically, we compress the activations and store only the h̃(l)v = Quant(RP(h (l) v )) in memory during the forward pass. During the backward pass, the node embedding is recovered as ĥ(l)v = IRP(Dequant(h̃ (l) v )). EXACT can achieve a superior compression ratio, e.g., the compression ratio of EXACT is roughly 128× if DR = 8 and the precision is INT2. However, one practical question is that, will the variance explode, leading to significantly worse performance when applying random projection and quantization sequentially? Here we try to answer the above question both theoretically and experimentally. Below we first theoretically show that after applying random projection, the quantization range of projected node embeddings is bounded.
Proposition 3 (Proof in Appendix E) For each projected node embedding h(l)vproj = RP(h (l) v ) =
h (l) v R, for ∀ > 0, by choosing s = ||h(l)v ||2
√ 2ln(2R/ )
R , we have P (||h (l) vproj ||∞ ≤ s) ≥ 1− .
From Proposition 3, after projection, the maximal absolute value among weights in h(l)vproj is bounded. Following the fact that the quantization range r(l)vproj = max{h (l) vproj} − min{h (l) vproj} ≤ 2||h (l) vproj ||∞,
r (l) vproj is also bounded. Recall that the variance of quantization scales with
[r(l)vproj ]2
B2 (Proposition 1), applying random projection and quantization sequentially will not lead to the variance explosion.
We also experimentally visualize the infinity norm of projected embeddings in Appendix E Figure 4. The infinity norm of projected embeddings may be even less than those of original embeddings when R is larger than a threshold (R = 0.5D in Figure 4). Thus, the extra variance is expected to be dominated by the random projection, and hence the loss in accuracy is largely determined by the DR ratio. We also experimentally investigate the effect of EXACT by plotting the test accuracy against each of quantization precision and the DR ratio.
Due to the page limit, we only present one representative result in Figure 2, namely, training GCNs on ogbn-arxiv with EXACT. More similar results can be found in Appendix I.1 Here we do not consider INT1 precision since its model performance drop is already near 0.5%. We summarize two main observations. First, the performance of EXACT is largely determined by the DR ratio of random projection. Second, when DR ≤ 8, the loss in accuracy generally ranges from ≈ 0.2% to ≈ 0.5%, regardless of the quantization precision. System Implementation. In EXACT, each operations (e.g., SPMM , ReLU, and BatchNorm) is reimplemented with different configurations using CUDA kernels. For example, BatchNorm only sup-
ports quantization, while the random projection is not applicable to it. Since Pytorch only supports precision down to INT8, we convert quantized tensors into bit-streams by CUDA kernels to maximize the memory saving. The configuration and implementation details are given in Appendix F.

heading: 4 RELATED WORK AND DISCUSSION
text: Due to the page limit, we briefly review and discuss the relationship between existing works and EXACT. A more comprehensive discussion can be found in Appendix B. Existing works can be roughly divided into scalable/efficient GNN inference and scalable GNN training according to the problem they try to solve. Almost all previous GNN quantization works try to solve the first problem, which is much simpler than the second problem that EXACT tries to address. Specifically, most of them try to enable the usage of low precision integer arithmetic during inference (Tailor et al., 2021; Zhao et al., 2020) by simulating the quantization effect, which may even increase the memory consumption during training. Feng et al. (2020) tries to address the second problem by proposing a heterogeneous quantization framework which assigns different bits to node embeddings in each layer. However, this framework is impractical on off-the-shell hardwares. For scalable GNN training methods, EXACT is orthogonal to most of them, including distributed training (Zheng et al., 2020b; Jia et al., 2020; Zhu et al., 2019; Wan et al., 2021), subgraph sampling (Hamilton et al., 2017; Zeng et al., 2020; Chiang et al., 2019), and historical embedding-based methods (Fey et al., 2021). We discuss the potential benefit of integrating EXACT with them in Appendix B.

heading: 5 EXPERIMENTS
text: The experiments are designed to answer the following research questions. RQ1: How effective is EXACT in terms of model performance at different compression rates (Section 5.1)? RQ2: Is the training process of deeper GNNs also robust to the noise introduced by EXACT (Section 5.1)? RQ3: How sensitive is EXACT to its key hyperparameters (Appendix I.4)? RQ4: What is the running time overhead of EXACT (Section 5.2)? RQ5: Is the convergence speed of GNNs impacted by EXACT (Appendix I.3)? RQ6: To what extent can EXACT reduce the hardware requirement for training GNNs on large graphs (Section 5.3)?
Datasets and Models: To evaluate the scalability of EXACT, we adopt five common large scale graph benchmark datasets from different domains. Namely, Reddit, Flickr, Yelp, ogbn-arxiv, and ogbn-products. We evaluate EXACT under both the mini-batch training and full-batch training setting. For the mini-batch training setting, we integrate EXACT with two state-of-the-art subgraph sampling methods, namely, Cluster-GCN (Chiang et al., 2019) and GraphSAINT (Zeng et al., 2020). For the full-batch training setting, we integrate EXACT with three popular models: two commonly used shallow models, namely, GCN (Kipf & Welling, 2017) and GraphSAGE (Hamilton et al., 2017), and one deep model GCNII (Chen et al., 2020b). To avoid creating confusions, GCN, GraphSAGE, and GCNII are both trained with the whole graph at each step. For a fair comparison, we use the mean aggregator for GraphSAGE, Cluster-GCN, and GraphSAINT throughput the paper. Details about the hyperparameters of models and datasets can be found in Appendix G.
Hyperparameter Settings: From Section 3.3, we show that the performance of EXACT is largely determined by the DR ratio of random projection, and the model performance drop with INT2 quantization is negligible. Thus, to balance the accuracy drop and memory saving ratio, we adopt INT2 precision with different DR ratio for EXACT. Specifically, EXACT (INT2) indicates it only applies 2- bit quantization to the activation maps of all applicable operations (see Table 8). EXACT(RP+INT2) indicates it applies the random projection followed by a 2-bit quantization to the activations of all applicable operations. We perform a grid search for DR ratio from {2, 4, 8}. Detailed D R configuration for EXACT(RP+INT2) can be found in Table 12.

heading: 5.1 TRADE-OFF BETWEEN SPACE AND MODEL PERFORMANCE
text: To answer RQ1 and RQ2, we first analyze the trade-off between the model performance and memory saving. Table 3 summarizes the model performance and the memory footprint of the activation maps. Besides the memory usage of activation maps, a detailed analysis about the overall memory saving ratio is provided in Appendix H. We make two main observations.
Table 3: Comparison on the test accuracy/F1-micro and memory saving on five datasets. The hardware here is a single RTX 3090 (24GB). “Act Mem.” is the memory (MB) occupied by activation maps. “OOM” indicates the out-of-memory error. Bold faces indicates the loss in accuracy is negligible (≈ 0.2%) or the result is better compared to the baseline.Underline numbers indicates the loss
in accuracy is moderate (≈ 0.5%). All reported results are averaged over ten random trials.
# nodes # edges
230K 11.6M
89K 450K
717K 7.9M
169K 1.2M
2.4M 61.9M
Model Methods Reddit Flickr Yelp ogbnarxiv
ogbnproducts
Acc. ActMem. Acc. Act Mem. F1-micro Act Mem. Acc. Act Mem. Acc. Act Mem.
ClusterGCN
Baseline 95.62±0.10 14.5 (1×) 49.61±0.47 16.5 (1×) 63.98±0.14 29.3 (1×) — — 78.62±0.26 35.2 (1×) EXACT(INT2) 95.58±0.09 2 (7.3×) 49.69±0.20 1.5 (11×) 63.90±0.15 4 (7.3×) — — 78.47±0.40 2.5 (14×) EXACT(RP+INT2) 95.32±0.07 1.4 (10.4×) 49.31±0.21 0.9 (18.3×) 63.61±0.21 3 (9.8×) — — 77.86±0.28 2.2 (16×)
GraphSAINT
Baseline 96.02±0.08 44.3 (1×) 51.11±0.28 88.7 (1×) 63.78±0.12 33.5 (1×) 71.49±0.20 270 (1×) 79.03±0.23 516 (1×) EXACT(INT2) 95.96±0.05 6.6 (6.7×) 50.86±0.32 7.8 (11.4×) 63.77±0.14 4.3 (7.8×) 71.76±0.15 20 (13.5×) 78.94±0.28 40.5 (12.7×) EXACT(RP+INT2) 95.69±0.06 3.6 (12.3×) 50.65±0.17 3.4 (26×) 63.41±0.19 3.3 (10.2×) 71.44±0.16 10.8 (25×) 78.50±0.41 29.5 (17.5×)
GCN Baseline 95.39±0.04 1029 (1×) 53.08±0.14 378.8 (1×) 40.22±0.47 6429 (1×) 72.07±0.16 729.4(1×) — —
EXACT(INT2) 95.36±0.03 122.8 (8.4×) 52.92±0.20 37 (10.2×) 40.20±0.38 640 (10×) 72.04±0.21 54.5 (13.4×) — — EXACT(RP+INT2) 95.30±0.03 67 (15.4×) 52.90±0.22 17.8 (21.3×) 39.89±0.56 427 (15×) 71.67±0.16 30.2 (24.1×) — —
GraphSAGE
Baseline 96.44±0.04 1527.4 (1×) 51.74±0.13 546.9 (1×) 62.05±0.14 6976 (1×) 71.85±0.24 786.2(1×) 78.782±0.19 OOM EXACT(INT2) 96.40±0.05 155.5 (9.8×) 51.97±0.22 49.3 (11.1×) 61.95±0.12 680 (10.3×) 71.71±0.38 60.8 (12.9×) 78.79±0.12 1144
EXACT(RP+INT2) 96.34±0.03 71.7 (21.3×) 51.83±0.21 20.4 (26.8×) 61.59±0.12 466.5 (15×) 71.32±0.26 30.5 (25.8×) 78.76±0.13 572
GCNII Baseline 96.71±0.07 5850 (1×) 54.23±0.77 4067 (1×) OOM OOM 72.85±0.27 14409 (1×) — —
EXACT(INT2) 96.65±0.06 388 (15×) 54.55±0.47 256.4 (15.9×) 64.79±0.09 2236 72.85±0.43 899 (16×) — — EXACT(RP+INT2) 96.51±0.09 197.8 (29.6×) 53.96±0.58 127.6 (31.9×) 64.01±0.17 1649 72.67±0.45 451.2 (32×) — —
First, EXACT aggressively reduces the memory footprint with ≤ 0.5% loss in accuracy. From Table 3, EXACT (INT2) reduces the memory footprint of activation maps from 7× to 16×, with negligible loss in accuracy (≈ 0.2%) in almost all experiments. EXACT (RP+INT2) generally reduces the memory footprint of activation maps from 10× to 32×, and the loss in accuracy tends to range from lossless (0.0%) to moderate (≈ 0.5%). Although EXACT can aggressively compress the memory usage of activations associated with MM and SPMM by up to 128×, the compression ratio for other operations usually cannot surpass 32× (e.g., the activation maps of ReLU and Dropout take exact one bit to store, see Appendix F). As a result, the overall compress ratio is pulled down. We show in Appendix H that the measured compression ratio is consistent with the theoretical one.
Second, the training process of deeper models is also robust to extremely compressed activations. As we analyzed in Appendix E.4, the vanilla quantization and random projection obtain unreasonable good performance might because the commonly used GNNs are shallow. Thus, one open question is that, to what extent can deeper GNNs robust to extremely compressed activation maps (RQ2)? The ablation study of GCNII in Table 3 tries to experimentally answer this question. Namely, the training process of deeper GNNs (up to 16-layer, see Table 17) is also robust to the noise introduced by the extremely compressed activations, with up to 32× less memory footprint of activation maps. We note that building deeper GNNs is still an open direction due to the oversmoothing and the extensive memory consumption problem (Oono & Suzuki, 2019; Li et al., 2021; Chen et al., 2020b). From the practical usage perspective, our GCNII ablation study indicates that EXACT may enable to train deeper GNNs on large graphs with minimal loss in performance.
To answer RQ3, we present the sensitivity study of EXACT (RP+INT2) to DR ratio in Appendix I.4. In general, the model performance drop of EXACT (RP+INT2) increases with the DR ratio. In summary, when DR = 8, the loss in accuracy of EXACT (RP+INT2) is ≤ 0.5% in roughly two-third of experiments and ≤ 1% in almost all experiments.

heading: 5.2 TRADE-OFFS BETWEEN SPACE AND SPEED
text: To answer RQ4, we compare the training throughput of EXACT with the baseline using a single RTX 3090 (24GB) GPU. Figure 3 shows the results among EXACT with different configurations and the baseline (see Table 9 for detailed package information). We make two main observations.
First, the overhead of EXACT (INT2) is roughly 12% ∼ 25%. The overhead of EXACT (INT2) comes entirely from quantization. We note that EXACT (INT2) adopts vanilla integer quantization with negligible loss in accuracy, which has the lowest overhead among all quantization methods. Sophisticated quantization strategies (e.g., mixed precision quantization and non-uniform quantization) further increase the time overhead, because they require an extra search process to find precision or
2The regular training raises OOM error. The result is obtained using the automated mixed-precision training (torch.cuda.amp).
quantization step for each layer/sample (Chakrabarti & Moseley, 2019; Fu et al., 2020; Chen et al., 2021a). From Figure 3, the overhead of EXACT is roughly 12% ∼ 25%. Second, the running speed of EXACT (RP+INT2) is comparable or even faster than quantization only. The above counter-intuitive observation can be explained by the following mechanism. Although random projection introduces two extra matrix multiplication operations, the total number of elements to be quantized is DR times less than directly quantizing the activations. Considering that quantization is the main bottleneck, compared to EXACT (INT2), EXACT (RP+INT2) achieves roughly twice overall compression ratios with a comparable or even smaller time overhead.
RQ4 focuses on the actual running speed measured by the hardware throughput. Here we investigate the convergence speed from the optimization perspective (RQ5). To answer RQ5, due to the page limit, we present the training curve of models trained on Reddit dataset with EXACT in Appendix I.3. The converge speed here is measured by the number of epochs it takes to reach a certain validation accuracy. In summary, the convergence speed of EXACT is slightly slower than the baseline, with limited impact on final accuracy.

heading: 5.3 ABLATION STUDIES
text: Train full-batch GraphSAGE on ogbn-products using a GPU with 11GB memory.
Table 3 shows that EXACT reduces the memory footprint of activation maps by up to 32×. However, as analyzed in Appendix H, besides the activation maps, input data (including feature matricesX , adjacency matrix A, and labels), activation gradients also occupy memory in the full-batch training. For ogbn-products, the input data occupies 4.7GB (see Appendix H). From the OGB leader-board, training a full-batch GraphSAGE re-
quires GPUs with at least 48GB memory. We note that EXACT can be integrated with AMP (see Appendix F) to further squeeze out the memory. To answer RQ6, by integrating AMP and EXACT (RP+INT2) with DR = 4, we successfully train a full-batch GraphSAGE on ogbn-products on a single GTX 1080 Ti with 11GB memory, with moderate loss (≈ 0.5%) in accuracy (Table 4).

heading: 6 CONCLUSION AND FUTURE WORK
text: In this paper, we propose EXACT, a simple-yet-effective framework for training GNNs with compressed activations. EXACT essentially applies random projection and quantization sequentially on activations, which can be plugged into any training method to save memory. We demonstrate the potential of EXACT for the real-world usage by systematically evaluating the trade-off among the memory-saving, time overhead, and accuracy drop. Specifically, EXACT achieves non-trivial memory saving with 0.2-0.5% accuracy drop and 10-25% time overhead across five models and five large scale graph datasets. We show that EXACT is orthogonal to almost all of existing scalable GNN training methods. We discuss in detail that how EXACT can be integrated with existing solutions in Appendix B. Future work includes (1) evaluating EXACT under multi-GPU settings for distributed training; (2) combining EXACT with historical embedding-based method; (3) combining EXACT with memory-efficient training system, such as gradient checkpointing and swapping.

heading: B RELATED WORK AND DISCUSSION
text: Due to the page limit, we discuss the relationship between EXACT and existing works in detail here. GNN is notoriously inefficient and non-scalable during both the training and inference phase. Previous works can be divided into two categories according to the problem they try to solve. Namely, scalable/efficient GNN inference and scalable GNN training. Below we introduce and discuss the relationship between them and EXACT.

heading: B.1 SCALABLE/EFFICIENT GNN INFERENCE
text: Scalable GNN inference is an orthogonal direction to the problem that EXACT tries to address. Below we introduces two popular directions for scalable GNN inference. Namely, GNN quantization and graph sparsification.
GNN Quantization. Most of the previous GNN quantization works focus on enabling the usage of low precision integer arithmetic during inference. Specifically, Tailor et al. (2021) proposes a Quantization Aware Training method tailored for GNNs, which emulates inference-time quantization during the training phase. Zhao et al. (2020) proposes to jointly search the quantization bit-width and GNN architectures. Note that both of them do not actually convert the node embeddings into lower numerical precision during training. Instead, they use full precision data type to simulate the effect of real quantization. Feng et al. (2020) proposes a heterogeneous quantization framework which assigns different bits to node embeddings in each layer while maintaining the weights at full precision. However, due to the mismatch in operands’ bit-width, it is impractical to use in general purpose hardwares, such as CPUs and GPUs.
Graph Sparsification. The adjacency matrices are usually stored as sparse matrices. When hardwares perform SPMM , the sparse matrix format leads to random memory accesses and limited data reuse due to its irregular structure. Thus, GNNs have higher inference latency than other neural networks. Graph sparsification can alleviate the issue by removing redundant edges from original graphs. Zheng et al. (2020a) proposes a learning-based graph sparsification method which removes
potentially task-irrelevant edges from input graphs Li et al. (2020) formulates the graph sparsification problem as an optimization objective which can be solved by alternating direction method of multipliers (ADMM). Chen et al. (2021b) proposes to co-simplify the input graph and GNN model by extending the iterative magnitude pruning to graph areas.

heading: B.2 SCALABLE GNN TRAINING
text: In this subsection, we introduce previous works that try to train GNNs on large graphs. We conclude that EXACT are parallel to most of the previous works.
Subgraph Sampling/Mini-Batch Training. As we introduced in the main body, most of the previous scalable GNN training methods fall into this category. The key idea is to train GNNs with sampled subgraphs instead of the whole graph at each step. In this way, only node embeddings that are present in the current subgraph will be retained in memory. Based on this idea, various sampling techniques have been proposed, including the node-wise sampling (Hamilton et al., 2017; Chen et al., 2017; Markowitz et al., 2021), layer-wise sampling (Zou et al., 2019; Chen et al., 2018; Huang et al., 2018), and subgraph sampling (Chiang et al., 2019; Zeng et al., 2020). Generally, methods in this category are orthogonal to EXACT, and they can certainly be combined. Similar to EXACT, subgraph sampling also introduce the gradient noise during the training process. However, the gradient noise does not necessarily result in the model performance drop. In contrast, many previous works experimentally show that subgraph sampling based methods may even yield better performance than the whole graph training (Zeng et al., 2020; Hu et al., 2020). Our experiment results also exhibit a similar phenomenon that EXACT may achieve better performance than the original one. This observation can be explained by the regularization effect introduced by the gradient noise (Hu et al., 2020).
In this paper, we show in experiments that EXACT can be combined with subgraph-sampling based methods with only limited accuracy drop. From the practical usage perspective, EXACT can enlarge the maximal subgraph size for subgraph training-based method. As a result, the error introduced by EXACT can be compensated by using larger subgraphs, which previously cannot fit into GPUs. We leave it as one future direction.
GNNAutoScale (Fey et al., 2021). GNNAutoScale follows the idea of utilizing historical embeddings from prior training iterations (Fey et al., 2021). Specifically, it stores a historical embedding for each node as an offline storage in CPU memory. At each training step, it first samples a minibatch of nodes. Then for out-of-mini-batch nodes, GNNAutoScale swaps their historical embeddings from CPU memory to GPU memory and utilizes their non-trainable historical embeddings for propagation. In this way, only the node embeddings inside the current mini-batch and those of their direct 1-hop neighbors are retained in memory. The time overhead of GNNAutoScale is mainly from the swapping step. EXACT is orthogonal to GNNAutoScale. It would be interesting to store node embeddings compressed by EXACT as historical embeddings for GNNAutoScale. Considering that quantization is more time-efficient compared to swapping, the time overhead of GNNAutoScale can be greatly reduced since the time cost of swapping scales with the total number of bits to be swapped.
Distributed Training. Unlike the data in other domains, the graph data cannot be trivially divided into mini-batches due to its connectivity between samples. The graph distributed training methods require to split the graph into mini-batches that minimizes the communication overhead. Specifically, AliGraph (Zhu et al., 2019) is a distributed GNN framework on CPU platforms, which does not support GPUs acceleration. ROC (Jia et al., 2020) learns to optimize the graph partitioning by predicting the execution time of performing a GNN operation on an input subgraph. Both BDSGCN (Wan et al., 2021) and DistDGL (Zheng et al., 2020b) leverage METIS (Karypis & Kumar, 1998) to performance the graph partitioning to minimize the communication cost. However, one main drawback of applying METIS is that the ahead-of-training overhead of METIS is relatively large on huge graphs. EXACT is also orthogonal to the distributed training methods and can be used to decrease the hardware requirements. It would be interesting to investigate the trade-off among the accuracy drop and the memory saving of EXACT under multi-GPU settings.
Memory efficient training system. Gradient checkpointing (Chen et al., 2016; Jain et al., 2019; Kirisame et al., 2020; Shah et al., 2020) trades computation for memory by dropping some of the activations in the forward pass and recomputing them in the backward pass. Swapping (Meng et al., 2017) utilizes the huge amount of available CPU memory by swapping tensors between CPU and
Algorithm 1: Forward Pass of the lth GCN layer Input: H(l), Θ(l), the total number of layers L. Output: H(l+1)
1 ctx(l) ← {} /* the context which saves tensors for backward */ 2 J (l) ← MM(H(l),Θ(l)) 3 Add H(l) and Θ(l) to ctx(l) /* used for MM.backward */ 4 H(l+1) ← SPMM(Â,J (l)) 5 Add Â (in CSR format) and J (l) to ctx(l) /* used for SPMM.backward */ 6 if l 6= L− 1 then 7 Add 1{H(l+1)>0} to ctx
(l) /* used for ReLU.backward */ 8 H(l+1) = ReLU(H(l+1)) 9 end
10 return H(l+1)
GPU. EXACT is also parallel to both the gradient checkpointing and swapping. As a result, EXACT can be combined with them to further squeeze out the memory, at the cost of larger time overhead.
Other Activation Compression Techniques. JPEG-ACT (Evans et al., 2020) extends the widelyused JPEG method to compress the activations. Cusz is a prediction-based lossy compression method with GPU support (Tian et al., 2020; Tao et al., 2017). Both of them can be applied to compress the activations of GNNs. However, one practical problem is that they have hyperparameters which are hard to tune. For example, in JPEG-ACT, there is a hyperparameter α controls the space/accuracy trade-off. This problem can be alleviated by incorporating one recent proposed method AC-GC (Evans & Aamodt, 2021). Given a preset allowable increase in loss, AC-GC can automatically find the hyperparameter such that the compression rate can adapt to the preset training conditions. We note that EXACT and AC-GC approaches the problem from different perspective. For EXACT, we value the compression rate more than the accuracy drop. Namely, given a preset memory budget, we try to fit the model with acceptable accuracy drop. For AC-GC, they value the accuracy change more than the compression rate in the sense that the compression rate is adapted to the given allowable accuracy drop.

heading: C ANALYZING THE MEMORY CONSUMPTION OF GCNS
text: Generally, each training step contains a forward pass phase and a backward pass phase. During the forward pass, activation maps of each operation are kept in memory for computing the gradients. During the backward pass, the weight gradients and activation gradients are calculated and stored in memory. For weight gradients, their memory usage is negligible since they share the same shape with model weights. For activation gradients, their memory usage is dynamic and will eventually be cleared to zero. In standard deep learning libraries, activation gradients will be deleted as soon as their reference counting becomes zero (Paszke et al., 2019; Abadi et al., 2016). Thus, storing activation maps take up most of the memory. The memory usage of activation maps in GCN layers are given in Algorithm 1.

heading: D RANDOM PROJECTION
text: In this paper, we use the normalized Rademacher random matrix (Achlioptas, 2001) for projecting the activation maps. Specifically, for a D × R normalized Rademacher random matrix R, each element in R is i.i.d. sampled from the following distribution:
Ri,j =
√ 1
R ×  +1, with probability 1 2 ,
−1, with probability 1 2 .
(7)
where √
1 R is the normalize factor such that E[RR >] = I . From above, we can see that the sampling cost of the normalized Rademacher random matrix is relatively low since sampling from Bernoulli distribution is much cheaper than sampling from Gaussian distribution.

heading: E THEORY

heading: E.1 PROOF OF PROPOSITION 1
text: Proposition 1 (Details in Appendix E) The above quantization and dequantization are unbiased operations, i.e., E[ĥ(l)v ] = E[Dequant(Quant(h(l)v ))] = h(l)v and Var(ĥ(l)v ) = D[r (l) v ] 2
6B2 .
Proposition 1 is adopted from the theoretical analysis in ActNN (Chen et al., 2021a). For completeness, we prove it here with more details. The conclusion of E[ĥ(l)v ] = h(l)v follows from the fact that the stochastic rounding is an unbiased operation (Courbariaux et al., 2015). Specifically, ∀ scalar h, the stochastic rounding can be expressed as
bhe = {dhe, with probability h− bhc, bhc, with probability 1− (h− bhc), (8)
where d·e is the ceil operation and b·c is the floor operation. First, we have E[bhe] = h because E[bhe] = dhe(h− bhc) + bhc(1− h+ bhc)
= h (following the fact that dhe − bhc = 1)
Hence,
E[ĥ(l)v ] = r (l) v
B E[bh (l) v − Z(l)v r (l) v Be] + Z(l)v = h(l)v .
Regarding the variance, let h̄ = h (l) v −Z (l) v
r (l) v B = (h1, · · · , hD). Suppose ∀i, hi − bhic = σ ∼ Uniform(0, 1), we have.
Var(ĥ(l)v ) = [r
(l) v ]2 B2 Var(bh̄e) = E[h̄>h̄]− (E[h̄])>(E[h̄])
= [r
(l) v ]2
B2 D∑ i=1 dhie2(hi − bhic) + bhic2(1− hi + bhic)− h2i
= [r
(l) v ]2D B2 ( 2bh1ch1 + h1 − bh1c2 − bh1c − h21 ) (substitute dh1e with bh1c+ 1)
= [r
(l) v ]2D B2 ( σ − σ2 ) (substitute bh1c with h1 − σ)
By taking expectation w.r.t. σ on the both side, we have Var(ĥ(l)v ) = [r(l)v ] 2D 6B2

heading: E.2 PROOF OF PROPOSITION 2
text: Proposition 2 (Proof in Appendix E) The above RP and IRP are unbiased operations, i.e., E[Ĥ(l)] = E[IRP(RP(H(l)))] = H(l). For each ĥ(l)v in Ĥ(l), we have Var(ĥ(l)v ) = D−1R ||h (l) v ||22.
We can get E[Ĥ(l)] = E[IRP(RP(H(l)))] = H(l) directly following the fact that E[IRP(RP(H(l)))] = E[H(l)RR>] = H(l)E[RR>] = H(l).
Regarding the variance, let P = RR> ∈ RD×D. For the sake of notation convenience, we ignore the subscript and superscript in this subsection. Namely, we use the notation ĥ to represent ĥ(l)v , and
use the notation h to represent h(l)v . Also, we assume the shapes of ĥ and h are both 1×D in this subsection. We have
Var(ĥ) = Var(hP )
= hE[PP>]h> − hh>
Next, we will characterize the distribution of PP>. Recall that P = RR>, thus,
Pi,j = R∑ k=1 ai,kaj,k,
where each ai,j = ± √ 1 R with equal probability. Thus,
[PP>]i,j = D∑ m=1 Pi,mPj,m
= D∑ m=1 ( ( R∑ k=1 ai,kam,k)( R∑ k=1 aj,kam,k) ) . (9)
From Equation 9, it is easy to show that E[PP>] is a diagonal matrix. Namely, ∀i 6= j, we have E[PP>]i,j = 0 following the fact that each ai,j are independent and E[ai,j ] = 0. For elements on the diagonal, we have
E[PP>]i,i = E[ D∑ m=1 Pi,mPi,m]
= E[ D∑ m=1 ( R∑ k=1 ai,kam,k) 2]
= E[( R∑ k=1 ai,kai,k) 2] + E[ ∑ m 6=i ( R∑ k=1 ai,kam,k) 2] = 1 + ∑ m6=i ( R∑ k=1 E[(ai,kam,k)2] + 2 ∑ p,q:p<q E[ai,pam,pai,qam,q] ) = 1 + ∑ m6=i ( R∑ k=1 E[(ai,kam,k)2] )
= 1 + D − 1 R . (10)
Let h = (h1, · · ·hD), we have
Var(ĥ) = hE[PP>]h> − hh>
= h 1 + D−1 R
. . . 1 + D−1R h> − hh> =
D∑ i=1 (1 + D − 1 R )h2i − D∑ i=1 h2i
= D − 1 R ||h||22. (11)

heading: E.3 PROOF OF PROPOSITION 3
text: Proposition 3 (Proof in Appendix E) For each projected node embedding h(l)vproj = RP(h (l) v ) =
h (l) v R, for ∀ > 0, by choosing s = ||h(l)v ||2
√ 2ln(2R/ )
R , we have P (||h (l) vproj ||∞ ≤ s) ≥ 1− .
For the normalized Rademacher random matrix R ∈ RD×R, we have Ri,j = ± 1√R with equal probability. Let h(l)v = (h1, · · · , hD) and h(l)vproj = (u1, · · · , uR). Since h (l) vproj = h (l) v R, we have
u1 = ∑D i=1 aihi, where ai = ± 1√ R with equal probability.
First, we have the following inequality:
E[etRu1 ] = D∏ i=1 E[etRaihi ]
= D∏ i=1 et √ Rhi + e−t √ Rhi 2 = D∏ i=1 cosh t √ Rhi
≤ D∏ i=1 e t2Rh2i 2 (following the fact that coshx ≤ e x 2 2 )
= e t2R||h(l)v || 2 2 2 . (12)
For notation convenience, we use the notation “C” to represent the node embedding norm ||h(l)v ||2. For ∀s, q > 0, according to Chernoff bound, we have:
P (|ui| ≥ s) = 2P (ui ≥ s) = 2P (equi ≥ eqs)
≤ 2E[e qui ]
eqs (by Chernoff bound)
≤ 2E[e sRui C2 ]
e s2R C2
(by setting q = sR
C2 )
≤ 2e− s2R 2C2 (by Equation 12) (13)
Note that Equation 13 holds for ∀s > 0. Given a specific , by setting s = C √
2 ln(2R/ ) R , we have
P (|ui| ≥ s) ≤ R . Also,
P (||h(l)vproj ||∞ ≤ s) = R∏ i=1 P (|ui| ≤ s)
= (1− P (|ui| ≥ s))R
≥ (1− R )R ≥ 1−
Here we also experimentally verify Proposition 3 by visualizing the infinity norm of projected node embeddings in Figure 4. To avoid creating confusions, DR = 1 in Figure 4 means we apply the random projection with R = D on the activation maps. We can observe that in general, the infinity norm increases with the compression ratio DR . However, we note that the infinity norm of projected embeddings may less than those of original embeddings when R is larger than a threshold. This suggests that when R is below a certain threshold, quantizing projected embeddings only have limited influence on model performance. This claim is also verified in the main body of this paper.

heading: E.4 THE COMPOUND EFFECT OF APPROXIMATION ERRORS
text: In this subsection, we analyze why GNNs can tolerate the extremely compressed activations. From the numerical analysis (Bauer, 1974), the jacobian matrix (gradient) between two tensors can be viewed as a sum over all ”paths” connecting these two tensors. Moreover, the architecture of GNNs is considerably shallower than CNNs and transformers. Informally, for GNNs, the key intuition is that these paths are significantly ”shorter” compared to those of CNNs. As a result, the approximation error (i.e., the variance from quantization) along the path is hard to get accumulated.
We adopt Theorem 3 in ActNN (Chen et al., 2021a) to make the above statement more rigorous. Let Ĉ(m) be the compressed context (either by quantization, random projection, or quantized random projection). Let ∇Θ(l) and ∇H(l) be the gradient of Θ(l) and H(l), respectively. ∇̂Θ(l) and ∇̂H(l) are the calculated gradient using the compressed context, respectively. Further, we use the notation G(l∼m)Θ (∇̂H(m), Ĉ(m)) to represent the variance introduced by utilizing the compressed context Ĉ(m). Specifically, for a L layer GNN, we have
Var(∇̂Θ(l)) = Var(∇Θ(l)) + L∑ m=l E [ Var ( G (l∼m) Θ (∇̂H (m), Ĉ(m))|∇̂H(m) )] , (14)
where Var(·|∇̂H(m)) is the conditional variance, and Var(∇Θ(l)) is the variance is from the minibatch sampling variance. The key insights from Equation 14 are two folds. First, since the variances introduced by compressed contexts at different layers will accumulate, it suggests that the noise introduced by the compressed context is relatively small for shallow models. Considering that most of GNNs are usually less than four layers, this may explain why the loss in accuracy is negligible when using the vanilla INT2 quantization and a relatively large DR ratio. Second, under the subgraph training setting, the extra variance introduced by quantization can be compensated when using larger batch size (i.e., the size of subgraph in the graph learning).
In Table 1, we evaluate GNNs under the full-batch setting. Therefore, in addition to the reason of mentioned shallower architecture, it is also possible that the activation compressed training benefits from the large batch size, which leads to a much smoother gradient. To quantitatively study the effect of the batch size, below we present the ablation studies of two representative sampling based methods with much smaller batch size. Namely, GraphSAINT (Table 6) and Cluster-GCN (Table 7). For the baseline, we use the same hyperparameters reported in the corresponding paper. We make two main observations. First, INT2 quantization works under a much smaller batch size. Namely, for both Cluster-GCN and GraphSAINT, the accuracy drop is negligible even when the sampled subgraph contains only ≈ 500 nodes. This observation implies that the activations of GNNs can be aggressively compressed regardless of the batch size. Second, using a much smaller batch size in general will lead to an accuracy drop. However, for Yelp, a much smaller batch size may even improve the performance.

heading: F SYSTEM IMPLEMENTATION OF EXACT
text: F.1 INDIVIDUAL LAYERS CONFIGURATIONS OF EXACT
The operation configurations are shown in Table 8. For all graph convolution operations, EXACT can apply both the quantization and random projection to their saved activation maps. In EXACT, the SPMM and its variants (i.e., SPMM MAX , SPMM MEAN , and SPMM MIN ) are implemented based
on those provided in Pytorch Sparse3, with extra supporting for compressing activation maps by quantization and random projection.
For BatchNorm layers, we found that quantizating its saved activation maps only impact model performance a little, which is experimentally verified in the experiments. However, we experimentally found that randomly projecting its saved activation maps will lead to divergence. We note that this observation is consistent with previous finding that BatchNorm is very sensitive to noise (Micikevicius et al., 2017). Hence, EXACT only apply the quantization to BatchNorm layers.
Regarding ReLU operations, we have y = ReLU(x) = x1x>0 and ∇y = ∇y1x>0. Hence, ReLU operations only need to store 1x>0 in the context for the backward pass, which takes a single bit per element to store, without introducing any errors. Since standard deep learning framework only support down to INT8 precision, here we convert the mask matrix into bit-stream and fuse this process into the CUDA kernel of ReLU.forward to minimize the time overhead.
Regarding Dropout operations, let p be the dropout probability. During the training process, we have
y = Dropout(x, p) = 1 1− p xg, ∇x = 1 1− p ∇yg, (15)
where g ∼ Bernoulli(1 − p) is a binary vector sharing the same shape as x. 11−p is the normalization factor such that E[y] = x. Similarly, Dropout Operations also only need to store g in the context for the backward pass, which takes a single bit per element to store, without introducing any errors. Similarly, we convert g into bit-stream and fuse this process into the CUDA kernel of Dropout.forward to minimize the time overhead..
To be clear, “fixed 1 bit” in Table 8 means the activation maps of ReLU and Dropout operations take only 1 bit per element to store. And the bit-width of all other operations can be adjusted.
F.2 IMPLEMENTATION DETAILS
We provide simple API for converting modules in Pytorch Geometric and Pytorch to its corresponding version in EXACT. For example, replacing torch geometric.nn.GCNConv with exact.GCNConv and replacing torch.dropout with exact.dropout. Currently, EXACT only support Pytorch Geometric and Pytorch. In future we will try to integrate EXACT with other popular graph learning packages, such as DGL.
Following ActNN (Chen et al., 2021a), to obtain the highest compression ratio, the quantization range r(l)v and the zero point Z (l) v are stored in the bfloat16 data type 4. The quantization and dequantization modules are both implemented using CUDA kernels. We note that Pytorch only support data types down to INT8. To obtain highest compression ratio, the quantized data is compressed into bit streams such that it can be decoded later during the dequantization process.
All CUDA kernels in EXACT support both full-precision and half-precision (i.e., bfloat16 and float16). Thus, EXACT can also be integrated with the automated mixed precision training, i.e., AMP5, to further decrease the memory consumption.

heading: G EXPERIMENT SETTINGS

heading: G.1 DATASETS, FRAMEWORKS, AND HARDWARES
text: We give the detailed statistics and the download URLs for all datasets used in our experiments in Table 10. We follow the standard data splits and all datasets are directly downloaded from Pytorch Geometric or the protocol of OGB (Hu et al., 2020). We implement all models based on Pytorch and Pytorch Geometric. Almost all experiments are done on a single NVIDIA GeForce RTX 3090 with 24GB GPU memory. During our experiments, we found that the version of Pytorch, Pytorch Sparse, and Pytorch Scatter can significantly impact the running speed of the baseline. Here we list the details of our used packages in all experiments in Table 9.
3https://github.com/rusty1s/pytorch sparse/blob/master/torch sparse/matmul.py 4https://pytorch.org/docs/stable/generated/torch.Tensor.bfloat16.html 5https://pytorch.org/docs/stable/amp.html

heading: G.2 MODEL HYPERPARAMETER CONFIGURATIONS OF TABLE 1 AND TABLE 2
text: We adopt three popular GNNs. Namely, GCN (Kipf & Welling, 2017), GraphSAGE (Hamilton et al., 2017), and GAT (Veličković et al., 2017) in Table 1 and Table 2. We follow the hyperparameter configurations and codebases provided on the OGB (Hu et al., 2020) leader-board. Specifically, the hyperparameter configuration is given in Table 11. The optimizer is Adam (Kingma & Ba, 2014) for GCN and GraphSAGE, while the optimizer is RMSprop for GAT. All methods terminate after a fixed number of epochs. We report the test accuracy associated with the highest validation score. The “Gradient Clipping” in Table 11 indicate the maximum norm for gradients. “Gradient Clipping= 0.0” means we do not clip the gradients in that experiment.

heading: G.3 HYPERPARAMETER CONFIGURATIONS OF EXACT (RP+INT2)
text: The EXACT (RP+INT2) has only one hyperparameter, namely, DR . The D R configuration of EXACT (RP+INT2) can be found in Table 12.

heading: G.4 MODEL HYPERPARAMETER CONFIGURATIONS IN SECTION 5
text: The optimizer used in all experiments in Section 5 is Adam (Kingma & Ba, 2014). We use the default hyperparameters for Adam optimizer, except for the learning rate. All methods terminate after a fixed number of epochs. We report the test accuracy/F1-micro associated with the highest validation score. Regarding Reddit, Flickr, and Yelp dataset, we follow the hyperparameter configurations reported in the respective papers as closely as possible. We clips the gradient during training. The “Gradient Clipping” in below tables indicate the maximum norm for gradients. “Gradient Clipping= 0.0” means we do not clip the gradients in that experiment.
6https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch geometric.datasets.Reddit 7https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch geometric.datasets.Flickr 8https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch geometric.datasets.Yelp 9https://ogb.stanford.edu/docs/nodeprop/#ogbn-arxiv
10https://ogb.stanford.edu/docs/nodeprop/#ogbn-products
Regarding ogbn-arxiv and ogbn-products dataset, we follow the hyperparameter configurations and codebases provided on the OGB (Hu et al., 2020) leader-board. Please refer to the OGB website for more details. Table 13 and Table 14 summarize the hyperparameter configuration of Cluster-GCN and GraphSAINT, respectively. Table 15, Table 16, and Table 17 summarize the hyperparameter configuration of full-Batch GCN, full-Batch GraphSAGE, and full-batch GCNII, respectively.

heading: H ANALYZING THE MEMORY USAGE

heading: H.1 ANALYZING THE MEMORY USAGE
text: We use torch.cuda.memory allocated for the memory measurement. As we mentioned in Table 3, “Act Mem.” is the memory occupied by activation maps. Besides activation maps, the model, optimizer, input data, weight gradients, and activation gradients also occupy GPU memory. We will analyze each of them below. First, the memory occupied by the model, optimizer, and weight gradients is negligible because the number of parameters in most of GNNs is very small. Second, the memory occupied by the input data depends on the graph size and often cannot be compressed. This part can take up a lot of memory when the graph is large. Third, the memory occupied by activation gradients is dynamic and hard to estimate, since these tensors are temporarily stored in GPUs. For standard deep learning library, they will be deleted as soon as their reference counting becomes zero (Abadi et al., 2016; Paszke et al., 2019). EXACT cannot compress activation gradients. However, the memory occupied by activation gradients can be compressed using AMP, because activation gradients are stored in float16 data type under AMP. As illustrated in Appendix F.2, we note that our EXACT framework can be integrated with AMP.
The memory usage of the model weights plus optimizer is about 2MB. Since activation gradients are dynamic and hard to analyze, here we report the peak memory usage during the backward pass, which encompass the activation gradients and other intermediate variables. For the memory usage of the input data, activation maps, and the peak memory usage during the backward pass, we provided a detailed analysis in Table 18. We can observe that the memory usage is mainly occupied by activation maps. For GCNII, the memory usage is dominated by the activation maps. This is because we need to store all layers’ activation maps during the forward pass. In contrast, during the backward pass, there is usually only one layer’ activation gradients are kept in memory.
Table 18: The detailed analysis about the memory usage of input data, activation maps, and peak memory usage during the backward pass. “Data Mem” is the memory usage of input data (including the input feature matrix X , adjacency matrices A, and labels ).“Act Mem” is the memory usage of activation maps. “Peak BWD Mem” is the peak memory usage during the backward pass. “Ratio
(%)” here equals Act MemData Mem+Act Mem+Peak BWD Mem .
Reddit Flickr Yelp ogbn-arxiv ogbnproducts
Data Mem Act Mem. Peak BWD Mem. Ratio (%) DataMem Act Mem Peak BWD Mem. Ratio (%) DataMem Act Mem Peak BWD Mem. Ratio (%) DataMem Act Mem Peak BWD Mem. Ratio (%) DataMem Act Mem Peak BWD Mem. Ratio (%)
Cluster-GCN 8.8 15 3 60.0 5.5 16.5 5.3 60.4 5.3 29.3 12 62.9 - - - - 4.9 35.2 11.5 68.2 GraphSAINT 27.6 44 9.3 54.4 31.2 88.7 29 59.6 5.8 33.5 13 64.0 26.7 270 62.4 75.1 57 516 157 70.7
GCN 1168 1029 316 40.9 208 379 86 56.3 1544 6429 1195 70.1 175.7 729.4 82.2 73.9 - - - - GraphSAGE 1168 1527 696 45.0 208 547 184 58.3 1544 6976 2881 61.2 175.7 786.2 192 68.1 4811 16555 6176 60.1
GCNII 1168 5850 239 80.6 208 4067 88 93.2 1544 33540 1197 92.4 175.7 14409 194 97.5 - - - -

heading: H.2 OVERALL COMPRESSION RATIO
text: The overall memory compression ratio is shown in Table 19. We observe that (1) for shallow GNN models, the overall memory compression ratio ranges from 1.5× to 4×. (2) for GCNII, the overall memory compression ratio ranges from 4× to 18×. We note that we store the graph structure data (e.g., node ID and edge ID) in torch.Long data type, which can safely cast to torch.Int data type to save the memory. Our implementation supports for using torch.Int as the data type for the graph structure data. However, for a fair comparison, we do not utilize this feature.

heading: H.3 ANALYZING THE MEMORY COMPRESSION RATIO
text: Below we analyze the compress ratio in Table 3. We take a three-layer, 128-dimensional GCNs trained on ogbn-arxiv for example. The computational graph of the baseline in Table 3 is:
Total bits =0(MM) + 32(SPMM) + 32(BN) + 0(ReLU) + 8(Dropout) + (the first layer) 32(MM) + 32(SPMM) + 32(BN) + 0(ReLU) + 8(Dropout) + (the second layer) 32(MM) + 32(SPMM) (the third layer)
Hence the baseline costs totally 240 bit per element. The first MM in the first layer does not cost extra bits because its activation map is exactly the input feature matrix X , which has been stored in GPU memory (recall that we need to first move the input data to GPU memory before training). Pytorch will save this part of memory via the “pass by reference” mechanism. The official ReLU operation in Pytorch does not need the extra space for saving activation maps (they can reuse the activation maps saved by the previous layer via passing by reference). Regarding Dropout operation, Pytorch stores the mask matrix using UINT8 data type, which costs 8 bit per element.
The computational graph of “EXACT (INT2)” in Table 3 is:
Total bits =2.25(MM) + 2.25(SPMM) + 2.25(BN) + 1(ReLU) + 1(Dropout) + (the first layer) 2.25(MM) + 2.25(SPMM) + 2.25(BN) + 1(ReLU) + 1(Dropout) + (the second layer) 2.25(MM) + 2.25(SPMM) (the third layer)
The 2.25 bit of MM , SPMM , and BN is from 2 (quantized activation maps) + 0.125 (the zero point tensor) + 0.125 (the range tensor), where 0.125 = 16 (the bit-width of bfloat16)128 (the dimension is 128) (See Appendix F.2 for details). Also for EXACT, we cannot leverage the “pass by reference” mechanism to save memory for ReLU since the exact activation maps of previous operation is replaced by the compressed one. Hence in EXACT, the 1-bit masks of ReLU and Dropout is converted into bit-streams using CUDA kernels, which cost 1 bit per element to store.
Hence the “EXACT (INT2)” costs totally 22 bit per element. And the theoretical compression ratio is 24022 = 10.9. The empirical compression ratio in Table 3 may have a small gap with the theoretical one. This is because the attributes of the sparse tensor (adjacency matrix) in pytorch sparse is lazily initialized, i.e., it may be generated during the forward pass and be account for the memory of activation maps.
The computational graph of “EXACT (RP+INT2)” with DR = 8 (see Table 12) in Table 3 is:
Total bits =0.28(MM) + 0.28(SPMM) + 2.25(BN) + 1(ReLU) + 1(Dropout) + (the first layer) 0.28(MM) + 0.28(SPMM) + 2.25(BN) + 1(ReLU) + 1(Dropout) + (the second layer) 0.28(MM) + 0.28(SPMM) (the third layer)
The 0.28 bit of MM and SPMM is from 28 + 16 8×128 + 16 8×128 = 0.28125. We note that EXACT does not apply random projection for BatchNorm layers (see Appendix F.2). Hence the “EXACT
(RP+INT2)” costs totally 10.18 bit per element. And the theoretical compression ratio is 24010.18 = 23.57. Again, there exists gap between the empirical compression ratio and the theoretical one because the lazy initialization mechanism, the existence of temporary tensors, and the neglect of the memory usage of these random projection matrices.
We emphasize that the compression ratio also depends on the number of input features of the dataset. (because the mentioned “input feature matrix X is passed by reference” mechanism). For the above example, the number of input features and the hidden dimension are both 128 and hence it is easy to be analyzed. Hence, the theoretical compression ratio may vary for different datasets.

heading: I ADDITIONAL EXPERIMENT RESULTS

heading: I.1 MORE RESULTS ON ACCURACY AGAINST THE PRECISION AND DR
text: To support the claim that “the performance of EXACT is mainly determined by the DR ratio of random projection”, we present more results on the test accuracy against the precision and DR ratio here. In Figure 5, we show the results of two models trained with EXACT using full-batch data on the ogbn-arxiv dataset. In Figure 6, we show the results of two models trained with EXACT using mini-batch data on Yelp dataset.

heading: I.2 COMPARISON AGAINST SWAPPING AND GRADIENT CHECKPOINTING
text: We also compare EXACT against the naive swapping and the naive gradient checkpointing. Both swapping and gradient checkpointing are lossless compression, so they do not have any accuracy drop. Below we present and discuss about their trade-off among the space and speed.
For swapping, we offload all activation maps to CPU memory, which offload all activation maps to CPU memory. Thus, it can achieve the highest compression ratio for activation maps. However, its running time overhead is roughly 50%-80%, which is often unacceptable.
For gradient checkpointing, we utilize torch.utils.checkpoint to insert checkpoints at each GNN layer. For the time overhead, as shown in Figure 3, it is comparable to EXACT. We present the memory usage of activations in Table 20. In summary, its time overhead is comparable to EXACT, however, the memory compression ratio of activations is 1.7 ∼ 2.3×, which is not large enough. Thus, we still cannot train the GCNII on Yelp dataset using a RTX 3090 (24GB) GPU.

heading: I.3 THE TRAINING CURVES OF EXACT
text: From the optimization theory, common convergence speed bounds (measured by the number of iterations) in stochastic optimization improve with smaller gradient variance (Bottou et al., 2018). EXACT essentially trades the gradient variance in return for reduced memory. Here we experimentally examine how EXACT affects the convergence speed. Figure 7 shows the training curves of GNNs trained with EXACT using different configurations on Reddit dataset. We make two main observations. First, EXACT with smaller DR typically converges faster per iteration. This is consistent with the mentioned optimization theory. Second, when increasing the DR ratio, the difference in the convergence speed is very small, where the convergence speed is measured by the gap in validation accuracy between consecutive epochs. In practical scenarios, EXACT hence only have limited impact on the model performance.

heading: I.4 HYPERPARAMETER SENSITIVITY EXPERIMENT
text: As we analyzed in the main body, INT2 is a suitable precision and we only vary the DR ratio for EXACT. In this section, we investigate the sensitivity of EXACT (RP+INT2) to the DR ratio. Table 12 shows the configuration of EXACT (RP+INT2) in the experiment in the main body. Here we present a comprehensive hyperparameter sensitivity study for EXACT. Specifically, the sensitivity studies of EXACT with GraphSAINT, ClusterGCN, GCN, GraphSAGE, and GCNII are shown in Figure 8, Figure 9, Figure 10, Figure 11, and Figure 12, respectively. Here we summarize some key observations. First, the model performance drop generally increases with the DR ratio. Second, when D R = 8, the loss in accuracy of EXACT (RP+INT2) is below 0.5% on two third of experiments. To be concrete, in Table 3, we totally adopt 22 combinations of different datasets and models. As shown in Table 3 and Table 8, for 15 of the 22 experiments, the loss of accuracy is below or near 0.5% when DR = 8. Third, when D R = 8, the loss in accuracy is below 1% in almost all experiments, excepting for two experiments. Namely, ClusterGCN with ogbn-products and GraphSAINT with ogbn-products.

heading: I.5 COMPARISON BETWEEN EXACT AND SAMPLING METHODS

heading: I.5.1 COMPARE EXACT TO SAMPLING METHODS UNDER A FIXED MEMORY BUDGET
text: As examined in Table 3, EXACT and subgraph sampling methods are orthogonal and can be applied over each others. Here we present an ablation study of comparing EXACT to subgraph sampling methods in a standalone way. We note that EXACT only focus on saving the memory for activations. In contrast, subgraph sampling methods can simultaneously reduce the memory of the input data, activations, and activation gradients since they directly reduce the number of node embeddings
retained in the memory. Hence compared to subgraph sampling methods, one limitation of EXACT is that the overall memory saving ratio of EXACT often cannot surpass that of sampling methods.
Under the full-batch setting, one way to further improve the overall memory saving ratio is to incorporate the graph compression methods with EXACT, such as the graph sparsification (Spielman & Srivastava, 2011) and graph coarsening (Cai et al., 2021a). However, it is beyond the scope of this paper. For a fair comparison, we control the batch size of subgraph subgraph sampling methods such that their activation memory usage equals to that of EXACT (INT2). For both EXACT(INT2) and EXACT(RP+INT2), we quote the “GraphSAGE” results from Table 3. For GraphSAINT and Cluster-GCN, we tune their batch size such that they have the same activation memory usage as EXACT (INT2). The results are shown in Table 21. We make two main observations. First, for Reddit and Flickr, full batch training with EXACT outperforms the two sampling methods. For Yelp, subgraph sampling methods are better than full batch training with EXACT. However, we note
that this gap is not from EXACT. As shown in Table 3, without EXACT, the F1-micro of full batch GraphSAGE is 62.05%, which is much lower than that of subgraph sampling methods. Also, from Table 6, one interesting observation is that for Yelp, a much smaller batch size (≈ 500) may even further improve the F1-micro of GraphSAINT from 63.20% to 64.05%.
In summary, subgraph sampling methods may outperform the full batch training on some datasets (e.g., Yelp). In this case, we think full batch training with EXACT cannot outperform subgraph sampling methods. If this is not the case, as shown in Table 21, full batch training with EXACT can outperform subgraph sampling methods under a fixed activation memory usage.

heading: I.5.2 ENLARGING THE BATCH SIZE OF SAMPLING METHODS WITH EXACT
text: Here we present a case study of scaling up the batch size of GraphSAINT with EXACT on the ogbn-products dataset. As shown in Table 14, the batch size of GraphSAINT (number of nodes in the sampled subgraphs) is controlled by the “Walk length” and “Roots”, and roughly equals “Walk length”×“Roots” (Zeng et al., 2020). For ogbn-products, we tripled the batch size by changing “Roots” from 20, 000 to 30, 000. All other hyperparameters are left unchanged. The results are shown in Table 22. We observe that EXACT may improve the accuracy when using larger batch size. We note that EXACT can scale up the batch size to more than 3× larger. However, we found that if we further scale up the batch size to 4×, there is an accuracy drop compared to the baseline
with the original batch size, regardless applying EXACT or not. This is consistent with previous finding that a larger batch size does not always lead to better performance for ogbn-products (Zeng et al., 2020; Hu et al., 2020). We quote the sentences from the OGB paper (Hu et al., 2020) to explain this counter-intuitive observation. “The recent mini-batch-based GNNs give promising results, even slightly outperforming the full-batch version of GraphSAGE that does not fit into ordinary GPU memory. The improved performance can be attributed to the regularization effects of mini-batch noise and edge dropout.” (Hu et al., 2020).
In summary, in the ogbn-products ablation study, we utilize EXACT to triple the batch size of GraphSAINT, which may even improve the accuracy over the original one. Moreover, from our experiments, the 3× larger batch size performs the best on ogbn-products. Thus, for some datasets, there exist an optimal batch size for subgraph sampling methods. Since this optimal batch size may be beyond the capacity of the hardware, the meaning of EXACT is to enlarge the “search space” for finding this optimal batch size.

Review summaries:
This work attempts to train GNN models with reduced memory requirements. Normally when training, we need to store activations for calculating the gradients on the backwards pass, and we typically keep these at full (32-bit) precision. This work argues that we don't need to do this, and instead we can significantly reduce the memory footprint at training time by keeping low precision activations for the backwards pass. They use two approaches for this: quantization, and random projections

This paper explores the training of GNN with compressed activation maps. It provides an optimized GPU implementation and comprehensively studies the trade-off among the memory saving, time overhead, and accuracy drop.  Experimental results show that the proposed framework can reduce the memory footprint of activations by up to 32x with only 0.2-0.5% accuracy drop and 10-25% time overhead.

This paper proposes EXACT, a framework for training GNNs with compressed activations with two methods: quantization and random projection. The primary design objective of EXACT is to trim down the memory consumption in GNN training while maintaining acceptable training speed and accuracy. EXACT achieves significant memory savings with 0.2-0.5% accuracy drop and 10-25% slowdown of training throughput across five models and five graph datasets.