RotatE

class RotatE(triples_factory, embedding_dim=200, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

Bases: pykeen.models.base.EntityRelationEmbeddingModel

An implementation of RotatE from [sun2019].

RotatE models relations as rotations from head to tail entities in complex space:

\[\textbf{e}_t= \textbf{e}_h \odot \textbf{r}_r\]

where \(\textbf{e}, \textbf{r} \in \mathbb{C}^{d}\) and the complex elements of \(\textbf{r}_r\) are restricted to have a modulus of one (\(\|\textbf{r}_r\| = 1\)). The interaction model is then defined as:

\[f(h,r,t) = -\|\textbf{e}_h \odot \textbf{r}_r - \textbf{e}_t\|\]

which allows to model symmetry, antisymmetry, inversion, and composition.

See also

Initialize the entity embedding model.

Parameters

relation_dim – The relation embedding dimensionality. If not given, defaults to same size as entity embedding dimension.

See also

Constructor of the base class pykeen.models.Model

See also

Constructor of the base class pykeen.models.EntityEmbeddingModel

Attributes Summary

hpo_default

The default strategy for optimizing the model’s hyper-parameters

Methods Summary

interaction_function(h, r, t)

Evaluate the interaction function of ComplEx for given embeddings.

score_h(rt_batch)

Forward pass using left side (head) prediction.

score_hrt(hrt_batch)

Forward pass.

score_t(hr_batch)

Forward pass using right side (tail) prediction.

Attributes Documentation

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 1024, 'low': 32, 'q': 16, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

Methods Documentation

static interaction_function(h, r, t)[source]

Evaluate the interaction function of ComplEx for given embeddings.

The embeddings have to be in a broadcastable shape.

WARNING: No forward constraints are applied.

Parameters
  • h (FloatTensor) – shape: (…, e, 2) Head embeddings. Last dimension corresponds to (real, imag).

  • r (FloatTensor) – shape: (…, e, 2) Relation embeddings. Last dimension corresponds to (real, imag).

  • t (FloatTensor) – shape: (…, e, 2) Tail embeddings. Last dimension corresponds to (real, imag).

Return type

FloatTensor

Returns

shape: (…) The scores.

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Raises

NotImplementedError – If the method was not implemented for this class.

Return type

FloatTensor

Returns

shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.