Repository: pyg-team/pytorch_geometric
Pull request #6898 (2023-03-12 00:29:02)

feat: Add `ComplEx` KGE
Hello! we are hoping to add the ComplEx knowledge graph embedding to PyG as part of our CS224W final project at Stanford.

Our implementation largely follows that of TransE. We implemented complex vectors by breaking the real and imaginary components into separate vectors, but are open to other approaches.

We'd also like to seek input on whether an example is necessary like for TransE [here](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/transe_fb15k_237.py), and if so, whether it would be preferred to be in the same file or a new one.

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

> :exclamation: Current head 1a154f2 differs from pull request most recent head 827f05a. Consider uploading reports for the commit 827f05a to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #6898      +/-   ##
==========================================
+ Coverage   91.56%   91.58%   +0.01%     
==========================================
  Files         427      428       +1     
  Lines       23177    23207      +30     
==========================================
+ Hits        21222    21254      +32     
+ Misses       1955     1953       -2     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6898?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/6898?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/base.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6898?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2tnZS9iYXNlLnB5) | `100.00% <100.00%> (+4.00%)` | :arrow_up: |
| [torch\_geometric/nn/kge/complex.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6898?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2tnZS9jb21wbGV4LnB5) | `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)

riyavsinha: Hi Matthias! I think this is ready for review again. The [PapersWithCode](https://paperswithcode.com/paper/complex-embeddings-for-simple-link-prediction) metric for ComplEx for Hits@10 on fb15k-237 is .42, and our current optimizer settings in the example are consistently about .36-37 (we didn't want to increase the hidden size when tuning though to keep it fast as an example). If you'd like us to find a better hyperparameter set please let us know! 

In the current set, we used Adagrad which ended up performing much better than Adam and which was the suggested optimizer in the paper, and we added some L2 weight regularization as also suggested by the paper.

We have also updated the loss function to use BCELogits since that is more numerically stable


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f05f78e46..a3230dd65 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `ComplEx` KGE model ([#6898](https://github.com/pyg-team/pytorch_geometric/pull/6898))
 - Added option to write benchmark results to csv ([#6888](https://github.com/pyg-team/pytorch_geometric/pull/6888))
 - Added `HeteroLayerNorm` and `HeteroBatchNorm` layers ([#6838](https://github.com/pyg-team/pytorch_geometric/pull/6838))
 - Added the `HeterophilousGraphDataset` suite ([#6846](https://github.com/pyg-team/pytorch_geometric/pull/6846))
diff --git a/README.md b/README.md
index 6a7fcb494..6ea8e5e20 100644
--- a/README.md
+++ b/README.md
@@ -300,7 +300,8 @@ Unlike simple stacking of GNN layers, these models could involve pre-processing,
 * **[LINKX](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.LINKX.html)** from Lim *et al.*: [Large Scale Learning on Non-Homophilous Graphs:
 New Benchmarks and Strong Simple Methods](https://arxiv.org/abs/2110.14446) (NeurIPS 2021) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/linkx.py)]
 * **[RevGNN](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.GroupAddRev.html)** from Li *et al.*: [Training Graph Neural with 1000 Layers](https://arxiv.org/abs/2106.07476) (ICML 2021) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/rev_gnn.py)]
-* **[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/transe_fb15k_237.py)]
+* **[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/pdf/1606.06357.pdf) (ICML 2016) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)]
 </details>
 
 **GNN operators and utilities:**
diff --git a/examples/transe_fb15k_237.py b/examples/kge_fb15k_237.py
similarity index 74%
rename from examples/transe_fb15k_237.py
rename to examples/kge_fb15k_237.py
index 062cb80f5..621fe0b44 100644
--- a/examples/transe_fb15k_237.py
+++ b/examples/kge_fb15k_237.py
@@ -1,9 +1,18 @@
+import argparse
 import os.path as osp
 
 import torch
+import torch.optim as optim
 
 from torch_geometric.datasets import FB15k_237
-from torch_geometric.nn import TransE
+from torch_geometric.nn import ComplEx, TransE
+
+model_map = {'transe': TransE, 'complex': ComplEx}
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', choices=model_map.keys(), type=str.lower,
+                    required=True)
+args = parser.parse_args()
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'FB15k')
@@ -12,7 +21,7 @@ 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 = TransE(
+model = model_map[args.model](
     num_nodes=train_data.num_nodes,
     num_relations=train_data.num_edge_types,
     hidden_channels=50,
@@ -22,11 +31,15 @@ loader = model.loader(
     head_index=train_data.edge_index[0],
     rel_type=train_data.edge_type,
     tail_index=train_data.edge_index[1],
-    batch_size=10000,
+    batch_size=1000,
     shuffle=True,
 )
 
-optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+optimizer_map = {
+    'transe': torch.optim.Adam(model.parameters(), lr=0.01),
+    'complex': optim.Adagrad(model.parameters(), lr=0.001, weight_decay=1e-6)
+}
+optimizer = optimizer_map[args.model]
 
 
 def train():
diff --git a/torch_geometric/nn/kge/__init__.py b/torch_geometric/nn/kge/__init__.py
index 0be296563..329167077 100644
--- a/torch_geometric/nn/kge/__init__.py
+++ b/torch_geometric/nn/kge/__init__.py
@@ -1,7 +1,9 @@
 from .base import KGEModel
 from .transe import TransE
+from .complex import ComplEx
 
 __all__ = classes = [
     'KGEModel',
     'TransE',
+    'ComplEx',
 ]
diff --git a/torch_geometric/nn/kge/base.py b/torch_geometric/nn/kge/base.py
index 9aff243ec..229fc79aa 100644
--- a/torch_geometric/nn/kge/base.py
+++ b/torch_geometric/nn/kge/base.py
@@ -34,15 +34,17 @@ class KGEModel(torch.nn.Module):
         self.node_emb = Embedding(num_nodes, hidden_channels, sparse=sparse)
         self.rel_emb = Embedding(num_relations, hidden_channels, sparse=sparse)
 
-        self.reset_parameters()
-
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         self.node_emb.reset_parameters()
         self.rel_emb.reset_parameters()
 
-    def forward(self, head_index: Tensor, rel_type: Tensor,
-                tail_index: Tensor) -> Tensor:
+    def forward(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tensor:
         r"""Returns the score for the given triplet.
 
         Args:
@@ -52,8 +54,12 @@ class KGEModel(torch.nn.Module):
         """
         raise NotImplementedError
 
-    def loss(self, head_index: Tensor, rel_type: Tensor,
-             tail_index: Tensor) -> Tensor:
+    def loss(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tensor:
         r"""Returns the loss value for the given triplet.
 
         Args:
@@ -63,8 +69,13 @@ class KGEModel(torch.nn.Module):
         """
         raise NotImplementedError
 
-    def loader(self, head_index: Tensor, rel_type: Tensor, tail_index: Tensor,
-               **kwargs) -> Tensor:
+    def loader(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+        **kwargs,
+    ) -> Tensor:
         r"""Returns a mini-batch loader that samples a subset of triplets.
 
         Args:
@@ -123,8 +134,12 @@ class KGEModel(torch.nn.Module):
         return mean_rank, hits_at_k
 
     @torch.no_grad()
-    def random_sample(self, head_index: Tensor, rel_type: Tensor,
-                      tail_index: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+    def random_sample(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
         r"""Randomly samples negative triplets by either replacing the head or
         the tail (but not both).
 
diff --git a/torch_geometric/nn/kge/complex.py b/torch_geometric/nn/kge/complex.py
new file mode 100644
index 000000000..2dc48b233
--- /dev/null
+++ b/torch_geometric/nn/kge/complex.py
@@ -0,0 +1,94 @@
+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 ComplEx(KGEModel):
+    r"""The ComplEx model from the `"Complex Embeddings for Simple Link
+    Prediction" <https://arxiv.org/pdf/1606.06357.pdf>`_ paper.
+
+    :class:`ComplEx` models relations as complex-valued bilinear mappings
+    between head and tail entities using the Hermetian dot product.
+    The entities and relations are embedded in different dimensional spaces,
+    resulting in the scoring function:
+
+    .. math::
+        d(h, r, t) = Re(< \mathbf{e}_h,  \mathbf{e}_r, \mathbf{e}_t>)
+
+    .. note::
+
+        For an example of using the :class:`ComplEx` 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.
+        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,
+        sparse: bool = False,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse)
+
+        self.node_emb_im = Embedding(num_nodes, hidden_channels, sparse=sparse)
+        self.rel_emb_im = Embedding(num_relations, 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.xavier_uniform_(self.rel_emb.weight)
+        torch.nn.init.xavier_uniform_(self.rel_emb_im.weight)
+
+    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)
+        rel_re = self.rel_emb(rel_type)
+        rel_im = self.rel_emb_im(rel_type)
+        tail_re = self.node_emb(tail_index)
+        tail_im = self.node_emb_im(tail_index)
+
+        return (triple_dot(head_re, rel_re, tail_re) +
+                triple_dot(head_im, rel_re, tail_im) +
+                triple_dot(head_re, rel_im, tail_im) -
+                triple_dot(head_im, rel_im, tail_re))
+
+    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)
+
+
+def triple_dot(x: Tensor, y: Tensor, z: Tensor) -> Tensor:
+    return (x * y * z).sum(dim=-1)
diff --git a/torch_geometric/nn/kge/transe.py b/torch_geometric/nn/kge/transe.py
index c2ed3d28d..084b7efe6 100644
--- a/torch_geometric/nn/kge/transe.py
+++ b/torch_geometric/nn/kge/transe.py
@@ -26,9 +26,9 @@ class TransE(KGEModel):
     .. note::
 
         For an example of using the :class:`TransE` model, see
-        `examples/transe_fb15k_237.py
+        `examples/kge_fb15k_237.py
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
-        transe_fb15k_237.py>`_.
+        kge_fb15k_237.py>`_.
 
     Args:
         num_nodes (int): The number of nodes/entities in the graph.
@@ -54,6 +54,8 @@ class TransE(KGEModel):
         self.margin = margin
         super().__init__(num_nodes, num_relations, hidden_channels, sparse)
 
+        self.reset_parameters()
+
     def reset_parameters(self):
         bound = 6. / math.sqrt(self.hidden_channels)
         torch.nn.init.uniform_(self.node_emb.weight, -bound, bound)
@@ -61,8 +63,13 @@ class TransE(KGEModel):
         F.normalize(self.rel_emb.weight.data, p=self.p_norm, dim=-1,
                     out=self.rel_emb.weight.data)
 
-    def forward(self, head_index: Tensor, rel_type: Tensor,
-                tail_index: Tensor) -> Tensor:
+    def forward(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+    ) -> Tensor:
+
         head = self.node_emb(head_index)
         rel = self.rel_emb(rel_type)
         tail = self.node_emb(tail_index)
@@ -73,8 +80,13 @@ class TransE(KGEModel):
         # Calculate *negative* TransE norm:
         return -((head + rel) - tail).norm(p=self.p_norm, dim=-1)
 
-    def loss(self, head_index: Tensor, rel_type: Tensor,
-             tail_index: Tensor) -> Tensor:
+    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))
 


Test patch:
diff --git a/test/nn/kge/test_complex.py b/test/nn/kge/test_complex.py
new file mode 100644
index 000000000..bc7af98cc
--- /dev/null
+++ b/test/nn/kge/test_complex.py
@@ -0,0 +1,58 @@
+import torch
+
+from torch_geometric.nn import ComplEx
+
+
+def test_complex_scoring():
+    model = ComplEx(num_nodes=5, num_relations=2, hidden_channels=1)
+
+    model.node_emb.weight.data = torch.tensor([
+        [2.],
+        [3.],
+        [5.],
+        [1.],
+        [2.],
+    ])
+    model.node_emb_im.weight.data = torch.tensor([
+        [4.],
+        [1.],
+        [3.],
+        [1.],
+        [2.],
+    ])
+    model.rel_emb.weight.data = torch.tensor([
+        [2.],
+        [3.],
+    ])
+    model.rel_emb_im.weight.data = torch.tensor([
+        [3.],
+        [1.],
+    ])
+
+    score = model(
+        head_index=torch.tensor([1, 3]),
+        rel_type=torch.tensor([1, 0]),
+        tail_index=torch.tensor([2, 4]),
+    )
+    assert score.tolist() == [58., 8.]
+
+
+def test_complex():
+    model = ComplEx(num_nodes=10, num_relations=5, hidden_channels=32)
+    assert str(model) == 'ComplEx(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)
+        assert mean_rank <= 10
+        assert hits_at_10 == 1.0