NBFNet — what it is, what we run, and how to size it for PrimeKG

Reading order: [paper PDF](NBF Networks Paper.pdf) → this doc → PDR2 fork README. Source: Zhu, Zhang, Xhonneux, Tang. Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction. NeurIPS 2021. arXiv:2106.06935. Code: DeepGraphLearning/NBFNet.

1. The one-paragraph version

NBFNet learns a pair representation h_q(u, v) for every candidate (head, query, tail) triple by parameterizing the generalized Bellman-Ford algorithm. The idea: any path-based link prediction method (Katz index, personalized PageRank, graph distance, widest path, most reliable path) can be written as h_q(u,v) = ⊕_{P ∈ P_uv} ⊗_{e ∈ P} w_q(e) for some semiring ⟨⊕, ⊗⟩. NBFNet replaces these handcrafted operators with three learned neural functions: INDICATOR (boundary condition — sets h_u^(0) = q), MESSAGE (the ⊗ — DistMult / TransE / RotatE), AGGREGATE (the ⊕ — sum / mean / max / PNA). T layers of Bellman-Ford iteration give a representation of v conditioned on u and q; an MLP on that vector produces the link score. The crucial structural difference vs. RGCN/CompGCN is that NBFNet computes a separate, source-conditioned node representation for each query — every other GNN computes one shared embedding per node and decodes pairs from it.

2. Architecture in numbers (paper-default, FB15k-237 / WN18RR)

Component	Value	Notes
Layers T	6	T=4 already gets 0.409 MRR vs 0.415 at T=6 (Table 6b) — diminishing returns past 6
Hidden dim d	32	All hidden layers same dim
MLP head	2 layers × 64 hidden	Final score f(h_q(u,v)) → ℝ
Message function	DistMult (best)	Hadamard product h ⊙ r. Also tried TransE (h + r) and RotatE (complex rotation). DistMult+PNA is best on FB15k-237 (Table 15).
Aggregate function	PNA (best)	Principal Neighborhood Aggregation: concatenates {mean, max, min, std} × {1, log(deg), 1/log(deg)} = 12 channels, projected to d. Output dim is 13d before linear (12 + boundary self-loop).
Optimizer	Adam, lr 5e-3	Self-adversarial neg sampling, temperature 0.5 (FB15k-237) / 1.0 (WN18RR)
Batch	256 / 128 positive	32 negatives per positive
Epochs	20
Augmentation	Reverse triples (v, q⁻¹, u)	Doubles `
Edge dropout	Yes — drop direct query edges	Forces the model to learn non-trivial paths
Layer norm	Yes	After AGGREGATE
Residual	short_cut: yes	Per layer if shapes match

Wall-clock on 4× V100: training 9.7 hr (FB15k-237), 4.4 hr (WN18RR); inference 4 min / 2.4 min.

3. Parameter count — the punchline

Table 8 in the paper:

Component	Formula	FB15k-237 (|R|=474, T=6, d=32, MLP m=64)
INDICATOR	|R| × d	15,168
MESSAGE	T × |R| × d × (d+1)	3,003,264
AGGREGATE	T × d × (13d+3)	80,448 (PNA)
MLP f(·)	m(2d+1) + m + 1	4,225
Total		3,103,105

Three properties to internalize:

1. No |V| term. NBFNet has no entity embedding table at all. Initial node states are zero except for the query source, which gets the query embedding. Every other node’s representation is derived by message passing.
2. No |E| term. Edge representations are reused per relation type, not per edge. (w_q(x, r, v) = W_r q + b_r.)
3. Linear in |R|. The MESSAGE block dominates and grows like T × |R| × d². So scaling pain is in the relations, not the graph.

This is why TransE needs 30M params on FB15k-237 (one d-dim vector per node) but NBFNet needs 3M (none).

4. Time and memory complexity

Per Bellman-Ford run:

• Time: O(T · (|E|d + |V|d²))
• Memory (naive): O(|E|d) — every edge holds a message vector
• Memory (fused rspmm): O(|V|d) — custom CUDA kernel computes message + scatter-add in one pass without materializing messages

Amortized per triple (one BF run scores |V| candidates simultaneously): O(|E|d/|V| + d²). Compared to RGCN’s O(|E|d²) and GraIL’s O(|E|d²), NBFNet is the only inductive method that’s not quadratic in d along the dominant |E| term.

5. ULTRA — what our fork actually is

Our lib/pdr2/code/ultra/ is not plain NBFNet. It’s ULTRA (Galkin et al. 2024), a strict superset that addresses NBFNet’s one remaining transfer limitation: per-relation parameters. ULTRA stacks two NBFNets:

                ┌─────────────────────┐
query relation ──►│   RelNBFNet (6L)    │── query-conditioned
                │ on graph of relations│   relation embeddings
                │  (4 fundamental rels:│   (one per |R|, per query)
                │  h2h, h2t, t2h, t2t) │
                └──────────┬──────────┘
                           │
                ┌──────────▼──────────┐
(h, ?, ?, KG) ─►│  EntityNBFNet (6L)  │── pair scores h_q(h, v)
                │   on the actual KG   │   for all v ∈ V
                └─────────────────────┘

The relation graph is universal: 4 fundamental interactions (head-to-head, head-to-tail, tail-to-head, tail-to-tail) regardless of how many real relations the KG has. So the only thing that changes when you swap in a new KG is the data — no new parameters. Pretrained ultra_3g.pth, ultra_4g.pth, ultra_50g.pth are ~2 MB each (≈168K parameters total) and transfer zero-shot to any KG. Local copies of these checkpoints live in ~/git/ULTRA/ckpts/ once the notebook downloads them.

Our lib/pdr2/code/ultra/models.py defines:

• Ultra — the wrapper that runs both NBFNets in sequence
• RelNBFNet — boundary uses torch.ones(d) instead of a query embedding (it operates on the relation graph; the query is a node there)
• EntityNBFNet — boundary uses the relation embeddings produced by RelNBFNet, fed in via self.query and layer.relation before forward

Both inherit from base_nbfnet.py:BaseNBFNet, which contains the bellmanford loop (lines 118–158) and the canonical NBFNet ops: remove_easy_edges, negative_sample_to_tail, beam-search visualization. The actual message-passing convolution is layers.py:GeneralizedRelationalConv and is essentially identical to upstream NBFNet — same DistMult/TransE/RotatE message functions, same PNA, same fused rspmm kernel.

6. Our config, side-by-side with the paper

lib/pdr2/code/config/transductive/inference.yaml and the PrimeKG-specific config inlined in lib/pdr2/notebooks/pdr2_run_primekg.ipynb:

Knob	Paper (FB15k-237)	Our PDR2 (PrimeKG)	Why we differ
T (layers)	6	6	Same
d (hidden)	32	64	ULTRA default — relation graph reps benefit from extra capacity
Message func	DistMult	DistMult	Same
Aggregate func	PNA	sum	ULTRA default — PNA’s [1, log_deg, 1/log_deg] scaling can blow up on PrimeKG super-hubs (avg deg 125, max deg ~10⁴ on protein-protein nodes); sum is safer until we measure
Optimizer	Adam	AdamW in the configs. The Vertex trainer runs Adam, hard-coded (training.py:1100), and ignores the YAML block (#861)	Weight-decay is helpful for the larger ULTRA
LR	5e-3	5e-4 in the configs. The Vertex trainer runs a flat 1e-4: its plateau scheduler is never stepped (#861)	10× lower — empirically necessary when training-set size grows past ~1M edges
Negative weighting	Self-adversarial, temperature 0.5 / 1.0	None on Vertex: plain mean BCE over the positive and its 10 strict negatives. adversarial_temperature is read only by pretrain.py and run_kgia.py (#861)	Inherited from the PDR2 release
Batch	256 pos × 32 neg = 8,192	512 × 10 = 5,120 (inference) / 64 (pretraining)	Memory-bound by `
Epochs	20	5 for PrimeKG (10 for GP-KG joint pretrain)	Compute budget — 1 Colab Pro+ session ≈ 6 hr
BPE	(full epoch)	500 batches/epoch	Subsampled — full epoch on PrimeKG would be ~16K batches
Edge dropout	Drop query edges	Same (remove_one_hop=False default, drops only the exact query triple)	Same
Reverse aug	Yes	Yes	Same — required by the bellmanford loop’s symmetry

7. Resizing for PrimeKG — concrete options

PrimeKG numbers (per `docs/KG_CATALOG.md`):

• |V| = 129,375 entities
• |E| = 8,100,498 edges (already bidirectional)
• |R| = 30 relation types (60 with reverse)
• avg degree ≈ 125 (≈6× FB15k-237’s avg deg 21, ≈2× Hetionet’s 96)

What is not a problem

• Parameter count. With |R|=60 and ULTRA’s relation-graph trick, our model is ~168K parameters total — smaller than NBFNet on FB15k-237 (3M). PrimeKG won’t grow the model.
• Inductivity. ULTRA + NBFNet are entity-agnostic. Adding/removing PrimeKG nodes between train and eval is fine.

What is a problem

• Per-batch wall-clock. Bellman-Ford is O(T × |E| × d) per query. PrimeKG has 6.5× the edges of GP-KG → expect ~6× slower per batch. Notebook predicts ~9 sec/batch (we measured 3.3s on GP-KG).
• Memory (without rspmm). Naive memory is O(|E|·d·batch_size). At d=64, batch=512: 8.1M × 64 × 512 × 4B = 1 TB. Fused rspmm drops this to O(|V|·d·batch_size) = 17 GB → tractable but tight on a single A100/H100.
• Validation cost. PrimeKG’s standard split is ~1.6M valid triples. Per-triple eval at batch 512 needs ~3,200 batches → ~8 hr. The notebook subsamples val_n_sample=5000 for fast iteration; full eval is run once at the end overnight (~7-8 hr).

Knobs worth tuning, ranked by leverage

1. Drop T from 6 → 4. Paper Table 6b: T=4 → 0.409 MRR, T=6 → 0.415 (1.5% delta) on FB15k-237. Per-batch time scales linearly with T. 33% wall-clock savings for ~1.5% metric loss is the best Pareto move if we’re compute-bound. (Try first.)

2. Aggregate: sum → PNA. Paper Table 6a: PNA gives +7% MRR on FB15k-237 vs sum (DistMult+PNA 0.415 vs DistMult+sum 0.388). Adds ~6.5× the linear-projection params per layer (output is 13d not 2d) but is essentially free at runtime — the rspmm kernel computes all four PNA stats in one pass. Try this on PrimeKG with the scale.clamp(min=1e-2) numerical guard already in our layers.py:171. PrimeKG’s high max-degree might still trip it; if loss goes NaN, fall back to sum.

3. Hidden dim 64 → 96 or 128. ULTRA’s d=64 is a guess inherited from upstream — the original NBFNet found d=32 was enough for FB15k-237. PrimeKG has more entity types (10 vs 1) and more semantic structure, so capacity should help. Cost: AGGREGATE linear is O(d²), MESSAGE is O(d) per edge, so doubling d quadruples the AGGREGATE cost but only doubles the dominant |E|d term. Test d=96 first; only go to 128 if metrics warrant.

4. Learning rate schedule. Every foundation trained on Vertex so far ran a flat Adam lr=1e-4, not the configs’ AdamW 5e-4 (#861). The paper used a flat Adam lr=5e-3, 50× higher, on a 10× smaller graph. Adding cosine decay or warmup is cheap. Whether Run 3 changes it is open (#861): it would break comparability with Run 2.

5. Bidirectional BFS subsampling (paper Appendix G). For WikiKG90M (87M nodes), the authors restricted each query to a 2-hop bidirectional BFS subgraph with degree clamp m=100. This gives a guarantee: any path of length ≤ 2k between head and any tail candidate is preserved. PrimeKG is small enough that we don’t need this (our 8M edges fit fine), but this is the playbook if we ever pre-train on a PrimeKG + DRKG + Hetionet union (which would push past 50M edges). The k=2, m=100 settings are a sane starting point.

Knobs not worth tuning

• Layers within RelNBFNet vs EntityNBFNet asymmetrically. Our config sets both to 6 — keep them in sync. RelNBFNet’s graph is tiny (4 fundamental relations × |R| nodes), so its T contributes negligibly to wall time.
• Negative count. 32 negatives is paper-validated on three datasets. Going lower trades MRR for speed; going higher hits memory.
• Removing layer norm or short-cut. Paper showed layer norm was needed for stability on FB15k-237. Don’t strip it unless we have a specific reason.

8. Open questions — things this audit can not answer

These need empirical runs, not paper-reading:

• Does PNA work on PrimeKG without numerical NaNs? Theoretically yes (the clamp protects it); practically need to A/B against sum on a 1-epoch run.
• What’s the right val_n_sample? Currently 5000; PDR1 used full eval. The 5000-sample MRR could be a biased estimator of full-eval MRR — needs a one-shot calibration check.
• Does increasing d from 64 to 96 actually help on PrimeKG, or does it just cost compute? No literature signal — needs an A/B.
• Are PrimeKG’s super-hubs (very high-degree proteins) hurting message passing? Paper doesn’t address this; we should compute the degree distribution and consider degree-aware sampling.
• Is sum-aggregate creating loss-of-information that propagates through 6 layers? The paper only ran sum vs PNA on FB15k-237 (avg deg 21). PrimeKG (avg deg 125) is well outside that calibration regime.

9. Citations to keep in scope

For our own PDR/MODS docs, when referring to NBFNet:

Zhu, Z., Zhang, Z., Xhonneux, L.-P., & Tang, J. (2021). Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction. Advances in Neural Information Processing Systems 34 (NeurIPS 2021). arXiv:2106.06935.

For ULTRA (which is what we actually run):

Galkin, M., Yuan, X., Mostafa, H., Tang, J., & Zhu, Z. (2024). Towards Foundation Models for Knowledge Graph Reasoning. International Conference on Learning Representations (ICLR 2024). github.com/DeepGraphLearning/ULTRA.