RDBLearn: Simple In-Context Prediction Over Relational Databases

TL;DR

RDBLearn extends in-context learning to relational databases by automatically featurizing linked records, enabling foundation models to perform relational prediction without extensive task-specific training.

Contribution

It introduces a simple, effective method for applying tabular in-context learning to relational data, packaged as an easy-to-use toolkit with flexible backend options.

Findings

RDBLearn outperforms existing foundation model approaches on multiple datasets.

It sometimes surpasses supervised models trained on the same data.

The approach simplifies relational prediction tasks with minimal tuning.

Abstract

Recent advances in tabular in-context learning (ICL) show that a single pretrained model can adapt to new prediction tasks from a small set of labeled examples, avoiding per-task training and heavy tuning. However, many real-world tasks live in relational databases, where predictive signal is spread across multiple linked tables rather than a single flat table. We show that tabular ICL can be extended to relational prediction with a simple recipe: automatically featurize each target row using relational aggregations over its linked records, materialize the resulting augmented table, and run an off-the-shelf tabular foundation model on it. We package this approach in \textit{RDBLearn} (https://github.com/HKUSHXLab/rdblearn), an easy-to-use toolkit with a scikit-learn-style estimator interface that makes it straightforward to swap different tabular ICL backends; a complementary agent-specific interface is provided as well. Across a broad collection of RelBench and 4DBInfer datasets, RDBLearn is the best-performing foundation model approach we evaluate, at times even outperforming strong supervised baselines trained or fine-tuned on each dataset.