Skip to content

OpenRFM v0.1.0

Latest

Choose a tag to compare

@guaiyoui guaiyoui released this 02 Jun 03:25

Model Card — OpenRFM pretrain (ctx96)

Summary

A relational foundation model pre-trained once on synthetic structural-causal-model
(SCM) relational tasks, intended for in-context (no fine-tuning) prediction on
relational databases (evaluated on RelBench). This is the best leak-free,
trick-free checkpoint from the OpenRFM reproduction.

  • Asset: openrfm-pretrain-ctx96.pt (released via GitHub Releases)
  • Source run: runs/pretrain_ctx96/checkpoint_step_1000.pt
  • Size: ~224 MB
  • Architecture: baseline 3-table hierarchical encoder (RelationalFoundationModel, --no-improved)

Headline result

Metric Value
RelBenchV1 classification, 12-task avg AUROC (ICL, leak-free, no tricks) ~69.0
Synthetic held-out validation (mixed-task mean score) 0.705
↳ binary AUROC / multiclass MRR / regression MAE 0.840 / 0.579 / 0.437
Paper target (KumoRFM-2) 79.60

Training configuration

Pre-trained on 2× RTX 3090 Ti with torchrun --nproc_per_node=2.

Hyperparameter Value
steps 2000 (release uses the step-1000 checkpoint)
context size 96
d_model / heads / layers 256 / 8 / 3
batch size × grad-accum (effective) 4 × 3 × 2 GPUs = 24
task type mixed (binary + regression + multiclass, max 8 classes)
optimizer AdamW, lr 2e-4, wd 1e-2, cosine decay, warmup 120, min-lr-ratio 0.02
precision FP32 (--no-amp)
rows_per_child / rows_per_aux 10 / 5
root/child/aux cols 8 / 8 / 6
entity_history_frac / entity_drift_scale 0.5 / 0.3
history_label_frac 0.75
root_shortcut_dropout 0.5
column_permutation on
typed_task_conditioning / typed_graph_attention on / on
schema_dropout / link_dropout / missing_prob 0.3 / 0.15 / 0.1
seed 42
throughput ~38 tasks/s, ~3.7k context-examples/s; ~21 min total
skipped non-finite steps 0 / 2000

Intended use

  • In-context prediction on relational databases using the OpenRFM
    icl_eval / icl_suite harness (strict temporal splits, no fine-tuning).
  • Research into synthetic-to-real transfer for relational foundation models.

Out of scope / limitations

  • Below the paper target. ~69.0 vs 79.60 AUROC. The dominant residual cause
    is the fixed 3-table-chain architecture, which cannot represent the multiple
    sibling tables that star-schema RelBench tasks need at once. See the repo
    README and progress_plan/reproduction_status.md.
  • Not a fine-tuned, task-specific model; do not compare to fine-tuned baselines.
  • Synthetic-to-real gap: weakest on data-rich / star-schema / domain-mismatched
    tasks (trial, stack-badge, amazon-item, event-repeat).
  • Trained on consumer hardware at a scale far below the original paper.

How to load

import torch
from kumorfm_repro.icl_eval import load_model_from_checkpoint
model = load_model_from_checkpoint("openrfm-pretrain-ctx96.pt", device="cuda")

Provenance

Trained from scratch in this repository. Contains no code or weights from the
original KumoRFM authors. The training data is fully synthetic (generated by
kumorfm_repro.data); no RelBench data was seen during pre-training.