Relational Transformer (RT)

The official implementation of the Relational Transformer (RT), an architecture for Relational Foundation Models (RFMs) that predict directly over relational databases (tables linked by foreign keys) and generalize zero-shot to new databases, tasks, and schemas.

This repository is RT-J, the current model and the one to start from. The two earlier papers are kept reproducible but are superseded; main tracks RT-J only.

Installation

Install the rt package (the model plus the native Rust data engine) from GitHub. It builds the extension from source, so you need a Rust toolchain and Python 3.12+:

pip install "git+https://github.com/stanford-star/relational-transformer.git"

Quickstart

The quickest way to try a released checkpoint is on a RelBench database already preprocessed into RT's tensor format on the Hub. The example below predicts whether an F1 driver fails to finish a race (driver-dnf) with a released RT-J checkpoint:

import os

# flex_attention's compiled kernel is CUDA-only; run it eager on CPU/MPS
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import torch
from huggingface_hub import snapshot_download

from rt import RelationalTransformer
from rt.eval import build_evaluator
from rt.data import get_tasks

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

# 1. download one RelBench database, already preprocessed into RT's tensor format
pre_dir = snapshot_download(
    "stanford-star/relbench-preprocessed",
    repo_type="dataset",
    allow_patterns="rel-f1/*",
)

# 2. load a pretrained checkpoint (RT-J here)
model = RelationalTransformer.from_pretrained(
    "stanford-star/rt-j", device=device
).to(torch.bfloat16)
cfg = model.config

# 3. build an evaluator for one task and predict zero-shot for 5 test rows.
#    the evaluator samples each row's context from the preprocessed DB;
tasks = get_tasks(pre_dir, [("rel-f1", "driver-dnf")], ("test",))
ev = build_evaluator(
    tasks, pre_dir,
    embedder=cfg["embedder"], d_text=cfg["d_text"], device=device,
    ctx_size_list=[128], local_ctx_size=64, bfs_width=32, prefer_latest=True,
    num_walks=10_000, walk_length=20, tokens_per_gpu=2**18,
    items_per_task=5, num_workers=0, prefetch_factor=2,
    shuffle_seed=0, context_seed=0, mmap_populate=True, vector_db_path=None,
    db_cutoff=None,
)

# evaluate_raw yields one (task, ctx, labels, preds, n) per task
results = ev.evaluate_raw([(model, "")], [128])
_task, _ctx, _labels, out, _n = next(iter(results))
preds = torch.sigmoid(torch.tensor(out[""], dtype=torch.float32))
print("driver-dnf probability:", [round(p, 3) for p in preds.tolist()])

Released checkpoints

Three checkpoints are public on the stanford-star HuggingFace org. All three were trained against the all-MiniLM-L12-v2 text embedder with d_text=384, which is the embedding published alongside the preprocessed data.

The embedder is not a free choice at inference time: the context is embedded with the vectors on disk in pre_dir, so evaluating a checkpoint under a different embedder scores garbage rather than failing. rt.eval.main asserts that the embedder and d_text you pass match the checkpoint's config.json.

The two legacy checkpoints also need differently preprocessed data (the legacy/ subdirectory of stanford-star/relbench-preprocessed) — see Inference. Fetch any of them with Downloads; from_pretrained resolves stanford-star/rt-j from the Hub on demand, so the quickstart above needs no manual download.

Each checkpoint's architecture, training recipe, evaluation protocol, licence and limitations are on its Hub card. The cards are kept in this repository under docs/cards/, which also holds the cards of the preprocessed datasets.

Bring your own database

Point RT at your own database, define a prediction task, and infer with a released checkpoint:

• Colab, no setup: the fully worked notebook (open in Colab) runs the whole flow end-to-end on your database (or the bundled demo).

Reproducing the paper

The released checkpoints and this library are enough to reproduce RT-J's inference numbers: every evaluation in the paper is rt.eval.main over a preprocessed RelBench directory, and both the checkpoints and the data are public. What the paper adds on top is a protocol — which context configuration each task won, how many ensemble members, which baselines were featurized how — and that protocol is reproduce/, a cleaned subset of the research directory the paper was produced from. It is research code with no stability promise, published so the numbers are checkable.

The reduced data behind every figure and table is committed, in reproduce/series/: one JSON per arm, with the aggregate metric at every context or ensemble size and the per-task value underneath it. So are the two most expensive stages' results — reproduce/tune/tuned_configs.json (the 120-point per-task context grid, roughly 200 GPU-hours) and reproduce/valtest/results.json — and every later stage reads them from there, so the grid never has to be re-run.

Expect agreement to a few tenths of a point, not to the digit. Two changes landed after the paper's runs: the BFS child sampler now draws without replacement, which leaves the distribution of contexts alone but changes the random stream; and input z-scoring statistics are now restricted to the train period, where they used to be computed over all rows, so the published preprocessed datasets carry the older normalization until they are regenerated. We measured the second one with the released checkpoint at its default context over seven tasks: mean −0.045 AUROC points and −0.035 nMAE points, every per-task move under 0.51 points, mixed in sign. reproduce/README.md has the table, the caveats, and which parts of reproduce/ have actually been run on a GPU versus which are untested.

The paper's own plotting scripts are not in this repository. They read our private Weights & Biases projects with no cache in between, so for anyone outside they would fail at the first API call — reproduce/series/ is the same data those plots were drawn from, committed instead. Two things stay genuinely out of reach: the pretraining-ablation figures, which compare five multi-day pretraining runs whose checkpoints are not released (the knobs and the recipe are written down, the runs are repeatable in principle), and six appendix schematics that never had a generating script.

reproduce/ carries no scheduler. The paper's runs went to a Slurm cluster, and the partitions, QoS names and measured wall-clock limits that made that work described one cluster on particular nights; they are gone. A stage instead enumerates its jobs and runs them sequentially in-process, and reproduce/launch.py is the one function to implement to put them on your own scheduler. reproduce/README.md is the full map of what produces which figure, what each stage costs, what was dropped from the research directory and why.

Development

We use pixi to manage one self-contained environment (Python, PyTorch + CUDA, Rust, and all dependencies), built on first use. There is nothing to build past pixi install: the native Rust extension is compiled as part of the project's own editable install.

git clone https://github.com/stanford-star/relational-transformer.git
cd relational-transformer
pixi run pytest                        # the test suite
pixi run python examples/train.py      # or eval.py, preprocess.py, ...

Documentation

There is no CLI: RT is a library, and a run is a script that calls it. Copy something from examples/ and edit it — every entry point takes its arguments explicitly, so nothing is hidden in a default you did not choose.

Citation

If you use RT, please cite the RT-J paper:

@article{ranjan2026rtj,
  title     = {RT-J: Large-Scale Pretraining of Relational Transformers for
               Context-Efficient Predictions},
  author    = {Ranjan, Rishabh and Kothapalli, Vignesh and Agarwal,
               Harshvardhan and Kanatsoulis, Charilaos and Upendra, Roshan and
               Palczewski, Tom and Guestrin, Carlos and Leskovec, Jure},
  year      = {2026},
}

CITATION.cff is the machine-readable form, and the source of truth for the author list. To cite one of the earlier models instead, use the paper linked in its row of the table at the top.

License

RT is released under the MIT License. The released checkpoints and the datasets on the stanford-star HuggingFace org carry their own licenses; "the Join" is assembled from third-party databases that keep theirs.