🦜 🏴‍☠️ Implement NodePiece representation and model

[mberr]

This is a first draft to add NodePiece representations to pykeen.

For now, it uses a simple variant of it, where each entity is represented by k randomly chosen incident relations.

• Official Code: https://github.com/migalkin/NodePiece
• Paper: https://arxiv.org/abs/2106.12144

[mberr]

@migalkin would be great to have your feedback, as you may be familiar with it 😉

[cthoyt]

Can we have a demo on how you would use this representation with a model? Like can we easily implement a TransE with NodePiece?

[mberr]

Can we have a demo on how you would use this representation with a model? Like can we easily implement a TransE with NodePiece?

Sure.

from typing import Optional

from class_resolver.api import HintOrType

from pykeen.models.nbase import ERModel
from pykeen.nn.emb import EmbeddingSpecification, NodePieceRepresentation
from pykeen.nn.modules import Interaction, TransEInteraction
from pykeen.pipeline import pipeline
from pykeen.triples.triples_factory import CoreTriplesFactory


class NodePieceModel(ERModel):
    def __init__(
        self,
        *,
        triples_factory: CoreTriplesFactory,
        embedding_specification: Optional[EmbeddingSpecification] = None,
        interaction: HintOrType[Interaction] = TransEInteraction,
        **kwargs,
    ) -> None:
        if embedding_specification is None:
            embedding_specification = EmbeddingSpecification(
                shape=(64,),
            )
        entity_representations = NodePieceRepresentation(
            triples_factory=triples_factory,
            token_representation=embedding_specification,
        )
        super().__init__(
            triples_factory=triples_factory,
            interaction=interaction,
            entity_representations=entity_representations,
            relation_representations=embedding_specification,
            **kwargs,
        )


result = pipeline(
    dataset="nations",
    model=NodePieceModel,
    model_kwargs=dict(
        interaction_kwargs=dict(
            p=2,
        ),
    ),
)
print(result.get_metric("hits_at_10"))

EDIT: added in 6b29c4e

[mberr]

fe90b42 - @cthoyt this is not really part of this PR

[mberr]

this typo is not really part of the PR

[mberr]

same here

[cthoyt]

Looks like the issues are now with ConvE's tests

[migalkin]

Looks like the issues are now with ConvE's tests

I was running the debugger for the ConvE test and for some reason after initialization of TestConvE(cases.ModelTestCase) it goes on to initialize NodePiece although it is not related anyhow to the ConvE test 🤔

[migalkin]

The implementation works 🎉
Exposing the ratio param from the MLP encoder sounds like a good idea (with the default value 2), otherwise everything looks ready!

[migalkin]

The current _ConcatMLP is not DeepSets :)

The idea of DeepSets is to project each set member independently through some encoder, then aggregate (like with mean) and then pass through another FF net. It would look like this:

enc1 = nn.Sequential(
    nn.Linear(embedding_dim, embedding_dim),
    nn.ReLU(),
    nn.Linear(embedding_dim, embedding_dim)
)

enc2 = nn.Sequential(nn.Linear(embedding_dim, embedding_dim), nn.ReLU(), nn.Linear(embedding_dim, output_dim))

and in forward pass:

# x: shape (bs, num_elements, embedding_dim)
x = enc1(x)               # the same shape (bs, num_elements, embedding_dim)
x = torch.mean(-2)  # here we do the aggregation to (bs, embedding_dim)
x = enc2(x)               # final projection keeping (bs, output_dim)

It can be added as an option along with mlp though

[mberr]

The correct docstring somehow got lost during the refactoring 😅

here it was still correct: #621 (comment)

[mberr]

Looks like the issues are now with ConvE's tests

I think the issue is that we did not yet think about what should be scored in score_r, if we have inverse relations. In the baseline implementation in _OldAbstractModel, we use relation_ids to create the individual hrt triples, which does not contain the inverse relations, but only the "real" ones.

We could either:

1. skip the score_r test for models with inverse triples
2. decide on how to handle this generally

To keep this PR focused on one thing, I tend towards option 1.

[cthoyt]

Yes let’s bump this. So let’s override the test in conve to be skipped and leave a todo for later

[cthoyt]