Degree-Quant: Quantization-Aware Training for Graph Neural Networks

1 INTRODUCTION . GNNs have received substantial attention in recent years due to their ability to model irregularly structured data . As a result , they are extensively used for applications as diverse as molecular interactions ( Duvenaud et al. , 2015 ; Wu et al. , 2017 ) , social networks ( Hamilton et al. , 2017 ) , recommendation systems ( van den Berg et al. , 2017 ) or program understanding ( Allamanis et al. , 2018 ) . Recent advancements have centered around building more sophisticated models including new types of layers ( Kipf & Welling , 2017 ; Velickovic et al. , 2018 ; Xu et al. , 2019 ) and better aggregation functions ( Corso et al. , 2020 ) . However , despite GNNs having few model parameters , the compute required for each application remains tightly coupled to the input graph size . A 2-layer Graph Convolutional Network ( GCN ) model with 32 hidden units would result in a model size of just 81KB but requires 19 GigaOPs to process the entire Reddit graph . We illustrate this growth in fig . 1 . One major challenge with graph architectures is therefore performing inference efficiently , which limits the applications they can be deployed for . For example , GNNs have been combined with CNNs for SLAM feature matching ( Sarlin et al. , 2019 ) , however it is not trivial to deploy this technique on smartphones , or even smaller devices , whose neural network accelerators often do not implement floating point arithmetic , and instead favour more efficient integer arithmetic . Integer quantization is one way to lower the compute , memory and energy budget required to perform inference , without requiring modifications to the model architecture ; this is also useful for model serving in data centers . Although quantization has been well studied for CNNs and language models ( Jacob et al. , 2017 ; Wang et al. , 2018 ; Zafrir et al. , 2019 ; Prato et al. , 2019 ) , there remains relatively little work addressing ∗Equal contribution . Correspondence to : Shyam Tailor < sat62 @ cam.ac.uk > GNN efficiency ( Mukkara et al. , 2018 ; Jia et al. , 2020 ; Zeng & Prasanna , 2020 ; Yan et al. , 2020 ) . To the best of our knowledge , there is no work explicitly characterising the issues that arise when quantizing GNNs or showing latency benefits of using low-precision arithmetic . The recent work of Wang et al . ( 2020 ) explores only binarized embeddings of a single graph type ( citation networks ) . In Feng et al . ( 2020 ) a heterogeneous quantization framework assigns different bits to embedding and attention coefficients in each layer while maintaining the weights at full precision ( FP32 ) . Due to the mismatch in operands ’ bit-width the majority of the operations are performed at FP32 after data casting , making it impractical to use in general purpose hardware such as CPUs or GPUs . In addition they do not demonstrate how to train networks which generalize to unseen input graphs . Our framework relies upon uniform quantization applied to all elements in the network and uses bit-widths ( 8-bit and 4-bit ) that are supported by off-the-shelf hardware such as CPUs and GPU for which efficient low-level operators for common operations found in GNNs exists . This work considers the motivations and problems associated with quantization of graph architectures , and provides the following contributions : • The explanation of the sources of degradation in GNNs when using lower precision arithmetic . We show how the choice of straight-through estimator ( STE ) implementation , node degree , and method for tracking quantization statistics significantly impacts performance . • An architecture-agnostic method for quantization-aware training on graphs , Degree-Quant ( DQ ) , which results in INT8 models often performing as well as their FP32 counterparts . At INT4 , models trained with DQ typically outperform quantized baselines by over 20 % . We show , unlike previous work , that models trained with DQ generalize to unseen graphs . We provide code at this URL : https : //github.com/camlsys/degree-quant . • We show that quantized networks achieve up to 4.7× speedups on CPU with INT8 arithmetic , relative to full precision floating point , with 4-8× reductions in runtime memory usage . 2 BACKGROUND . 2.1 MESSAGE PASSING NEURAL NETWORKS ( MPNNS ) . Many popular GNN architectures may be viewed as generalizations of CNN architectures to an irregular domain : at a high level , graph architectures attempt to build representations based on a node ’ s neighborhood ( see fig . 2 ) . Unlike CNNs , however , this neighborhood does not have a fixed ordering or size . This work considers GNN architectures conforming to the MPNN paradigm ( Gilmer et al. , 2017 ) . A graph G = ( V , E ) has node features X ∈ RN×F , an incidence matrix I ∈ N2×E , and optionallyDdimensional edge features E ∈ RE×D . The forward pass through an MPNN layer consists of message passing , aggregation and update phases : h ( i ) l+1 = γ ( h ( i ) l , ∧ j∈N ( i ) [ φ ( h ( j ) l , h ( i ) l , eij ) ] ) . Messages from node u to node v are calculated using function φ , and are aggregated using a permutation-invariant function ∧ . The features at v are subsequently updated using γ . We focus on three architectures with corresponding update rules : 1 . Graph Convolution Network ( GCN ) : h ( i ) l+1 = ∑ j∈N ( i ) ∪ { i } ( 1√ didj Wh ( j ) l ) ( Kipf & Welling , 2017 ) , where di refers to the degree of node i . 2 . Graph Attention Network ( GAT ) : h ( i ) l+1 = αi , iWh ( i ) l + ∑ j∈N ( i ) ( αi , jWh ( j ) l ) , where α represent attention coefficients ( Velickovic et al. , 2018 ) . 3 . Graph Isomorphism Network ( GIN ) : h ( i ) l+1 = fΘ [ ( 1 + ) h ( i ) l + ∑ j∈N ( i ) h ( j ) l ] , where f is a learnable function ( e.g . a MLP ) and is a learnable constant ( Xu et al. , 2019 ) . 2.2 QUANTIZATION FOR NON-GRAPH NEURAL NETWORKS . Quantization allows for model size reduction and inference speedup without changing the model architecture . While there exists extensive studies of the impact of quantization at different bitwidths ( Courbariaux et al. , 2015 ; Han et al. , 2015 ; Louizos et al. , 2017 ) and data formats ( Micikevicius et al. , 2017 ; Carmichael et al. , 2018 ; Kalamkar et al. , 2019 ) , it is 8-bit integer ( INT8 ) quantization that has attracted the most attention . This is due to INT8 models reaching comparable accuracy levels to FP32 models ( Krishnamoorthi , 2018 ; Jacob et al. , 2017 ) , offer a 4× model compression , and result in inference speedups on off-the-shelf hardware as 8-bit arithmetic is widely supported . Quantization-aware training ( QAT ) has become the de facto approach towards designing robust quantized models with low error ( Wang et al. , 2018 ; Zafrir et al. , 2019 ; Wang et al. , 2018 ) . In their simplest forms , QAT schemes involve exposing the numerical errors introduced by quantization by simulating it on the forward pass Jacob et al . ( 2017 ) and make use of STE ( Bengio et al. , 2013 ) to compute the gradients—as if no quantization had been applied . For integer QAT , the quantization of a tensor x during the forward pass is often implemented as : xq = min ( qmax , max ( qmin , bx/s+ zc ) ) , where qmin and qmax are the minimum and maximum representable values at a given bit-width and signedness , s is the scaling factor making x span the [ qmin , qmax ] range and , z is the zero-point , which allows for the real value 0 to be representable in xq . Both s and z are scalars obtained at training time . hen , the tensor is dequantized as : x̂ = ( xq − z ) s , where the resulting tensor x̂ ∼ x for a high enough bit-width . This similarity degrades at lower bit-widths . Other variants of integer QAT are presented in Jacob et al . ( 2017 ) and Krishnamoorthi ( 2018 ) . To reach performance comparable to FP32 models , QAT schemes often rely on other techniques such as gradient clipping , to mask gradient updates based on the largest representable value at a given bit-width ; stochastic , or noisy , QAT , which stochastically applies QAT to a portion of the weights at each training step ( Fan et al. , 2020 ; Dong et al. , 2017 ) ; or the re-ordering of layers ( Sheng et al. , 2018 ; Alizadeh et al. , 2019 ) . 3 QUANTIZATION FOR GNNS . In this section , we build an intuition for why GNNs would fail with low precision arithmetic by identifying the sources of error that will disproportionately affect the accuracy of a low precision model . Using this insight , we propose our technique for QAT with GNNs , Degree-Quant . Our analysis focuses on three models : GCN , GAT and GIN . This choice was made as we believe that these are among the most popular graph architectures , with strong performance on a variety of tasks ( Dwivedi et al. , 2020 ) , while also being representative of different trends in the literature . 3.1 SOURCES OF ERROR . QAT relies upon the STE to make an estimate of the gradient despite the non-differentiable rounding operation in the forward pass . If this approximation is inaccurate , however , then poor performance will be obtained . In GNN layers , we identify the aggregation phase , where nodes combine messages from a varying number of neighbors in a permutation-invariant fashion , as a source of substantial numerical error , especially at nodes with high in-degree . Outputs from aggregation have magnitudes that vary significantly depending on a node ’ s in-degree : as it increases , the variance of aggregation values will increase.1 Over the course of training qmin and qmax , the quantization range statistics , become severely distorted by infrequent outliers , reducing the resolution for the vast majority of values observed . This reults in increased rounding error for nodes with smaller in-degrees . Controlling qmin and qmax hence becomes a trade-off balancing truncation error and rounding error . We can derive how the mean and variance of the aggregation output values vary as node in-degree , n , increases for each of the three GNN layers . Suppose we model incoming message values for a single output dimension with random variables Xi , without making assumptions on their exact distribution or independence . Further , we use Yn as the random variable representing the value of node output after the aggregation step . With GIN layers , we have Yn = ( 1 + ) X0 + ∑n i=1Xi . It is trivial to prove that E ( Yn ) = O ( n ) . The variance of the aggregation output is also O ( n ) in the case that that ∑ i 6=j Cov ( Xi , Xj ) ∑ i Var ( Xi ) . We note that if ∑ i 6=j Cov ( Xi , Xj ) is large then it implies that the network has learned highly redundant features , and may be a sign of over-fitting . Similar arguments can be made for GCN and GAT layers ; we would expect GCN aggregation values to grow like O ( √ n ) , and GAT aggregation values to remain constant ( O ( 1 ) ) due to the attention coefficients . We empirically validate these predictions on GNNs trained on Cora ; results are plotted in fig . 3 . We see that the aggregation values do follow the trends predicted , and that for the values of in-degree in the plot ( up to 168 ) the covariance terms can be neglected . As expected , the variance and mean of the aggregated output grow fastest for GIN , and are roughly constant for GAT as in-degree increases . From this empirical evidence , it would be expected that GIN layers are most affected by quantization . By using GIN and GCN as examples , we can see how aggregation error causes error in weight updates . Suppose we consider a GIN layer incorporating one weight matrix in the update function i.e . h ( i ) l+1 = f ( Wy ( i ) GIN ) , where f is an activation function , y ( i ) GIN = ( 1 + ) h ( i ) l + ∑ j∈N ( i ) h ( j ) l , and N ( i ) denotes the in-neighbors of node i . Writing y ( i ) GCN = ∑ k∈N ( i ) ( 1√ didk Wh ( j ) l ) , we see that the derivatives of the loss with respect to the weights for GCN and GIN are : GIN ∂L ∂W = |V |∑ i=1 ( ∂L ∂h ( i ) l+1 ◦ f ′ ( Wy ( i ) GIN ) ) y ( i ) > GIN GCN ∂L ∂W = |V |∑ i=1 ∑ j∈N ( i ) 1√ didj ( ∂L ∂h ( i ) l+1 ◦ f ′ ( y ( i ) GCN ) ) h ( j ) > l The larger the error in y ( i ) GIN—caused by aggregation error—the greater the error in the weight gradients for GIN , which results in poorly performing models being obtained . The same argument applies to GCN , with the h ( j ) > l and y ( i ) GCN terms introducing aggregation error into the weight updates .

Summary: The authors propose a new technique for quantization aware training of neural networks that is specially suited for graph neural networks. They do a good job of motivating the problem by demonstrating that the large variation of input degree in GNNs can lead to unique challenges for numerical precision, forcing a compromise between truncation error and rounding error. Th proposed technique incorporates stochastic masking and quantization proportional to the input degree to allow higher input-degree nodes to operate at higher resolution on average.