Towards a Relationship-Aware Transformer for Tabular Data

TL;DR

This paper introduces a relationship-aware transformer model for tabular data that incorporates external dependencies between samples via a modified attention mechanism, enhancing tasks like treatment effect estimation.

Contribution

It proposes a novel attention-based approach to encode relationships between data points in tabular data, extending the capabilities of transformers beyond standard methods.

Findings

Outperforms gradient boosting decision trees on regression tasks

Effective in modeling sample relationships for treatment effect estimation

Demonstrates advantages on synthetic and real-world datasets

Abstract

Deep learning models for tabular data typically do not allow for imposing a graph of external dependencies between samples, which can be useful for accounting for relatedness in tasks such as treatment effect estimation. Graph neural networks only consider adjacent nodes, making them difficult to apply to sparse graphs. This paper proposes several solutions based on a modified attention mechanism, which accounts for possible relationships between data points by adding a term to the attention matrix. Our models are compared with each other and the gradient boosting decision trees in a regression task on synthetic and real-world datasets, as well as in a treatment effect estimation task on the IHDP dataset.