Knowledge Graph Embedding (KGE)

K3-node provides a multi-backend implementation of Knowledge Graph Embedding (KGE) models and negative sampling triplet iterators under k3_node.layers.kge.

KGE models learn low-dimensional vector representations for entities $\mathcal{E}$ and relations $\mathcal{R}$ in multi-relational knowledge graphs, scoring triplets $(h, r, t)$ according to score functions $f(h, r, t)$.

Models

Base Class: KGEModel

KGEModel defines the base architecture for knowledge graph embeddings, handling entity/relation lookup, margin ranking loss, and negative sampling scoring.

from k3_node.layers.kge import KGEModel

`k3_node.layers.kge.KGEModel`

Bases: Layer

An abstract base class for implementing custom KGE models.

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes/entities in the graph.	required
`num_relations`	`int`	The number of relations in the graph.	required
`hidden_channels`	`int`	The hidden embedding size.	required
`sparse`	`bool`	Kept for API compatibility with PyG; has no effect since Keras optimizers do not distinguish sparse embedding gradients the way PyTorch does. (default: :obj:`False`)	`False`

`call(head_index, rel_type, tail_index)`

Returns the score for the given triplet.

Parameters:

Name	Description	Default
`head_index`	The head indices.	required
`rel_type`	The relation type.	required
`tail_index`	The tail indices.	required

`compile(optimizer)`

Sets the optimizer used by :meth:fit.

`evaluate(data, batch_size=20000, k=10)`

Ranks the true tail of every triplet in data among all entities and returns the mean rank, the mean reciprocal rank (MRR) and Hits@k.

`fit(data, epochs=1, batch_size=1000, validation_data=None, validation_batch_size=20000, verbose=1)`

Trains on the triplets of data (edge_index holds heads and tails, edge_type the relations), in shuffled mini-batches.

Parameters:

Name	Type	Description	Default
`data`	`Data`	The training knowledge graph.	required
`epochs`	`int`	Passes over all triplets. (default: `1`)	`1`
`batch_size`	`int`	Triplets per gradient step. (default: `1000`)	`1000`
`validation_data`	`Data`	Evaluated with :meth:`evaluate` after every epoch (this ranks all entities for every triplet, so keep it small).	`None`
`validation_batch_size`	`int`	Entities scored at once during validation.	`20000`
`verbose`	`int`	`0` is silent, otherwise one line is printed per epoch.	`1`

Returns:

Name	Type	Description
`dict`		The mean loss (and validation metrics) of every epoch.

`loader(head_index, rel_type, tail_index, **kwargs)`

Returns a mini-batch loader that samples a subset of triplets.

Parameters:

Name	Type	Description	Default
`head_index`		The head indices.	required
`rel_type`		The relation type.	required
`tail_index`		The tail indices.	required
`**kwargs`	`optional`	Additional arguments of :class:`k3_node.layers.kge.KGTripletLoader`, such as `batch_size`, `shuffle` or `drop_last`.	`{}`

`loss(head_index, rel_type, tail_index)`

Returns the loss value for the given triplet.

`random_sample(head_index, rel_type, tail_index)`

Randomly samples negative triplets by either replacing the head or the tail (but not both).

Parameters:

Name	Description	Default
`head_index`	The head indices.	required
`rel_type`	The relation type.	required
`tail_index`	The tail indices.	required

`reset_parameters()`

Resets all learnable parameters of the module.

`test(head_index, rel_type, tail_index, batch_size, k=10, log=True)`

Evaluates the model quality by computing Mean Rank, MRR and Hits@:math:k across all possible tail entities.

Parameters:

Name	Type	Description	Default
`head_index`		The head indices.	required
`rel_type`		The relation type.	required
`tail_index`		The tail indices.	required
`batch_size`	`int`	The batch size to use for evaluating.	required
`k`	`int`	The :math:`k` in Hits @ :math:`k`. (default: :obj:`10`)	`10`
`log`	`bool`	If set to :obj:`False`, will not print a progress bar to the console. (default: :obj:`True`)	`True`

TransE

Translational distance model ($h + r \approx t$) mapping entities and relations to a real vector space: $$d(h + r, t) = -\| \mathbf{e}_h + \mathbf{e}_r - \mathbf{e}_t \|_p$$

from k3_node.layers.kge import TransE

model = TransE(num_nodes=1000, num_relations=50, hidden_channels=64, margin=1.0, p_norm=1.0)

`k3_node.layers.kge.TransE`

Bases: KGEModel

The TransE model from the "Translating Embeddings for Modeling Multi-Relational Data" <https://proceedings.neurips.cc/paper/2013/file/ 1cecc7a77928ca8133fa24680a88d2f9-Paper.pdf>_ paper.

:class:TransE models relations as a translation from head to tail entities such that

.. math:: \mathbf{e}_h + \mathbf{e}_r \approx \mathbf{e}_t,

resulting in the scoring function:

.. math:: d(h, r, t) = - {| \mathbf{e}_h + \mathbf{e}_r - \mathbf{e}_t |}_p

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes/entities in the graph.	required
`num_relations`	`int`	The number of relations in the graph.	required
`hidden_channels`	`int`	The hidden embedding size.	required
`margin`	`float`	The margin of the ranking loss. (default: :obj:`1.0`)	`1.0`
`p_norm`	`float`	The order embedding and distance normalization. (default: :obj:`1.0`)	`1.0`
`sparse`	`bool`	Kept for API compatibility. (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.layers import TransE

head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
rel = np.random.randint(0, 5, size=(10,))
tail = np.random.randint(0, 20, size=(10,))

model = TransE(num_nodes=20, num_relations=5, hidden_channels=8)
score = model(head, rel, tail)  # plausibility score of every triple
print(tuple(score.shape))  # (10,)
loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
print(tuple(loss.shape))  # (): a scalar

RotatE

Knowledge graph embedding by relational rotation in complex space ($\mathbf{e}_t = \mathbf{e}_h \circ \mathbf{r}$ where $|\mathbf{r}_i| = 1$): $$d(h \circ r, t) = -\| \mathbf{e}_h \circ \mathbf{e}_r - \mathbf{e}_t \|$$

from k3_node.layers.kge import RotatE

model = RotatE(num_nodes=1000, num_relations=50, hidden_channels=64, margin=6.0)

`k3_node.layers.kge.RotatE`

Bases: KGEModel

The RotatE model from the "RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space" <https://arxiv.org/abs/ 1902.10197>_ paper.

:class:RotatE models relations as a rotation in complex space from head to tail such that

.. math:: \mathbf{e}_t = \mathbf{e}_h \circ \mathbf{e}_r,

resulting in the scoring function

.. math:: d(h, r, t) = - {| \mathbf{e}_h \circ \mathbf{e}_r - \mathbf{e}_t |}_p

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes/entities in the graph.	required
`num_relations`	`int`	The number of relations in the graph.	required
`hidden_channels`	`int`	The hidden embedding size.	required
`margin`	`float`	The margin of the ranking loss. (default: :obj:`1.0`)	`1.0`
`sparse`	`bool`	Kept for API compatibility. (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.layers import RotatE

head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
rel = np.random.randint(0, 5, size=(10,))
tail = np.random.randint(0, 20, size=(10,))

model = RotatE(num_nodes=20, num_relations=5, hidden_channels=8)
score = model(head, rel, tail)  # plausibility score of every triple
print(tuple(score.shape))  # (10,)
loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
print(tuple(loss.shape))  # (): a scalar

DistMult

Bilinear diagonal model capturing symmetric relation interactions: $$f(h, r, t) = \langle \mathbf{e}_h, \mathbf{e}_r, \mathbf{e}_t \rangle = \sum_{i} (\mathbf{e}_h)_i (\mathbf{e}_r)_i (\mathbf{e}_t)_i$$

from k3_node.layers.kge import DistMult

model = DistMult(num_nodes=1000, num_relations=50, hidden_channels=64, margin=1.0)

`k3_node.layers.kge.DistMult`

Bases: KGEModel

The DistMult model from the "Embedding Entities and Relations for Learning and Inference in Knowledge Bases" <https://arxiv.org/abs/1412.6575>_ paper.

:class:DistMult models relations as diagonal matrices, which simplifies the bi-linear interaction between the head and tail entities to the score function:

.. math:: d(h, r, t) = < \mathbf{e}_h, \mathbf{e}_r, \mathbf{e}_t >

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes/entities in the graph.	required
`num_relations`	`int`	The number of relations in the graph.	required
`hidden_channels`	`int`	The hidden embedding size.	required
`margin`	`float`	The margin of the ranking loss. (default: :obj:`1.0`)	`1.0`
`sparse`	`bool`	Kept for API compatibility. (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.layers import DistMult

head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
rel = np.random.randint(0, 5, size=(10,))
tail = np.random.randint(0, 20, size=(10,))

model = DistMult(num_nodes=20, num_relations=5, hidden_channels=8)
score = model(head, rel, tail)  # plausibility score of every triple
print(tuple(score.shape))  # (10,)
loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
print(tuple(loss.shape))  # (): a scalar

ComplEx

Complex embeddings for simple link prediction, handling asymmetric relations via Hermitian dot product: $$f(h, r, t) = \text{Re}(\langle \mathbf{e}_h, \mathbf{e}_r, \bar{\mathbf{e}}_t \rangle)$$

from k3_node.layers.kge import ComplEx

model = ComplEx(num_nodes=1000, num_relations=50, hidden_channels=64, margin=1.0)

`k3_node.layers.kge.ComplEx`

Bases: KGEModel

The ComplEx model from the "Complex Embeddings for Simple Link Prediction" <https://arxiv.org/abs/1606.06357>_ paper.

:class:ComplEx models relations as complex-valued bilinear mappings between head and tail entities using the Hermetian dot product. The entities and relations are embedded in different dimensional spaces, resulting in the scoring function:

.. math:: d(h, r, t) = Re(< \mathbf{e}_h, \mathbf{e}_r, \mathbf{e}_t>)

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes/entities in the graph.	required
`num_relations`	`int`	The number of relations in the graph.	required
`hidden_channels`	`int`	The hidden embedding size.	required
`sparse`	`bool`	Kept for API compatibility. (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.layers import ComplEx

head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
rel = np.random.randint(0, 5, size=(10,))
tail = np.random.randint(0, 20, size=(10,))

model = ComplEx(num_nodes=20, num_relations=5, hidden_channels=8)
score = model(head, rel, tail)  # plausibility score of every triple
print(tuple(score.shape))  # (10,)
loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
print(tuple(loss.shape))  # (): a scalar

Data Loading & Triplet Batching

KGTripletLoader

Framework-agnostic batch iterator for knowledge graph triplets $(h, r, t)$ with corrupt head / tail negative sampling:

from k3_node.layers.kge import KGTripletLoader

loader = KGTripletLoader(
    head_index=head_indices,
    rel_type=rel_types,
    tail_index=tail_indices,
    batch_size=256,
    shuffle=True
)

for head, rel, tail in loader:
    loss = model.loss(head, rel, tail)

`k3_node.layers.kge.KGTripletLoader`

A minimal, framework-agnostic batching iterator over knowledge-graph triplets, mirroring :class:torch.utils.data.DataLoader usage in :meth:k3_node.layers.kge.KGEModel.loader.

Parameters:

Name	Type	Description	Default
`head_index`		The head indices.	required
`rel_type`		The relation type.	required
`tail_index`		The tail indices.	required
`batch_size`	`int`	The batch size. (default: :obj:`1`)	`1`
`shuffle`	`bool`	If set to :obj:`True`, shuffles the triplets at every epoch. (default: :obj:`False`)	`False`
`drop_last`	`bool`	If set to :obj:`True`, drops the last incomplete batch. (default: :obj:`False`)	`False`