Repository: pykeen/pykeen
Issue: How do I use external representations in KGEMs?
GitHub labels: question

Suppose I work on a problem of link prediction where each entity may have some external source of embedding (a vector of a certain dimension, like a chemical fingerprint). I would like to add this external embedding to the existing entity embedding of a model (say TransE) via simple vector addition. What's the best way to do this? Is there a way to get this external embedding into a triples class?

This is definitely more of a research topic than one we can build for you, so there are two things you should check out:

1. The tutorial on building your own interaction function (https://pykeen.readthedocs.io/en/stable/byo/interaction.html) and building your own model (https://pykeen.readthedocs.io/en/stable/extending/models_new.html) which shows how general you can actually make the models and the various representations they wrap

2. Check how we did this in the LiteralE model. This model does a simple concatenation between the literal vectors loaded for each entity with the ones from any given interaction function - this might inspire you: https://github.com/pykeen/pykeen/blob/master/src/pykeen/models/multimodal/base.py. Rather than literal vectors, you can bring your own representation.

Maybe you also want to see #470 that's about a similar idea

Thank you very much for the pointers. This is very helpful. I did not know LiteralE before.

Btw another big philosophical issue we've faced when looking at stuff like this is what to do when you have different kinds of external representations (i.e., vectors) for different nodes in the graph. For example, when you have a biomedical KG, there are chemicals, proteins, diseases, etc. Each of them might have a different kind of external representation like a chemical fingerprint from HTS screens, a protein fingerprint from recent applications of language models, and some disease phewas profiles. Maybe you could get them all to fit into vectors of the same size, or have a much more complicated interaction function that makes different combinations based on types.

One more alternative:

You can also write your own subclass of [RepresentationModule](https://github.com/pykeen/pykeen/blob/b8659b97016bc8c2120c498d50675be69a607e6f/src/pykeen/nn/emb.py#L46). 
A representation module is basically a mapping from integer IDs to some vector representation, cf.

https://github.com/pykeen/pykeen/blob/b8659b97016bc8c2120c498d50675be69a607e6f/src/pykeen/nn/emb.py#L89-L94

For example you could have
```python
class CombinedRepresentation(RepresentationModule):
    def forward(self, indices):
        return self.embedding(indices) + self.external(indices)
```

where both, `self.embeddings` and `self.external` are `pykeen.nn.Embedding`, both only the first is trainable. The second one only serves to store the external features (and offer some utility methods for reshaping).

Now you should be able to use this representation everywhere where Embeddings are used inside the models, in particular also in [ERModel](https://github.com/pykeen/pykeen/blob/b8659b97016bc8c2120c498d50675be69a607e6f/src/pykeen/models/nbase.py#L290)


```python
combined_representation = CombinedRepresentations(...)
model = ERModel(
    interaction="transe",
    entity_representations=combined_representation,
    relation_representations=EmbeddingSpecification(shape=combined_representation.shape)
)
```

@pitehu I saw you've starred Himmelstein's DrugBank repo. If you ever want to discuss biomedical applications of KGs, please let me know!

> ```python
> CombinedRepresentation
> ```

Would this allow me to train my external embedding model jointly also? @mberr 
If I were to apply this to an EntityRelationEmbeddingModel, then it appears that I will need to rewrite the model (e.g TuckER). In other words, there is no drop-in replacement that's universally applicable if I want to test the performance of adding textual literal on multiple different models, correct?

> > ```python
> > CombinedRepresentation
> > ```
> 
> Would this allow me to train my external embedding model jointly also? @mberr

If you want to train the external representation, too, i.e. use the chemical fingerprint as *initialization*, but then allow the model to change this representation, you as well just use a "normal" embedding model and only set a custom embedding initializer. The more interesting case (from my point of view), is to use the external information as a feature vector (which is *not* trained), maybe optionally transform it, e.g., by some MLP, and then combine it, or replace the entity embedding.

In any case, everything which contributes in computing the final loss, and which requires a gradient, will receive updates from the optimizer.


> If I were to apply this to an EntityRelationEmbeddingModel, then it appears that I will need to rewrite the model (e.g TuckER). In other words, there is no drop-in replacement that's universally applicable if I want to test the performance of adding textual literal on multiple different models, correct?


`ERModel` is our "new-style" model, where we clearly separate entity and relation representations (i.e. components which produce one vector / matrix / tensors ... for each entity / relation ID) from interaction functions, which operate on vector representations, and produce scores for triples.

This makes it very flexible to exchange single components. For TuckER, you can simply change the interaction function as follows
```python
combined_representation = CombinedRepresentations(...)
model = ERModel(
    interaction="tucker",
    entity_representations=combined_representation,
    relation_representations=EmbeddingSpecification(shape=combined_representation.shape)
)
```

This works out of the box for every interaction function, which trains a single vector for each entity and relation.

Some other interactions, e.g., NTN, require you to adapt the relation representations to be a list of embedding specifications, since NTN uses 5 different relation representations for each relation, cf. https://github.com/pykeen/pykeen/blob/dc20e024fff419087ec0e4eeccd5de3cbfc0c032/src/pykeen/models/unimodal/ntn.py#L87-L115

As of https://github.com/pykeen/pykeen/pull/964, it's much more ergonomic to include external representations in a KGEM.