Large-Scale Pretraining unlocks Few-Shot Prediction for Relational Data
Abstract
The ability to learn a new task from a few examples has been crucial to the success of foundation models such as large language models (LLMs). However, current foundation models for structured data like tables and relational databases still require tens of thousands of labeled examples to perform well on new tasks. Here we show that when pretrained at scale with the right recipe, Relational Transformers (RTs) can make state-of-the-art (SoTA) predictions with only hundreds of labels. This capability arises from the following ingredients: First, we assemble THE JOIN, the largest pretraining corpus of relational data to date, comprising 6k forecasting tasks across 650 real-world databases from diverse domains. Second, we develop a pretraining recipe that combines (1) mixed context sizes for multi-scale learning, (2) multi-cell masking for dense supervision, and (3) a novel random-walk-based retriever to efficiently gather relevant context. Third, we identify context ensembling and context tuning as complementary axes for scaling test-time compute to further improve few-shot performance. Pretrained with our recipe on THE JOIN, RT achieves parity with LLM Agent + TabICLv2 and RDBLearn + TabICLv2 pipelines with 32–23× fewer labels respectively, and even surpasses the prior SoTA of full task-specific training, on average nMAE for RelBench regression tasks. Context ensembling and tuning improves this further by up to 3%, and 4% respectively. Our ablations highlight the importance of schema semantics, multi-cell masking, and random-walk retrieval. Overall, our work paves the way for developing relational foundation models with strong few-shot capabilities.