自带交互

这是一个关于如何实现自己的交互模块(也称为评分函数)作为pykeen.nn.modules.Interaction子类的教程,以便在PyKEEN中使用。

实现你的第一个交互模块

想象一下,你乘坐时光机回到了2013年,你刚刚发明了TransE,定义为:

\[f(h, r, t) = -\| \mathbf{e}_h + \mathbf{r}_r - \mathbf{e}_t \|_2\]

其中 \(\mathbf{e}_i\) 是实体 \(i\)\(d\) 维表示, \(\mathbf{r}_j\) 是关系 \(j\)\(d\) 维表示,而 \(\|...\|_2\)\(L_2\) 范数。

要在PyKEEN中实现TransE,你需要子类化 pykeen.nn.modules.Interaction。这个类本身是 torch.nn.Module的子类,这意味着你需要提供 torch.nn.Module.forward()的实现。然而,参数 被预定义为hrt,它们分别对应于头、关系和尾的表示。

from pykeen.nn.modules import Interaction

class TransEInteraction(Interaction):
    def forward(self, h, r, t):
        return -(h + r - t).norm(p=2, dim=-1)

注意 dim=-1 因为这个操作实际上是定义在整个批次的头、关系和尾表示上的。

另请参阅

pykeen.nn.modules.TransEInteraction中提供了一个参考实现

作为一名刚刚发明了TransE的研究人员,你可能会想知道如果将加法+替换为乘法*会发生什么。你可能会得到一个新的交互方式(这恰好是DistMult,它在TransE发布后仅一年就被发表了):

\[f(h, r, t) = \mathbf{e}_h^T diag(\mathbf{r}_r) \mathbf{e}_t\]

其中 \(\mathbf{e}_i\) 是实体 \(i\)\(d\) 维表示, \(\mathbf{r}_j\) 是关系 \(j\)\(d\) 维表示。

from pykeen.nn.modules import Interaction

class DistMultInteraction(Interaction):
    def forward(self, h, r, t):
        return (h * r * t).sum(dim=-1)

另请参阅

pykeen.nn.modules.DistMultInteraction中提供了一个参考实现

与超参数的交互

虽然我们之前用\(L_2\)范数定义了TransE,但它可以用\(p\)的不同值来计算:

\[f(h, r, t) = -\| \mathbf{e}_h + \mathbf{r}_r - \mathbf{e}_t \|_p\]

这可以通过使用__init__()将其纳入交互定义中,将\(p\)的值存储在实例中,然后在forward()中访问它。

from pykeen.nn.modules import Interaction

class TransEInteraction(Interaction):
    def __init__(self, p: int):
        super().__init__()
        self.p = p

    def forward(self, h, r, t):
        return -(h + r - t).norm(p=self.p, dim=-1)

一般来说,你可以在__init__()中放入任何你想要的内容来支持分数的计算。

与可训练参数的交互

在ER-MLP中,多层感知器由一个具有\(3 \times d\)神经元的输入层、一个具有\(d\)神经元的隐藏层和一个具有一个神经元的输出层组成。输入由头部、关系和尾部的嵌入连接表示。它被定义为:

\[f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2\]

具有隐藏维度 \(y\), \(W_1 \in \mathcal{R}^{3d \times y}\), \(W_2\ \in \mathcal{R}^y\), 以及 偏置 \(b_1 \in \mathcal{R}^y\)\(b_2 \in \mathcal{R}\).

\(W_1\), \(W_1\), \(b_1\), 和 \(b_2\)全局参数,意味着它们是可训练的,但既不附加到实体也不附加到关系。与TransE中的\(p\)不同,这些全局可训练参数不被视为超参数。然而,像超参数一样,它们也可以在pykeen.nn.modules.Interaction类的__init__函数中定义。它们在训练过程中与实体和关系嵌入一起训练。

import torch.nn
from pykeen.nn.modules import Interaction
from pykeen.utils import broadcast_cat

class ERMLPInteraction(Interaction):
    def __init__(self, embedding_dim: int, hidden_dim: int):
        super().__init__()
        # The weights of this MLP will be learned.
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(in_features=3 * embedding_dim, out_features=hidden_dim, bias=True),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=hidden_dim, out_features=1, bias=True),
        )

    def forward(self, h, r, t):
        x = broadcast_cat([h, r, t], dim=-1)
        return self.mlp(x)

请注意,使用了pykeen.utils.broadcast_cat()而不是标准的torch.cat(),这是因为头、关系和尾向量的形状标准化。

另请参阅

pykeen.nn.modules.ERMLPInteraction中提供了一个参考实现

与不同形状向量的交互

结构化嵌入使用2-张量来表示每个关系,其交互定义为:

\[f(h, r, t) = - \|\textbf{M}_{r}^{head} \textbf{e}_h - \textbf{M}_{r}^{tail} \textbf{e}_t\|_p\]

其中 \(\mathbf{e}_i\) 是实体 \(i\)\(d\) 维表示, \(\mathbf{M}^{head}_j\) 是关系 \(j\) 对于头实体的 \(d \times d\) 维表示, \(\mathbf{M}^{tail}_j\) 是关系 \(j\) 对于尾实体的 \(d \times d\) 维表示,并且 \(\|...\|_2\)\(L_p\) 范数。

在本教程中,我们将提出对结构化嵌入(类似于TransR)的简化,其中使用相同的关系2-张量来投影头部和尾部实体,如下所示:

\[f(h, r, t) = - \|\textbf{M}_{r} \textbf{e}_h - \textbf{M}_{r} \textbf{e}_t\|_2\]

其中 \(\mathbf{e}_i\) 是实体 \(i\)\(d\) 维表示, \(\mathbf{M}_j\) 是关系 \(j\)\(d \times d\) 维表示,而 \(\|...\|_2\)\(L_2\) 范数。

from pykeen.nn.modules import Interaction

class SimplifiedStructuredEmbeddingInteraction(Interaction):
    relation_shape = ('dd',)

    def forward(self, h, r, t):
        h_proj = r @ h.unsqueeze(dim=-1)
        t_proj = r @ t.unsqueeze(dim=-1)
        return -(h_proj - t_proj).squeeze(dim=-1).norm(p=2, dim=-1)

注意relation_shape的定义。默认情况下,entity_shaperelation_shape都等于('d', ),这使用特征符号来表示 它们都是具有相同形状的1-张量。在这个简化版本的 结构化嵌入中,我们需要表示关系的形状是\(d \times d\), 所以它被写成dd

另请参阅

参考实现可以在pykeen.nn.modules.StructuredEmbeddingInteractionpykeen.nn.modules.TransRInteraction中找到。

与多种表示的交互

有时,就像在结构化嵌入的规范版本中一样,您需要为实体和/或关系提供多个表示。要指定这一点,您只需扩展relation_shape的元组,添加更多条目,每个条目对应表示序列。

from pykeen.nn.modules import Interaction

class StructuredEmbeddingInteraction(Interaction):
    relation_shape = (
        'dd',  # Corresponds to $\mathbf{M}^{head}_j$
        'dd',  # Corresponds to $\mathbf{M}^{tail}_j$
    )

    def forward(self, h, r, t):
        # Since the relation_shape is more than length 1, the r value is given as a sequence
        # of the representations defined there. You can use tuple unpacking to get them out
        r_h, r_t = r
        h_proj = r_h @ h.unsqueeze(dim=-1)
        t_proj = r_t @ t.unsqueeze(dim=-1)
        return -(h_proj - t_proj).squeeze(dim=-1).norm(p=2, dim=-1)

与不同维度向量的交互

TransD 是一个交互模块的示例,它不仅为每个实体使用两种不同的表示,为每个关系使用两种表示,而且它们的维度也不同。

可以通过在entity_shape和/或relation_shape字典中选择不同的字母来实现。最终,使用的字母是任意的,但在使用pykeen.models.make_model()pykeen.models.make_model_cls()pykeen.pipeline.interaction_pipeline()函数来实例化模型、创建模型类或使用自定义交互模块运行管道时,需要记住它们是什么。

from pykeen.nn.modules import Interaction
from pykeen.utils import project_entity

class TransDInteraction(Interaction):
    entity_shape = ("d", "d")
    relation_shape = ("e", "e")

    def forward(self, h, r, t):
        h, h_proj = h
        r, r_proj = r
        t, t_proj = t
        h_bot = project_entity(
            e=h,
            e_p=h_p,
            r_p=r_p,
        )
        t_bot = project_entity(
            e=t,
            e_p=t_p,
            r_p=r_p,
        )
        return -(h_bot + r - t_bot).norm(p=2, dim=-1)

注意

在这个实现中使用了pykeen.utils.project_entity()函数来降低复杂性。到目前为止,所有使用多种不同表示维度的模型都非常复杂,不符合展示简单示例的范式。

另请参阅

pykeen.nn.modules.TransDInteraction中提供了一个参考实现

pykeen.nn.modules.Interactionpykeen.models.Model 之间的区别

高级的 pipeline() 函数允许你传递预定义的子类,例如 pykeen.models.Model 的子类 pykeen.models.TransEpykeen.models.DistMult。这些类是围绕交互函数 pykeen.nn.modules.TransEInteractionnn.modules.DistMultInteraction 的高级封装,更适合用于运行基准测试实验或知识图谱嵌入的实际应用,这些应用包含大量关于默认超参数、推荐的超参数优化策略以及更复杂的正则化方案应用的信息。

作为一名研究人员,pykeen.nn.modules.Interaction 是一种快速将想法转化为新模型的方式,这些模型可以在不需要定义pykeen.models.Model的所有开销的情况下使用。这些组件在PyKEEN中也是完全可重用的(例如,在自定义的训练循环中),并且可以在PyKEEN之外作为独立组件使用。

如果您对您的交互模块感到满意,并希望进一步使其具有通用可重用性,请查看“扩展模型”教程。

Ad hoc 交互模型

一个 pykeen.models.ERModel 可以从 pykeen.nn.modules.Interaction 构建。

新的样式类,pykeen.models.ERModel 将交互从表示中抽象出来,使得不同的交互可以互换使用。可以直接从交互模块构建一个新模型,给定一个dimensions映射。在每个pykeen.nn.modules.Interaction中,有一个名为entity_shaperelation_shape的字段,允许使用特征符号来定义模型的不同维度。大多数模型共享d维度用于实体和关系向量。一些(但不是全部)例外是:

考虑到这一点,您需要通过PyKEEN文档研究向量的维度。 如果您正在实现自己的向量,您可以控制这一点,并且会知道要指定哪些维度(尽管 d 对于实体和关系都是标准的)。作为 {'d': value} 的简写,您可以直接 传递 value 作为维度,它将被自动解释为 {'d': value}

从交互模块类的查找中创建一个模型类:

>>> from pykeen.nn.modules import TransEInteraction
>>> from pykeen.models import make_model_cls
>>> embedding_dim = 3
>>> model_cls = make_model_cls(
...     dimensions={"d": embedding_dim},
...     interaction='TransE',
...     interaction_kwargs={'p': 2},
... )

如果entity_shapesrelation_shapes中只有一个维度,可以直接将其作为整数给出作为快捷方式。

>>> # Implicitly can also be written as:
>>> model_cls_alt = make_model_cls(
...     dimensions=embedding_dim,
...     interaciton='TransE',
...     interaction_kwargs={'p': 2},
... )

从交互模块类创建一个模型类:

>>> from pykeen.nn.modules import TransEInteraction
>>> from pykeen.models import make_model_cls
>>> embedding_dim = 3
>>> model_cls = make_model_cls({"d": embedding_dim}, TransEInteraction, {'p': 2})

从实例化的交互模块中创建一个模型类:

>>> from pykeen.nn.modules import TransEInteraction
>>> from pykeen.models import make_model_cls
>>> embedding_dim = 3
>>> model_cls = make_model_cls({"d": embedding_dim}, TransEInteraction(p=2))

所有这些模型类都可以直接传递给pykeen.pipeline.pipeline()model参数。

交互管道

pykeen.pipeline.pipeline() 还允许传递一个交互,因此以下代码块可以被压缩:

from pykeen.pipeline import pipeline
from pykeen.nn.modules import TransEInteraction

model = make_model_cls(
    interaction=TransEInteraction,
    interaction_kwargs={'p': 2},
    dimensions={'d': 100},
)
results = pipeline(
    dataset='Nations',
    model=model,
    ...
)

进入:

from pykeen.pipeline import pipeline
from pykeen.nn.modules import TransEInteraction

results = pipeline(
    dataset='Nations',
    interaction=TransEInteraction,
    interaction_kwargs={'p': 2},
    dimensions={'d': 100},
    ...
)

这可以与pykeen.nn.modules.Interaction的任何子类一起使用,不仅限于PyKEEN包中实现的子类。