Knowledge Graph Reasoning

Overview

Knowledge graphs represent structured information as entities and relations in a graph format. TorchDrug provides comprehensive support for knowledge graph completion (link prediction) using embedding-based models and neural reasoning approaches.

Available Datasets

General Knowledge Graphs

FB15k (Freebase subset): - 14,951 entities - 1,345 relation types - 592,213 triples - General world knowledge from Freebase

FB15k-237: - 14,541 entities - 237 relation types - 310,116 triples - Filtered version removing inverse relations - More challenging benchmark

WN18 (WordNet): - 40,943 entities (word senses) - 18 relation types (lexical relations) - 151,442 triples - Linguistic knowledge graph

WN18RR: - 40,943 entities - 11 relation types - 93,003 triples - Filtered WordNet removing easy inverse patterns

Biomedical Knowledge Graphs

Hetionet: - 45,158 entities (genes, compounds, diseases, pathways, etc.) - 24 relation types (treats, causes, binds, etc.) - 2,250,197 edges - Integrates 29 public biomedical databases - Designed for drug repurposing and disease understanding

Task: KnowledgeGraphCompletion

The primary task for knowledge graphs is link prediction - given a head entity and relation, predict the tail entity (or vice versa).

Task Modes

Head Prediction: - Given (?, relation, tail), predict head entity - "What can cause Disease X?"

Tail Prediction: - Given (head, relation, ?), predict tail entity - "What diseases does Gene X cause?"

Both: - Predict both head and tail - Standard evaluation protocol

Evaluation Metrics

Ranking Metrics: - Mean Rank (MR): Average rank of correct entity - Mean Reciprocal Rank (MRR): Average of 1/rank - Hits@K: Percentage of correct entities in top K predictions - Typically reported for K=1, 3, 10

Filtered vs Raw: - Filtered: Remove other known true triples from ranking - Raw: Rank among all possible entities - Filtered is standard for evaluation

Embedding Models

Translational Models

TransE (Translation Embedding): - Represents relations as translations in embedding space - h + r ≈ t (head + relation ≈ tail) - Simple and effective baseline - Works well for 1-to-1 relations - Struggles with N-to-N relations

RotatE (Rotation Embedding): - Relations as rotations in complex space - Better handles symmetric and inverse relations - State-of-the-art on many benchmarks - Can model composition patterns

Semantic Matching Models

DistMult: - Bilinear scoring function - Handles symmetric relations naturally - Cannot model asymmetric relations - Fast and memory efficient

ComplEx: - Complex-valued embeddings - Models asymmetric and inverse relations - Better than DistMult for most graphs - Balances expressiveness and efficiency

SimplE: - Extends DistMult with inverse relations - Fully expressive (can represent any relation pattern) - Two embeddings per entity (canonical and inverse)

Neural Logic Models

NeuralLP (Neural Logic Programming): - Learns logical rules through differentiable operations - Interprets predictions via learned rules - Good for sparse knowledge graphs - Computationally more expensive

KBGAT (Knowledge Base Graph Attention): - Graph attention networks for KG completion - Learns entity representations from neighborhood - Handles unseen entities through inductive learning - Better for incomplete graphs

Training Workflow

Basic Pipeline

from torchdrug import datasets, models, tasks, core

# Load dataset
dataset = datasets.FB15k237("~/kg-datasets/")

# Define model
model = models.RotatE(
    num_entity=dataset.num_entity,
    num_relation=dataset.num_relation,
    embedding_dim=2000,
    max_score=9
)

# Define task
task = tasks.KnowledgeGraphCompletion(
    model,
    num_negative=128,
    adversarial_temperature=2,
    criterion="bce"
)

# Train with PyTorch Lightning or custom loop

Negative Sampling

Strategies: - Uniform: Sample entities uniformly at random - Self-Adversarial: Weight samples by current model's scores - Type-Constrained: Sample only valid entity types for relation

Parameters: - num_negative: Number of negative samples per positive triple - adversarial_temperature: Temperature for self-adversarial weighting - Higher temperature = more focus on hard negatives

Loss Functions

Binary Cross-Entropy (BCE): - Treats each triple independently - Balanced classification between positive and negative

Margin Loss: - Ensures positive scores higher than negative by margin - max(0, margin + score_neg - score_pos)

Logistic Loss: - Smooth version of margin loss - Better gradient properties

Model Selection Guide

By Relation Patterns

1-to-1 Relations: - TransE works well - Any model will likely succeed

1-to-N Relations: - DistMult, ComplEx, SimplE - Avoid TransE

N-to-1 Relations: - DistMult, ComplEx, SimplE - Avoid TransE

N-to-N Relations: - ComplEx, SimplE, RotatE - Most challenging pattern

Symmetric Relations: - DistMult, ComplEx - RotatE with proper initialization

Antisymmetric Relations: - ComplEx, SimplE, RotatE - Avoid DistMult

Inverse Relations: - ComplEx, SimplE, RotatE - Important for bidirectional reasoning

Composition: - RotatE (best) - TransE (reasonable) - Captures multi-hop paths

By Dataset Characteristics

Small Graphs (< 50k entities): - ComplEx or SimplE - Lower embedding dimensions (200-500)

Large Graphs (> 100k entities): - DistMult for efficiency - RotatE for accuracy - Higher dimensions (500-2000)

Sparse Graphs: - NeuralLP (learns rules from limited data) - Pre-train embeddings on larger graphs

Dense, Complete Graphs: - Any embedding model works well - Choose based on relation patterns

Biomedical/Domain Graphs: - Consider type constraints in sampling - Use domain-specific negative sampling - Hetionet benefits from relation-specific models

Advanced Techniques

Multi-Hop Reasoning

Chain multiple relations to answer complex queries: - "What drugs treat diseases caused by gene X?" - Requires path-based or rule-based reasoning - NeuralLP naturally supports this

Temporal Knowledge Graphs

Extend to time-varying facts: - Add temporal information to triples - Predict future facts - Requires temporal encoding in models

Few-Shot Learning

Handle relations with few examples: - Meta-learning approaches - Transfer from related relations - Important for emerging knowledge

Inductive Learning

Generalize to unseen entities: - KBGAT and other GNN-based methods - Use entity features/descriptions - Critical for evolving knowledge graphs

Biomedical Applications

Drug Repurposing

Predict "drug treats disease" links in Hetionet: 1. Train on known drug-disease associations 2. Predict new treatment candidates 3. Filter by mechanism (gene, pathway involvement) 4. Validate predictions experimentally

Disease Gene Discovery

Identify genes associated with diseases: 1. Model gene-disease-pathway networks 2. Predict missing gene-disease links 3. Incorporate protein interactions, expression data 4. Prioritize candidates for validation

Protein Function Prediction

Link proteins to biological processes: 1. Integrate protein interactions, GO terms 2. Predict missing GO annotations 3. Transfer function from similar proteins

Common Issues and Solutions

Issue: Poor performance on specific relation types - Solution: Analyze relation patterns, choose appropriate model, or use relation-specific models

Issue: Overfitting on small graphs - Solution: Reduce embedding dimension, increase regularization, or use simpler models

Issue: Slow training on large graphs - Solution: Reduce negative samples, use DistMult for efficiency, or implement mini-batch training

Issue: Cannot handle new entities - Solution: Use inductive models (KBGAT), incorporate entity features, or pre-compute embeddings for new entities based on their neighbors

Best Practices

1. Start with ComplEx or RotatE for most tasks
2. Use self-adversarial negative sampling
3. Tune embedding dimension (typically 500-2000)
4. Apply regularization to prevent overfitting
5. Use filtered evaluation metrics
6. Analyze performance per relation type
7. Consider relation-specific models for heterogeneous graphs
8. Validate predictions with domain experts