RT-J: Large-Scale Pretraining of Relational Transformers for Context-Efficient Predictions

AI summary

Overview

• Research area: Machine learning for relational databases — specifically relational foundation models (RFMs) that perform in-context learning over multi-table data.
• Technical level: Advanced.
• Scope: The paper introduces RT-J, an 85M-parameter relational foundation model pretrained on a new 639-database corpus (the Join) with long-context training and random-walk retrieval, and shows it matches strong feature-engineering pipelines on RelBench regression with roughly 10x fewer in-context examples.

What This Paper Is About

Foundation models for tables and relational databases currently need tens of thousands of labeled examples to adapt to an unseen dataset, and relational architectures that operate natively on multi-table databases usually must be retrained per schema and task. The authors argue this is not an architectural limitation but a mismatch between existing pretraining recipes — which target zero-shot or fine-tuning regimes — and the goal of predicting from only a handful of in-context examples. Their goal is to build a pretrained relational transformer whose training data, context retrieval, and test-time compute are all designed for a few-example regime.

Key Contributions

1. The Join, a real-world relational pretraining corpus. A collection of 639 web-crawled relational databases across domains such as e-commerce, sports, media, finance, and healthcare, with 6,023 LLM-generated forecasting tasks (and 11K autocomplete tasks). It is described as the largest open pretraining corpus of real-world relational data to date, curated through collection, data-type standardization, LLM task generation, and XGBoost-based task filtering.

2. A pretraining recipe for context-efficient in-context learning. Long-context training with randomly sampled context sizes up to 8,192 database cells, plus a novel lightweight random-walk retriever that ranks candidate in-context rows by visit count before a bounded-width BFS expansion around each seed.

3. Test-time compute scaling via context tuning and context ensembling. Per-task tuning of the context hyperparameters (L, ℓ, b, ρ) and averaging predictions over multiple independently sampled contexts (n_ens in {1,2,4,8,16}).

4. RT-J and RT-PluRel models plus released artifacts. RT-PluRel is an 86M-parameter RT pretrained on synthetic PluRel data that ranked first among all system and model submissions on the RelArena-α leaderboard with an Elo of 1861. Continued pretraining of RT-PluRel on the Join gives RT-J. The authors release the models, the Join data, the LLM agent task curation pipeline, and the pretraining, fine-tuning, and inference code.

Main Findings

• Fewer examples for the same accuracy. On in-context regression tasks from RelBench, RT-J requires 10x fewer examples on average to match strong feature engineering pipelines built from an LLM-agent or RDBLearn featurizer combined with a TabICLv2 or LightGBM predictor. With only 15 in-context examples, RT-J already beats the best error any baseline reaches with 225 examples, and the strongest pipeline needs around 2,300 examples — about 10x more — to surpass it.

• Consistent lead at matched example counts. RT-J leads at every context size by 2.3 to 5.0 nMAE points on regression (reported as up to 5 nMAE points in Figure 3, and as more than 2 nMAE points better than the best baseline at all example counts in Figure 2). On classification, RT-J leads by 3.1 AUROC points when examples are scarcest, though the strongest baseline catches up at the largest context size.

• Context budget. At its maximum training context size of 8,192 database cells, RT-J fits around 225 examples on average.

• Real pretraining data matters most in context. Synthetic PluRel pretraining improves fine-tuning over a randomly initialized RT by 1.7 nMAE and 0.8 AUROC points; continued pretraining on the Join adds 0.5 nMAE and 0.5 AUROC points on top. In the in-context setting, RT-J improves over RT-PluRel by 0.6 nMAE and 1.5 AUROC points. Full averaged numbers: RT (random init) fine-tuned 29.3 nMAE / 77.8 AUROC; RT-PluRel fine-tuned 27.6 nMAE / 78.5 AUROC, in-context 32.9 nMAE / 73.0 AUROC; RT-J fine-tuned 27.1 nMAE / 79.0 AUROC, in-context 32.4 nMAE / 74.5 AUROC.

• Schema semantics are used, not just cell values. Shuffling table and column names at evaluation degrades classification by up to 2.1 AUROC points and regression by up to 0.8 nMAE points, and the gap widens with more in-context examples — suggesting schema text provides prior information the examples do not replace.

• Random-walk retrieval is label-efficient. Random walk ties RDBLearn-feature similarity at L=8192 (33.4 versus 33.5 nMAE and 72.9 versus 73.0 AUROC), leads regression at L=256 (by 0.7 nMAE) and L=512 and again at L=8192, while the similarity retriever is ahead by 0.4 to 0.8 nMAE at L=1024 to 4096 and by at most 0.9 AUROC. RDBLearn-feature similarity packs 394 labels into the context versus 223 for random walk at L=8192, while uniform random selection retrieves a similar count (208) yet trails by 4.1 nMAE and 6.1 AUROC. RT-J-embedding similarity stays 2.3 to 3.2 nMAE and 3.6 to 5.4 AUROC behind random walk at every L. The BFS variants come only within 1.6 to 1.8 nMAE of random walk with 2.7 to 5.6x fewer labels (84 and 40).

• Context tuning helps regression, hurts classification on average. With a single context seed, tuning reduces mean nMAE from 33.4 to 32.5 and improves eight of nine regression tasks, with the largest gains on rel-trial/site-success (91.0 to 88.6) and rel-hm/item-sales (14.0 to 11.8, via a smaller 512-cell context). For classification, tuning hurts a single seed on average (72.9 to 72.4 AUROC) and transfers to test on only five of twelve tasks.

• Context ensembling improves both metrics. With the per-task-tuned configuration, increasing n_ens from 1 to 16 yields +1.3 AUROC for classification, slightly lower at +1.1 AUROC at the default context. Regression shows smaller consistent gains (0.2 nMAE on average at the default context). Largest default-context gains: rel-trial/study-outcome +5.3, rel-event/user-repeat +2.6, and rel-hm/user-churn +2.4 AUROC; rel-avito/user-clicks degrades by 3.0 AUROC. Regression gains top out at 0.7 nMAE points on rel-event/user-attendance.

• Summary of ablation gains (Table 2). RT-J versus RT-PluRel: +0.6 nMAE, +1.5 AUROC. The Join-only versus PluRel-only: +1.8 nMAE, +2.8 AUROC. Random walk versus width-256 BFS: +1.6 nMAE, +3.9 AUROC. Random walk versus uniform random: +4.1 nMAE, +6.1 AUROC. Tuned versus default context: +0.8 nMAE, −0.5 AUROC. 16 context ensemble versus 1 context: +0.2 nMAE, +1.1 AUROC.

• Scale relative to the original RT recipe. Compared to the original RT pretraining, this recipe used 11x more total tokens, up to 8x larger contexts, and a 4x larger model, amounting to roughly 70x more pretraining FLOPs.

• The predictor mattered more than the featurizer among baselines. TabICLv2 beat LightGBM at every context size on both metrics, while LLM Agent and RDBLearn features performed similarly.

Methodology in Plain English

The authors keep the existing relational transformer (RT) architecture, which treats every database cell as a token and uses three attention layers in sequence — column attention (tokens attend within the same column), feature attention (tokens within the same row), and neighbor attention (tokens across rows connected by primary–foreign key links). Each token embedding combines the cell value, datatype information, and frozen text embeddings of the table and column names. Instead of sampling labeled examples separately, the model receives a single query-centered context built by expanding from the target row along primary–foreign key links, with the target cell masked.

For data, they crawled databases from five disjoint source classes (curated relational repositories, public APIs, domain-specific portals, community dumps, and benchmarks), kept only databases with at least two tables joined by a foreign key and at least one timestamp-parseable column, then applied five automatic fixes: boolean coercion, numeric coercion, foreign-key inference by name match plus an at least 80 percent value-containment test, temporal normalization, and entity-axis subsampling. An LLM (Claude Opus 4.6) proposed forecasting tasks from 25 registered template families across eight reasoning regimes, a deterministic checker removed invalid proposals, and an XGBoost model over window-aggregation features filtered tasks for learnability.

For context construction, they run W = 10,000 random walks of K = 20 steps from the target row, count how often each target-table row is visited, and use the top-ranked rows as seeds for a local breadth-first search with context size ℓ and width bound b until the total context reaches L cells. During pretraining, L is sampled uniformly from {512, ..., 8192}, ℓ from {128, ..., 8192}, b from {8, ..., 256} (all powers of two), and the recency parameter ρ from {0, 1}, so the model sees long and diverse contexts. Only the target cell is masked; the model predicts all masked cells and is trained against ground-truth values.

Pretraining happens in two phases: first on the synthetic PluRel corpus (86,211 tasks over 1,900 databases), then on 13,243 tasks from a 523-database subset of the Join (excluding the largest databases due to memory bottlenecks). Optimization uses Muon at a constant learning rate after warmup, with the best checkpoint selected after stochastic weight averaging. Each phase takes about a day on a single node with 8x A100 GPUs. Evaluation is on RelBench's 21 forecasting tasks across 7 databases (9 regression, 12 binary classification), all held out from pretraining, using Z-score normalized mean absolute error (nMAE) and AUROC. Test-time context size is swept over L in {256, 512, 1024, 2048, 4096, 8192} cells.

Why This Matters

• Impact on research. The paper reframes context efficiency in relational learning as a recipe problem — data, context engineering, and test-time compute — rather than an architecture problem. It provides the first large open corpus of real-world relational databases with generated forecasting tasks, which the authors position as an alternative to the synthetic pretraining corpora common for structured-data foundation models, and it quantifies the value of real data, schema semantics, and retrieval quality through targeted ablations.

• Real-world applications.

  • Flagging likely churn or fraud for a small enterprise without a large customer base.
  • Predicting outcomes of novel clinical trials with small cohorts.
  • Answering one-off predictive questions from an analyst, without an expensive labeling step.
  • Serving as a tool that LLM agents can call when a question requires predicting an unknown value rather than transforming existing ones with SQL.

• Industry relevance. Strong relational predictors today typically require task-specific training or substantial labeled data, which limits predictive modeling precisely where it is most valuable. RT-J's ability to make accurate predictions from hundreds of examples or fewer in a single forward pass, and the possibility of exposing it directly inside a database system for ad-hoc forecasts, points toward deployed enterprise tooling rather than bespoke feature-engineering pipelines.

Future Directions

• Longer contexts. RT-J was trained up to 8,192 context cells; efficient longer-context training would allow both more in-context examples and larger relational contexts for them.
• Beyond RelBench. Evaluation is limited to RelBench, with only 21 tasks on 7 databases, so the authors call for expanding the scope of real-world relational benchmarks, inspired by the Join.
• Unsupported task types. RT-J cannot solve link prediction and recommendation tasks, a limitation shared with the original RT.
• Cheaper retrieval and tuning at inference. Sensitivity analysis in the appendix indicates that lower random-walk counts W and K reduce latency overhead significantly while largely retaining accuracy, and the compute, time, and label overheads of tuning versus ensembling are compared — leaving room for more efficient test-time scaling.

Target Audience

Researchers and practitioners working on tabular and relational foundation models, in-context learning, and database-native machine learning, as well as applied scientists who need predictions over multi-table enterprise or scientific data with few labels. Engineers building feature-engineering pipelines or LLM-agent-based data systems will find the baseline comparisons and the released corpus, models, and code most directly useful. The paper assumes familiarity with transformer attention mechanisms, pretraining objectives, and standard regression and classification metrics.

Authors’ abstract

Foundation models for tables and relational databases are increasingly used for predictive tasks across various scientific and industrial domains. However, adapting existing models to unseen datasets requires tens of thousands of examples. Here we introduce RT-J, an 85M-parameter relational foundation model (RFM) that makes accurate predictions even with hundreds of in-context examples or fewer. To achieve this, we first assemble the Join, a large collection of 639 web-crawled databases with 6,023 LLM-generated forecasting tasks. We pretrain with long contexts of up to 8,192 database cells, using in-context examples from a novel lightweight random-walk retriever. Finally, we leverage context tuning and ensembling to scale up test-time compute. On in-context regression tasks from RelBench, RT-J requires 10x fewer examples on average to match strong feature engineering pipelines with LLM-agent or RDBLearn featurizers and TabICLv2 or LightGBM predictors. Overall, our work pushes the frontier for context-efficient predictions on relational data and highlights the importance of real-world semantics, context engineering and large-scale pretraining for RFMs.