SlowGuess/NeurIPS_2023_OCR huggingface dataset:
GeKCs also offer an unexpected opportunity to better scale the computation of the PLL objective (Eq. (1)) on very large knowledge graphs. This is because computing the PLL for a batch of  $|B|$  triples with GeKCs obtained via non-negative restriction and by squaring (Sections 4.1 and 4.2) does not require storing a matrix of size  $\mathcal{O}(|\mathcal{E}| \cdot |B|)$  to fully exploit GPU parallelism [36]. For instance, in Appendix C.4.2 we show that computing the PLL for CP [40] with embedding size  $d$  requires time  $\mathcal{O}(|\mathcal{E}| \cdot |B| \cdot d)$  and additional space  $\mathcal{O}(|\mathcal{E}| \cdot |B|)$ . On the other hand, for  $\mathrm{CP}^2$  (resp.  $\mathrm{CP}^+$ ) it requires time  $\mathcal{O}((|\mathcal{E}| + |B|) \cdot d^2)$  (resp.  $\mathcal{O}((|\mathcal{E}| + |B|) \cdot d)$ ) and space  $\mathcal{O}(|B| \cdot d)$ . Table C.1 summarises similar reduced complexities for other instances of GeKCs, such as COMPLEX $^+$  and COMPLEX $^2$ . The reduced time and memory requirements with GeKCs allow us to use larger batch sizes and better scale to large knowledge graphs. Fig. 3 clearly highlights this when measuring the time and GPU memory required to train these models on a KG with millions of entities such as ogbI-wikikg2 [32].
4.4 Sampling new triples with GeKCs
GeKCs only allowing non-negative parameters, such as  $\mathrm{CP}^+$ , RESCAL $^+$  and TUCKER $^+$ , support ancestral sampling as sum units can be interpreted as marginalised discrete latent variables, similarly to the latent variable interpretation in mixture models [55, 53] (see Appendix C.3 for details). This is however not possible in general for COMPLEX $^+$  and GeKCs obtained by squaring, as negative parameters break this interpretation. Luckily, as these circuits still support efficient marginalisation (Prop. 2 and Thm. 1) and hence also conditioning, we can perform inverse transform sampling. That is, to generate a triple  $(s, r, o)$ , we can sample in an autoregressive fashion, e.g., first  $s \sim p(S)$ , then  $r \sim p(R \mid s)$  and  $o \sim p(O \mid s, r)$ , hence requiring only three marginalization steps.
5 Injection of Logical Constraints
Converting KGE models to PCs provides the opportunity to "embed" logical constraints in the neural link predictor such that (i) predictions are always guaranteed to satisfy the constraints at test time, and (ii) training can still be done by efficient MLE (or PLL). This is in stark contrast with previous approaches for KGEs, which relax the constraints or enforce them only at training time (see Section 6). Consider, as an example, the problem of integrating the logical constraints induced by a schema of a KG, i.e., enforcing that triples not satisfying a domain constraint have probability zero.
Definition 4 (Domain constraint). Given a predicate  $r \in \mathcal{R}$  and  $\kappa_S(r), \kappa_O(r) \subseteq \mathcal{E}$  the sets of all subjects and objects that are semantically coherent with respect to  $r$ , a domain constraint is a propositional logic formula defined as
$$
K _ {r} \equiv S \in \kappa_ {S} (r) \wedge R = r \wedge O \in \kappa_ {O} (r) \equiv (\vee_ {u \in \kappa_ {S} (r)} S = u) \wedge R = r \wedge (\vee_ {v \in \kappa_ {O} (r)} O = v). \tag {3}
$$
Given  $\mathcal{R} = \{r_1, \ldots, r_m\}$  a set of predicates, the disjunction  $K \equiv K_{r_1} \lor \ldots \lor K_{r_m}$  encodes all the domain constraints that are defined in a KG. An input triple  $(s, r, o)$  satisfies  $K$ , written as  $(s, r, o) \models K$ , if  $s \in \kappa_S(r)$  and  $o \in \kappa_O(r)$ . To design GeKCs such as their predictions always satisfy logical constraints (which might not be necessarily domain constraints), we follow Ahmed et al. [1] and define a score function to represent a probability distribution  $p_K$  that assigns probability mass only to triples that satisfy the constraint  $K$ , i.e.,  $\phi_{\mathrm{pc}}(s, r, o) \cdot c_K(s, r, o) \propto p_K(s, r, o)$ . Here,  $\phi_{\mathrm{pc}}$  is a GeKC and  $c_K(s, r, o) = 1 \{(s, r, o) \models K\}$  is an indicator function that ensures that zero mass is assigned to triples violating  $K$ . In words, we are "cutting" the support of  $\phi_{\mathrm{pc}}$ , as illustrated in Fig. 4.
Figure 4: Injection of domain constraints. Given a circuit  $c_{K}$  encoding domain constraints and a GeKC  $\phi_{\mathrm{pc}}$ , the probability assigned by the product circuit  $\phi_{\mathrm{pc}} \cdot c_{K}$  to the inconsistent triple showed in Section 1 is 0, and a positive probability is assigned to consistent triples only, e.g., for the interacts predicate those involving drugs (Ds) as subjects and proteins (Ps) as objects. Best viewed in colours.
Computing  $p_K(s, r, o)$  exactly but naively would require computing a new partition function  $Z_K = \sum_{s' \in \mathcal{E}} \sum_{r' \in \mathcal{R}} \sum_{o' \in \mathcal{E}} (\phi_{\mathrm{pc}}(s', r', o') \cdot c_K(s', r', o'))$ , which is impractical as previously discussed (Section 2). Instead, we compile  $c_K$  as a smooth and decomposable circuit, sometimes called a constraint or logical circuit [19, 1], by leveraging compilers from the knowledge compilation literature [18, 52]. In a nutshell,  $c_K$  is another circuit over variables  $S, R, O$  that outputs 1 if an input triple satisfies the encoded logical constraint  $K$  and 0 otherwise. See Def. A.2 for a formal definition of such circuits. Then, similarly to what we have showed for computing squared circuits that enable efficient marginalisation (Section 4.2), the satisfaction of compatibility between a GeKC  $\phi_{\mathrm{pc}}$  and a constraint circuit  $c_K$  enables us to compute  $Z_K$  efficiently, as certified by the following theorem.
Theorem 2 (Tractable integration of constraints in GeKCs). Let  $c_{K}$  be a constraint circuit encoding a logical constraint  $K$  over variables  $\{S, R, O\}$ . Then exactly computing the partition function  $Z_{K}$  of the product  $\phi_{\mathrm{pc}}(s, r, o) \cdot c_{K}(s, r, o) \propto p_{K}(s, r, o)$  for any GeKC  $\phi_{\mathrm{pc}}$  derived from CP, RESCAL, TUCKER or COMPLEX (Section 4) can be done in time  $\mathcal{O}((|\mathcal{E}| + |\mathcal{R}|) \cdot |\phi_{\mathrm{pc}}| \cdot |c_{K}|)$ .
In Prop. A.1 we show that the compilation of domain constraints  $K$  (Def. 4) is straightforward and results in a constraint circuit  $c_{K}$  having compact size. For example, the size of the constraint circuit encoding the domain constraints of ogbI-biokg is approximately  $|c_{K}| = 307\cdot 10^{3}$ . To put this number in perspective, the size of the circuit for COMPLEX with embedding size 1000 is the much larger  $375\cdot 10^{6}$ . Since  $|c_{K}|$  is much smaller, by the same argument on the efficiency of GeKCs obtained via squaring (Section 4.3) it results that the integration of logical constraints adds a negligible overhead.
SOTA KGEs and current limitations. A plethora of ways to represent and learn KGEs has been proposed, see [31] for a review. KGEs such as CP and COMPLEX are still the de facto go-to choices in many applications [5, 56, 12]. Works performing density estimation in embedding space [78, 11] can sample embeddings, but to sample triples one would need to train a decoder. Several works try to modify training for KGEs as to introduce a penalty for triples that do not satisfy given logical constraints [8, 39, 34, 42, 23, 30], or casting it as a min-max game [44]. Unlike our GeKCs (Section 5), none of these approaches guarantee that test-time predictions satisfy the constraints. Moreover, several heuristics have been proposed to calibrate the probabilistic predictions of KGE models ex-post [61, 79]. As showed in [64], the triple distribution can be modelled autoregressively as  $p(S,R,O) = p(S)\cdot p(O\mid S)\cdot p(R\mid S,O)$  where each conditional distribution is encoded by a neural network. However, differently from our GeKCs, integrating constraints exactly or computing any marginal (thus conditional) probability is inefficient. KGE models based on non-negative tensor decompositions [63] are equivalent to GeKCs obtained by non-negative restriction (Section 4.1), but are generally trained by minimizing different non-probabilistic losses.
Circuits. Circuits provide a unifying framework for several tractable probabilistic models such as sum-product networks (SPNs) and hidden Markov models, which are smooth and decomposable
Table 1: GeKCs are competitive with their energy-based counterparts. Best average test MRRs of CP, COMPLEX and GeKCs trained with the PLL and MLE objectives (Eqs. (1) and (2)). For standard deviations and training times see Table F.2.
Model	FB15k-237	WN18RR	ogbl-biokg
PLL	MLE	PLL	MLE	PLL	MLE
CP	0.310	—	0.105	—	0.831	—
CP+	0.237	0.230	0.027	0.026	0.496	0.501
CP2	0.315	0.282	0.104	0.091	0.848	0.829
COMPLEX	0.342	—	0.471	—	0.829	—
COMPLEX+	0.214	0.205	0.030	0.029	0.503	0.516
COMPLEX2	0.334	0.300	0.420	0.391	0.858	0.840
PCs [13], as well as compact logical representations [19, 1]. See [75, 13, 17] for an overview. PCs with negative parameters are also called non-monotonic [58], but are surprisingly not as well investigated as their monotonic counterparts, i.e., PCs with only non-negative parameters, at least from the learning perspective. Similarly to our construction for  $\mathrm{COMPLEX}^+$  (Appendix C.2), Dennis [20] constrains the output of the non-monotonic sub-circuits of a larger PC to be less than their monotonic counterparts. Squaring a circuit has been investigating for tractably computing several divergences [76] and is related to the Born-rule of quantum mechanics [51].
Circuits for relational data. Logical circuits to compile formulas in first-order logic (FOL) [25] have been used to reason over relational data, e.g. via exchangeability [68, 49]. Other formalisms such as tractable Markov Logic [77], probabilistic logic bases [50], relational SPNs [46, 45] and generative clausal networks [71] use underlying circuit-like structures to represent probabilistic models over a tractable fragment of FOL formulas. These works assume that every atom in a grounded formula is associated to a random variable, also called the possible world semantics in probabilistic logic programs [57] and databases (PDBs) [14]. In this semantics, TractOR [26] casts answering complex queries over KGEs as to performing inference in PDBs. Differently from these works, our GeKCs are models defined over only three variables (Section 2). In Appendix E we reconcile these two semantics by interpreting the probability of a triple to be proportional to that of all KGs containing it.
7 Empirical Evaluation
We aim to answer the following research questions: RQ1) are GeKCs competitive with commonly used KGEs for link prediction? RQ2) Does integrating domain constraints in GeKCs benefit training and prediction?: RQ3) how good are the triples sampled from GeKCs?
7.1 Link Prediction (RQ1)
Experimental setting. We evaluate GeKCs on standard KG benchmarks for link prediction<sup>4</sup>: FB15k-237 [62], WN18RR [21] and ogbl-biokg [32], whose statistics can be found in Appendix F.1. As usual [48, 56, 54], we assess the models for predicting objects (queries  $(s, r, ?)$ ) and subjects (queries  $(? , r , o )$  ), and report their mean reciprocal rank (MRR) and fraction of hits at  $k$  (Hits@  $k$  ) (see Appendix F.2). We remark that our aim in this Section is not to score the new state-of-the-art link prediction performance on these benchmarks. Instead, we aim to rigorously assess how close GeKCs can be to commonly used and reasonably tuned KGE models. We focus on CP and COMPLEX as they currently are the go-to models of choice for link prediction [40, 56, 12]. We compare them against our GeKCs  $\mathrm{CP}^{+}$ , COMPLEX $^{+}$ ,  $\mathrm{CP}^{2}$  and COMPLEX $^{2}$  (Section 4). Appendix F.4 collects all the details about the model hyperparameters and training for reproducibility.
Link prediction results. Table 1 reports the MRR and times for all benchmarks and models when trained by PLL or MLE. First,  $\mathrm{CP}^2$  and COMPLEX achieve competitive scores when compared to CP and COMPLEX. Moreover,  $\mathrm{CP}^2$  (resp. COMPLEX) always outperforms  $\mathrm{CP}^+$  (resp. COMPLEX), thus providing empirical evidence that negative embedding values are crucial for model expressiveness. Concerning times, Table F.2 shows that squared GeKCs can train much faster on large KGs (see Section 4.3):  $\mathrm{CP}^2$  and COMPLEX require less than half the training time of CP and COMPLEX on ogbl-biokg, while also unexpectedly scoring the current SOTA MRR on it. We experiment also on the much larger ogbl-wikikg2 KG [32], comprising millions of entities. Even more remarkably, we are
$^{4}$ Code is available at https://github.com/april-tools/gekcs.
5Across non-ensemble methods and according to the OGB leaderboard, updated at the time of this writing.
Model	k	Embedding size
10	50	200	1000
COMPLEX	1	99.68	99.90	99.93	99.94
20	99.81	99.79	99.85	99.91
100	99.60	99.44	99.60	99.77
COMPLEX2	1	82.50	94.22	99.30	99.50
20	86.50	96.70	99.42	99.64
100	90.66	97.71	99.23	98.78
D-COMPLEX2	1	100.00	100.00	100.00	100.00
20	100.00	100.00	100.00	100.00
100	100.00	100.00	100.00	100.00
Figure 5: GeKCs with domain constraints guarantee domain-consistent predictions. Semantic consistency scores (Sem@k) [33] on ogbl-biokg achieved by COMPLEX, COMPLEX $^2$  and its integration with domain constraints (D-COMPLEX $^2$ ) (left), and MRRs computed on test queries (right). COMPLEX infers  $200+$  triples violating constraints as the highest scoring completions ( $k = 1$ ).
able to score an MRR of 0.572 after just  $\sim 3$  hours with COMPLEX $^2$  trained by PLL with a batch size of  $10^{4}$  and embedding size  $d = 100$ . To put this in context, we were able to score 0.562 with the best configuration of COMPLEX but after  $\sim 3$  days, as we could not fit in memory more than a batch size 500. $^6$  The same trends are shown for the Hits@  $k$  (Table F.3) and likelihood (Table F.4) metrics.
Distilling GeKCs. Table F.5 reports the results achieved by  $\mathrm{CP}^2$  and COMPLEX $^2$  initialised with the parameters of learned CP and COMPLEX (see Section 4.2) and confirms we can quickly turn an EBM into GeKC, thus inheriting all the perks of being a tractable generative model.
Calibration study. We also measure how well calibrated the predictions of the models in Table 1 are, which is essential to ensure trustworthiness in critical tasks. For example, given a perfectly calibrated model, for all the triples predicted with a probability of  $80\%$ , exactly  $80\%$  of them would actually exist [79]. On all KGs but WN18RR, GeKCs achieve lower empirical calibration errors [29] and better calibrated curves than their counterparts, as we report in Appendix F.5.3. The worse performance of all models on WN18RR can be explained by the distribution shift that exists between its training and test split, which we better confirm in Section 7.3.
7.2 Integrating Domain Constraints (RQ2)
We focus on ogbl-biokg [32], as it contains the domain metadata for each entity (i.e., disease, drug, function, protein, or side effect). Given the entity domains allowed for each predicate, we formulate domain constraints as in Def. 4. First, we want to estimate how likely are the models to predict triples that do not satisfy the domain constraints. We focus on COMPLEX and COMPLEX $^2$ , as they have been shown to achieve the best results in Section 7.1 and introduce D-COMPLEX $^2$  as the constraints-aware version of COMPLEX $^2$  (Section 5). For each test query  $(s,r,?)$  (resp.  $(?,r,o))$ , we compute the Sem@k score [33] as the average percentage of triples in the first  $k$  positions of the rankings of potential object (resp. subject) completions that satisfy the domain constraints (see Appendix F.2).
Fig. 5 highlights how both COMPLEX and  $\mathrm{COMPLEX}^2$  systematically predict object (or subject) completions that violate domain constraints even for large embedding sizes. For instance, a Sem@1 score of  $99\%$  (resp.  $99.9\%$ ) means that  $\sim 3200$  (resp.  $\sim 320$ ) predicted test triples violate domain constraints. While for COMPLEX and  $\mathrm{COMPLEX}^2$  there is no theoretical guarantee of consistent predictions with respect to the domain constraints, D-COMPLEX always guarantee consistent predictions by design. Furthermore, we observe a significant improvement in terms of MRR when integrating constraints for smaller embedding sizes, as reported in Fig. 5.
7.3 Quality of sampled triples (RQ3)
Inspired by the literature on evaluating deep generative models for images, we propose a metric akin to the kernel Inception distance [4] to evaluate the quality of the triples we can sample with GeKCs.
Definition 5 (Kernel triple distance (KTD)). Given  $\mathbb{P}$ ,  $\mathbb{Q}$  two probability distributions over triples, and a positive definite kernel  $k\colon \mathbb{R}^h\times \mathbb{R}^h\to \mathbb{R}$ , we define the kernel triple distance KTD  $(\mathbb{P},\mathbb{Q})$  as
<sup>6</sup>A smaller ( $d = 50$ ) and highly tuned version of COMPLEX achieves 0.639 MRR but still after days [12].
Table 2: GeKCs trained by MLE generate new likely triples. Empirical KTD scores between test triples and triples generated by baselines and GeKCs trained with the PLL objective or by MLE (Eqs. (1) and (2)). Lower is better. For standard deviations see Table F.6.
Model	FB15k-237	WN18RR	ogbl-biokg
Training set	0.055	0.260	0.029
Uniform	0.589	0.766	1.822
NNMFAug	0.414	0.607	0.518
PLL	MLE	PLL	MLE	PLL	MLE
CP+	0.404	0.433	0.633	0.578	0.966	0.738
CP2	0.253	0.070	0.768	0.768	0.039	0.017
COMPLEX+	0.336	0.323	0.456	0.478	0.175	0.097
COMPLEX2	0.326	0.102	0.338	0.278	0.104	0.034
the squared maximum mean discrepancy [28] between triple latent representations obtained via a map  $\psi \colon \mathcal{E} \times \mathcal{R} \times \mathcal{E} \to \mathbb{R}^h$  that projects triples to an  $h$ -dimensional embedding, i.e.,
$$
\operatorname {K T D} (\mathbb {P}, \mathbb {Q}) = \mathbb {E} _ {x, x ^ {\prime} \sim \mathbb {P}} [ k (\psi (x), \psi (x ^ {\prime})) ] + \mathbb {E} _ {y, y ^ {\prime} \sim \mathbb {Q}} [ k (\psi (y), \psi (y ^ {\prime})) ] - 2 \cdot \mathbb {E} _ {x \sim \mathbb {P}, y \sim \mathbb {Q}} [ k (\psi (x), \psi (y)) ].
$$
An empirical estimate of the KTD score close to zero indicates that there is little difference between the two triple distributions  $\mathbb{P}$  and  $\mathbb{Q}$  (see Appendix F.3). For images,  $\psi$  is typically chosen as the last embedding of a SOTA neural classifier. We choose  $\psi$  to be the  $L_{2}$ -normed outputs of the product units of a circuit [73, 74], specifically the SOTA COMPLEX learned by Chen et al. [12] with  $h = 4000$ . We choose  $k$  as the polynomial kernel  $k(\mathbf{x},\mathbf{y}) = (\mathbf{x}^{\top}\mathbf{y} + 1)^{3}$ , following Binkowski et al. [4].
Table 2 shows the empirical KTD scores computed between the test triples and the generated ones, and Fig. F.1 visualises triple embeddings. We employ two baselines: a uniform probability distribution over all possible triples and NNMFAug [9], the only work to address triple sampling to the best of our knowledge. We also report the KTD scores for training triples as an empirical lower bound. Squared GeKCs achieve lower KTD scores with respect to the ones obtained by non-negative restriction, confirming again a better estimation of the joint distribution. In addition, they achieve far lower KTD scores than all competitors when learning by MLE (Eq. (2)), which justifies its usage as an objective. Lastly, we confirm the distribution shift on WN18RR: training set KTD scores are far from zero, but even in this challenging scenario, COMPLEX<sup>2</sup> scores KTD values that are closer to the training ones.
8 Conclusions and Future Work
We proposed to re-interpret the representation and learning of widely used KGE models such as CP, RESCAL, TUCKER and COMPLEX, as generative models, overcoming some of the classical limitation of their usual EBM interpretation (see Sections 1 and 2). GeKC-variants for other KGE models whose scores are multilinear maps can be readily devised in the same way. Moreover, we conjecture that other KGE models defining score functions having a distance-based semantics such as TransE [7] and RotatE [60] can be reinterpreted to be GeKCs as well. Our GeKCs open up a number of interesting future directions. First, we plan to investigate how the enhanced efficiency and calibration of GeKCs can help in complex reasoning tasks beyond link prediction [2]. Second, we can leverage the rich literature on learning the structure of circuits [72, 75] to devise smaller and sparser KGE circuit architectures that better capture the triple distribution or sporting structural properties that can make reasoning tasks other than marginalisation efficient [22, 16, 71].
We would like to acknowledge Iain Murray for his thoughtful feedback and suggestions on a draft of this work. In addition, we thank Xiong Bo and Ricky Zhu for pointing out related works in the knowledge graphs literature. AV was supported by the "UNREAL: Unified Reasoning Layer for Trustworthy ML" project (EP/Y023838/1) selected by the ERC and funded by UKRI EPSRC. RP was supported by the Graz Center for Machine Learning (GraML).
[1] Kareem Ahmed, Stefano Teso, Kai-Wei Chang, Guy Van den Broeck, and Antonio Vergari. Semantic probabilistic layers for neuro-symbolic learning. In Advances in Neural Information
[76] Antonio Vergari, YooJung Choi, Anji Liu, Stefano Teso, and Guy Van den Broeck. A compositional atlas of tractable circuit operations for probabilistic inference. In Advances in Neural Information Processing Systems 34 (NeurIPS), pages 13189-13201. Curran Associates, Inc., 2021.
[77] W Austin Webb and Pedro Domingos. Tractable probabilistic knowledge bases with existence uncertainty. In *Workshops at the Twenty-Seventh AAAI Conference on Artificial Intelligence*, 2013.
[78] Han Xiao, Minlie Huang, and Xiaoyan Zhu. Transg: A generative model for knowledge graph embedding. In Annual Meeting of the Association for Computational Linguistics, 2016.
[79] Ruiqi Zhu, Fangrong Wang, Alan Bundy, Xue Li, Kwabena Nuamah, Lei Xu, Stefano Mauceri, and Jeff Z. Pan. A closer look at probability calibration of knowledge graph embedding. In 11th International Joint Conference on Knowledge Graphs, page 104–109, 2023.
Figure A.1: Evaluation of circuit representations of score functions as in neural networks. Feed-forward evaluation of the TUCKER score function as a circuit over 2-dimensional embeddings and parameterised by the core tensor  $\mathcal{T}$  (see proof of Prop. 1 below) (a). Given a triple (loxoprofen, interacts, COX2), the input units (Def. 1) map subject, predicate and object to their embedding entries (in violet boxes). Then, the circuit is evaluated similarly to neural networks: the products (in orange) are evaluated before the weighted sum (in blue), which is parameterised by the core tensor values (in green) (b). The output of the circuit is the score of the input triple.
A.1 KGE Models as Circuits
Proof. We present the proof by construction for TUCKER [3], as CP [40], RESCAL [47] and COMPLEX [65] define score functions that are a specialisation of it [3] (see below). Given a triple  $(s, r, o) \in \mathcal{E} \times \mathcal{R} \times \mathcal{E}$ , the TUCKER score function computes
$$
\phi_ {\mathrm {T U C K E R}} (s, r, o) = \mathcal {T} \times_ {1} \mathbf {e} _ {s} \times_ {2} \mathbf {w} _ {r} \times_ {3} \mathbf {e} _ {o} = \sum_ {i = 1} ^ {d _ {e}} \sum_ {j = 1} ^ {d _ {r}} \sum_ {k = 1} ^ {d _ {e}} \tau_ {i j k} e _ {s i} w _ {r j} e _ {o k} \tag {4}
$$