Repository: pykeen/pykeen
Issue: Investigate _get_representations
Root cause: Algorithmic Error

Pushed a draft on the Inductive NodePiece with a GNN encoder with a surprising finding - I got lost in all the Representation Modules and reshaping happening in the forward pass 😞  and just applied index_select + unsqueezing to GNN outputs, and found that it actually works >2x faster than plain InductiveNodePiece w/o GNN 👀 

On a M1 laptop, both models 64d, 32 negative samples, 100 epochs:
* `InductiveNodePiece`: 16 min
* `InductiveNodePieceGNN`: 7 min

It seems that this part in ERModel `_get_representations` could introduce dramatic slowdowns:
```python
h, r, t = [
            [representation.get_in_more_canonical_shape(dim=dim, indices=indices) for representation in representations]
            for dim, indices, representations in (
                ("h", h_indices, entity_representations),
                ("r", r_indices, self.relation_representations),
                ("t", t_indices, entity_representations),
            )
        ]
```

For the complete reproduction of paper experiments, I'd need to add the Attention aggregator to the CompGCN layer - but this can actually be a separate PR 🤔

_Originally posted by @migalkin in https://github.com/pykeen/pykeen/issues/722#issuecomment-1016000935_

Relevant commit: https://github.com/pykeen/pykeen/pull/722/commits/b0d8c6dd164426009b48cfc59ddd65dc5b7f217b

Current `master` code
https://github.com/pykeen/pykeen/blob/9958de424c9354ae21646af4c8a078de17346dc2/src/pykeen/models/nbase.py#L501-L520

Updated code by @migalkin 
https://github.com/pykeen/pykeen/blob/b0d8c6dd164426009b48cfc59ddd65dc5b7f217b/src/pykeen/models/unimodal/inductive_nodepiece_gnn.py#L93-L132

The relevant part should be to compare
```python
x_e.index_select(dim=0, index=h_indices).unsqueeze(1) if h_indices is not None else x_e.unsqueeze(0)
```
against
```python
representation.get_in_more_canonical_shape(dim=dim, indices=indices)
```

@migalkin what was the `representation` you compared against?

Here some measurements on `Embedding` representation on GPU

<summary>Code</summary>
<details>

```python
import torch
from torch.utils.benchmark import Compare, Timer

from pykeen.nn.emb import EmbeddingSpecification

emb = EmbeddingSpecification(shape=(256,)).make(num_embeddings=15_000)
print(
    Compare(
        [
            Timer(
                stmt=stmt,
                setup=f"e=emb.to(device='cuda'); indices = torch.randint(15_000, size={index_shape}, device='cuda')",
                globals=dict(torch=torch, emb=emb),
                description=str(index_shape),
            ).blocked_autorange()
            for stmt in [f"e.get_in_more_canonical_shape(dim='{dim}', indices=indices)" for dim in "hrt"]
            + [
                "e(indices=indices)",
                "e._embeddings(indices)",
                "e._embeddings.weight.index_select(dim=0, index=indices.view(-1)).view(*indices.shape, *e.shape)",
                "e._embeddings.weight.index_select(dim=0, index=indices.view(-1))",
            ]
            for index_shape in ((256,), (4096,), (128, 256), (128 * 256,))
        ]
    )
)
```
</details>

```
indices.shape                                                                                          |  (256,)  |  (4096,)  |  (128, 256)  |  (32768,)
1 threads: ---------------------------------------------------------------------------------------------------------------------------------------------
      e.get_in_more_canonical_shape(dim='h', indices=indices)                                          |   34.9   |    35.1   |    140.0     |   140.0  
      e.get_in_more_canonical_shape(dim='r', indices=indices)                                          |   34.9   |    35.6   |    140.2     |   140.0  
      e.get_in_more_canonical_shape(dim='t', indices=indices)                                          |   35.0   |    35.2   |    140.1     |   140.0  
      e(indices=indices)                                                                               |   23.6   |    24.2   |    140.1     |   140.7  
      e._embeddings(indices)                                                                           |   15.0   |    20.3   |    140.2     |   140.4  
      e._embeddings.weight.index_select(dim=0, index=indices.view(-1)).view(*indices.shape, *e.shape)  |   18.5   |    20.2   |    140.4     |   140.5  
      e._embeddings.weight.index_select(dim=0, index=indices.view(-1))                                 |   14.3   |    20.3   |    140.5     |   140.7  

Times are in microseconds (us).
```

Ah so there is indeed a big performance gap on small (but most widely used) batch sizes!

> @migalkin what was the `representation` you compared against?

It was `NodePieceRepresentation` used in the default NodePiece and Inductive NodePiece


The `_get_representations` part seems to have been accelerated. The overall training speed does not seem to change much, though.

<summary>Code</summary>
<details>

```python
import torch
from torch.utils.benchmark import Compare, Timer
​
from pykeen.datasets import get_dataset
from pykeen.models.unimodal.mure import MuRE
​
​
def main():
    dataset = get_dataset(dataset="fb15k237", dataset_kwargs=dict(create_inverse_triples=True))
    model = MuRE(
        triples_factory=dataset.training,
        embedding_dim=64,
    )
    print(
        Compare(
            [
                Timer(
                    stmt=stmt,
                    setup=f"h=torch.randint({dataset.num_entities}, size={index_shape}, device='cuda'); r=torch.randint({dataset.num_relations}, size={index_shape}, device='cuda')",
                    globals=dict(torch=torch, model=model),
                    description=str(index_shape),
                ).blocked_autorange()
                for index_shape in ((256,), (4096,), (128, 256), (128 * 256,))
                for stmt in [
                    "model._get_representations(h, r, None)",
                    "model._get_representations(h, r, h)",
                ]
            ]
        )
    )
​

if __name__ == "__main__":
    main()
```
</details>

<summary>Times</summary>
<details>

```
(venv) user@server:/tmp/pykeen$ git status
On branch master
Your branch is up to date with 'origin/master'.
Untracked files:
  (use "git add <file>..." to include in what will be committed)
        measure_speed.py
nothing added to commit but untracked files present (use "git add" to track)
(venv) user@server:/tmp/pykeen$ python3 measure_speed.py 
You're trying to map triples with 30 entities and 0 relations that are not in the training set. These triples will be excluded from the mapping.
In total 28 from 20466 triples were filtered out
No random seed is specified. This may lead to non-reproducible results.
[----------------------------------------------  ---------------------------------------------]
                                              |  (256,)  |  (4096,)  |  (128, 256)  |  (32768,)
1 threads: ------------------------------------------------------------------------------------
      model._get_representations(h, r, None)  |  232.1   |   234.3   |    250.2     |   235.6  
      model._get_representations(h, r, h)     |  301.6   |   312.0   |    328.8     |   308.8  
Times are in microseconds (us).
(venv) user@server:/tmp/pykeen$ git checkout fix-performance-regression 
Switched to branch 'fix-performance-regression'
Your branch is up to date with 'origin/fix-performance-regression'.
(venv) user@server:/tmp/pykeen$ python3 measure_speed.py 
You're trying to map triples with 30 entities and 0 relations that are not in the training set. These triples will be excluded from the mapping.
In total 28 from 20466 triples were filtered out
No random seed is specified. This may lead to non-reproducible results.
[----------------------------------------------  ---------------------------------------------]
                                              |  (256,)  |  (4096,)  |  (128, 256)  |  (32768,)
1 threads: ------------------------------------------------------------------------------------
      model._get_representations(h, r, None)  |  150.3   |   151.7   |    160.3     |   152.3  
      model._get_representations(h, r, h)     |  205.5   |   203.8   |    219.5     |   212.4  
Times are in microseconds (us).
```
</details>

More discoveries on the speed pertaining particularly to NodePiece: NodePiece passes to the aggregator the whole batch of entities `entity_size` which is in the sLCWA case of size `batch_size x num_negatives` , so in the case of 256*32 it encodes 8192 entities. If `bs x num_negs` > `total_num_entities` this leads to performance issues as the batch contains **many** duplicates and we do unnecessary work computing the same representations for duplicates.

In the gnn case, we first materialize all entities, then run message passing, and then perform `index_select`, so in the case of 256 * 32 entities we don’t actually run the encoder for 8192 entities, but we just extract their indices from the updated nodes matrix.

Therefore, NodePieceGNN wins when:
* `batch_size x num_negatives` > `total_num_entities` + O(GNN) - some complexity of the GNN forward pass, eg training with batch size 256 and 32 negatives on a graph with 1k unique entities

Otherwise vanilla NodePiece wins.

Max proposed a simple patch for the [NodePieceRepresentation](https://github.com/pykeen/pykeen/blob/ad3183cda470c6ecd017089596fdb8a7614f84d1/src/pykeen/nn/emb.py#L1077) to aggregate only unique indices:

```
indices, inverse = indices.unique(return_inverse=True)

# the old forward...
x = self.aggregation(...)

return x[inverse]
```
I tested that and it indeed gives 8x performance increase on the inductive pipeline 🚀 

I believe though it can be even further improved when joining both `head` and `tail` indices and getting unique values among those two, such that we won't repeat the NodePiece aggregation pass twice - this could potentially lead to 2x more improvements.