Source-linked AI summary
KumoRFM-2: Scaling Foundation Models for Relational Learning
Valter Hudovernik, Federico López, Vid Kocijan, Akihiro Nitta, Jan Eric Lenssen, Jure Leskovec, Matthias Fey
TL;DR
Relational predictive modeling needs foundation models that operate directly on connected tables rather than flattened representations. KumoRFM-2 combines task-conditioned relational processing with in-context learning and fine-tuning, outperforming supervised and foundation-model approaches across 41 tasks while scaling to production databases.
Problem
Existing tabular extensions flatten relational data, requiring feature engineering or task-agnostic aggregation that misses task-specific relational structure.
Method
KumoRFM-2 pre-trains on synthetic and real relational data and processes temporal multi-table contexts through task-conditioned relational attention, supporting in-context learning and fine-tuning.
Results
Across 41 predictive tasks, KumoRFM-2 outperforms supervised and few-shot foundation models, with reported gains of ≈5% over RelGNN and ≈8% over AutoGluon on specified benchmarks.
Takeaways & Limitations
A single pre-trained model can rival or exceed task-specific pipelines while providing training-free, ad-hoc predictions on relational databases at production scale.
Takeaways & Limitations
On some highly volatile regression tasks, data scientist-engineered pipelines can outperform automated graph-based approaches because long-term coarse-grained patterns are difficult to capture from fine-grained transaction subgraphs.
Abstract
from arXiv · showhide
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.
1 Introduction
Foundation models enable few-shot prediction on structured data, but existing tabular extensions inadequately represent multi-table relational structure. KumoRFM-2 addresses this gap and outperforms supervised and foundation-model baselines across 41 predictive tasks.
- Motivation: In-context learning adapts a pre-trained model to unseen predictive tasks from a few labeled examples, reducing data requirements and deployment time.The supplied passages describe deployment changing from hours or days to seconds for novel tasks.
- Problem: Flattening relational data requires feature engineering and target generation or task-agnostic aggregation that misses task-specific relational structure.These limitations undermine training-free, ad-hoc prediction on real-world multi-table data.
- Approach: KumoRFM-2 performs task-conditioned feature extraction across separate intra-table and cross-table attention stages.The design learns to filter, correlate, and aggregate information across multiple tables.
- Results: 41 predictive tasks show KumoRFM-2 outperforming supervised and few-shot foundation models across tabular and relational settings.The evaluation spans real-world databases from RelBenchV1, RelBenchV2, SALT, and 4DBInfer.
- Implication: KumoRFM-2 supports accurate on-demand relational predictions that can integrate into agentic workflows for real-time decision-making.The conclusion frames few-shot foundation models as able to rival carefully tuned task-specific approaches.
2 Relational Foundation Models and KumoRFM
Relational foundation models treat databases as temporal graphs and use context examples to predict new targets without task-specific training. KumoRFM couples database-native integration, declarative task specification, and automatic temporal context construction for scalable inference.
- Relational Foundation Models: Relational Deep Learning represents records as nodes, primary–foreign key links as edges, and time-consistent database snapshots as G≤t.Each training example combines an entity, timestamp, and ground-truth target.
- Relational Foundation Models: Relational tasks require extracting relevant signals from local subgraphs across multiple connected tables rather than from the task table alone.The relevant subgraph is restricted to records available up to the entity’s timestamp.
- KumoRFM System: KumoRFM integrates with relational databases, warehouses, and graph infrastructure, supporting scalable execution and stateless multi-tenant serving.The system also provides embeddings, evaluation metrics, and explanations alongside predictions.
- KumoRFM System: KumoRFM uses a declarative Predictive Query Language or natural language to specify tasks and constructs context examples and input subgraphs directly from the database.Processing is pushed to the backend, and queries support aggregations, filters, binary operations, static tasks, and temporal predictions.
- KumoRFM System: Automatic context generation replays historical database states for temporal tasks while preserving time consistency and preventing leakage.Context targets and entity-centric subgraphs are computed at earlier timestamps according to the query definition.
3 The KumoRFM-2 Architecture
KumoRFM-2 combines task-conditioned table processing with relational and cross-sample attention, addressing expressivity limits of flattened representations. It also supports fine-tuning and robustness-oriented ensembling for large, diverse relational tasks.
- Architecture: KumoRFM-2 processes row/column information, foreign-key structure, and cross-sample patterns at distinct database scales.Context targets enter early so processing remains task-conditioned throughout the model.
- Architecture: A hierarchical attention scheme first builds row representations through alternating column and row attention, then relates them across tables and examples.Foreign-key and cross-sample attention provide the second, larger stage.
- Robustness: Randomized column and class shuffles, hop ensembling, and post-processing reduce variance and sensitivity to arbitrary input orderings.These mechanisms provide more fine-grained ensembling control than KumoRFM-1.
- Expressivity: Fixed-function column-wise encoders fail when labels depend on row-level co-occurrence, whereas task-conditioned extraction captures the required alignment.The adversarial example defines a positive label when child features A and B co-occur in at least one related row.
- Expressivity: AUROC = 1.0 for KumoRFM-2 versus AUROC = 0.5 for column-wise encoders on the co-occurrence example.The result demonstrates the expressivity gap induced by task-agnostic, fixed-function relational flattening.
- Fine-Tuning: Fine-tuning specializes KumoRFM-2 to a dataset and task, reducing reliance on in-context learning and context selection when training data exceeds context limits.The model weights then encode domain- and task-specific information.
4 Experimental Results
Across binary classification, multi-class classification, and regression benchmarks, KumoRFM-2 matches or surpasses strong supervised and foundation-model baselines. Its performance remains robust with limited context, benefits from relational information, and improves further with fine-tuning.
- Binary Classification: 79.60 average AUROC gives KumoRFM-2 the best overall rank on 12 RelBenchV1 binary classification tasks.It exceeds the best tabular model by 3.98 points, KumoRFM-1 by 2.89 points, and RelGNN by 1.54 points.
- Binary Classification: 0.2% of available data can suffice for KumoRFM-2 to surpass the strongest supervised relational baseline by 1.54 points on average.This advantage is achieved without task-specific training, although supervised models retain a slight edge on some tasks with larger training sets.
- Regression: µn = 0.822 is KumoRFM-2’s lowest normalized MAE on RelBenchV1 regression, beating the strongest baseline by 4.7% on average.It also improves over KumoRFM-1 by 10.5% and achieves the top result on 5 of 9 tasks.
- Robustness and Scaling: 1k context examples, equal to 0.05% of training examples, still yield competitive performance, demonstrating strong sample efficiency.Performance is evaluated under varying context sizes on a SALT subset against supervised XGBoost and in-context TabICLv2.
- Robustness and Scaling: Increasing relational context and neighborhood size generally improves performance, but the optimal neighborhood size depends on the task.Broader subgraphs provide more informative signals on average, while some tasks peak with localized neighborhoods such as n = 4.
5 KumoRFM in Practice
KumoRFM-2 provides a unified workflow for constructing relational graphs, specifying predictive tasks, and issuing predictions directly over database data. Its abstractions automate schema handling and support both high-level queries and explicit context control without requiring traditional training-data pipelines.
- Workflow and interfaces: KumoRFM-2 is available through a Python SDK with unified interfaces for graph construction, task specification, and model inference.The SDK is available from Python 3.10 onwards, while in-context inference is offered through a public endpoint and fine-tuning through enterprise deployments.
- Workflow and interfaces: The SDK automatically infers primary keys, time columns, and inter-table links from database schema information while allowing declarative overrides.This automatic graph construction reduces manual schema engineering without removing user control.
- Task specification: Users can express relational and temporal prediction tasks with a predictive query language or provide explicit context and prediction tables through TaskTable.Runtime parameters allow users to trade off performance and efficiency after the query is established.
- Task specification: KumoRFM-2 uses one underlying relational graph for context and prediction examples rather than requiring a separate training-data pipeline.Feature construction and context selection are handled automatically when practitioners specify prediction tasks over the database.
- Practical guidance: The recommended workflow starts with a focused query over relevant tables and expands the graph iteratively as needed.The system is designed to benefit from relational and temporal structure, especially meaningful entity relationships and histories.
6 Conclusion
KumoRFM-2 provides training-free, ad-hoc predictions on relational databases and scales to production environments with 500B+ rows.
- 500B+ rows: KumoRFM-2 scales to production environments while providing training-free, ad-hoc predictions on relational databases.