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

TL;DR

This work addresses the challenge of learning generalizable pre-training for relational databases under task heterogeneity and temporal dynamics.It introduces Task Vector Estimation (TVE), a pre-training framework that constructs predictive supervisory signals from next-window dynamics via set-based aggregation over schema traversal graphs, producing task-aware representations.The authors provide information-theoretic justification showing that task-informed representations retain more downstream-relevant information than traditional self-supervised methods, and demonstrate TVE's empirical advantages on RelBench, especially in low-data regimes, with complementary use alongside standard SSL losses.Overall, the paper argues for pre-training objectives that encode task heterogeneity and temporal structure as design principles for predictive modeling on relational data, and it shares code to enable reproducibility.

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.

Figures (8)

• Figure 1: Overview of pre‐training frameworks. Input graph is omitted for brevity.
• Figure 2: Construction of task vectors from relational schema. From schema graph, we derive a schema traversal graph, where edges represent valid join keys. For each reachable table $R$ in $k$-hop from the source node $v$, we apply a set function $S$ to the set of next-window values ${}_R^k\mathbf{X}_v^{(t_v, t_v + \Delta t]}$ reachable via the traversal path, where only paths contributing to the prediction task are retained. The task vector concatenates outputs from these set functions. A null indicator (purple cell) is activated when all other entries are absent (i.e., zero-valued).
• Figure 3: Markov Chain of Labels.
• Figure 4: Test performance curves averaged over runs on data sufficient tasks.
• Figure 5: Linear probing results averaged over runs on test set.

Theorems & Definitions (12)

• Remark 4.1
• Remark 4.2
• Definition 4.1: Sufficient representation given inputs and side-channel information
• Theorem 4.1
• Proposition A.1: federici2020
• Proposition A.2
• proof
• Corollary A.1
• proof
• proof