wikimovies-transe-rgcn (v1 — experimental)

⚠️ Experimental baseline — part of an active research PoC.
This checkpoint is the first public v1 of the Graph-Embedding Alignment experiment (graph_int_LLM). It is not final: link-prediction quality is below our internal target (MRR > 0.2), retrieval coverage is only ~56% at N_max=500, and we expect to publish improved v2+ weights after stronger KGE init, longer R-GCN training, and larger subgraph budgets.

Safe to use for reproduction and Phase 3 projection experiments; do not treat these metrics as SOTA on WikiMovies.

Frozen TransE-initialized 2-layer R-GCN encoder for the WikiMovies knowledge graph, trained for link prediction and used as a frozen subgraph encoder in the Graph-Align pipeline (graph tokens → projection MLP → frozen Gemma 4).

Property	Value
Entities	43,234
Relations	9
Triples	134,741
Embedding dim	128
Init	TransE (200 epochs, PyKEEN)
R-GCN train	50 epochs, DistMult link-pred, frozen after train
QA dataset	MetaQA 3-hop (vanilla)

Metrics (recorded 2026-06-05)

Link prediction (held-out triple ranking)

Model	Split	MRR	Hits@10	Hits@1
TransE	val	0.069	17.2%	0.6%
TransE	test	0.068	17.0%	0.4%
R-GCN (this checkpoint)	val	0.073	14.4%	3.9%
R-GCN (this checkpoint)	test	0.079	15.4%	4.4%

Internal Phase 2 target: MRR > 0.2 — not met (test MRR ≈ 0.079). This does not block the MVP path: QA accuracy depends more on retrieval coverage and projection alignment than global link-pred MRR.

Subgraph retrieval (MetaQA train, hop_max=3, N_max=500)

Metric	Value
answer_in_khop	100.0%
answer_in_retrieved (capped subgraph)	56.2%

~44% of training QA pairs have the gold answer outside the capped 500-node subgraph.

Encode latency (1000 subgraphs, sequential forward)

Device	Mode	ms/subgraph
CPU	node	2.98
CPU	graph	2.91
MPS	node	40.32
MPS	graph	30.86

Sequential per-subgraph forward; MPS slower than CPU on small subgraphs (kernel launch overhead). Use CPU or batched encode_subgraphs for inference.

How we plan to improve (roadmap)

1. (link-pred) Stronger KGE init: train ComplEx or RotatE (PyKEEN) on same split; use as R-GCN node init — expected: Often +0.05–0.15 MRR on multi-rel KGs vs TransE L1
2. (link-pred) R-GCN hyperparam sweep: 100–200 epochs, lr 5e-4, dropout 0.1, num_bases=9, unfreeze/finetune init — expected: Modest MRR gain; current 50-epoch run under-trained
3. (link-pred) Better negative sampling: type-constrained / degree-aware (PyKEEN Bernoulli or LCWA+filtered) — expected: Better ranking tail prediction
4. (link-pred) Add inverse relations (+9 rels) and self-loops in message-passing graph — expected: Improves undirected reachability in conv layers
5. (link-pred) Calibrate eval: run PyKEEN TransE/ComplEx baseline on identical split; verify our eval matches — expected: Rule out eval implementation gap
6. (retrieval) Raise N_max budget (1000 / 2000) and rebuild subgraphs; re-measure answer_in_retrieved_rate (python scripts/build_subgraphs.py --n-max-budget 2000) — expected: Direct lift toward 80%+ retrieved; linear cost in subgraph size
7. (retrieval) Smarter truncation: hop-first then degree (already); add answer-aware cap ablation for upper bound — expected: Diagnostic only — sets ceiling for retrieval

MVP path (current): use this frozen v1 encoder + train only the projection MLP into Gemma; retrain encoder separately and publish v2 when link-pred / retrieval metrics improve.

Usage

From Hugging Face Hub

from graph_align.encoders import load_encoder
from graph_align.data.subgraph import extract_subgraph
from graph_align.data.kg import load_kg_cache

encoder = load_encoder("shikhragimov/wikimovies-transe-rgcn", device="cpu")
kg = load_kg_cache()  # build locally: python scripts/build_kg_cache.py

sg = extract_subgraph(kg, topic_entity="The Matrix (1999 film)", hop_max=3, n_max=500)
node_emb = encoder.encode_subgraph(sg, mode="node")   # (N, 128)
graph_emb = encoder.encode_subgraph(sg, mode="graph")  # (1, 128)

From local checkpoint

encoder = load_encoder("checkpoints/wikimovies_rgcn.pt", device="cpu")

Batch encoding (DataLoader)

from graph_align.encoders import load_encoder

encoded = encoder.encode_dataloader_batch(batch, mode="node")
# encoded.embeddings: (B, N_max, 128), encoded.mask: (B, N_max)

Install the project from source:

git clone https://github.com/riverlab/graph_int_LLM.git
cd graph_int_LLM
pip install -e ".[dev]"
python scripts/build_kg_cache.py

Files in this repo

File	Description
wikimovies_rgcn.pt	Frozen R-GCN weights + embedded entity/relation maps
encoder_config.json	Architecture hyperparameters and metadata
metrics.json	Link-pred + rollup metrics JSON
entity_to_id.json / relation_to_id.json	String ↔ id mappings
id_to_entity.json / id_to_relation.json	Reverse lookups

Limitations

• Trained only on WikiMovies; not a general-purpose KG encoder.
• v1 experimental — expect breaking changes when v2 is published.
• Subgraph encoding requires the same KG vocabulary (entity strings from WikiMovies KB).
• MPS can be slower than CPU for small subgraphs (per-forward overhead).

Citation / links

• Code: graph_int_LLM
• Dataset: camazlucas/MetaQA
• Paper: preprint TBD (Stage 1 MVP)