Table of Contents
Fetching ...

How to Turn Your Knowledge Graph Embeddings into Generative Models

Lorenzo Loconte, Nicola Di Mauro, Robert Peharz, Antonio Vergari

TL;DR

This work reframes widely used knowledge graph embedding scores (e.g., CP, RESCAL, TuckER, ComplEx) as circuits and introduces generative KG circuits (GeKC) that yield tractable probabilistic models over triples. By enforcing non-negativity or squaring score outputs, GeKC becomes smooth, decomposable probabilistic circuits enabling exact maximum-likelihood estimation, efficient sampling, and guaranteed satisfaction of logical constraints via constraint circuits. The approach provides substantial training efficiency advantages and scales to graphs with millions of entities, while preserving competitive link prediction performance and offering richer calibration and sampling capabilities. It opens avenues for principled probabilistic reasoning, exact normalization, and constraint-aware generation in large relational knowledge graphs.

Abstract

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.

How to Turn Your Knowledge Graph Embeddings into Generative Models

TL;DR

This work reframes widely used knowledge graph embedding scores (e.g., CP, RESCAL, TuckER, ComplEx) as circuits and introduces generative KG circuits (GeKC) that yield tractable probabilistic models over triples. By enforcing non-negativity or squaring score outputs, GeKC becomes smooth, decomposable probabilistic circuits enabling exact maximum-likelihood estimation, efficient sampling, and guaranteed satisfaction of logical constraints via constraint circuits. The approach provides substantial training efficiency advantages and scales to graphs with millions of entities, while preserving competitive link prediction performance and offering richer calibration and sampling capabilities. It opens avenues for principled probabilistic reasoning, exact normalization, and constraint-aware generation in large relational knowledge graphs.

Abstract

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.
Paper Structure (75 sections, 7 theorems, 29 equations, 11 figures, 10 tables)

This paper contains 75 sections, 7 theorems, 29 equations, 11 figures, 10 tables.

Key Result

Proposition 1

The computational graphs of the score functions $\phi$ of CP, RESCAL, TuckER and ComplEx are smooth and decomposable circuits over $\bm{\mathrm{X}}=\{S,R,O\}$, whose evaluation cost is $\mathop{\mathrm{cost}}\limits(\phi)\in\Theta(|\phi|)$, where $|\phi|$ denotes the number of edges in the circuit,

Figures (11)

  • 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 (GeKC s) we need to either restrict its activations to be non-negative (in blue, see \ref{['sec:non-negative-restriction']}) or square it (in red, see \ref{['sec:non-monotonic-squaring']}).
  • 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 \ref{['app:kge-as-circuits']}. 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.
  • Figure 3: GeKC s 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 \ref{['app:training-benchmark-details']} for details.
  • Figure 4: Injection of domain constraints. Given a circuit $c_K$ encoding domain constraints and a GeKC $\phi_\mathsf{pc}$, the probability assigned by the product circuit $\phi_\mathsf{pc} \cdot c_K$ to the inconsistent triple showed in \ref{['sec:introduction']} 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.
  • Figure 5: GeKC s with domain constraints guarantee domain-consistent predictions. Semantic consistency scores (Sem@$k$) hubert2022new-strategies on ogbl-biokg achieved by ComplEx, ComplEx2 and its integration with domain constraints (d-ComplEx2) (left), and MRRs computed on test queries (right). ComplEx infers 200+ triples violating constraints as the highest scoring completions ($k=1$).
  • ...and 6 more figures

Theorems & Definitions (22)

  • Definition 1: Circuit choi2020pcvergari2021compositional
  • Definition 2: Smoothness and Decomposability
  • Proposition 1: Score functions of KGE models as circuits
  • Proposition 2: Efficient summations
  • Definition 3: Generative KGE circuit
  • Theorem 1: Efficient summations on squared GeKC s
  • Definition 4: Domain constraint
  • Theorem 2: Tractable integration of constraints in GeKC s
  • Definition 5: Kernel triple distance (KTD)
  • proof
  • ...and 12 more