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 complex tasks. We propose a novel, theoretically sound 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 our approach CInA effectively generalizes to out-of-distribution datasets and various real-world datasets, matching or even surpassing traditional per-dataset causal inference methodologies.

基于理论结果我们提出了一种新的方法Causal Inference with Attention (CInA)，它基于多个无标签数据集进行自监督因果学习，在训练完成的transformer-type架构的最后一层通过理论上的最优协变量平衡与自注意机制之间的原始-对偶关系来实现对未知任务的零样本因果推断。实证结果表明，我们的方法CInA能够在分布外数据集和各种真实世界数据集上有效推广，并且与传统的基于每个数据集的因果推断方法相匹配甚至超越。

朝向因果基础模型：因果推断与注意力之间的二元性