Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

A Pre-training Framework for Relational Data with Information-theoretic Principles

About

Relational databases underpin critical infrastructure across a wide range of domains, yet the design of generalizable pre-training strategies for learning from relational databases remains an open challenge due to task heterogeneity. Specifically, there exist many possible downstream tasks, as tasks are defined based on relational schema graphs, temporal dependencies, and SQL-defined label logics. An effective pre-training framework is desired to take these factors into account in order to obtain task-aware representations. By incorporating knowledge of the underlying distribution that drives label generation, downstream tasks can benefit from relevant side-channel information. To bridge this gap, we introduce Task Vector Estimation (TVE), a novel pre-training framework that constructs predictive supervisory signals via set-based aggregation over schema traversal graphs, explicitly modeling next-window relational dynamics. We formalize our approach through an information-theoretic lens, demonstrating that task-informed representations retain more relevant signals than those obtained without task priors. Extensive experiments on the RelBench benchmark show that TVE consistently outperforms traditional pre-training baselines. Our findings advocate for pre-training objectives that encode task heterogeneity and temporal structure as design principles for predictive modeling on relational databases. Our code is publicly available at https://github.com/quang-truong/task-vector-estimation.

Quang Truong, Zhikai Chen, Mingxuan Ju, Tong Zhao, Neil Shah, Jiliang Tang• 2025

Related benchmarks

TaskDatasetResultRank
Driver Top 3 Predictionrel-f1
ROC-AUC74.5
70
Driver DNF Predictionrel-f1
ROC-AUC0.569
67
User Churn PredictionAmazon Rel
ROC-AUC0.624
64
Item Churn Predictionrel-amazon
ROC-AUC63.5
64
User Churn Predictionrel-hm
ROC-AUC64.8
62
User Badge PredictionRel Stack User Badge
ROC-AUC75.6
37
Entity Regression (study-adverse)rel (trial)
MAE0.212
22
Entity Regression (post-votes)rel-stack
MAE0.148
19
Entity Classification (user-engagement)rel-stack
ROC-AUC73.9
17
Entity Regression (item-sales)rel-hm
MAE0.189
16
Showing 10 of 12 rows

Other info

Follow for update