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: In this paper, the authors present a general and scalable approach to obtain shallow embeddings of many entities. The authors achieve this by first mapping each entity into a code vector using LSH, and then learning a neural decoder on top of the code vectors to reconstruct the original embeddings/adjacency matrix. The authors then apply their approach to node classification problems, where node features are unavailable (the authors ended up artificially removing node features from existing OGB datasets). The superior performance is demonstrated over the random code vector baseline.