↔️🎱 Merge QuatE (#1472)

Merge QuatE functional form into module, and re-organize documentation.

See also: #1102

---------

Co-authored-by: Max Berrendorf <[email protected]>

Author: cthoyt

docs/source/reference/nn/utils.rst

@@ -3,5 +3,5 @@ Utilities
 .. automodule:: pykeen.nn.utils
     :members:
 
-.. automodule:: pykeen.nn.algebra
+.. automodule:: pykeen.nn.quaternion
     :members:

src/pykeen/models/unimodal/quate.py

@@ -4,50 +4,22 @@
 from typing import Any, ClassVar, Optional
 
 import torch
-from torch.nn import functional
 
 from ..nbase import ERModel
 from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE
 from ...losses import BCEWithLogitsLoss, Loss
+from ...nn import quaternion
 from ...nn.init import init_quaternions
 from ...nn.modules import QuatEInteraction
 from ...regularizers import LpRegularizer, Regularizer
-from ...typing import Constrainer, FloatTensor, Hint, Initializer
+from ...typing import Constrainer, Hint, Initializer
 from ...utils import get_expected_norm
 
 __all__ = [
     "QuatE",
 ]
 
 
-def quaternion_normalizer(x: FloatTensor) -> FloatTensor:
-    r"""
-    Normalize the length of relation vectors, if the forward constraint has not been applied yet.
-
-    Absolute value of a quaternion
-
-    .. math::
-
-        |a + bi + cj + dk| = \sqrt{a^2 + b^2 + c^2 + d^2}
-
-    L2 norm of quaternion vector:
-
-    .. math::
-        \|x\|^2 = \sum_{i=1}^d |x_i|^2
-                 = \sum_{i=1}^d (x_i.re^2 + x_i.im_1^2 + x_i.im_2^2 + x_i.im_3^2)
-    :param x:
-        The vector.
-
-    :return:
-        The normalized vector.
-    """
-    # Normalize relation embeddings
-    shape = x.shape
-    x = x.view(*shape[:-1], -1, 4)
-    x = functional.normalize(x, p=2, dim=-1)
-    return x.view(*shape)
-
-
 class QuatE(ERModel):
     r"""An implementation of QuatE from [zhang2019]_.
 
@@ -56,13 +28,17 @@ class QuatE(ERModel):
     $\textbf{e}_i, \textbf{r}_i \in \mathbb{H}^d$, and the plausibility score is computed using the
     quaternion inner product.
 
+    The representations are stored in an :class:`~pykeen.nn.representation.Embedding`.
+    Scores are calculated with :class:`~pykeen.nn.modules.QuatEInteraction`.
+
     .. seealso ::
 
         Official implementation: https://github.com/cheungdaven/QuatE/blob/master/models/QuatE.py
     ---
     citation:
         author: Zhang
         year: 2019
+        arxiv: 1904.10281
         link: https://arxiv.org/abs/1904.10281
         github: cheungdaven/quate
     """
@@ -92,7 +68,7 @@ def __init__(
         relation_initializer: Hint[Initializer] = init_quaternions,
         relation_regularizer: Hint[Regularizer] = LpRegularizer,
         relation_regularizer_kwargs: Optional[Mapping[str, Any]] = None,
-        relation_normalizer: Hint[Constrainer] = quaternion_normalizer,
+        relation_normalizer: Hint[Constrainer] = quaternion.normalize,
         **kwargs,
     ) -> None:
         """Initialize QuatE.

src/pykeen/nn/algebra.py

@@ -9,11 +9,9 @@
 import torch
 
 from ..typing import FloatTensor
-from ..utils import einsum
 
 __all__ = [
     "circular_correlation",
-    "quat_e_interaction",
 ]
 
 
@@ -44,29 +42,3 @@ def circular_correlation(
     p_fft = a_fft * b_fft
     # inverse real FFT
     return torch.fft.irfft(p_fft, n=a.shape[-1], dim=-1)
-
-
-def quat_e_interaction(
-    h: FloatTensor,
-    r: FloatTensor,
-    t: FloatTensor,
-    table: FloatTensor,
-):
-    """Evaluate the interaction function of QuatE for given embeddings.
-
-    The embeddings have to be in a broadcastable shape.
-
-    :param h: shape: (`*batch_dims`, dim, 4)
-        The head representations.
-    :param r: shape: (`*batch_dims`, dim, 4)
-        The head representations.
-    :param t: shape: (`*batch_dims`, dim, 4)
-        The tail representations.
-    :param table:
-        the quaternion multiplication table.
-
-    :return: shape: (...)
-        The scores.
-    """
-    # TODO: this sign is in the official code, too, but why do we need it?
-    return -einsum("...di, ...dj, ...dk, ijk -> ...", h, r, t, table)

src/pykeen/nn/functional.py

@@ -30,8 +30,7 @@
 from typing_extensions import Self
 
 from . import functional as pkf
-from . import init
-from .algebra import quaterion_multiplication_table
+from . import init, quaternion
 from .compute_kernel import batched_dot
 from .sim import KG2ESimilarity, kg2e_similarity_resolver
 from .utils import apply_optional_bn
@@ -2782,16 +2781,23 @@ def forward(self, h: FloatTensor, r: tuple[FloatTensor, FloatTensor], t: FloatTe
 
 
 @parse_docdata
-class QuatEInteraction(
-    FunctionalInteraction[
-        FloatTensor,
-        FloatTensor,
-        FloatTensor,
-    ],
-):
-    """A module wrapper for the QuatE interaction function.
+class QuatEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]):
+    r"""The state-less QuatE interaction function.
+
+    It is given as
+
+    .. math ::
+        \langle \mathbf{h} \otimes \mathbf{r}, \mathbf{t} \rangle
+
+    where $\mathbf{h}, \mathbf{r}, \mathbf{t} \in \mathbb{H}^d$ are quanternion representations,
+    $\otimes$ denotes the Hamilton product, and $\langle \cdot, \cdot \rangle$ the inner product.
 
-    .. seealso:: :func:`pykeen.nn.functional.quat_e_interaction`
+    .. warning ::
+        In order to representation a rotation, $\mathbf{r}$ must be normalized to unit length,
+        cf. :func:`pykeen.nn.quaternion.normalize`.
+
+    .. seealso::
+        - https://en.wikipedia.org/wiki/Quaternion
 
     ---
     citation:
@@ -2805,15 +2811,34 @@ class QuatEInteraction(
     # with k=4
     entity_shape: Sequence[str] = ("dk",)
     relation_shape: Sequence[str] = ("dk",)
-    func = pkf.quat_e_interaction
 
     def __init__(self) -> None:
         """Initialize the interaction module."""
         super().__init__()
-        self.register_buffer(name="table", tensor=quaterion_multiplication_table())
+        self.register_buffer(name="table", tensor=quaternion.multiplication_table())
 
-    def _prepare_state_for_functional(self) -> MutableMapping[str, Any]:
-        return dict(table=self.table)
+    def forward(self, h: FloatTensor, r: tuple[FloatTensor, FloatTensor], t: FloatTensor) -> FloatTensor:
+        """Evaluate the interaction function of QuatE for given embeddings.
+
+        The embeddings have to be in a broadcastable shape.
+
+        .. seealso::
+            :meth:`Interaction.forward <pykeen.nn.modules.Interaction.forward>` for a detailed description about
+            the generic batched form of the interaction function.
+
+        :param h: shape: (`*batch_dims`, dim, 4)
+            The head representations.
+        :param r: shape: (`*batch_dims`, dim, 4)
+            The head representations.
+        :param t: shape: (`*batch_dims`, dim, 4)
+            The tail representations.
+
+        :return: shape: (...)
+            The scores.
+        """
+        # TODO: this sign is in the official code, too, but why do we need it?
+        # note: this is a fused kernel for computing the Hamilton product and the inner product at once
+        return -einsum("...di, ...dj, ...dk, ijk -> ...", h, r, t, self.table)
 
 
 class MonotonicAffineTransformationInteraction(

src/pykeen/nn/modules.py

@@ -0,0 +1,94 @@
+"""Utilities for quaternions."""
+
+from functools import lru_cache
+
+import torch
+
+from ..typing import FloatTensor
+
+__all__ = [
+    "normalize",
+    "hamiltonian_product",
+    "multiplication_table",
+]
+
+
+def normalize(x: FloatTensor) -> FloatTensor:
+    r"""
+    Normalize the length of relation vectors, if the forward constraint has not been applied yet.
+
+    Absolute value of a quaternion
+
+    .. math::
+
+        |a + bi + cj + dk| = \sqrt{a^2 + b^2 + c^2 + d^2}
+
+    L2 norm of quaternion vector:
+
+    .. math::
+        \|x\|^2 = \sum_{i=1}^d |x_i|^2
+                 = \sum_{i=1}^d (x_i.re^2 + x_i.im_1^2 + x_i.im_2^2 + x_i.im_3^2)
+
+    :param x: shape: ``(*batch_dims, 4 \cdot d)``
+        The vector in flat form.
+
+    :return: shape: ``(*batch_dims, 4 \cdot d)``
+        The normalized vector.
+    """
+    # Normalize relation embeddings
+    shape = x.shape
+    x = x.view(*shape[:-1], -1, 4)
+    x = torch.nn.functional.normalize(x, p=2, dim=-1)
+    return x.view(*shape)
+
+
+def hamiltonian_product(qa: FloatTensor, qb: FloatTensor) -> FloatTensor:
+    """Compute the hamiltonian product of two quaternions (which enables rotation)."""
+    return torch.stack(
+        [
+            qa[0] * qb[0] - qa[1] * qb[1] - qa[2] * qb[2] - qa[3] * qb[3],
+            qa[0] * qb[1] + qa[1] * qb[0] + qa[2] * qb[3] - qa[3] * qb[2],
+            qa[0] * qb[2] - qa[1] * qb[3] + qa[2] * qb[0] + qa[3] * qb[1],
+            qa[0] * qb[3] + qa[1] * qb[2] - qa[2] * qb[1] + qa[3] * qb[0],
+        ],
+        dim=-1,
+    )
+
+
+@lru_cache(1)
+def multiplication_table() -> FloatTensor:
+    """
+    Create the quaternion basis multiplication table.
+
+    :return: shape: (4, 4, 4)
+        the table of products of basis elements.
+
+    ..seealso:: https://en.wikipedia.org/wiki/Quaternion#Multiplication_of_basis_elements
+    """
+    _1, _i, _j, _k = 0, 1, 2, 3
+    table = torch.zeros(4, 4, 4)
+    for i, j, k, v in [
+        # 1 * ? = ?; ? * 1 = ?
+        (_1, _1, _1, 1),
+        (_1, _i, _i, 1),
+        (_1, _j, _j, 1),
+        (_1, _k, _k, 1),
+        (_i, _1, _i, 1),
+        (_j, _1, _j, 1),
+        (_k, _1, _k, 1),
+        # i**2 = j**2 = k**2 = -1
+        (_i, _i, _1, -1),
+        (_j, _j, _1, -1),
+        (_k, _k, _1, -1),
+        # i * j = k; i * k = -j
+        (_i, _j, _k, 1),
+        (_i, _k, _j, -1),
+        # j * i = -k, j * k = i
+        (_j, _i, _k, -1),
+        (_j, _k, _i, 1),
+        # k * i = j; k * j = -i
+        (_k, _i, _j, 1),
+        (_k, _j, _i, -1),
+    ]:
+        table[i, j, k] = v
+    return table

src/pykeen/nn/quaternion.py

@@ -14,9 +14,15 @@
 import pykeen.nn.modules
 import pykeen.nn.sim
 import pykeen.utils
-from pykeen.models.unimodal.quate import quaternion_normalizer
+from pykeen.nn import quaternion
 from pykeen.typing import Representation, Sign
-from pykeen.utils import clamp_norm, complex_normalize, einsum, ensure_tuple, project_entity
+from pykeen.utils import (
+    clamp_norm,
+    complex_normalize,
+    einsum,
+    ensure_tuple,
+    project_entity,
+)
 from tests import cases
 
 logger = logging.getLogger(__name__)
@@ -226,34 +232,21 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor:
         )
 
 
-def _rotate_quaternion(qa: torch.FloatTensor, qb: torch.FloatTensor) -> torch.FloatTensor:
-    # Rotate (=Hamilton product in quaternion space).
-    return torch.stack(
-        [
-            qa[0] * qb[0] - qa[1] * qb[1] - qa[2] * qb[2] - qa[3] * qb[3],
-            qa[0] * qb[1] + qa[1] * qb[0] + qa[2] * qb[3] - qa[3] * qb[2],
-            qa[0] * qb[2] - qa[1] * qb[3] + qa[2] * qb[0] + qa[3] * qb[1],
-            qa[0] * qb[3] + qa[1] * qb[2] - qa[2] * qb[1] + qa[3] * qb[0],
-        ],
-        dim=-1,
-    )
-
-
 class QuatETests(cases.InteractionTestCase):
     """Tests for QuatE interaction."""
 
     cls = pykeen.nn.modules.QuatEInteraction
     shape_kwargs = dict(k=4)  # quaternions
     atol = 1.0e-06
 
-    def _exp_score(self, h: torch.Tensor, r: torch.Tensor, t: torch.Tensor, table: torch.Tensor) -> torch.FloatTensor:  # noqa: D102
+    def _exp_score(self, h: torch.Tensor, r: torch.Tensor, t: torch.Tensor) -> torch.FloatTensor:  # noqa: D102
         # we calculate the scores using the hard-coded formula, instead of utilizing table + einsum
-        x = _rotate_quaternion(*(x.unbind(dim=-1) for x in [h, r]))
+        x = quaternion.hamiltonian_product(*(x.unbind(dim=-1) for x in [h, r]))
         return -(x * t).sum()
 
     def _get_hrt(self, *shapes):
         h, r, t = super()._get_hrt(*shapes)
-        r = quaternion_normalizer(r)
+        r = quaternion.normalize(r)
         return h, r, t
 
 

tests/test_nn/test_modules.py

@@ -2,7 +2,7 @@
 
 import torch
 
-import pykeen.nn.algebra
+from pykeen.nn import quaternion
 
 
 def _test_multiplication_table(t: torch.Tensor):
@@ -22,4 +22,4 @@ def _test_multiplication_table(t: torch.Tensor):
 
 def test_quaternion_multiplication_table():
     """Test quaternion multiplication table."""
-    _test_multiplication_table(pykeen.nn.algebra.quaterion_multiplication_table())
+    _test_multiplication_table(quaternion.multiplication_table())