Repository: pyg-team/pytorch_geometric
Pull request #7026 (2023-03-24 01:34:06)

feat: add `RotatE` KGE model
[No longer contributing affiliated with CS224W]

Paper reported FB15k-237 H@10: .533
Implementation Example H@10: .431

Of note, the paper only reports their score using their Adversarial Negative Sampling (ANS), which is not implemented as a part of this PR. Table 7 from the paper comparing the use of ANS vs. uniform random sampling shows a 5% increase in H@10 for FB15k-237, which may be a factor to the disparity in performance, combined with the small hidden size in the example, I believe. The margin=9 argument is from the paper's reported best performance in Table 12 on FB15k-237 as well.

Happy to try to implement the ANS as part of this PR if the current performance is below standard, let me know!



codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7026?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7026](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7026?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ac93c6f) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/f0c72186286f257778c1d9293cfd0d35472d30bb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (f0c7218) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head ac93c6f differs from pull request most recent head bbd5d44. Consider uploading reports for the commit bbd5d44 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7026      +/-   ##
==========================================
+ Coverage   91.51%   91.52%   +0.01%     
==========================================
  Files         435      436       +1     
  Lines       23773    23810      +37     
==========================================
+ Hits        21756    21793      +37     
  Misses       2017     2017              
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7026?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/kge/\_\_init\_\_.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7026?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2tnZS9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/kge/rotate.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7026?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2tnZS9yb3RhdGUucHk=) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ec7477948..5d782d9ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [2.4.0] - 2023-MM-DD
 
 ### Added
+- Added the `RotatE` KGE model ([#7026](https://github.com/pyg-team/pytorch_geometric/pull/7026))
 
 ### Changed
 
diff --git a/README.md b/README.md
index 22d6c63d9..e9807f77c 100644
--- a/README.md
+++ b/README.md
@@ -304,6 +304,7 @@ New Benchmarks and Strong Simple Methods](https://arxiv.org/abs/2110.14446) (Neu
 * **[TransE](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.kge.TransE.html)** from Bordes *et al.*: [Translating Embeddings for Modeling Multi-Relational Data](https://proceedings.neurips.cc/paper/2013/file/1cecc7a77928ca8133fa24680a88d2f9-Paper.pdf) (NIPS 2013) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)]
 * **[ComplEx](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.kge.ComplEx.html)** from Trouillon *et al.*: [Complex Embeddings for Simple Link Prediction](https://arxiv.org/abs/1606.06357) (ICML 2016) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)]
 * **[DistMult](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.kge.DistMult.html)** from Yang *et al.*: [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://arxiv.org/abs/1412.6575) (ICLR 2015) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)]
+* **[RotatE](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.kge.RotatE.html)** from Sun *et al.*: [RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space](https://arxiv.org/abs/1902.10197) (ICLR 2019) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)]
 </details>
 
 **GNN operators and utilities:**
diff --git a/examples/kge_fb15k_237.py b/examples/kge_fb15k_237.py
index b4d1ac7f0..a8fe23dae 100644
--- a/examples/kge_fb15k_237.py
+++ b/examples/kge_fb15k_237.py
@@ -5,12 +5,13 @@ import torch
 import torch.optim as optim
 
 from torch_geometric.datasets import FB15k_237
-from torch_geometric.nn import ComplEx, DistMult, TransE
+from torch_geometric.nn import ComplEx, DistMult, RotatE, TransE
 
 model_map = {
     'transe': TransE,
     'complex': ComplEx,
     'distmult': DistMult,
+    'rotate': RotatE,
 }
 
 parser = argparse.ArgumentParser()
@@ -25,11 +26,11 @@ train_data = FB15k_237(path, split='train')[0].to(device)
 val_data = FB15k_237(path, split='val')[0].to(device)
 test_data = FB15k_237(path, split='test')[0].to(device)
 
-model = model_map[args.model](
-    num_nodes=train_data.num_nodes,
-    num_relations=train_data.num_edge_types,
-    hidden_channels=50,
-).to(device)
+model_arg_map = {'rotate': {'margin': 9.0}}
+model = model_map[args.model](num_nodes=train_data.num_nodes,
+                              num_relations=train_data.num_edge_types,
+                              hidden_channels=50,
+                              **model_arg_map.get(args.model, {})).to(device)
 
 loader = model.loader(
     head_index=train_data.edge_index[0],
@@ -43,6 +44,7 @@ optimizer_map = {
     'transe': optim.Adam(model.parameters(), lr=0.01),
     'complex': optim.Adagrad(model.parameters(), lr=0.001, weight_decay=1e-6),
     'distmult': optim.Adam(model.parameters(), lr=0.0001, weight_decay=1e-6),
+    'rotate': optim.Adam(model.parameters(), lr=1e-3),
 }
 optimizer = optimizer_map[args.model]
 
diff --git a/torch_geometric/nn/kge/__init__.py b/torch_geometric/nn/kge/__init__.py
index 35e285218..cd0c302ba 100644
--- a/torch_geometric/nn/kge/__init__.py
+++ b/torch_geometric/nn/kge/__init__.py
@@ -2,10 +2,12 @@ from .base import KGEModel
 from .transe import TransE
 from .complex import ComplEx
 from .distmult import DistMult
+from .rotate import RotatE
 
 __all__ = classes = [
     'KGEModel',
     'TransE',
     'ComplEx',
     'DistMult',
+    'RotatE',
 ]
diff --git a/torch_geometric/nn/kge/rotate.py b/torch_geometric/nn/kge/rotate.py
new file mode 100644
index 000000000..a4724c5e9
--- /dev/null
+++ b/torch_geometric/nn/kge/rotate.py
@@ -0,0 +1,97 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Embedding
+
+from torch_geometric.nn.kge import KGEModel
+
+
+class RotatE(KGEModel):
+    r"""The RotatE model from the `"RotatE: Knowledge Graph Embedding by
+    Relational Rotation in Complex Space" <https://arxiv.org/abs/
+    1902.10197>`_ paper.
+
+    :class:`RotatE` models relations as a rotation in complex space
+    from head to tail such that:
+
+    .. math::
+        \mathbf{e}_t = \mathbf{e}_h \circ \mathbf{e}_r
+
+    Resulting in the scoring function:
+
+    .. math::
+        d(h, r, t) = - {\| \mathbf{e}_h \circ \mathbf{e}_r - \mathbf{e}_t \|}_p
+
+    .. note::
+
+        For an example of using the :class:`RotatE` model, see
+        `examples/kge_fb15k_237.py
+        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
+        kge_fb15k_237.py>`_.
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        margin (float, optional): The margin of the ranking loss.
+        sparse (bool, optional): If set to :obj:`True`, gradients w.r.t. to
+            the embedding matrices will be sparse. (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        margin: float = 1.0,
+        sparse: bool = False,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse)
+        self.register_buffer('margin', torch.Tensor([margin]))
+        self.node_emb_im = Embedding(num_nodes, hidden_channels, sparse=sparse)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.node_emb.weight)
+        torch.nn.init.xavier_uniform_(self.node_emb_im.weight)
+        torch.nn.init.uniform_(self.rel_emb.weight, 0, 2 * math.pi)
+
+    def forward(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tensor:
+
+        head_re = self.node_emb(head_index)
+        head_im = self.node_emb_im(head_index)
+        tail_re = self.node_emb(tail_index)
+        tail_im = self.node_emb_im(tail_index)
+
+        rel_theta = self.rel_emb(rel_type)
+        rel_re, rel_im = torch.cos(rel_theta), torch.sin(rel_theta)
+
+        re_score = (rel_re * head_re - rel_im * head_im) - tail_re
+        im_score = (rel_re * head_im + rel_im * head_re) - tail_im
+        complex_score = torch.stack([re_score, im_score], dim=2)
+        score = torch.linalg.vector_norm(complex_score, dim=(1, 2))
+
+        return self.margin - score
+
+    def loss(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tensor:
+
+        pos_score = self(head_index, rel_type, tail_index)
+        neg_score = self(*self.random_sample(head_index, rel_type, tail_index))
+        scores = torch.cat([pos_score, neg_score], dim=0)
+
+        pos_target = torch.ones_like(pos_score)
+        neg_target = torch.zeros_like(neg_score)
+        target = torch.cat([pos_target, neg_target], dim=0)
+
+        return F.binary_cross_entropy_with_logits(scores, target)


Test patch:
diff --git a/test/nn/kge/test_rotate.py b/test/nn/kge/test_rotate.py
new file mode 100644
index 000000000..87640f386
--- /dev/null
+++ b/test/nn/kge/test_rotate.py
@@ -0,0 +1,24 @@
+import torch
+
+from torch_geometric.nn import RotatE
+
+
+def test_rotate():
+    model = RotatE(num_nodes=10, num_relations=5, hidden_channels=32)
+    assert str(model) == 'RotatE(10, num_relations=5, hidden_channels=32)'
+
+    head_index = torch.tensor([0, 2, 4, 6, 8])
+    rel_type = torch.tensor([0, 1, 2, 3, 4])
+    tail_index = torch.tensor([1, 3, 5, 7, 9])
+
+    loader = model.loader(head_index, rel_type, tail_index, batch_size=5)
+    for h, r, t in loader:
+        out = model(h, r, t)
+        assert out.size() == (5, )
+
+        loss = model.loss(h, r, t)
+        assert loss >= 0.
+
+        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
+        assert mean_rank <= 10
+        assert hits_at_10 == 1.0