---
language: en
license: apache-2.0
tags:
  - knowledge-graph
  - rgcn
  - transe
  - metaqa
  - wikimovies
  - graph-embedding
  - experimental
library_name: pytorch
pipeline_tag: feature-extraction
---

# 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](https://github.com/riverlab/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

```python
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

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

### Batch encoding (DataLoader)

```python
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:

```bash
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](https://github.com/riverlab/graph_int_LLM)
- Dataset: [camazlucas/MetaQA](https://huggingface.co/datasets/camazlucas/MetaQA)
- Paper: preprint TBD (Stage 1 MVP)