Towards Causal Foundation Model: on Duality between Causal Inference and Attention

TL;DR

This paper introduces CInA, a novel self-supervised attention-based method for causal inference that enables zero-shot treatment effect estimation, demonstrating strong generalization across diverse datasets and paving the way for causal foundation models.

Contribution

The paper presents a theoretically justified approach linking attention mechanisms to causal inference, allowing zero-shot causal reasoning with foundation models.

Findings

CInA effectively generalizes to out-of-distribution datasets.

It matches or surpasses traditional methods on real-world data.

Theoretical connection between attention and covariate balancing is established.

Abstract

Foundation models have brought changes to the landscape of machine learning, demonstrating sparks of human-level intelligence across a diverse array of tasks. However, a gap persists in complex tasks such as causal inference, primarily due to challenges associated with intricate reasoning steps and high numerical precision requirements. In this work, we take a first step towards building causally-aware foundation models for treatment effect estimations. We propose a novel, theoretically justified method called Causal Inference with Attention (CInA), which utilizes multiple unlabeled datasets to perform self-supervised causal learning, and subsequently enables zero-shot causal inference on unseen tasks with new data. This is based on our theoretical results that demonstrate the primal-dual connection between optimal covariate balancing and self-attention, facilitating zero-shot causal inference through the final layer of a trained transformer-type architecture. We demonstrate empirically that CInA effectively generalizes to out-of-distribution datasets and various real-world datasets, matching or even surpassing traditional per-dataset methodologies. These results provide compelling evidence that our method has the potential to serve as a stepping stone for the development of causal foundation models.