Knowledge Graph Embedding Compression (2020)

Abstract: Knowledge graph (KG) representation learning techniques that learn continuous embeddings of entities and relations in the KG have become popular in many AI applications. With a large KG, the embeddings consume a large amount of storage and memory. This is problematic and prohibits the deployment of these techniques in many real world settings. Thus, we propose an approach that compresses the KG embedding layer by representing each entity in the KG as a vector of discrete codes and then composes the embeddings from these codes. The approach can be trained end-to-end with simple modifications to any existing KG embedding technique. We evaluate the approach on various standard KG embedding evaluations and show that it achieves 50-1000x compression of embeddings with a minor loss in performance. The compressed embeddings also retain the ability to perform various reasoning tasks such as KG inference.

## introduction knowledge graphs (kgs) are a popular way of storing world knowledge, lending support to a number of ai applications such as search @xcite , question answering @xcite and dialog systems @xcite . typical kgs are huge, consisting of millions of entities and relations. with the growth in use of kgs, researchers have explored ways to learn better representations of kgs in order to improve generalization and robustness in downstream tasks. in particular, there has been interest in learning embeddings of kgs in continuous vector spaces @xcite @xcite . kg embedding approaches represent entities as learnable continuous vectors while each relation is modeled as an operation in the same space such as translation, projection, etc. @xcite @xcite . these approaches give us a way to perform reasoning in kgs with simple numerical computation in continuous spaces. despite the simplicity and wide-applicability of kg embedding approaches, they have a few key issues. a major issue is that the number of embedding parameters grow linearly with the number of entities. this is challenging when we have millions or billions of entities in the kg, especially when there are a lot of sparse entities or relations in the kg. there is a clear redundancy in the continuous parameterization of embeddings given that many entities are actually similar to each other. this overparameterization can lead to a drop in performance due to overfitting in downstream models. the large memory requirement of continuous representations also prevents models that rely on them from being deployed on modest user-facing computing devices such as mobile phones. to address this issue, we propose a coding scheme that replaces the traditional kg embedding layer by representing each entity in the kg with a k-way d dimensional code (kd code) @xcite @xcite . each entity in the kg is represented as a sequence of d codes where each code can take values in {1 . . . k}. the codes for each entity are learnt in such a way that they capture the semantics and the relational structure of the kg -i.e., the codes that represent similar or related entities are typically also similar foot_0 . the coding scheme is much more compact than traditional kg embedding schemes. we learn the discrete codes for entities using an autoencoder style model which learns a discretization function that maps continuous entity representations to discrete codes and a reversediscretization function that maps the discrete codes back to continuous entity representations. the discretization and reverse-discretization functions are jointly learnt end-to-end. the inherent discreteness of the representation learning problem poses several learning issues. we tackle these issues by resorting to the straight-through estimator @xcite or the tempering softmax @xcite and using guidance from existing kg embeddings to smoothly guide learning of the discrete representations. we evaluate our approach on various standard kg embedding evaluations and we find that we can massively reduce the size of the kg embedding layer while suffering only a minimal loss in performance (if at all). we show that the proposed approach for learning discrete kg representations leads to a good performance in the task of link prediction (cloze entity prediction) as well as in the task of kg reasoning and inference. typical kg embedding approaches are multilayer neural networks which consist of an embedding component and a scoring component. the embedding component maps each entity to its corresponding embedding. the scoring component learns a scoring function f : e × r × @xmath0, r, e j ) defines the score of the triplet (e i , r, e j ). kg embeddings are learnt by defining a loss function l and solving the following optimization problem: ## discrete kg representation learning in order to learn discrete kg representations, we define a quantization function q : r @xmath1, which (during training) takes raw kg embeddings and produces their quantized representations. @xmath2: during training, both d and r are learned. then, every entity in the kg is represented by a kd code via applying the discretization function d to save space (compression). the continuous embeddings and the parameters of the discretization function are then no longer needed. in the test/inference stage, the reverse-discretization function r is used to decode the kd codes into regular embedding vectors for every entity. we use vector quantization @xcite and codebook learning @xcite to define the discretization and reverse-discretization functions d and r. we describe them below. ## experiments we compare the baseline continuous representations described earlier in ## related work deep learning model compression has attracted many research efforts in the last few years @xcite . these efforts include network pruning @xcite , weight sharing @xcite , quantization @xcite , low-precision computation @xcite and knowledge distillation @xcite these techniques can also be used for embedding compression. press and wolf (2016) and @xcite propose weight-tying approaches that learn input and output representations jointly. matrix factorization-based methods @xcite @xcite have also been proposed which approximate an embedding matrix with smaller matrices or clusters. closest to our work are @xcite @xcite who present similar approaches to learn discrete codings for word embeddings using multiple codebooks, i.e. product quantization @xcite . similar techniques have used been used by @xcite who extend vaes to learn discrete representations using vector quantization in the image domain. this allows the vae model to circumvent its well known issues of "posterior collapse". all these previous works have been applied to the image domain, and sometimes in language to learn discrete word embeddings. in this work, we present the first results on compressing kg embeddings and also show how the compressed embeddings can be used to support various knowledge based applications such as kg inference. ## conclusion the embedding layer contains majority of the parameters in any representation learning approach on knowledge graphs. this is a barrier in successful deployment of models using knowledge graphs at scale on user-facing computing devices. in this work, we proposed novel and general approaches for kg embedding compression. our approaches learn to represent entities in a kg as a vector of discrete codes in an end-to-end fashion. at test time, the discrete kg representation can be cheaply and efficiently converted to a dense embedding and then used in any downstream application requiring the use of a knowledge graph. we evaluated our proposed methods on different link prediction and kg inference tasks and show that the proposed methods for kg embedding compression can effectively compress the kg embedding table without suffering any significant loss in performance. in this work, we only considered the problem of learning discrete entity representations. in the future, we would like to jointly learn discrete representations of entities as well as relations.