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

TL;DR

This paper introduces Task Vector Estimation (TVE), a pre-training framework for relational databases that leverages information-theoretic principles to create task-aware representations, improving performance on diverse downstream tasks.

Contribution

The paper proposes TVE, a novel pre-training method that models relational dynamics and task heterogeneity using set-based aggregation and information theory, advancing relational data learning.

Findings

TVE outperforms traditional pre-training baselines on RelBench.

Task-informed representations retain more relevant signals.

Encoding task heterogeneity improves predictive modeling on relational databases.

Abstract

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.