Repository: pyg-team/pytorch_geometric
Pull request #6958 (2023-03-18 20:33:40)

feat: add `DistMult` KGE model
Achieves .40 Hits@10 on FB15K in example (.41 on validation), which is comparable to the [PapersWithCode](https://paperswithcode.com/paper/embedding-entities-and-relations-for-learning) metric value of .419

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

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

```diff
@@           Coverage Diff           @@
##           master    #6958   +/-   ##
=======================================
  Coverage   91.45%   91.46%           
=======================================
  Files         429      430    +1     
  Lines       23357    23381   +24     
=======================================
+ Hits        21361    21385   +24     
  Misses       1996     1996           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6958?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/6958?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/distmult.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6958?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2tnZS9kaXN0bXVsdC5weQ==) | `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 1333c5bc2..4b26a4634 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 `DistMult` KGE model ([#6958](https://github.com/pyg-team/pytorch_geometric/pull/6958))
 - Added `HeteroData.set_value_dict` functionality ([#6961](https://github.com/pyg-team/pytorch_geometric/pull/6961))
 - Added PyTorch >= 2.0 support ([#6934](https://github.com/pyg-team/pytorch_geometric/pull/6934))
 - Added PyTorch Lightning >= 2.0 support ([#6929](https://github.com/pyg-team/pytorch_geometric/pull/6929))
diff --git a/README.md b/README.md
index 6ea8e5e20..bce039d45 100644
--- a/README.md
+++ b/README.md
@@ -301,7 +301,8 @@ Unlike simple stacking of GNN layers, these models could involve pre-processing,
 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/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)]
+* **[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)]
 </details>
 
 **GNN operators and utilities:**
diff --git a/examples/kge_fb15k_237.py b/examples/kge_fb15k_237.py
index 621fe0b44..b4d1ac7f0 100644
--- a/examples/kge_fb15k_237.py
+++ b/examples/kge_fb15k_237.py
@@ -5,9 +5,13 @@ import torch
 import torch.optim as optim
 
 from torch_geometric.datasets import FB15k_237
-from torch_geometric.nn import ComplEx, TransE
+from torch_geometric.nn import ComplEx, DistMult, TransE
 
-model_map = {'transe': TransE, 'complex': ComplEx}
+model_map = {
+    'transe': TransE,
+    'complex': ComplEx,
+    'distmult': DistMult,
+}
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--model', choices=model_map.keys(), type=str.lower,
@@ -36,8 +40,9 @@ loader = model.loader(
 )
 
 optimizer_map = {
-    'transe': torch.optim.Adam(model.parameters(), lr=0.01),
-    'complex': optim.Adagrad(model.parameters(), lr=0.001, weight_decay=1e-6)
+    '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),
 }
 optimizer = optimizer_map[args.model]
 
diff --git a/torch_geometric/nn/kge/__init__.py b/torch_geometric/nn/kge/__init__.py
index 329167077..35e285218 100644
--- a/torch_geometric/nn/kge/__init__.py
+++ b/torch_geometric/nn/kge/__init__.py
@@ -1,9 +1,11 @@
 from .base import KGEModel
 from .transe import TransE
 from .complex import ComplEx
+from .distmult import DistMult
 
 __all__ = classes = [
     'KGEModel',
     'TransE',
     'ComplEx',
+    'DistMult',
 ]
diff --git a/torch_geometric/nn/kge/complex.py b/torch_geometric/nn/kge/complex.py
index 2dc48b233..7e9711a13 100644
--- a/torch_geometric/nn/kge/complex.py
+++ b/torch_geometric/nn/kge/complex.py
@@ -8,7 +8,7 @@ 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.
+    Prediction" <https://arxiv.org/abs/1606.06357>`_ paper.
 
     :class:`ComplEx` models relations as complex-valued bilinear mappings
     between head and tail entities using the Hermetian dot product.
diff --git a/torch_geometric/nn/kge/distmult.py b/torch_geometric/nn/kge/distmult.py
new file mode 100644
index 000000000..8063d32d6
--- /dev/null
+++ b/torch_geometric/nn/kge/distmult.py
@@ -0,0 +1,82 @@
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.nn.kge import KGEModel
+
+
+class DistMult(KGEModel):
+    r"""The DistMult model from the `"Embedding Entities and Relations for
+    Learning and Inference in Knowledge Bases"
+    <https://arxiv.org/abs/1412.6575>`_ paper.
+
+    :class:`DistMult` models relations as diagonal matrices, which simplifies
+    the bi-linear interaction between the head and tail entities to the score
+    function:
+
+    .. math::
+        d(h, r, t) = < \mathbf{e}_h,  \mathbf{e}_r, \mathbf{e}_t >
+
+    .. note::
+
+        For an example of using the :class:`DistMult` 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.
+            (default: :obj:`1.0`)
+        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.margin = margin
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.node_emb.weight)
+        torch.nn.init.xavier_uniform_(self.rel_emb.weight)
+
+    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)
+
+        return (head * rel * tail).sum(dim=-1)
+
+    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))
+
+        return F.margin_ranking_loss(
+            pos_score,
+            neg_score,
+            target=torch.ones_like(pos_score),
+            margin=self.margin,
+        )


Test patch:
diff --git a/test/nn/kge/test_distmult.py b/test/nn/kge/test_distmult.py
new file mode 100644
index 000000000..b9f02f139
--- /dev/null
+++ b/test/nn/kge/test_distmult.py
@@ -0,0 +1,24 @@
+import torch
+
+from torch_geometric.nn import DistMult
+
+
+def test_distmult():
+    model = DistMult(num_nodes=10, num_relations=5, hidden_channels=32)
+    assert str(model) == 'DistMult(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