SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lorenzo Lo Conte
University of Edinburgh, UK  
1.loconte@sms.ed.ac.uk
Nicola Di Mauro
University of Bari, Italy nicola.dimauro@uniba.it
Robert Peharz
TU Graz, Austria robert. peharz@tugraz.at
Antonio Vergari
University of Edinburgh, UK  
avergari@ed.ac.uk
Some of the most successful knowledge graph embedding (KGE) models for link prediction – CP, RESCAL, TUCKER, COMPLEX – can be interpreted as energy-based models. Under this perspective they are not amenable for exact maximum-likelihood estimation (MLE), sampling and struggle to integrate logical constraints. This work re-interprets the score functions of these KGEs as circuits – constrained computational graphs allowing efficient marginalisation. Then, we design two recipes to obtain efficient generative circuit models by either restricting their activations to be non-negative or squaring their outputs. Our interpretation comes with little or no loss of performance for link prediction, while the circuits framework unlocks exact learning by MLE, efficient sampling of new triples, and guarantee that logical constraints are satisfied by design. Furthermore, our models scale more gracefully than the original KGEs on graphs with millions of entities.
Knowledge graphs (KGs) are a popular way to represent structured domain information as directed graphs encoded as collections of triples (subject, predicate, object), where subjects and objects (entities) are nodes in the graph, and predicates their edge labels. For example, the information that the drug "loxoprofen" interacts with the protein "COX2" is represented as a triple (loxoprofen, interacts, COX2) in the biomedical KG ogbl-biokg [32]. As real-world KGs are often incomplete, we are interested in performing reasoning tasks over them while dealing with missing information. The simplest reasoning task is link prediction [48], i.e., querying for the entities that are connected in a KG by an edge labelled with a certain predicate. For instance, we can retrieve all proteins that the drug "loxoprofen" interacts with by asking the query (loxoprofen, interacts,?).
Knowledge graph embedding (KGE) models are state-of-the-art (SOTA) models for link prediction [40, 56, 12] that map entities and predicates to sub-symbolic representations, which are used to assign a real-valued degree of existence to triples in order to rank them. For example, the SOTA KGE model COMPLEX [65] assigns (loxoprofen, interacts, phosphoric-acid) and (loxoprofen, interacts, COX2) scores 2.3 and 1.3, hence ranking the first higher than the second in our link prediction example.
This simple example, however, also highlights some opportunities that are missed by KGE models. First, triple scores cannot be directly compared across different queries and across different KGE models, as they can be seen as negative energies and not normalised probabilities over triples [41, 7, 6, 27, 43]. To establish a sound probabilistic interpretation and therefore have probabilities instead of scores that can be easily interpreted and compared [79], we would need to compute the normalisation
Figure 1: Which KGE models can be used as efficient generative models of triples? The score functions of popular KGE models such as COMPLEX, CP, RESCAL and TUCKER can be easily represented as circuits (lilac). However, to retrieve a valid probabilistic circuit (PC, in orange) that encodes a probability distribution over triples (GeKCs) we need to either restrict its activations to be non-negative (in blue, see Section 4.1) or square it (in red, see Section 4.2).
constant (or partition function), which is impractical for real-world KGs due to their considerable size (see Section 2). Therefore, learning KGE models by maximum-likelihood estimation (MLE) would be computationally infeasible, which is the canonical probabilistic way to learn a generative model over triples. A generative model would enable us to sample new triples efficiently, e.g., to generate a surrogate KG whose statistics are consistent with the original one or to do data augmentation [9]. Furthermore, traditional KGE models do not provide a principled way to guarantee the satisfaction of hard constraints, which are crucial to ensure trustworthy predictions in safety-critical contexts such as biomedical applications. The result is that predictions of these models can blatantly violate simple constraints such as KG schema definitions. For instance, the triple that the SOTA COMPLEX ranks higher in our example above violates the semantics of "interacts", i.e., such predicate can only hold between drugs (e.g., loxoprofen) and proteins (e.g., COX2) but phosphoric-acid is not a protein.
Contributions. We show that KGE models that have become a de facto standard, such as COMPLEX and alternatives based on multilinear score functions [47, 40, 66], can be represented as structured computational graphs, named circuits [13]. Under this light, i) we propose a different interpretation of these computational graphs and their parameters, to retrieve efficient and yet expressive probabilistic models over triples in a KG, which we name generative KGE circuits (GeKCs) (Section 4). We show that ii) not only GeKCs can be efficiently trained by exact MLE, but learning them with widely-used discriminative objectives [37, 40, 56, 12] also scales far better over large KGs with millions of entities. In addition, iii) we are able to sample new triples exactly and efficiently from GeKCs, and propose a novel metric to measure their quality (Section 7.3). Furthermore, by leveraging recent theoretical advances in representing circuits [76], iv) we guarantee that predictions at test time will never violate logical constraints such as domain schema definitions by design (Section 5). Finally, our experimental results show that these advantages come with no or minimal loss in terms of link prediction accuracy.
2 From KGs and embedding models...
KGs and embedding models. A KG  $\mathcal{G}$  is a directed multi-graph where nodes are entities and edges are labelled with predicates, i.e., elements of two sets  $\mathcal{E}$  and  $\mathcal{R}$ , respectively. We define  $\mathcal{G}$  as a collection of triples  $(s,r,o)\subseteq \mathcal{E}\times \mathcal{R}\times \mathcal{E}$ , where  $s$ ,  $r$ ,  $o$  denote the subject, predicate and object, respectively. A KG embedding (KGE) model maps a triple  $(s,r,o)$  to a real scalar via a score function  $\phi \colon \mathcal{E}\times \mathcal{R}\times \mathcal{E}\to \mathbb{R}$ . A common recipe to construct differentiable score functions for many state-of-the-art KGE models [56] is to (i) map entities and predicates to embeddings of rank  $d$ , i.e., elements of normed vector spaces (e.g.,  $\mathbb{R}^d$ ), and (ii) combine the embeddings of subject, predicate and object via multilinear maps. This is the case for KGE models such as CP [40], RESCAL [47], TUCKER [3], and COMPLEX [65] (see Fig. 2). For instance, the score function of COMPLEX [65] is defined as  $\phi_{\mathrm{COMPLEX}}(s,r,o) = \mathrm{Re}(\langle \mathbf{e}_s,\mathbf{w}_r,\overline{\mathbf{e}_o}\rangle)$  where  $\mathbf{e}_s,\mathbf{w}_r,\mathbf{e}_o\in \mathbb{C}^d$  are the complex embeddings of subject, predicate and object,  $\langle \cdot ,\cdot ,\cdot \rangle$  denotes a trilinear product,  $\overline{\cdot}$  denotes the complex conjugate operator and  $\mathrm{Re}(\cdot)$  the real part of complex numbers.
Probabilistic loss-derived interpretation. KGE models have been traditionally interpreted as energy-based models (EBMs) [41, 7, 6, 27, 43]: their score function is assumed to compute the negative energy of a triple. This interpretation induces a distribution over possible KGs by associating a Bernoulli variable, whose parameter is determined by the score function, to every triple [48]. Learning EBMs under this perspective requires using contrastive objectives [7, 48, 56], but several recent works observed that to achieve SOTA link prediction results one needs only to predict
subjects, objects [37, 40, 56] or more recently also predicates of triples [12], i.e., to treat KGEs as discriminative classifiers. Specifically, they are learned by minimising a categorical cross-entropy, e.g., by maximising  $\log p(o\mid s,r) = \phi (s,r,o) - \log \sum_{o^{\prime}\in \mathcal{E}}\exp \phi (s,r,o^{\prime})$  for object prediction. From this perspective, we observe that we can recover an energy-based interpretation if we assume there exist a joint probability distribution  $p$  over three variables  $S,R,O$ , denoting respectively subjects, predicates and objects. Written as a Boltzmann distribution, we have that  $p(S = s,R = r,O = o) = (\exp \phi (s,r,o)) / Z$ , where  $Z = \sum_{s^{\prime}\in \mathcal{E}}\sum_{r^{\prime}\in \mathcal{R}}\sum_{o^{\prime}\in \mathcal{E}}\exp \phi (s^{\prime},r^{\prime},o^{\prime})$  denotes the partition function. This interpretation is apparently in contrast with the traditional one over possible KGs [48]. We reconcile it with ours in Appendix E. Under this view, we can reinterpret and generalise the recently introduced discriminative objectives [12] as a weighted pseudo-log-likelihood (PLL) [70]
$$
\mathcal {L} _ {\mathrm {P L L}} := \sum_ {(s, r, o) \in \mathcal {G}} \omega_ {s} \log p (s \mid r, o) + \omega_ {o} \log p (o \mid s, r) + \omega_ {r} \log p (r \mid s, o) \tag {1}
$$
where  $\omega_s, \omega_o, \omega_r \in \mathbb{R}_+$  can differently weigh each term, which is a conditional log-probability that can be computed by summing out either  $s, r$  or  $o$ , e.g., to compute  $\log p(o \mid s, r)$  above. Optimisation is usually carried out by mini-batch gradient ascent [56] and, given a batch of triples  $B \subset \mathcal{G}$ , we have that exactly computing the PLL objective requires time  $\mathcal{O}(|\mathcal{E}| \cdot |B| \cdot \mathrm{cost}(\phi))$  and space  $\mathcal{O}(|\mathcal{E}| \cdot |B|)$  to exploit GPU parallelism [36],<sup>1</sup> where  $\mathrm{cost}(\phi)$  denotes the complexity of evaluating the  $\phi$  once.
Note that the PLL objective (Eq. (1)) is a traditional proxy for learning generative models for which it is infeasible to evaluate the maximum-likelihood estimation (MLE) objective<sup>2</sup>
$$
\mathcal {L} _ {\mathrm {M L E}} := \sum_ {(s, r, o) \in \mathcal {G}} \log p (s, r, o) = - | \mathcal {G} | \log Z + \sum_ {(s, r, o) \in \mathcal {G}} \phi (s, r, o). \tag {2}
$$
In theory, evaluating  $\log p(s,r,o)$  exactly can be done in polynomial time under our three-variable interpretation, as computing  $Z$  requires  $\mathcal{O}(|\mathcal{E}|^2\cdot |\mathcal{R}|\cdot \mathrm{cost}(\phi))$  time, but in practice this cost is still prohibitive for real-world KGs. In fact, it would require summing over  $|\mathcal{E}\times \mathcal{R}\times \mathcal{E}|$  evaluations of the score function  $\phi$ , which for FB15k-237 [62], the small fragment of Freebase [48], translates to  ${\sim}10^{11}$  evaluations of  $\phi$ . This practical bottleneck hinders the generative capabilities of these models and their ability to yield normalised and interpretable probabilities. Next, we show how we can reinterpret KGE score functions as to retrieve a generative model over triples for which computing  $Z$  exactly can be done in time  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|)\cdot \mathrm{cost}(\phi))$ , making renormalisation feasible.
3 ...to Circuits...
In this section, we show that popular and successful KGE models such as CP, RESCAL, TUCKER and COMPLEX (see Fig. 2 and Section 2), can be viewed as structured computational graphs that can, in principle, enable summing over all possible triples efficiently. Later, to exploit this efficient summation for marginalisation over triple probabilities, we reinterpret the semantics of these computational graphs as to yield circuits that output valid probabilities. We start with the needed background about circuits and show that some score functions can be readily represented as circuits.
Definition 1 (Circuit [13, 76]). A circuit  $\phi$  is a parametrized computational graph over variables  $\mathbf{X}$  encoding a function  $\phi(\mathbf{X})$  and comprising three kinds of computational units: input, product, and sum. Each product or sum unit  $n$  receives as inputs the outputs of other units, denoted with the set  $\operatorname{in}(n)$ . Each unit  $n$  encodes a function  $\phi_n$  defined as: (i)  $l_n(\operatorname{sc}(n))$  if  $n$  is an input unit, where  $l_n$  is a function over variables  $\operatorname{sc}(n) \subseteq \mathbf{X}$ , called its scope, (ii)  $\prod_{i \in \operatorname{in}(n)} \phi_i(\operatorname{sc}(\phi_i))$  if  $n$  is a product unit, and (iii)  $\sum_{i \in \operatorname{in}(n)} \theta_i \phi_i(\operatorname{sc}(\phi_i))$  if  $n$  is a sum unit, with  $\theta_i \in \mathbb{R}$  denoting the weighted sum parameters. The scope of a product or sum unit  $n$  is the union of the scopes of its inputs, i.e.,  $\operatorname{sc}(n) = \bigcup_{i \in \operatorname{in}(n)} \operatorname{sc}(i)$ .
Fig. 2 and Fig. A.1 show examples of circuits. Next, we introduce the two structural properties that enable efficient summation and Prop. 1 certifies that the aforementioned KGEs have these properties.
Definition 2 (Smoothness and Decomposability). A circuit is smooth if for every sum unit  $n$ , its input units depend all on the same variables, i.e.,  $\forall i,j \in \mathsf{in}(n) \colon \mathsf{sc}(i) = \mathsf{sc}(j)$ . A circuit is decomposable if the inputs of every product unit  $n$  depend on disjoint sets of variables, i.e.,  $\forall i,j \in \mathsf{in}(n) \, i \neq j \colon \mathsf{sc}(i) \cap \mathsf{sc}(j) = \emptyset$ .
For large real-world KGs it is reasonable to assume that  $|\mathcal{E}|\gg |\mathcal{R}|$
In general, the PLL is recognised as a proxy for MLE because, under certain assumptions, it is possible to show it can retrieve the MLE solution asymptotically with enough data [35].
Figure 2: Interpreting the score functions of CP, RESCAL, TUCKER, COMPLEX as circuits over 2-dimensional embeddings. Input, product and sum units are coloured in purple, orange and blue, respectively. Output sum units are labelled with the score functions, and their parameters are assumed to be 1, if not specified. The detailed construction is presented in Appendix A.1. Given a triple  $(s,r,o)$ , the input units map subject  $s$ , predicate  $r$  and object  $o$  to their embedding entries. Then, the products are evaluated before the weighted sum, which outputs the score of the input triple.
Proposition 1 (Score functions of KGE models as circuits). The computational graphs of the score functions  $\phi$  of CP, RESCAL, TUCKER and COMPLEX are smooth and decomposable circuits over  $\mathbf{X} = \{S,R,O\}$ , whose evaluation cost is  $\mathrm{cost}(\phi)\in \Theta (|\phi |)$ , where  $|\phi |$  denotes the number of edges in the circuit, also called its size. For example, the size of the circuit for CP is  $|\phi_{\mathrm{CP}}|\in \mathcal{O}(d)$
Appendix A.1 reports the complete proof by construction for the score functions of these models and the circuit sizes, while Fig. 2 illustrates them. Intuitively, since the score functions of the cited KGE models are based on products and sums as operators, they can be represented as circuits where input units map entities and predicates into the corresponding embedding components (similarly to look-up tables). As the inputs of each sum unit are product units that share the same scope  $\{S,R,O\}$  and fully decompose it in their input units, they satisfy smoothness and decomposability (Def. 2).
Smooth and decomposable circuits enable summing over all possible (partial) assignments to  $\mathbf{X}$  by (i) performing summations at input units over values in their domains, and (ii) evaluating the circuit once in a feed-forward way [13, 76]. This re-interpretation of score functions allows to "push" summations to the input units of a circuit, greatly reducing complexity, as detailed in the following proposition.
Proposition 2 (Efficient summations). Let  $\phi$  be a smooth and decomposable circuit over  $\mathbf{X} = \{S, R, O\}$  that encodes the score function of a KGE model. The sum  $\sum_{s \in \mathcal{E}} \sum_{r \in \mathcal{R}} \sum_{o \in \mathcal{E}} \phi(s, r, o)$  or any other summation over subjects, predicates or objects can be computed in time  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|) \cdot |\phi|)$ .
However, these sums are in logarithm space, as we have that  $p(s,r,o) \propto \exp \phi (s,r,o)$  (see Section 2). As a consequence, summation in this context does not correspond to marginalising variables in probability space. This drives us to reinterpret the semantics of these circuit structures as to operate directly in probability space, rather than in logarithm space, i.e., by encoding non-negative functions.
4 ...to Probabilistic Circuits
We now present how to reinterpret the semantics of the computational graphs of KGE score functions to directly output non-negative values for any input. That is, we cast them as probabilistic circuits (PCs) [13, 76, 19]. First, we define our subclass of PCs that encodes a possibly unnormalised probability distribution over triples in a KG, but allows for efficient marginalisation.
Definition 3 (Generative KGE circuit). A generative KGE circuit (GeKC) is a smooth and decomposable PC  $\phi_{\mathrm{pc}}$  over variables  $\mathbf{X} = \{S,R,O\}$  that encodes a probability distribution over triples, i.e.,  $\phi_{\mathrm{pc}}(s,r,o)\propto p(s,r,o)$  for any  $(s,r,o)\in \mathcal{E}\times \mathcal{R}\times \mathcal{E}$
Since our GeKCs are smooth and decomposable (Def. 2) they guarantee the efficient computation of  $Z$  in time  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|) \cdot |\phi_{\mathrm{pc}}|)$  (Prop. 2). This is in contrast with existing KGE models, for which it would require the evaluation of the whole score function on each possible triple (Section 2). For example, assume a non-negative CP score function  $\phi_{\mathrm{CP}}^{+}(s,r,o) = \langle \mathbf{e}_{s},\mathbf{w}_{r},\mathbf{e}_{o}\rangle \in \mathbb{R}_{+}$  for some embeddings  $\mathbf{e}_s,\mathbf{w}_r,\mathbf{e}_o\in \mathbb{R}^d$ . Then, we can compute  $Z$  by pushing the outer summations inside the trilinear product, i.e.,  $Z = \langle \sum_{s\in \mathcal{E}}\mathbf{e}_s,\sum_{r\in \mathcal{R}}\mathbf{w}_r,\sum_{o\in \mathcal{E}}\mathbf{e}_o\rangle$ , which can be done in time  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|)\cdot d)$ . In the following sections, we propose two ways to turn the computational graphs of CP, RESCAL, TUCKER and COMPLEX into GeKCs without additional space requirements.
Figure 3: GeKCs scale better. Time (in seconds) and peak GPU memory (in GiB as bubble sizes) required for computing the PLL objective and back-propagating through it for a single batch on ogbl-wikikg2, by increasing the batch size and number of entities. See Appendix C.4.3 for details.
4.1 Non-negative restriction of a score function
The most natural way of casting existing KGE models to GeKCs is by constraining the computational units of their circuit structures (Section 3) to output non-negative values only. We will refer to this conversion method as the non-negative restriction of a model. To achieve the non-negative restriction of CP, RESCAL and TUCKER we can simply restrict the embedding values and additional parameters in their score functions to be non-negative, as products and sums are operators closed in  $\mathbb{R}_+$ . Thus, each input unit  $n$  over variables  $S$  or  $O$  (resp.  $R$ ) in a non-negative GeKC encodes a function  $l_n$  (Def. 1) modelling an unnormalised categorical distribution over entities (resp. predicates). For example, each  $i$ -th entry  $e_{si}$  of the embedding  $\mathbf{e}_s \in \mathbb{R}_+^d$  associated to an entity  $s \in \mathcal{E}$  becomes a parameter of the  $i$ -th unnormalised categorical distribution over  $S$ . See Fig. C.1 for an example. We denote the non-negative restriction of these KGEs as  $\mathrm{CP}^+$ , RESCAL $^+$  and TUCKER $^+$ , respectively.
However, deriving COMPLEX $^{+}$  [65] by restricting the embedding values of COMPLEX to be non-negative is not sufficient, because its score function includes a subtraction as it operates on complex numbers. To overcome this, we re-parameterise the imaginary part of each complex embedding to be always greater than or equal to its real part. Appendix C.2 details this procedure. Even if this reparametrisation allows for more flexibility, imposing non-negativity on GeKCs can restrict their ability to capture intricate interactions over subjects, predicates and objects given a fixed number of learnable parameters [67]. We empirically confirm this in our experiments in Section 7. Therefore, we propose an alternative way of representing KGEs as PCs via squaring.
4.2 Squaring the score function
Squaring works by taking the score function of a KGE model  $\phi$ , and multiplying it with itself to obtain  $\phi^2 = \phi \cdot \phi$ . Note that  $\phi^2$  would be guaranteed to be a PC, as it always outputs non-negative values. The challenge is to represent the product of two circuits as a smooth and decomposable PC as to guarantee efficient marginalisation (Prop. 2). In general, this task is known to be #P-hard [76].
However, it can be done efficiently if the two circuits are compatible [76], as we further detail in Appendix B.1. Intuitively, the circuit representations of the score functions  $\phi$  of CP, RESCAL, TUCKER and COMPLEX (see Fig. 2) are simple enough that every product unit is defined over the same scope  $\{S,R,O\}$  and fully decomposes it on its input units. As such, these circuits can be easily multiplied with any other smooth and decomposable circuit, a property also known as omni-compatibility [76]. This property enables us to build the squared version of these KGE models, which we denote as  $\mathrm{CP}^2$ , RESCAL $^2$ , TUCKER $^2$  and COMPLEX $^2$ , as PCs that are still smooth and decomposable. Note that these squared GeKCs do allow for negative parameters, and hence can be more expressive. The next theorem, instead, guarantees that we can normalize them efficiently.
Theorem 1 (Efficient summations on squared GeKCs). Performing summations as stated in Prop. 2 on  $\mathrm{CP}^2$ , RESCAL $^2$ , TUCKER $^2$  and COMPLEX $^2$  can be done in time  $\mathcal{O}(|\mathcal{E}| + |\mathcal{R}|) \cdot |\phi|^2$ .
For instance, the partition function  $Z$  of  $\mathrm{CP}^2$  with embedding size  $d$  would require  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|)\cdot d^2)$  operations to be computed, while a simple feed-forward pass for a batch of  $|B|$  triples is still  $\mathcal{O}(|B|\cdot d)$ . While in this case marginalisation requires an increase in complexity that is quadratic in  $d$ , it is still faster than the brute force approach to compute  $Z$  (see Section 2 and Appendix C.4.1).
<sup>3</sup>In fact, even though  $\phi^2$  can be easily built by introducing a product unit over two copies of  $\phi$  as sub-circuits, the resulting circuit would be not decomposable (Def. 2), as the sub-circuits are defined on the same scope.
Quickly distilling KGEs to squared GeKCs. Consider a squared GeKC obtained by initialising its parameters with those of its energy-based KGE counterpart. If the score function of the original KGE model always assigns non-negative scores to triples, then the "distilled" squared GeKC will output the same exact ranking of the original model for the answers to any link prediction queries. Although the premise of the non-negativity of the scores might not be guaranteed, we observe that, in practice, learned KGE models do assign positive scores to all or most of the triples of common KGs (see Appendix D). Therefore, we can use this observation to either instantly distil a GeKC or provide a good heuristic to initialise its parameters and fine-tune them (by MLE or PLL maximisation). In both cases, the result is that we can convert the original KGE model into a GeKC that provides comparable probabilities, enable efficient marginalisation, sampling, and the integration of logical constraints with little or no loss of performance for link prediction (Section 7.1).
4.3 On the Training Efficiency of GeKCs