Enhancing causal effect foundation model with task-specific model

Causal effect estimation aims to evaluate the effects of a treatment and plays a vital role in decision-making. Benefiting from in-context learning, tabular foundation models have shown promising performance for causal effect estimation without fine-tuning on task-specific context data. Despite their strong generalization ability to new tasks, foundation models cannot fully address the distribution shift between treatment groups caused by confounding bias, which limits the performance of causal effect estimation. One natural approach is to learn an encoder to explicitly align the representations of groups. However, enforcing representation alignment could lead to over-balancing, which means that predictive information for outcome prediction is partially removed. Although minimizing the prediction loss of the foundation model can provide guidance for learning the encoder to some extent, the complex architecture of the foundation model makes it difficult to effectively propagate task-relevant gradient signals from the output back to the encoder. To address this issue, we incorporate a lightweight task-specific predictor after the encoder to preserve outcome-relevant predictive information in the learned representations. By doing this, the foundation model is enhanced by the task-specific encoder and predictor trained on context data. Moreover, we design a gated aggregation mechanism to adaptively fuse the predictions of the foundation model and the task-specific model. We theoretically analyze the estimation error of our method, and conduct extensive experiments based on multiple backbones of foundation models. The experimental results on benchmark datasets demonstrate the effectiveness of our method and its general applicability across different backbones.