SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Model	Complexity of log p(o | s, r)	Complexity of Z
Time	Space	Time	Space
CP	O(|E| · |B| · d)	O(|E| · |B|)	O(|E|^2 · |R| · d)	O(d)
RESCAL	O(|E| · |B| · d + |B| · d^2)	O(|E| · |B|)	O(|E|^2 · |R| · d^2)	O(d^2)
TUCKER	O(|E| · |B| · de + |B| · de^2 · dr)	O(|E| · |B|)	O(|E|^2 · |R| · de^2 · dr)	O(de^2 · dr)
COMPLEX	O(|E| · |B| · d)	O(|E| · |B|)	O(|E|^2 · |R| · d)	O(d)
CP+	O((|E| + |B|) · d)	O(|B| · d)	O((|E| + |R|) · d)	O(d)
RESCAL+	O((|E| + |B| · d) · d)	O(|B| · d^2)	O((|E| + |R| · d) · d)	O(d^2)
TUCKER+	O((|E| + |B| · de · dr) · de)	O(|B| · de · dr)	O(|E| · de + |R| · dr + de^2 · dr)	O(de^2 · dr)
COMPLEX+	O((|E| + |B|) · d)	O(|B| · d)	O((|E| + |R|) · d)	O(d)
CP^2	O((|E| + |B|) · d^2)	O(|B| · d)	O((|E| + |R|) · d^2)	O(d^2)
RESCAL^2	O((|E| + |B|) · d^2)	O(|B| · d^2)	O((|E| + |R| · d) · d^2)	O(|R| · d^2)
TUCKER^2	O((|E| + |B| · dr) · de^2)	O(|B| · de · dr)	O(|E| · de^2 + |R| · dr^2 + de^2 · dr)	O(d^2·dr)
COMPLEX^2	O((|E| + |B|) · d^2)	O(|B| · d)	O((|E| + |R|) · d^2)	O(d^2)
C.4.1 Computing the Partition Function
In this section we derive the computational complexity of computing the partition function for GeKCs obtained via squaring (Section 4.2). For a summary of these complexities, see Table C.1.
$\mathbf{CP}^2$  and  $\mathbf{COMPLEX}^2$ . Here we derive the partition function of  $\mathbf{CP}^2$ . For  $\mathbf{COMPLEX}^2$  the derivation is similar, as the score function of  $\mathbf{COMPLEX}$  can be written in terms of trilinear products just like  $\mathbf{CP}$  (see Eq. (5)). The score function  $\phi_{\mathrm{CP}^2}$  encoded by  $\mathbf{CP}^2$  can be written as
$$
\phi_ {\mathrm {C P} ^ {2}} (s, r, o) = \left\langle \mathbf {e} _ {s}, \mathbf {w} _ {r}, \mathbf {e} _ {o} \right\rangle^ {2} = \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {d} e _ {s i} e _ {s j} w _ {r i} w _ {r j} e _ {o i} e _ {o j}
$$
where  $\mathbf{e}_s, \mathbf{e}_o \in \mathbb{R}^d$  (resp.  $\mathbf{w}_r \in \mathbb{R}^d$ ) are rows of the matrices  $\mathbf{U}, \mathbf{V} \in \mathbb{R}^{|\mathcal{E}| \times d}$  (resp.  $\mathbf{W} \in \mathbb{R}^{|\mathcal{R}| \times d}$ ), which associate to each entity (resp. predicate) a vector. By leveraging the einsum notation for
brevity, the partition function of  $\phi_{\mathrm{CP}^2}$  can be written as
$$
\begin{array}{l} Z = \sum_ {s \in \mathcal {E}} \sum_ {r \in \mathcal {R}} \sum_ {o \in \mathcal {E}} \phi_ {\mathrm {C P} ^ {2}} (s, r, o) = \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {d} \left(\sum_ {s \in \mathcal {E}} e _ {s i} e _ {s j}\right) \left(\sum_ {r \in \mathcal {R}} w _ {r i} w _ {r j}\right) \left(\sum_ {o \in \mathcal {E}} e _ {o i} e _ {o j}\right) \\ = \mathbf {U} _ {i j} ^ {\prime} \mathbf {W} _ {i j} ^ {\prime} \mathbf {V} _ {i j} ^ {\prime} \\ \end{array}
$$
where  $\mathbf{U}' = \mathbf{U}^\top \mathbf{U}$ ,  $\mathbf{W}' = \mathbf{W}^\top \mathbf{W}$  and  $\mathbf{V}' = \mathbf{V}^\top \mathbf{V}$  are  $d \times d$  matrices. With the simplest algorithm for matrix multiplication, we recover that computing  $Z$  requires time  $\mathcal{O}(|\mathcal{E}| \cdot d^2 + |\mathcal{R}| \cdot d^2)$  and additional space  $\mathcal{O}(d^2)$ .
RESCAL<sup>2</sup>. The score function  $\phi_{\mathrm{RESCAL}^2}$  encoded by RESCAL<sup>2</sup> can be written as
$$
\phi_ {\text {R E S C A L}} ^ {2} (s, r, o) = \left(\mathbf {e} _ {s} ^ {\top} \mathbf {W} _ {r} \mathbf {e} _ {o}\right) ^ {2} = \sum_ {(i, j, k, l) \in [ d ] ^ {4}} e _ {s i} e _ {s k} w _ {r i j} w _ {r k l} e _ {o j} e _ {o l}
$$
where  $[d]$  denotes the set  $\{1, \ldots, d\}$ ,  $\mathbf{e}_s, \mathbf{e}_o \in \mathbb{R}^d$  are rows of the matrix  $\mathbf{E} \in \mathbb{R}^{|\mathcal{E}| \times d}$  and  $\mathbf{W}_r \in \mathbb{R}^{d \times d}$  are slices along the first mode of the tensor  $\mathcal{W} \in \mathbb{R}^{|\mathcal{R}| \times d \times d}$ , which consists of stacked matrix embeddings associated to predicates. The partition function of  $\phi_{\mathrm{RESCAL}^2}$  can be written as
$$
\begin{array}{l} Z = \sum_ {s \in \mathcal {E}} \sum_ {r \in \mathcal {R}} \sum_ {o \in \mathcal {E}} \phi_ {\text {R E S C A L} ^ {2}} (s, r, o) = \sum_ {(i, j, k, l) \in [ d ] ^ {4}} \left(\sum_ {s \in \mathcal {E}} e _ {s i} e _ {s k}\right) \left(\sum_ {r \in \mathcal {R}} w _ {r i j} w _ {r k l}\right) \left(\sum_ {o \in \mathcal {E}} e _ {o j} e _ {o l}\right) \\ = \mathbf {E} _ {i k} ^ {\prime} \mathcal {W} _ {r i j} \mathcal {W} _ {r k l} \mathbf {E} _ {j l} ^ {\prime} \tag {9} \\ \end{array}
$$
where  $\mathbf{E}^{\prime} = \mathbf{E}^{\top}\mathbf{E}\in \mathbb{R}^{d\times d}$ . The complexity of computing  $Z$  depends on the order of tensor contractions in the einsum operation showed in Eq. (9). By optimising the order of tensor contractions (e.g., by using software libraries like opt_einsum), we retrieve that computing  $Z$  requires time  $\mathcal{O}(|\mathcal{E}|\cdot d^2 +|\mathcal{R}|\cdot d^3)$  and additional space  $\mathcal{O}(|\mathcal{R}|\cdot d^2)$ . Notice that the time complexity here is slightly lower than the theoretical upper bound given in Thm. 1, which would be  $\mathcal{O}(|\mathcal{E}|\cdot d^2 +|\mathcal{R}|\cdot d^4)$ .
$\mathbf{T}\mathbf{U}\mathbf{C}\mathbf{K}\mathbf{E}\mathbf{R}^{2}$ . Lastly, we present the derivation of the partition function for  $\mathrm{T}\cup \mathrm{C}\cup \mathrm{K}\mathrm{E}\mathrm{R}^{2}$ . The score function  $\phi_{\mathrm{T}\cup \mathrm{C}\cup \mathrm{K}\mathrm{E}\mathrm{R}^{2}}$  encoded by  $\mathrm{T}\cup \mathrm{C}\cup \mathrm{K}\mathrm{E}\mathrm{R}^{2}$  can be written as
$$
\begin{array}{l} \phi_ {\mathrm {T U C K E R} ^ {2}} (s, r, o) = \left(\mathcal {T} \times_ {1} \mathbf {e} _ {s} \times_ {2} \mathbf {w} _ {r} \times_ {3} \mathbf {e} _ {o}\right) ^ {2} \\ = \sum_{\substack{(i,j,k)\in \\ [d_{e}]\times [d_{r}]\times [d_{e}]}}\sum_{\substack{(l,m,n)\in \\ [d_{e}]\times [d_{r}]\times [d_{e}]}}\tau_{ijk}\tau_{lmn}e_{si}e_{sl}w_{rj}w_{rm}e_{ok}e_{on} \\ \end{array}
$$
where  $\mathbf{e}_s, \mathbf{e}_o \in \mathbb{R}^{d_e}$  are rows of the matrix  $\mathbf{E} \in \mathbb{R}^{|\mathcal{E}| \times d_e}$ ,  $\mathbf{w}_r$  is a row of the matrix  $\mathbf{W} \in \mathbb{R}^{|\mathcal{R}| \times d_r}$ , and  $\mathcal{T} \in \mathbb{R}^{d_e \times d_r \times d_e}$  denotes the core tensor. The partition function of  $\phi_{\mathrm{TUCKER}^2}$  can be written as
$$
\begin{array}{l} Z = \sum_ {s \in \mathcal {E}} \sum_ {r \in \mathcal {R}} \sum_ {o \in \mathcal {E}} \phi_ {\mathrm {T U C K E R} ^ {2}} (s, r, o) \\ = \sum_{\substack{(i,j,k)\in \\ [d_{e}]\times [d_{r}]\times [d_{e}]}}\sum_{\substack{(l,m,n)\in \\ [d_{e}]\times [d_{r}]\times [d_{e}]}}\tau_{ijk}\tau_{lmn}\left(\sum_{s\in \mathcal{E}}e_{si}e_{sl}\right)\left(\sum_{r\in \mathcal{R}}w_{rj}w_{rm}\right)\left(\sum_{o\in \mathcal{E}}e_{ok}e_{on}\right) \\ = \mathcal {T} _ {i j k} \mathcal {T} _ {l m n} \mathbf {E} _ {i l} ^ {\prime} \mathbf {W} _ {j m} ^ {\prime} \mathbf {E} _ {k n} ^ {\prime} \tag {10} \\ \end{array}
$$
where  $\mathbf{E}' = \mathbf{E}^\top \mathbf{E} \in \mathbb{R}^{d_e \times d_e}$  and  $\mathbf{W}' = \mathbf{W}^\top \mathbf{W} \in \mathbb{R}^{d_r \times d_r}$ . Similarly to RESCAL<sup>2</sup>, by optimising the order of tensor contractions in the einsum operation showed in Eq. (10), we retrieve that computing  $Z$  requires time  $\mathcal{O}(|\mathcal{E}| \cdot d_e^2 + |\mathcal{R}| \cdot d_r^2 + d_e^2 \cdot d_r)$  and additional space  $\mathcal{O}(d_e^2 \cdot d_r)$ . Similarly to RESCAL<sup>2</sup>, the time complexity is lower than the theoretical upper bound given in Thm. 1, which would be  $\mathcal{O}(|\mathcal{E}| \cdot d_e^2 + |\mathcal{R}| \cdot d_r^2 + d_e^4 \cdot d_r^2)$ .
C.4.2 Complexity of Computing the PLL Objective
In this section, we show that GeKCs enable to better scale the computation of the PLL objective (Eq. (1)) with respect to energy-based KGE models (see Section 2). We present this concept for CP and GeKCs derived from it (Section 4), as for the other score functions it is similar.
Complexity of the PLL objective on CP. Let  $\phi_{\mathrm{CP}}(s,r,o) = \langle \mathbf{e}_s,\mathbf{w}_r,\mathbf{e}_o\rangle = \sum_{i = 1}^d e_{si}w_{ri}e_{oi}$  be the score function of CP [40], where  $\mathbf{e}_s,\mathbf{e}_o\in \mathbb{R}^d$  (resp.  $\mathbf{w}_r\in \mathbb{R}^d$ ) are rows of the matrices  $\mathbf{U},\mathbf{V}\in \mathbb{R}^{|\mathcal{E}|\times d}$  (resp.  $\mathbf{W}\in \mathbb{R}^{|\mathcal{R}|\times d}$ ), which associate to each entity (resp. predicate) a vector. Given a training triple  $(s,r,o)$ , the computation of the term  $\log p(o\mid s,r) = \phi (s,r,o) - \log \sum_{o'\in \mathcal{E}}\exp \phi (s,r,o')$  requires evaluating  $\phi_{\mathrm{CP}}(s,r,o')$  for all objects  $o'\in \mathcal{E}$ . In order to fully exploit GPU parallelism [36], this is usually done with the matrix-vector multiplication  $\mathbf{V}(\mathbf{e}_s\odot \mathbf{w}_r)\in \mathbb{R}^{|\mathcal{E}|}$ , where  $\odot$  denotes the Hadamard product [40, 12]. Therefore, computing  $\log p(o\mid s,r)$  for each triple  $(s,r,o)$  in a mini-batch  $B\subset \mathcal{E}\times \mathcal{R}\times \mathcal{E}$  such that  $|\mathcal{E}|\gg |\mathcal{B}|$  requires time  $\mathcal{O}(|\mathcal{E}|\cdot |B|\cdot d)$  and space  $\mathcal{O}(|\mathcal{E}|\cdot |B|)$ . For the other terms of the PLL objective (i.e.,  $\log p(s\mid r,o)$  and  $\log p(r\mid s,o)$ ) the derivation is similar. Moreover, for real-world large KGs it is reasonable to assume that  $|\mathcal{E}|\gg |\mathcal{R}|$  and therefore the cost of computing  $\log p(r\mid s,o)$  is negligible.
Complexity of the PLL objective on GeKCs. GeKCs obtained from CP either by non-negative restriction (Section 4.1) or by squaring (Section 4.2) encode  $\phi_{\mathrm{pc}}(s,r,o)\propto p(s,r,o)$  for any input triple. As such, the component  $\log p(o\mid s,r)$  of the PLL objective can be written as
$$
\log p (o \mid s, r) = \log \phi_ {\mathrm {p c}} (s, r, o) - \log \sum_ {o ^ {\prime} \in \mathcal {E}} \phi_ {\mathrm {p c}} (s, r, o ^ {\prime}). \tag {11}
$$
The absence of the exponential function in the summed terms in Eq. (11) allows us to push the outer summation inside the circuit computing  $\phi_{\mathrm{pc}}(s,r,o)$ , and to sum over the input units relative to objects. For instance, for  $\mathrm{CP}^+$  we can write
$$
\sum_ {o ^ {\prime} \in \mathcal {E}} \phi_ {\mathrm {C P} ^ {+}} (s, r, o ^ {\prime}) = \sum_ {o ^ {\prime} \in \mathcal {E}} \sum_ {i = 1} ^ {d} e _ {s i} w _ {r i} e _ {o i} = \sum_ {i = 1} ^ {d} e _ {s i} w _ {r i} \left(\sum_ {o \in \mathcal {E}} e _ {o ^ {\prime} i}\right) = (\mathbf {e} _ {s} \odot \mathbf {w} _ {r}) ^ {\top} (\mathbf {V} ^ {\top} \mathbf {1} _ {\mathcal {E}})
$$
where  $\mathbf{e}_s, \mathbf{w}_r, \mathbf{e}_o \in \mathbb{R}_+^d$ ,  $\mathbf{V} \in \mathbb{R}_+^{|\mathcal{E}| \times d}$  denotes the matrix whose rows are object embeddings, and  $\mathbf{1}_{\mathcal{E}} = [1 \ldots 1]^{|\mathcal{E}|}$  is a vector of ones. Note that  $\mathbf{V}^\top \mathbf{1}_{\mathcal{E}} \in \mathbb{R}_+^d$  does not depend on the input triple. Therefore, given a mini-batch of triples  $B$ , computing  $\log p(o \mid s, r)$  requires time  $\mathcal{O}((| \mathcal{E}| + |B|) \cdot d)$  and space  $\mathcal{O}(|B| \cdot d)$ , which is much lower than the complexity on CP showed above, and we can still leverage GPU parallelism. For  $\mathrm{CP}^2$ , the complexity is similar to the derivation of the partition function complexity showed in Appendix C.4.1. That is, for  $\mathrm{CP}^2$  we can write
$$
\begin{array}{l} \sum_ {o ^ {\prime} \in \mathcal {E}} \phi_ {\mathrm {C P} ^ {2}} (s, r, o ^ {\prime}) = \sum_ {o ^ {\prime} \in \mathcal {E}} \left(\sum_ {i = 1} ^ {d} e _ {s i} w _ {r i} e _ {o i}\right) ^ {2} = \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {d} e _ {s i} e _ {s j} w _ {r i} w _ {r j} \left(\sum_ {o ^ {\prime} \in \mathcal {E}} e _ {o ^ {\prime} i} e _ {o ^ {\prime} j}\right) \\ \mathbf {\Sigma} = \left(\mathbf {e} _ {s} \odot \mathbf {w} _ {r}\right) ^ {\top} (\mathbf {V} ^ {\top} \mathbf {V}) \left(\mathbf {e} _ {s} \odot \mathbf {w} _ {r}\right) \\ \end{array}
$$
where  $\mathbf{e}_s, \mathbf{w}_r, \mathbf{e}_o \in \mathbb{R}^d$ ,  $\mathbf{V} \in \mathbb{R}^{|\mathcal{E}| \times d}$ . Note that the matrix  $\mathbf{V}^\top \mathbf{V} \in \mathbb{R}^{d \times d}$  does not depend on the input triple. Therefore, given a mini-batch of triples  $B$ , computing  $\log p(o \mid s, r)$  requires time  $\mathcal{O}((|E| + |B|) \cdot d^2)$  and space  $\mathcal{O}(|B| \cdot d)$ . While the time complexity is quadratic in the embedding size  $d$ , it is still much lower than the time complexity on CP. A similar discussion can also be carried out for the other KGE models and the corresponding GeKCs, which retrieves the complexities showed in Table C.1.
C.4.3 Training Speed-up Benchmark Details
In this section we report the details about the training benchmark on COMPLEX, COMPLEX $^+$  and COMPLEX $^2$ , whose results are showed in Fig. 3. We measure time and peak GPU memory usage required for computing the PLL objective (Eq. (1)) and to do an optimisation step for a single batch on ogbl-wikikg2 [32], a large knowledge graph with millions of entities (see Table F.1). We fix the embedding size to  $d = 100$  for the three models. For the benchmark with increasing batch size, we keep all the entities and increase the batch size from 100 to 5000. For the benchmark with increasing number of entities, we keep the batch size fixed to 500 (the maximum allowed for COMPLEX by our GPUs) and progressively increase the number of entities, from about  $3 \cdot 10^5$  to  $2.5 \cdot 10^6$ . We report the average time over 25 independent runs on a single Nvidia RTX A6000 with 48 GiB of memory.
D Distribution of Scores
Figure D.1: Scores are mostly non-negative. Histograms of the scores assigned by COMPLEX to existing validation triples (left) and their perturbation (right) on three data sets. The vast majority of triple scores are non-negative, suggesting that squaring them has minimal effect on the rankings.
In Fig. D.1 we show the histograms of the scores assigned to the validation triples and their perturbations of three data sets (see Appendix F.1). Following Socher et al. [59], we generate triple perturbations that are challenging for link prediction. That is, for each validation triple  $(s,r,o)$ , the corresponding perturbation  $(s,r,\hat{o})$  is obtained by replacing the object with a random entity that has appeared at least once as an object in a training triple with predicate  $r$ . The bottom line is that scores are mostly non-negative, and hence can be used as a heuristic to effectively initialise GeKCs or quickly distil them (e.g., on FB15K-237), as we further discuss in Appendix F.5.1.
E Reconciling Knowledge Graph Embeddings Interpretations
Triples as boolean variables. KGE models such as CP, RESCAL, TUCKER and COMPLEX have been historically introduced as factorizations of a tensor-representation of a KG, which we discuss next. In fact, a KG  $\mathcal{G}$  can be represented as a 3-way binary tensor  $\mathbf{Y} \in \{0,1\}^{|\mathcal{E}| \times |\mathcal{R}| \times |\mathcal{E}|}$  in which every entry  $Y_{sro}$  is 1 if  $(s,r,o) \in \mathcal{G}$  and 0 otherwise [48]. Under this light, a KGE model like RESCAL factorises every slice  $\mathbf{Y}_r \in \{0,1\}^{|\mathcal{E}| \times |\mathcal{E}|}$ , corresponding to the predicate  $r$  as  $\mathbf{EW}_r\mathbf{E}^\top$  where  $\mathbf{E}$  is an  $|\mathcal{E}| \times d$  matrix comprising the entity embeddings and  $\mathbf{W}_r$  is an  $d \times d$  matrix containing the embeddings for the predicate  $r$ . To deal with uncertainty and incomplete KGs,  $Y_{sro}$  can be interpreted as a Bernoulli random variable. As such, its distribution becomes  $p(Y_{sro} = 1 \mid s,r,o)$  which is usually modelled as  $\sigma(\phi(s,r,o))$ , where  $\sigma$  denotes the logistic function and  $\phi$  is the score function of a KGE model. Note that this distribution of triple introduces  $|\mathcal{E} \times \mathcal{R} \times \mathcal{E}|$  random variables, one for each possible triple.
KGE models as estimators of a distribution over KGs. At the same time, the interpretation of triples as boolean variables induces a distribution over possible KGs,  $q(\mathcal{G})$ , which is the distribution over all possible binary tensors  $p(\mathbf{Y})$ . The probability of a KG  $\mathcal{G}$  can therefore be computed as the product of the likelihoods of all variables  $Y_{sro}$ , i.e.,  $q(\mathcal{G}) = \prod_{(s,r,o)\in \mathcal{G}}p(Y_{sro} = 1\mid s,r,o) \cdot \prod_{(s,r,o)\notin \mathcal{G}}p(Y_{sro} = 0\mid s,r,o)$ . Note that (re-)normalising this distribution is intractable in general, as it would require summing over all possible  $2^{|\mathcal{E}\times \mathcal{R}\times \mathcal{E}|}$  binary tensors. This is why historically KGE models have been interpreted as energy-based models, by directly optimising for  $\phi (s,r,o)$ , interpreted as the negated energy associated to every triple, and not  $p(Y_{sro} = 1\mid s,r,o)$  (see Section 2). This has been done via negative sampling or other contrastive learning objectives [6, 7]. We point out that this very same interpretation can be found in the literature of probabilistic logic programming [25], probabilistic databases (PDBs) [14] and statistical relational learning (see Section 6) where the distribution over possible "worlds" is over sets of boolean assignments to ground atoms or facts, or tuples in a PDB, each interpreted as Bernoulli random variables.
Estimating a distribution over triples. In this work, instead, we interpret existing KGE models and our GeKCs as models that encode a possibly unnormalised probability distribution over three random variables,  $S, R, O$ , which induces a distribution over triples that is tractable to renormalise.<sup>7</sup>
The polynomial cost of renormalising an energy-based KGE is unfortunately infeasible for real-world KGs, see Section 2.
To reconcile these two perspectives, we interpret the probability of a triple  $p(s,r,o)$  to be proportional to the probability of all KGs  $\mathcal{G}$  where  $(s,r,o)$  holds, i.e., those  $\mathcal{G}$  such that  $(s,r,o) \in \mathcal{G}$ . Intuitively, a triple will be more probable to exist if it does exist in highly probable KGs. More formally, given  $q$  a probability distribution over KGs, we define  $p$  as an unnormalised probability distribution over triples, i.e.,  $\mu(s,r,o) \propto p(s,r,o)$ , where
$$
\mu (s, r, o) = \sum_ {\substack {\mathcal {G} \in \mathcal {H} \\ (s, r, o) \in \mathcal {G}}} q (\mathcal {G}) = \sum_ {\mathcal {G} \in \mathcal {H}} q (\mathcal {G}) \cdot \mathbb {1} \{(s, r, o) \in \mathcal {G} \} = \mathbb {E} _ {\mathcal {G} \sim q} [ \mathbb {1} \{(s, r, o) \in \mathcal {G} \} ] \tag{12}
$$
and  $\mathcal{H} = 2^{\mathcal{E}\times \mathcal{R}\times \mathcal{E}}$  denotes the set of all possible KGs. Computing the expectation in Eq. (12) exactly is equivalent to solving a weighted model counting (WMC) problem [10], where we sum the probabilities of all possible KGs containing  $(s,r,o)$ . Alternatively, it is equivalent to computing the probability of the simplest possible query in a PDB (i.e., asking for a single tuple), where each stored tuple is interpreted as an independent Bernoulli random variable  $Y_{sro}$ . Therefore, we have that  $\mu (s,r,o)$  is simply the likelihood that  $Y_{sro}$  is true, i.e.,  $p(Y_{sro} = 1\mid s,r,o)$ . Furthermore, the normalisation constant of  $\mu$  (Eq. (12)) can be written as
$$
Z = \sum_{\substack{(s,r,o)\in \\ \mathcal{E}\times \mathcal{R}\times \mathcal{E}}}\mu (s,r,o) = \sum_{\mathcal{G}\in \mathcal{H}}q(\mathcal{G})\cdot \sum_{\substack{(s,r,o)\in \\ \mathcal{E}\times \mathcal{R}\times \mathcal{E}}}\mathbb{1}\{(s,r,o)\in \mathcal{G}\} = \sum_{\mathcal{G}\in \mathcal{H}}q(\mathcal{G})\cdot |\mathcal{G}| = \mathbb{E}_{\mathcal{G}\sim q}[|\mathcal{G}|]
$$
which is the expected size of a KG according to the probability distribution  $q$ . Written in this way, however, computing  $Z$  through  $q(\mathcal{G})$  is intractable. For this reason, we directly encode  $\mu$  with GeKCs and compute  $Z$  by summing over all triples, and therefore without modelling  $q(\mathcal{G})$ .
Further interpretations in related works. Under the interpretation of a KG as a PDB, Friedman and Van den Broeck [26] further decompose the likelihood that  $Y_{sro}$  is true as
$$
p \left(Y _ {s r o} = 1 \mid s, r, o\right) = p \left(E _ {s} = 1 \mid s\right) \cdot p \left(T _ {r} = 1 \mid r\right) \cdot p \left(E _ {o} = 1 \mid o\right)
$$
where  $E_{s}, E_{o}, T_{r}$  are new Bernoulli variables that are assumed to be conditionally independent given the parameters of the PDB. That is, instead of introducing one random variable per triple, they introduce one random variable per entity and predicate. In this framework, they reinterpret the score function of DISTMULT, a simplified variant of CP, as an implicit circuit that models an unnormalized distribution over the collection of variables  $\mathbf{Z} = \{E_{u}\}_{u\in \mathcal{E}} \cup \{T_{r}\}_{r\in \mathcal{R}}$ , trained by negative sampling. This decomposition permits to compute the probability of any database query efficiently, which otherwise is known to be either a PTIME or a #P-hard problem, depending on the query type [15]. If we were to interpret our distribution  $\mu(S = s, R = r, O = o)$  as the unnormalized marginal distribution  $p(E_{s} = 1, T_{r} = 1, E_{o} = 1) = \sum_{\mathbf{z}'} p(E_{s} = 1, T_{r} = 1, E_{o} = 1, \mathbf{Z}' = \mathbf{z}')$ , where  $\mathbf{Z}' = \mathbf{Z} \setminus \{E_{s}, T_{r}, E_{o}\}$ , we could equivalently compute any probabilistic query efficiently. Note that under this interpretation, training our GeKCs by MLE over  $S, R, O$  would be equivalent to maximise a composite marginal log-likelihood [70] over  $\mathbf{Z}$ .
F Empirical Evaluation
F.1 Datasets Statistics
Table F.1 shows statistics of commonly-used datasets to benchmark KGE models for link prediction. We employ standard benchmark datasets [62, 21, 32] whose number of entities (resp. predicates) ranges from  $\approx 14\mathrm{k}$  to  $\approx 2.5\mathrm{M}$  (resp. from 11 to  $\approx 500$ ).
Table F.1: Dataset statistics. Statistics of multi-relational knowledge graphs: number of entities  $\left( {\left| \mathcal{E}\right| }\right)$  ,number of predicates  $\left( {\left| \mathcal{R}\right| }\right)$  ,number of training/validation/test triples.
Dataset	|E|	|R|	# Train	# Valid	# Test
FB15k-237	[62]	14,541	237	272,115	17,535	20,466
WN18RR	[21]	40,943	11	86,835	3,034	3,134
ogbl-biokg	[32]	93,773	51	4,763 · 103	163 · 103	163 · 103
ogbl-wikikg2	[32]	2.5 · 106	535	16,109 · 103	429 · 103	598 · 103
F.2 Metrics
Mean reciprocal rank and hits at  $k$ . Given a test triple  $(s, r, o)$ , we rank the possible object  $o'$  (resp. subject  $s'$ ) completions to link prediction queries  $(s, r, ?)$  (resp.  $(?, r, o)$ ) based on their scores in descending order. The position of the test triple  $(s, r, o)$  in the ranking of object (resp. subject) completed queries  $(s, r, o')$  (resp.  $(s', r, o)$ ) is then used to compute the mean reciprocal rank (MRR)
$$
\mathrm {M R R} = \frac {1}{2 | \mathcal {G} _ {\text {t e s t}} |} \sum_ {(s, r, o) \in \mathcal {G} _ {\text {t e s t}}} \left(\frac {1}{\operatorname {r a n k} (o \mid s , r)} + \frac {1}{\operatorname {r a n k} (s \mid r , o)}\right)
$$
where  $\mathcal{G}_{\mathrm{test}}$  denotes the set of test triples, and  $\mathrm{rank}(o\mid s,r),\mathrm{rank}(s\mid r,o)$  denote respectively the positions of the true completion  $(s,r,o)$  in the rankings of object and subject completed queries. The fraction of hits at  $k$  (Hits@  $k$  ) for  $k > 0$  is computed as
$$
\text {H i t s} @ k = \frac {1}{2 | \mathcal {G} _ {\text {t e s t}} |} \sum_ {(s, r, o) \in \mathcal {G} _ {\text {t e s t}}} \left(\mathbb {1} \left\{\operatorname {r a n k} (o \mid s, r) \leq k \right\} + \mathbb {1} \left\{\operatorname {r a n k} (s \mid r, o) \leq k \right\}\right).
$$
Consistently with existing works on link prediction [56, 12], the MRRs and Hits@  $k$  metrics are computed under the filtered setting, i.e., we rank true completed triples against potential ones that do not appear in the union of training, validation and test splits.