Learning Causal Orderings for In-Context Tabular Prediction

TL;DR

This paper introduces TabOrder, a model that learns and enforces causal variable orderings in tabular data to improve prediction and imputation, especially under distribution shifts.

Contribution

It proposes a novel causal order-constrained attention mechanism integrated into predictive models, learned via an unsupervised likelihood-based objective.

Findings

TabOrder accurately recovers causal variable orderings.

It improves prediction and imputation in tabular data with missingness.

Provides insights into biological data under intervention.

Abstract

In-context learning for tabular data sets strong predictive standards in observational settings; it however primarily relies on correlational structure, which becomes unreliable under distribution shift or intervention. While established methods to discover causal structure exist, they are often focused on structure identifiability and decoupled from the predictive architectures that could benefit from them. To bridge these perspectives, we study how to simultaneously infer and enforce causal structure in the form of topological variable orderings into tabular prediction. Unlike standard architectures, our model TabOrder uses causal order-constrained attention, basing predictions only on features that precede a target under a learned causal order. Similar to causal discovery methods, TabOrder learns the optimal variable ordering in an unsupervised manner through a likelihood-based objective. We justify this choice under standard functional model classes and also study how sample missingness, a common challenge in tabular data, interacts with causal direction identification. Empirically, we confirm that TabOrder recovers accurate variable orderings while addressing prediction and imputation tasks, as well as gives insight into real-world biological data under intervention.