SlowGuess/ACL_2020_OCR huggingface dataset:
Mrinmaya Sachan  
Toyota Technological Institute at Chicago  
mrinmaya@ttic.edu
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.
Knowledge graphs (KGs) are a popular way of storing world knowledge, lending support to a number of AI applications such as search (Singhal, 2012), question answering (Lopez et al., 2013; Berant et al., 2013) and dialog systems (He et al., 2017; Young et al., 2018). 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 (Bordes et al., 2011, 2013; Socher et al., 2013). 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.
(Bordes et al., 2013; Wang et al., 2014; Lin et al., 2015; Ji et al., 2015). 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) (van den Oord et al., 2017; Chen et al., 2018; Chen and Sun, 2019). Each entity in the KG is represented as a sequence of  $D$  codes where each code can take values in  $\{1 \ldots 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<sup>1</sup>. 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 reversed-discretization function that maps the discrete codes
For example, Barack Obama = "2-1-3-3" and Michelle Obama = "2-1-3-2" (for  $D = 4$  and  $K = 3$ )
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2681-2691 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
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 (Bengio et al., 2013) or the tempering softmax (Maddison et al., 2016; Jang et al., 2016) 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.
2.1 Knowledge Graph Embeddings
A knowledge graph (KG)  $\mathcal{G} \subseteq \mathcal{E} \times \mathcal{R} \times \mathcal{E}$  can be formalized as a set of triplets  $(e_i, r, e_j)$  composed of head and tail entities  $e_i$  and  $e_j$  ( $e_i, e_j \in \mathcal{E}, \mathcal{E}$  being the set of entities) and a relation  $r \in \mathcal{R}$  ( $\mathcal{R}$  being the set of relations)  $-n_e = |\mathcal{E}|, n_r = |\mathcal{R}|$ . The goal of learning KG embeddings is to learn vector embeddings  $\mathbf{e} \in \mathbb{R}^{d_e}$  for each entity  $e \in \mathcal{E}$  (and possibly also relation embeddings  $\mathbf{r} \in \mathbb{R}^{d_r}$ ).
Typical KG embedding approaches are multi-layer 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  $\mathbf{f}:\mathcal{E}\times \mathcal{R}\times \mathcal{E}\to \mathbb{R}$  where  $f(e_i,r,e_j)$  defines the score of the triplet  $(e_i,r,e_j)$ . KG embeddings are learnt by defining a loss function  $\mathcal{L}$  and solving the following optimization problem:
$$
\min  _ {\Theta} \sum_ {(e _ {i}, r, e _ {j}) \in \mathcal {G}} \mathcal {L} _ {\Theta} (e _ {i}, r, e _ {j}) \tag {1}
$$
Here  $\Theta$  includes all embedding parameters and any other neural network parameters. The loss function typically encourages the score of a positive triplet  $(e_i, r, e_j)$  to be higher than that of a (corrupted) negative triplet. In Table 1, we summarize the scoring function for several existing KG embedding approaches as well as their corresponding entity (and relation) representation parameters.
In all the KG embedding models, the number of parameters grow super-linearly with the number of entities and relations in the KG as well as the size of their representations. This number can be very large and learning KG embeddings can be a challenge for large, sparse KGs. In this paper, we present a novel coding scheme that significantly reduces the number of embedding parameters. We do so by leveraging recent advances in discrete representation learning. We summarize them below.
2.2 Discrete Representation Learning
Typical deep learning methods define an embedding function as  $F: \mathcal{V} \to \mathbb{R}^d$ , where  $\mathcal{V}$  denotes the vocabulary such as words, sub-words, entities, relations, etc. and each symbol in the vocabulary is mapped to a continuous vector in  $\mathbb{R}^d$ . The embedding function can be trained separate from the task in a completely unsupervised manner or jointly with other neural net parameters to optimize the target loss function. A common specification of the embedding function in NLP is a lookup table  $L \in \mathbb{R}^{n \times d}$  with  $n = |\mathcal{V}|$ . The total number of bits used to represent this table is  $O(nd)$  (32nd if each real number is represented by 32-bit floating point). This is problematic for large  $n$  and/or  $d$ .
Thus, various approaches have been proposed to compress embedding layers in neural networks. These include weight-tying (Press and Wolf, 2016; Inan et al., 2016; Li et al., 2018), matrix-factorization based approaches (Acharya et al., 2019), and approaches that rely on gumbel softmax (Baevski and Auli, 2018), vector quantization (Chen and Sun, 2019) and codebook learning (Shu and Nakayama, 2017). In this work, we build on discrete representation learning approaches (van den Oord et al., 2017; Chen et al., 2018; Chen and Sun, 2019). Discrete representation learning gives us a way to mitigate this issue by representing each symbol  $v$  in the vocabulary as a discrete vector  $\mathbf{z}_v = [\mathbf{z}_v^{(1)},\dots,\mathbf{z}_v^{(D)}]$ . Discrete representations have another clear benefit that they are interpretable and are a natural fit for complex reasoning, planning and predictive learning tasks.
Learning discrete representations is challenging due to the inherent non-differentiability in the embedding layer. Thus, a number of solutions such as the gumbel softmax trick (Maddison et al., 2016; Jang et al., 2016) and the straight-through estimator (Bengio et al., 2013) have been proposed to tackle this issue. Making the discrete representation learn
KG Embedding Model	Scoring function f	# Ent. Rel. Params
SE (Bordes et al., 2011)	||Wr(L)e_i - Wr(R)e_j||p	nede + 2nrde dr
NTN (Socher et al., 2013)	urTf(e_iWr[1...k]e_j + V_r[e_i] + br)	nede + nr(kd^2 + 2kde + k)
TransE (Bordes et al., 2013)	||ei + r - ej||p	(ne + nr)d
TransH (Wang et al., 2014)	||(ei - wr)e_iwr) + dr - (ej - wr^Tejwr)||2	nede + 2nrdr
TransR (Lin et al., 2015)	||eiWr + r - ejWr||2	nede + nr(dr + dr^2)
TransD (Ji et al., 2015)	-||(rpep_i + I)e_i + r - (rpep_j + I)e_j||2	2nede + 2nrdr
DistMult (Yang et al., 2014)	{ei, r, ej}	(ne + nr)d
ComplEx (Trouillon et al., 2016)	Re{ei, r, ej}	2(ne + nr)d
HolE (Nickel et al., 2016)	{r, ei ⊗ ej}	(ne + nr)d
SimpleE (Kazemi and Poole, 2018)	1/2{⟨ei(h), r, ej(t)⟩ + ⟨ej(h), r^(inv), ei(t))}	2(ne + nr)d
ConvE (Dettmers et al., 2018)	{σ(VEC(σ([ei; r] ⊘ ω))W, ej}	nede + nrdr
RotatE (Sun et al., 2019)	-||ei • r - ej||2	(ne + nr)d
HypER (Balažević et al., 2019a)	{σ(VEC(σ(e_i * vec^-1(w_rH))))W, ej}	nede + nrdr
TuckER (Balažević et al., 2019b)	W ×_1 ei ×_2 wr ×_3 ej	nede + nrdr
Table 1: Scoring functions  $f$  of some popular knowledge graph embedding approaches in the literature and the number of entity and relation specific parameters. Here,  $n_e = |\mathcal{E}|$  and  $n_r = |\mathcal{R}|$  respectively denote the number of entities and relation types and  $d_e$  and  $d_r$  respectively denote the dimension of entity and relation representations.  $d$  is defined when (as)  $d = d_e = d_r$ .  $\langle x^1, \ldots, x^k \rangle = \sum_i x_i^1 \ldots x_i^k$  denotes the generalized dot product,  $\bar{\sigma}$  denotes the conjugate of a complex number and  $\otimes$  denotes circular correlation,  $\sigma$  denotes an activation function,  $\circ$  denotes the convolution operator,  $\bullet$  denotes the hadamard product and  $\times_k$  denotes tensor product along the  $k^{th}$  mode.
ing process differentiable enables end-to-end learning of discrete representations via optimizing some task-specific objectives from language modelling and machine translation. In this work, we use discrete representation learning to compress KG embeddings. We describe it below.
3 Discrete KG Representation Learning
In order to learn discrete KG representations, we define a quantization function  $\mathcal{Q}:\mathbb{R}^d\to \mathbb{R}^d$  , which (during training) takes raw KG embeddings and produces their quantized representations.  $\mathcal{Q} =$ $\mathcal{D}\circ \mathcal{R}$  is composed of two functions:
1. A discretization function  $\mathcal{D}:\mathbb{R}^{d_e}\to \mathbb{Z}^D$  that maps the continuous KG embedding into a  $K$ -way  $D$ -dimensional discrete code with cardinality  $|\mathbb{Z}| = K$  (we call this KD code)
2. A reverse-discretization function  $\mathcal{R}$  :  $\mathbb{Z}^D\to \mathbb{R}^{d_e}$  that maps the KD code back to the continuous embedding.
During training, both  $\mathcal{D}$  and  $\mathcal{R}$  are learned. Then, every entity in the KG is represented by a KD code via applying the discretization function  $\mathcal{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  $\mathcal{R}$  is used to decode the KD codes into regular embedding vectors for every entity. We use vector quantization (Chen et al., 2018; Chen and Sun, 2019) and
codebook learning (Cai et al., 2010) to define the discretization and reverse-discretization functions  $\mathcal{D}$  and  $\mathcal{R}$ . We describe them below.
3.1 Discretization Function  $\mathcal{D}$
The goal of the discretization function is to map continuous KG embedding vectors into KD codes. We model the discretization function using nearest neighbor search (Cayton, 2008). Given continuous KG embeddings  $\{\mathbf{e}_i|i = 1\dots n_e\}$  as query vectors, we define a set of  $K$  key vectors  $\{\mathbf{k}_k|k = 1\dots K\}$  where  $\mathbf{k}_k\in \mathbb{R}^{d_e}$
In order to learn  $D$ -dimensional discrete codes, we partition the query and key vectors into  $D$  partitions where each partition corresponds to one of the  $D$  discrete codes  $-\mathbf{e}_i^{(j)} \in \mathbb{R}^{n \times d_e / D}$  and  $\mathbf{k}_k^{(j)} \in \mathbb{R}^{K \times d_e / D}, j = 1 \dots D$ .
Vector Quantization (VQ): Our first alternative for discretization is vector-quantization (Ballard, 1997), a classical quantization technique for data compression. We assume that the  $j^{th}$  discrete code of the  $i^{th}$  entity  $z_{i}^{(j)}$  can be computed by calculating distances between the corresponding query vector partition  $\mathbf{e}_i^{(j)}$  and various corresponding key vector partitions  $\{\mathbf{k}_k^{(j)}\}$ , and choosing the one with the minimum distance:
$$
z _ {i} ^ {(j)} = \underset {k} {\arg \min } d i s t \left(\mathbf {e} _ {i} ^ {(j)}, \mathbf {k} _ {k} ^ {(j)}\right) \tag {2}
$$
We use the Euclidean distance function:
dist(a,b) = ||a - b||2 in our experiments. Note that the argmin operation is inherently non-differentiable. The resulting quantization function  $\mathcal{Q}$  has no gradient towards the input query vectors. Thus, we use the straight-through estimator (Bengio et al., 2013) to compute a pseudo gradient. This means that during the forward pass, we compute  $\mathcal{Q}$  as defined here, but during the backward pass, we use the gradient of the query vectors.
Tempering Softmax (TS): Vector quantization is a popular method for learning discrete representations. Yet another popular approach is continuous relaxation of (2) via the tempering softmax (Maddison et al., 2016; Jang et al., 2016). We again use dot product and softmax for computing the proximity between query and key vectors:
$$
z _ {i} ^ {(j)} = \underset {k} {\arg \max} \frac {\exp \left(\langle \mathbf {e} _ {i} ^ {(j)} , \mathbf {k} _ {k} ^ {(j)} \rangle / \tau\right)}{\sum_ {k ^ {\prime}} \exp \left(\langle \mathbf {e} _ {i} ^ {(j)} , \mathbf {k} _ {k ^ {\prime}} ^ {(j)} \rangle / \tau\right)}
$$
Here,  $\tau$  is the temperature and  $\langle \mathbf{a},\mathbf{b}\rangle = \mathbf{a}^T\mathbf{b}$  denotes the dot product operation. Note that this function still carries an inherent non-differentiability. Hence, we relax the above and compute probability vectors  $\bar{z}_i^{(j)}$  which represent the probability distribution of the  $j^{th}$  dimension of the discrete code for the  $i^{th}$  entity taking a particular value (say  $k$ ). Given probabilistic vectors  $\bar{z}_i^{(j)}$ , we can compute the discrete codes  $z_i^{(j)}$  simply by taking the argmax. To compute discrete KD codes, we set a small value of  $\tau$ . As  $\tau \to 0$ , the softmax becomes spiky concentrated on the true  $z_i^{(j)}$ -th dimension. We again estimate pseudo gradients by setting a very small  $\tau$  in the forward pass (i.e. close to the discrete case (eq. 1)) and  $\tau = 1$  in the backward pass.
3.2 Reverse-discretization Function  $\mathcal{R}$
The goal of the reverse-discretization function is to map discrete KD codes into continuous KG embedding vectors. We model the reverse-discretization process first by a simple linear model which maps the discrete codes to continuous vectors by looking up a learnt codebook. Then, we present an alternative - a non-linear model for reverse-discretization based on recurrent neural networks.
Codebook Lookup (CL): We first define the reverse-discretization function in a simple manner where we substitute every discrete code with a continuous vector from a codebook. Let  $\mathcal{C}$  be a
set of codebooks.  $\mathcal{C}$  consists of a number of codebooks - a separate codebook  $\mathcal{C}^{(j)}$  for each position  $j = 1\dots D$  in the KD code. We model each codebook simply as a set of vectors:  $\mathcal{C}^{(j)} = \{\mathbf{c}_i^{(j)}|i = 1\dots K\}$  where  $\mathbf{c}_i^{(j)}\in \mathbb{R}^{d_e / D}$ . We simply compute the embedding vector for the  $j^{th}$  dimension of the  $i^{th}$  entity as:
$$
\mathbf {e} _ {i} ^ {(j)} = \mathbf {c} _ {i} ^ {(j)}
$$
The final entity embedding vector  $\mathbf{e}_i$  is achieved by the concatenation of the embedding vectors for each dimension:  $\mathbf{e}_i = [\mathbf{e}_i^{(1)}\dots \mathbf{e}_i^{(D)}]$ .
Non-linear Reconstruction (NL): While the codebook lookup approach is simple and efficient, due to its linear nature, the capacity of the generated KG embedding may be limited. Thus, we also employ neural network based non-linear approaches for embedding reconstruction. We propose a non-linear embedding reconstruction approach based on the Bi-LSTM network.
Given the KD code  $\mathbf{z}_i$  as a sequence of codes  $z_i^{(1)},\ldots ,z_i^{(D)}$ , we map the KD code to a continuous embedding vector by feeding the code to a Bi-LSTM followed by mean pooling.
$$
\text {L e t} \quad \left(\mathbf {h} _ {i} ^ {(1)}, \dots , \mathbf {h} _ {i} ^ {(D)}\right) \quad =
$$
Bi-LSTM  $\left(z_{i}^{(1)},\ldots ,z_{i}^{(D)}\right)$  be the hidden state representations for the various Bi-LSTM cells. Finally, we reconstruct the entity embedding  $\hat{\mathbf{e}}_i$  by mean-pooling the code embedding vectors followed by a linear transformation:  $\mathbf{e}_i = \mathbf{W}_{rev}^T\left(\sum_j\mathbf{h}_i^{(j)}\right)$ .
We also tried to map the KD code to a continuous embedding vector by feeding the code to variations of a character level CNN (Kim et al., 2016). However, the Char CNN model always performed worse than the Bi-LSTM model in our experiments. This was because our discretization function which discretizes contiguous partitions of the continuous representation better suits the Bi-LSTM reconstruction model. In the future, we would like to consider more complex discretization functions with other complex non-linear reconstruction models.
Storage Efficiency: A key motivation of learning discrete representations is that we can significantly compress the embedding layer at test time. The size of the embedding layer for typical KG representations is  $32n_{e}d_{e}$  (assuming a 32 bit representation) – this can be very large. In contrast, with discrete representation learning, we only need to store code embeddings  $\{\mathbf{z}_i\}$  and the parameters used in the reverse-discretization function such as
the codebooks  $\mathcal{C}$  or the parameters of the embedding reconstruction Bi-LSTM  $\{\Theta_{\mathrm{LSTM}},\mathbf{W}_{rev}\}$ .
The entity codes require  $n_e D \log_2 K$  bits.
The codebook lookup approach needs to also maintain codebooks which require  $32Kd_{e}$  parameters and the non-linear reconstruction approach requires  $Dd^{\prime}\times 6$  parameters (two set of parameter matrices each for the input, output and forget gates) for the Bi-LSTM and  $d_{e}d^{\prime}$  parameters for storing  $\mathbf{W}_{rev}$  - a total of  $(6D + d_{e})d^{\prime}$  parameters. Here,  $d^{\prime}$  is the size of the code embedding vectors.
In both codebook lookup and non-linear reconstruction formulations, discrete representation learning neatly decouples the KG size (number of entities) and dimensionality of the continuous embeddings. Thus, the discrete embedding layer can be compactly stored as typically  $D$  and  $\log_2 K$  are smaller than  $32d_e$  (considering only the dominating term  $n_e$ ).
Test Time Inference of Embeddings: At test time, we retrieve continuous embeddings for an entity by looking up the codebook or running inference on the reconstruction model using its discrete representation. For codebook lookup, the steps involved are (a) looking up a simple index for each code, and (b) concatenation. Since only index lookups and concatenation are needed, the extra computation complexity and memory footprint are very small -  $\mathcal{O}(D)$  time and memory. In the nonlinear reconstruction setting, we need to run inference on the Bi-LSTM model. This requires  $\mathcal{O}(D)$  matrix vector multiplications (to compute various LSTM gates) which takes  $\mathcal{O}(Dd_{e}d^{\prime})$  time. Finally, we have another linear transformation  $\mathbf{W}_{rev}$  - this takes  $\mathcal{O}(d_{e}d^{\prime})$  time.
We can further cache the embedding lookups and various intermediate results such as matrix vector products to improve performance. We show in our results that the test time inference overhead is typically very small.
Learning: Similar to previous continuous KG representation learning methods, we learn discrete entity representations by minimizing the triplet loss function. We extend equation 1 as:
$$
\min  _ {\left\{\mathbf {z} _ {e} \right\}, \theta , \Theta} \sum_ {(e _ {i}, r, e _ {j}) \in \mathcal {G}} \mathcal {L} _ {\left\{\mathbf {z} _ {e} \right\}, \theta , \Theta} (e _ {i}, r, e _ {j} | \theta , \Theta) \tag {3}
$$
Here,  $\mathbf{z}_e$  are code embeddings,  $\theta$  are the parameters of the reverse-discretization function (C or  $\{\theta_{LSTM},\mathbf{W}_{rev}\}$ ) and  $\Theta$  denotes parameters of the KG embedding approaches (listed in Table 1). The
aforementioned loss function (eq 3) is differentiable w.r.t. the embedding parameters and parameters of entity representation learning methods. However, the discrete codes introduce a non-differentiability. Thus, we use straight-through (Bengio et al., 2013) or the tempering softmax (Maddison et al., 2016; Jang et al., 2016) to estimate pseudo-gradients as described before (section 3.1).
Guidance from KG embeddings: We find that even with sophisticated discrete representation learning methods, solving the above optimization problem can be challenging in practice. Due to discreteness of the problem, this can lead to a suboptimal solution where discrete codes are not as good. Therefore, we also use guidance from continuous KG embeddings to solve (3) when provided<sup>2</sup>. The key idea is that in addition to optimizing (3), we can encourage the reconstructed embeddings from the learnt discrete codes to mimic continuous embeddings.
In order to provide this guidance from continuous embeddings, during the training, instead of using the reconstructed embedding vector generated from the discrete code, we use a weighted average of the reconstructed embeddings and continuous embeddings obtained using methods described in Table 1:  $(1 - \lambda)\mathcal{D} \circ \mathcal{R}(\mathbf{e}) + \lambda \mathbf{e}$ . Here  $\lambda \in (0,1)$  is a linear interpolant for selecting between reconstructed embeddings and pre-learnt continuous embeddings. We initialize  $\lambda$  to 1 and gradually decrease  $\lambda$  as training proceeds. This enables the method to gradually rely more and more on reconstruction from discrete embeddings. We also add a regularization term  $||\mathcal{D} \circ \mathcal{R}(\mathbf{e}) - \mathbf{e}||_2^2$  during the training to encourage the reconstructed embeddings to match the pre-learnt continuous embeddings. This procedure is similar to knowledge-distillation guidance (Hinton et al., 2015) in previous discrete representation learning works (Chen et al., 2018).