Embedding Compression with Hashing for Efficient Representation Learning in Graph

Graph neural networks ( GNNs ) are deep learning models designed specifically for graph data , and they typically rely on node features as the input node representation to the first layer . When applying such type of networks on graph without node feature , one can extract simple graph-based node features ( e.g. , number of degrees ) or learn the input node representation ( i.e. , embeddings ) when training the network . While the latter approach , which trains node embeddings , more likely leads to better performance , the number of parameters associated with the embeddings grows linearly with the number of nodes . It is therefore impractical to train the input node embeddings together with GNNs within graphics processing unit ( GPU ) memory in an end-to-end fashion when dealing with industrial scale graph data . Inspired by the embedding compression methods developed for natural language processing ( NLP ) models , we develop a node embedding compression method where each node is compactly represented with a bit vector instead of a float-point vector . The parameters utilized in the compression method can be trained together with GNNs . We show that the proposed node embedding compression method achieves superior performance compared to the alternatives . 1 INTRODUCTION . Graph neural networks ( GNNs ) are representation learning methods for graph data . When a GNN model is applied on node classification problems , the model typically learns the node representation from input node features X and its graph G where the node features X are used as the input node representation to the first layer of the model and the graph G dictates the propagation of information ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Zhou et al. , 2020 ) . However , the input node features X may not always be available for certain datasets . In order to apply such type of model on graph without node features X , we could either 1 ) extract simple graph based node features ( e.g. , number of degrees ) from the graph G or 2 ) use embedding learning methods to learn the node embeddings as features X ( Duong et al. , 2019 ) . While both approaches are valid , it has been shown that the second approach constantly outperforms the first one with a noticeable margin ( Duong et al. , 2019 ) , and most recent methods learn the node embeddings jointly with the parameters of GNNs ( He et al. , 2017 ; 2020 ; Wang et al. , 2019 ) . Learning node features ( or embedding ) X for graph with small number of nodes may not be much of a problem for common computer system . But , as the size of the embedding matrix X grows linearly with the number of nodes , scalability quickly becomes a problem , especially when attempting to apply such method on industrial grade graph data . For example , if a given graph has 1 billion nodes , we set the dimension of the learned embedding to 64 , and store the embedding X using singleprecision floating-point format , the memory cost for the embedding layer alone is 238 gigabytes , which is beyond the capability of common graphics processing unit ( GPU ) . To solve the scalability issue , we adopt the embedding compression idea originally developed for natural language processing ( NLP ) models ( Suzuki & Nagata , 2016 ; Shu & Nakayama , 2017 ; Svenstrup et al. , 2017 ; Takase & Kobayashi , 2020 ) . Particularly , we study the ALONE method proposed by Takase & Kobayashi ( 2020 ) as it only requires single stage of training unlike other methods . ALONE represents each word using a randomly generated compositional code vector ; then a decoder model , which can be trained end-to-end with the downstream model1 , uncompresses the compositional code vector into a floating-point vector . The bit size of the compositional code vector is parametrization by c andmwhere c is the cardinality of each element in the code vector and m is the length of the code vector . For example , if we set c = 4 andm = 6 , one valid code vector is [ 2 , 0 , 3 , 1 , 0 , 1 ] where the length of the vector is 6 and each element in the vector is within the set { 0 , 1 , 2 , 3 } . The code vector can be converted to a bit vector of length m log2 c by representing each element in the code vector as a binary number and concatenating the resulting binary numbers2 . Continuing the example , the code vector [ 2 , 0 , 3 , 1 , 0 , 1 ] can be compactly stored as [ 10 00 11 01 00 01 ] . Using the same conversion trick , it only requires 48 bits to store each word with the parametrization ( c = 64 , m = 8 , 8 log2 64 = 48 bits ) used by Takase & Kobayashi ( 2020 ) in their experiments . The coding scheme can uniquely represent up to 248 words , which is way beyond the number of words ( or sub-words ) used in conventional NLP models ( Vaswani et al. , 2017 ; Takase & Okazaki , 2019 ) . However , generating the code vectors in a random fashion hinders the quality of the uncompressed embedding . One way to quickly benchmark an embedding compression method ’ s capability is by evaluating the performance of reconstructed ( or uncompressed ) pre-trained embeddings . As shown in Figure 1 , when the model compresses more and more embeddings , the performance of the reconstructed embeddings drops considerably when each entity is represented with a randomly generated code vector following Takase & Kobayashi ( 2020 ) ( see lines labeled as “ random ” ) . The phenomenon is observed in our experiments with GloVe word embeddings ( Pennington et al. , 2014b ) on both word analogy/similarity task and metapath2vec/metapath2vec++ node embeddings ( Dong et al. , 2017a ) on node clustering task . In order to solve this performance degradation problem , instead of using a randomly generated code vector to represent each entity , we adopt an efficient random projection hashing method to generate a code vector for each entity using auxiliary information such as the adjacency matrix associated with the graph G or the pre-trained embedding3 . The adopted random projection hashing method is 1For the experiments conducted by Takase & Kobayashi ( 2020 ) , ALONE ’ s decoder model is trained together with transformer ( Vaswani et al. , 2017 ; Takase & Okazaki , 2019 ) on machine translation and summarization . 2The conversion mechanism is more space efficient when c is set to a power of 2 . 3We only apply the hashing-based coding method on the pre-trained embedding in experiments where the pre-trained embeddings are available ( i.e. , Section 5.1 ) . For other experiments , we use the adjacency matrix provided by the associated graph G. a locality-sensitive hashing ( LSH ) method ( Charikar , 2002 ) because it hashes entities with similar auxiliary information into similar code vectors . By referring once again to Figure 1 , the proposed method with hashing-based coding scheme ( see lines labeled as “ hashing ” ) outperformed the random coding scheme in all scenarios . Similar to ALONE , the proposed method does not introduce additional training stage since it is based on random projection . On top of that , the memory footprint is identical to ALONE as the proposed method only replaces the coding scheme . In additional to the proxy tasks of pre-trained embedding reconstruction , we also compare the effectiveness of different coding schemes where the GNN model and the decoder model are trained together in an end-to-end fashion . Particularly , we trained the GraphSAGE model ( Hamilton et al. , 2017 ) on three different node classification datasets . We choose to use the GraphSAGE model ( Hamilton et al. , 2017 ) because it is one of the most scalable GNNs in the literature ( Ying et al. , 2018 ) . To show that the proposed method also benefits other GNNs/tasks , we present additional experiment results in Appendix B.4 on more GNNs with link prediction task . The experimental results have confirmed the superb performance of the proposed hashing-based coding scheme comparing to the random coding scheme used in ALONE under the intended use scenario . 2 RELATED WORK . The embedding compression problem is extensively studied for NLP models because of the memory cost associated with storing the embedding vectors , and one of the most popular strategy is parameter sharing ( Suzuki & Nagata , 2016 ; Shu & Nakayama , 2017 ; Svenstrup et al. , 2017 ; Takase & Kobayashi , 2020 ) . For example , Suzuki & Nagata ( 2016 ) trains a small set of sub-vectors shared by all words called “ reference vectors ” where each word embedding is constructed by concatenating different sub-vectors together . Both the shared sub-vectors and sub-vector assignments are optimized during the training time . The resulting compressed representation is capable of representing each word compactly , but the training process is memory costly as it still needs to train the full embedding matrix to solve the sub-vector assignment problem . Therefore , the method proposed by Suzuki & Nagata ( 2016 ) is not suitable for large-scale data . Shu & Nakayama ( 2017 ) trains a encoder-decoder model ( i.e. , autoencoder ) where the encoder converts a pre-trained embedding into the corresponding compositional code representation and the decoder reconstructs the pre-trained embedding from the compositional code . Once the encoder and decoder are trained , all the pre-trained embeddings are converted to the compact compositional code representation using the encoder ; then the decoder can be trained together with the downstream models . Because the memory cost associated with the compositional code is much smaller than the raw embedding and the decoder is shared by all words , the method reduces the overall memory consumption associated with representing words . However , since it also requires training the embeddings prior to training the encoder-decoder , the training process still has high memory cost associated with the conventional embedding training similar to the previous work . In other words , the method proposed by Shu & Nakayama ( 2017 ) is also not applicable to extreme scale data . Svenstrup et al . ( 2017 ) represents each word compactly with a unique integer and k floating-point values where k is much smaller than the dimension of the embedding . To obtain the embedding from the compact representation of a word , k hash functions4are used to hash the word ’ s associated unique integer to an integer in [ 0 , c ) where c is much smaller than the number of words . Next , k vectors are extracted from a set of c learnable “ component vectors ” based on the output of hash function . The final embedding of the word is generated by computing the weighted sum of the k vectors where the weights are based on the k learnable floating-point values associated with the word . Similar to our work , Svenstrup et al . ( 2017 ) also uses hash functions in their proposed method . But , the role of the hash function is different : Svenstrup et al . ( 2017 ) uses hash functions for reducing cardinality while we use hash functions to perform LSH . On top of that , as the k learnable floating-point values are associated with each word in the method proposed by Svenstrup et al . ( 2017 ) , their method has its parameter size grown linearly with respect to the vocabulary size which makes the method not ideal for our application . The ALONE method proposed by Takase & Kobayashi ( 2020 ) represents each word with a randomly generated compositional code , and the embedding is obtained by inputting the compositional 4The hash function used by Svenstrup et al . ( 2017 ) is the hash function proposed by Carter & Wegman ( 1979 ) for hashing integers . code to a decoder where the number of learnable parameters in the decoder model is independent of the vocabulary size . The ALONE method checks all the requirements for our application ; however , the performance suffers when the vocabulary size increases comparing to autoencoder-based approach ( Shu & Nakayama , 2017 ) as demonstrated in Figure 1 ( labeled as “ learn ” ) . In contrast , our proposed method has similar performance comparing to autoencoder-based approach ( Shu & Nakayama , 2017 ) but does not require the additional training phases required by the method proposed by Shu & Nakayama ( 2017 ) .

Summary: This paper studies the embedding compression problem related to GNNs and graph representation. A two-stage method is proposed to generate the compressed embeddings: firstly, it encodes each node into its composite code with hashing; secondly, it uses a MLP module to decode the embedding for the node. Experiments are performed to evaluate the compression effect with both pretrained graph embeddings and node classification task with GraphSage.