ContextQMD
Back to Claude Scientific Skills

Knowledge Graph Reasoning

skills/torchdrug/references/knowledge_graphs.md

2.54.05.3 KB
Original Source

Knowledge Graph Reasoning

The official TorchDrug 0.2.1 reasoning tutorial covers two workflows:

  • knowledge graph embeddings with RotatE,
  • neural inductive logic programming with NeuralLP.

Both use tasks.KnowledgeGraphCompletion.

Datasets

Documented knowledge graph datasets:

  • FB15k: 14,951 entities, 1,345 relations, 592,213 triplets
  • FB15k237: 14,541 entities, 237 relations, 310,116 triplets
  • WN18: 40,943 entities, 18 relations, 151,442 triplets
  • WN18RR: 40,943 entities, 11 relations, 93,003 triplets
  • Hetionet: 45,158 entities, 24 relations, 2,025,177 triplets

Use predefined splits:

python
from torchdrug import datasets

dataset = datasets.FB15k237("~/kg-datasets/")
train_set, valid_set, test_set = dataset.split()

RotatE embedding workflow

Model

python
import torch
from torchdrug import core, models, tasks

model = models.RotatE(
    num_entity=dataset.num_entity,
    num_relation=dataset.num_relation,
    embedding_dim=2048,
    max_score=9,
)

embedding_dim=2048 follows the tutorial and may be reduced for memory or speed.

Task

python
task = tasks.KnowledgeGraphCompletion(
    model,
    num_negative=256,
    adversarial_temperature=1,
)
  • num_negative controls negative samples per positive.
  • adversarial_temperature enables score-weighted negative sampling.

Train and evaluate

python
optimizer = torch.optim.Adam(task.parameters(), lr=2e-5)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=1024,
)
solver.train(num_epoch=200)
solver.evaluate("valid")

Add gpus=[0] for a supported CUDA device. Reduce the epoch count for smoke tests.

NeuralLP workflow

NeuralLP learns weighted chain-like rules up to a configured maximum length.

python
model = models.NeuralLP(
    num_relation=dataset.num_relation,
    hidden_dim=128,
    num_step=3,
    num_lstm_layer=2,
)

task = tasks.KnowledgeGraphCompletion(
    model,
    fact_ratio=0.75,
    num_negative=256,
    sample_weight=False,
)

fact_ratio=0.75 reserves 75% of training facts for the background graph used for reasoning.

python
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=64,
)
solver.train(num_epoch=10)
solver.evaluate("valid")

Other documented models

Embedding models:

  • models.TransE
  • models.DistMult
  • models.ComplEx
  • models.SimplE
  • models.RotatE

Graph-attention model:

  • models.KBGAT

Verify each constructor in the model API. Do not transfer argument names from PyKEEN, DGL-KE, or PyTorch Geometric.

Task behavior

KnowledgeGraphCompletion owns:

  • negative sampling,
  • fact-graph construction,
  • loss computation,
  • head and tail prediction,
  • filtered ranking evaluation.

Important constructor options include:

  • criterion
  • metric
  • num_negative
  • margin
  • adversarial_temperature
  • strict_negative
  • fact_ratio
  • sample_weight
  • full_batch_eval

TorchDrug 0.2.1 added full-batch evaluation support. Choose it according to graph size and available memory.

Evaluation

Use filtered ranking metrics:

  • mean rank (MR)
  • mean reciprocal rank (MRR)
  • Hits@1
  • Hits@3
  • Hits@10

Filtered evaluation removes other known true triples before ranking. Preserve training, validation, and test facts exactly as the task expects to avoid leakage or incorrect filtering.

Also report:

  • results by relation,
  • head vs tail prediction,
  • variance across seeds,
  • memory/runtime settings,
  • whether reciprocal relations were added.

Biomedical use

Hetionet supports biomedical link-prediction experiments, but a high model score does not establish a new treatment, causal mechanism, or validated association.

For drug-repurposing analysis:

  1. define the exact relation being predicted,
  2. preserve entity and relation type constraints,
  3. exclude known positives correctly,
  4. check for train/test leakage through inverse or duplicate relations,
  5. calibrate or rank model scores,
  6. validate candidates against independent evidence and domain experts.

TorchDrug's generic KnowledgeGraphCompletion API does not automatically apply biomedical type constraints or causal interpretation.

Common failures

Entity/relation mismatch

Build model sizes from dataset.num_entity and dataset.num_relation.

Evaluation out of memory

Lower batch size or disable full-batch evaluation. Reducing negative samples mainly affects training, not the size of all-entity ranking.

NeuralLP produces invalid shapes

Use num_relation=dataset.num_relation and let KnowledgeGraphCompletion.preprocess() construct the fact graph.

Inflated metrics

Check for inverse-relation leakage, duplicate triples, accidental use of test facts, and raw rather than filtered ranking.