Task-Aligned Temporal Context for Relational Graph Transformers

Relational deep learning enables predictive modeling over relational databases with temporal information by representing rows and their relationships as heterogeneous graphs. Recently, graph transformers have been shown to achieve strong performance on entity-level tasks, but their context construction is less suited to event-level prediction, where an interaction connects entities with unequal histories. We examine this gap on anti-money laundering (AML) and on RelBench prediction tasks, finding that uniform neighborhood sampling can waste the token budget on padding or overrepresent one entity's history while omitting the participating entities themselves. Motivated by these observations, we introduce task-aligned temporal context construction for relational graph transformers. Our tokenization (i) balances context across an event's participating entities, (ii) prioritizes recent events, and (iii) summarizes repeated interactions with counterparties. In addition, a complementary profile branch attends to each entity's recent history to provide context for identifying deviations from its usual behavior. On TransXion, an AML benchmark, our approach improves RelGT's average precision from 0.296 with recency sampling to 0.389, a 31.4% relative gain. Across three additional tasks, adding attention to entity histories improves RelGT, increasing AUROC by 2.0 percentage points on rel-avito and reducing mean absolute error by 5.8–16.2% on rel-f1 and rel-event. The sampling changes provide further gains on rel-avito, while rel-f1 and rel-event perform best with standard sampling and entity histories. Together, these results highlight the importance of matching temporal context construction to the prediction task and show that explicit entity histories can complement standard neighborhood sampling.