KumoRFM-2: Scaling Foundation Models for Relational Learning

Valter Hudovernik; Federico López; Vid Kocijan; Akihiro Nitta; Jan Eric Lenssen; Jure Leskovec; Matthias Fey

cs.LGApr 14, 2026

KumoRFM-2: Scaling Foundation Models for Relational Learning

Authors: Valter Hudovernik, Federico López, Vid Kocijan, Akihiro Nitta, Jan Eric Lenssen, Jure Leskovec, Matthias Fey

Abstract

We introduce KumoRFM-2, the next iteration of a pre-trained foundation model for relational data. KumoRFM-2 supports in-context learning as well as fine-tuning and is applicable to a wide range of predictive tasks. In contrast to tabular foundation models, KumoRFM-2 natively operates on relational data, processing one or more connected tables simultaneously without manual table flattening or target variable generation, all while preserving temporal consistency. KumoRFM-2 leverages a large corpus of synthetic and real-world data to pre-train across four axes: the row and column dimensions at the individual table level, and the foreign key and cross-sample dimensions at the database level. In contrast to its predecessor, KumoRFM-2 injects task information as early as possible, enabling sharper selection of task-relevant columns and improved robustness to noisy data. Through extensive experiments on 41 challenging benchmarks and analysis around expressivity and sensitivity, we demonstrate that KumoRFM-2 outperforms supervised and foundational approaches by up to 8%, while maintaining strong performance under extreme settings of cold start and noisy data. To our knowledge, this is the first time a few-shot foundation model has been shown to surpass supervised approaches on common benchmark tasks, with performance further improving upon fine-tuning. Finally, while KumoRFM-1 was limited to small-scale in-memory datasets, KumoRFM-2 scales to billion-scale relational datasets.

Figures & tables

Figure 1: Overview of the KumoRFM system. (a) KumoRFM operates directly on a relational database, where the schema defines the underlying graph structure. (b) Users specify predictive tasks via a declarative Predictive Query Language or natural language. (c) KumoRFM constructs context examples $(v_{i},t_{i},y_{i})$ and corresponding input subgraphs $\mathcal{G}^{\leq t_{i}}[v_{i}]$ , where targets $y_{i}$ are derived according to the query definition. All data processing is pushed down to the underlying database for scalable execution, after which the constructed context is passed to the model for inference.

Figure 2: Automatic context generation. The query ( Kocijan et al., 2026 ) defines how context input/output pairs can be automatically constructed from historical database states while preventing data leakage.

Figure 3: Overview of the KumoRFM-2 architecture. (a) Given a task table with connected tables as context examples, (b) KumoRFM-2 first computes representations for all rows in the context via alternating column- and row-wise attention. These representations are made task-conditioned by incorporating context targets directly into the input tables. (c) The representations are then enriched via graph attention over primary–foreign key relationships and cross-sample attention over context examples, enabling in-context learning beyond single tables. Finally, the target is predicted from the resulting representation of the prediction example.

Figure 4: An adversarial example for fixed-function column-wise encoders. The label is positive iff both child features co-occur within at least one connected row. Row-level alignment is required to determine the label, which column-wise encoders ( Kanter and Veeramachaneni, 2015 ) fail to capture.

Figure 5: Performance overview of KumoRFM-2 , where scores are normalized relative to the strongest (supervised) model on each benchmark suite. KumoRFM-2 shows a significant improvement over KumoRFM-1 and consistently outperforms prior supervised or foundational few-shot approaches. Further details are provided in Tables 3 , 4 , 6 , 5 , 7 , and 8 .

	Dataset	Domain	Tasks #Tasks		#Tables #Rows #Cols			#Train	ICL
	Dataset	Domain	Tasks #Tasks		#Tables #Rows #Cols			#Train	Coverage
RelBenchV1	f1	Sports	Race outcome	3	9	74,063	67	$\leq$ 11,977	83.5%
	avito	E-commerce	User engagement, CTR	3	8	20,679,117	42	$\leq$ 116,598	8.6%
	event	Social	Event participation	3	5	41,328,337	128	$\leq$ 23,424	42.7%
	trial	Medical	Clinical trial outcome	3	15	5,434,924	140	$\leq$ 12,954	77.2%
	amazon	E-commerce	Churn, LTV	4	3	15,000,713	15	$\leq$ 5,142,347	0.2%

Table 1: Database summary and statistics across RelBenchV1, RelBenchV2, SALT and 4DBInfer. Datasets vary in the number of tables, rows, columns, and available training examples (#Train). ICL Coverage denotes the fraction of training data used for in-context learning, reaching as low as 0.2%.

	Category	Model	Description
Supervised Tabular ML	Gradient Boosted	XGBoost ( Chen and Guestrin, 2016 )	Ensemble of decision trees trained via gradient boosting, where each tree corrects the residual errors of the previous ones, forming a competitive baseline on tabular data.
	Decision Tree	LightGBM ( Ke et al., 2017 )
		CatBoost ( Prokhorenkova et al., 2018 )
	Deep Tabular ML	CARTE ( Kim et al., 2024 )	Deep representation learning on tabular data, trained in a (self-)supervised fashion.
	AutoML	AutoGluon ( Erickson et al., 2020 )	Ensemble of diverse learners ( e.g. , Neural Networks, GBDTs) with automated hyperparameter optimization. Generally provides best performance among tabular models.
	Manual Feature Engineering	DS+LightGBM DS+AutoGluon DS+TabPFN-2.5 … ( Robinson et al., 2024 )	Current gold-standard. A data scientist expert that engineers task-specific features ( e.g. , temporal aggregations, multi-hop joins) to capture relational dependencies, followed by a tabular model.

Table 2: Overview of utilized baselines . We compare against supervised tabular models (including hand-crafted and automatic feature engineering pipelines for flattening multi-table data), state-of-the-art relational deep learning models, and language-based, tabular, and relational foundation models.

	Method	f1		avito		event		trial	amazon		stack		hm	Avg	Rank
	Method	dnf	top3	click	visit	repeat	ignore	out	user	item	eng	badge	churn	$\uparrow$	$\downarrow$
Sup. TabML	LightGBM	68.56	73.92	53.60	53.05	68.04	79.93	70.09	52.22	62.54	63.39	63.43	55.21	63.67	16.75
	DS+LightGBM	69.80	82.40	64.27	64.46	75.42	84.23	72.00	67.60	81.80	90.30	86.20	69.00	75.62	7.75
	DS+AutoGluon ∗	70.01	82.53	65.66	65.87	74.85	82.31	70.51	68.12	80.45	89.94	85.38	68.71	75.36	8.50
	DS+TabPFN-2.5 ∗	71.16	77.70	64.21	64.53	76.43	85.70	70.75	66.28	79.83	89.07	85.08	68.31	74.92	9.50
Sup. RelationalML	GraphSAGE	72.62	75.54	65.90	66.20	76.89	81.62	68.60	70.42	82.81	90.59	88.86	69.88	75.83	6.42

Table 3: Test results on the binary classification tasks in RelBenchV1. Higher is better (AUROC). denotes the best, and the second-best among all models and tasks.

Method	mimic	ratebeer			arxiv	Avg	Rank
Method	stay	beer-churn	user-churn	dormant	citation	$\uparrow$	$\downarrow$
LightGBM	51.81	76.21	83.92	75.79	71.21	71.79	3.80
GraphSAGE	55.01	78.67	94.27	80.51	82.50	78.19	2.00
KumoRFM-1	56.33	75.06	91.38	77.10	80.62	76.22	2.80
KumoRFM-2	55.28	83.84	97.43	80.65	81.71	79.96	1.75

Table 4: Test results on the binary classification tasks in RelBenchV2. Higher is better (AUROC). denotes the best, and the second-best among all models and tasks.

	Method	AB	OB	RR	SE		Avg	Rank
	Method	churn	ctr	cvr	upvote	churn	$\uparrow$	$\downarrow$
Sup. TabML	XGBoost	50.00	50.00	50.00	49.68	50.84	50.10	10.6
	AutoGluon	50.00	49.69	50.96	50.81	50.00	50.29	10.4
	DFS+XGBoost	69.22	54.21	79.06	86.75	82.51	74.35	9.0
	DFS+AutoGluon	72.91	54.94	80.08	88.49	83.96	76.08	7.6
Sup. RelML	GAT	76.22	61.46	82.84	88.53	86.45	79.10	5.8

Table 5: Test results on the binary classification tasks in 4DBInfer. Higher is better (AUROC). denotes the best, and the second-best among all models and tasks.

	Method	Plant	Shipping	Item	Header	Sales	Sales	Payment	Shipping	Avg	Rank
	Method	Plant	Point	Incoterm	Incoterm	Office	Group	Terms	Condition	$\uparrow$	$\downarrow$
Base	Random	0.59	0.11	0.62	0.62	0.99	0.02	0.02	0.12	0.39	10.38
Base	Majority	0.59	0.54	0.62	0.62	0.99	0.05	0.23	0.41	0.51	9.63
	DS+XGBoost	0.99	0.95	0.70	0.70	0.99	0.51	0.57	0.68	0.76	5.25
Sup. TabML	DS+LightGBM	0.61	0.28	0.73	0.73	0.99	0.02	0.10	0.51	0.50	9.00
	DS+CatBoost	0.99	0.80	0.80	0.81	0.99	0.16	0.44	0.71	0.71	4.63

Table 6: Test results on the multi-class classification tasks in SALT. Higher is better (MRR). denotes the best, and the second-best among all models and tasks.

	Method	f1	avito	event	trial		amazon		stack	hm	$\mu$ n	Rank
	Method	pos	ctr	attend	adverse	success	user	item	votes	sales	$\downarrow$	$\downarrow$
Base	Global Median	4.399	0.043	0.264	57.533	0.462	16.783	64.234	0.068	0.076	1.062	12.78
Base	Entity Median	8.519	0.046	0.269	57.930	0.441	17.423	66.436	0.069	0.078	1.190	14.00
Sup. TabML	LightGBM	4.170	0.041	0.264	44.011	0.425	16.783	60.569	0.068	0.076	1.000	10.22
	DS+LightGBM	3.963	0.044	0.284	40.581	0.407	13.928	41.122	0.065	0.036	0.879	5.11
	DS+AutoGluon ∗	4.251	0.045	0.256	44.706	0.430	14.399	45.390	0.068	0.043	0.918	7.11

Table 7: Test results on the regression tasks in RelBenchV1. Lower is better (MAE). Average performance

\bm{\mu}_{n}

is normalized relative to the LightGBM baseline. denotes the best, and the second-best among all models and tasks.

Method	ratebeer	arxiv	$\mu$ n	Rank
Method	user-count	publication	$\downarrow$	$\downarrow$
Global Median	15.124	0.577	0.872	4.5
Entity Median	13.079	0.874	1.079	5.0
LightGBM	20.350	0.577	1.000	5.0
GraphSAGE	7.374	0.513	0.626	2.0
KumoRFM-1	11.063	0.518	0.712	3.0

Table 8: Test results on the regression tasks in RelBenchV2. Lower is better (MAE). Average performance

\bm{\mu}_{n}

is normalized relative to the LightGBM baseline. denotes the best, and the second-best among all models and tasks.

Panel of Figure 6 — Figure 6: Robustness of KumoRFM-2 with respect to context construction . KumoRFM-2 exhibits strong sample efficiency, and generally benefits from increased relational subgraph breadth and depth. We identify subgraph breadth as a key factor for optimizing task-specific performance.

Panel of Figure 7 — Figure 7: Robustness of KumoRFM-2 to sparse, incomplete and noisy data . KumoRFM-2 maintains strong resilience to data scarcity, and remains robust to noise compared to its predecessor.

Task	DS+Auto	DS+	KumoRFM-2
Task	Gluon	CARTE	ICL	Tuned
Plant	0.99	0.99	0.99	1.00
Shipping Point	0.98	0.97	0.98	0.99
Item Incoterm	0.78	0.75	0.78	0.82
Header Incoterm	0.78	0.77	0.81	0.88
Sales Office	0.99	0.99	1.00	1.00

Table 9: Test results on SALT when fine-tuning KumoRFM-2. Higher is better (MRR). denotes the best, and the second-best among all models and tasks.

Explore similar work

Relational In-Context Learning via Synthetic Pre-training with Structural Prior
Mar 4, 2026Yanbo Wang, Jiaxuan You, Chuan Shi +1Relational Foundation Models Relational Databases
Institute for Artificial Intelligence, Peking University · University of Illinois at Urbana-Champaign · Institute of Computing Technology, Beijing University of Post and Telecommunication +1
RT-J: Large-Scale Pretraining of Relational Transformers for Context-Efficient Predictions
Oct 13, 2026Rishabh Ranjan, Vignesh Kothapalli, Harshvardhan Agarwal +5Relational Foundation Models Relational Deep Learning
Stanford University · SAP
OpenRFM: Dissecting Relational In-Context Learning
Jun 3, 2026Zhikai Chen, Junyu Yin, Jialiang Gu +5Relational Foundation Models Synthetic Data Pretraining
Michigan State University · George Mason University · Georgia Institute of Technology +1

Mar 4, 2026J/K move · Enter open · S save

Relational In-Context Learning via Synthetic Pre-training with Structural Prior

Yanbo Wang, Jiaxuan You, Chuan Shi, Muhan Zhang

Institute for Artificial Intelligence, Peking University · University of Illinois at Urbana-Champaign · Institute of Computing Technology, Beijing University of Post and Telecommunication+1

Relational Foundation Models Relational Databases In-Context Learning Synthetic Data Pretraining Synthetic Data Generation

Open paper AskPDF arXiv

Relational Databases (RDBs) are the backbone of modern business, yet they lack foundation models comparable to those in text or vision. A key obstacle is that high-quality RDBs are private, scarce, and structurally heterogeneous, making internet-scale pre-training infeasible. To overcome this data scarcity, we introduce RDB-PFN, the first relational foundation model trained purely via synthetic data. Inspired by Prior-Data Fitted Networks (PFNs), where synthetic data generated from Structural Causal Models (SCMs) enables reasoning on single tables, we design a Relational Prior Generator to create an infinite stream of diverse RDBs from scratch. Pre-training on over 2 million synthetic single-table and relational tasks, RDB-PFN learns to adapt to any new database instantly via genuine in-context learning. Experiments show that RDB-PFN achieves strong few-shot performance on 19 real-world relational prediction tasks, outperforming state-of-the-art tabular foundation models evaluated on the same DFS-linearized inputs, while using a lightweight architecture and fast inference. The code is available at https://github.com/MuLabPKU/RDBPFN.

Abstract

Explore similar work

Relational In-Context Learning via Synthetic Pre-training with Structural Prior

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

OpenRFM: Dissecting Relational In-Context Learning