Paper: Towards Causal Foundation Model: on Duality between Causal Inference and Attention (ICLR-2024)

Reviewer comment (Theory):
The paper needs to provide an intuition of why this model works. What kind of information is learned by this model that enables zero-shot causal inference. Though it is relatively intuitive when it comes to text and image based foundation models, it needs further intuition and justification for causal effect estimation.

Author response:
In essence, we use attention to learn to produce optimal balancing weights that **debias** any given dataset. Then, treatment effect estimation can be estimated easily. We expect this to generalize well, since our theory suggests that attentions (under the correct loss function and assumptions) are **exact** optimal balancing solvers. 

In particular, our proposed method applies the attention mechanism to different units and tries to balance among the units. The optimal balancing weights are learned using the proposed loss in Eq. (7). In the zero-shot setting, we amortize the learning procedure and encode the procedure that produces the optimal balancing weight using a neural network (its parameters are denoted by $\phi$). Intuitively, once this is learned using multiple datasets, the model learns how to solve the optimal balancing problem. This is similar to learning how to optimize `[1]` or learning how to solve combinatorial optimization problems `[2]`, where the model acquires the ability to solve a new task (i.e., balance in a new dataset in our context). We discussed this briefly after introducing Algorithm 2 and 3, but we have edited it to further reflect and clarify the intuitions.

Reviewer comment (Reproducibility):
As a followup to the previous point, the experiments are not comprehensive. ALso more information on the model implementation and experimental set is needed.

Author response:
Regarding comprehensiveness: as causal inference tasks necessitate counterfactual information for validation and evaluation, experiments and benchmarks are often done on synthetic datasets or carefully curated real-world datasets. We included almost all common real-world causal inference datasets used by various previous works (c.f., benchmarks in `[3,4]`) and also created different simulations of varying difficulties. In addition, in revision (Appendix I.1), we included comparisons to more baselines: Dragonnet `[5]`, and RieszNet `[6]`, both considered as state-of-the-art neural estimation methods. We have also added larger scale ACIC experiments in Appendix I.2. We would be happy to add additional experiments if there are specific ones that we are missing. 

Regarding the model implementation and the experimental setup: apart from model architecture, hyperparameters, evaluations, and etc, which are provided in Appendix G and H, we have added the details of the data-generating process of simulation B to Appendix G. The data generating process of simulation A provided in Section 5.1 of the main text. We will also release the source code once this work becomes public.

Reviewer comment (Novelty):
This paper purports to be an initial step towards creating causal foundation models. But how does it compare to paper [1]? And could you elucidate the link between the causal transformer and the methods they proposed?  
[1] Causal Transformer for Estimating Counterfactual Outcomes.

Author response:
> **“This paper purports to be an initial step towards creating causal foundation models. But how does it compare to paper [1]? And could you elucidate the link between the causal transformer and the methods they proposed?”**

We thank the reviewer for this pointer and have added this discussion to the related work discussion of the revised manuscript (Appendix B).  

The method proposed by `[1]` is different from our work on a few aspects.  First, `[1]` trains their model on a **per-dataset** basis; and whenever a new dataset comes in, their model needs to be re-trained. Thus, this still falls into the realm of traditional, non-foundational causal models. While in our proposed framework towards building causal foundation models, we aim to train a single model that can perform zero-shot inference to a range of different contexts/datasets. 

Second, although both works make use of transformer structures, in their method, the transformer attends to the **time dimension for the same unit**, and produces a conditional estimand of the outcome of a single unit at the next step. While in our method, the attention is applied across different units to produce an estimand of the average effect across different units.

Reviewer comment (Presentation):
The methods you've presented can be challenging to understand, particularly for those not well-versed in the subject. It would be helpful to include an introductory section that provides background information. After that, you can highlight your contributions.

Author response:
> **“The methods you've presented can be challenging to understand, particularly for those not well-versed in the subject. It would be helpful to include an introductory section that provides background information. After that, you can highlight your contributions.”**

We thank the reviewer for this suggestion, we have modified the introduction to include a short description of the problem under study (highlighted in contributions), as well as the related work section to include pointers to background information in Section 3.1 (highlighted in “Causal Inference via Optimal Balancing”).  

The problem formulation of estimating average treatment effect is summarized in the first paragraph of Section 3.1. The method we proposed is following the literature on optimal covariance balancing, which produces an estimand for the treatment effect based on re-weighting different units in a dataset (see our response to the above question of how attention is used to produce balancing weights). We have provided a discussion on optimal covariance balancing in Section 2.

Reviewer comment (Reproducibility):
The source code is not available, but the authors say they will release it upon publication.

Author response:
We will open source the code upon publication.

Reviewer comment (Theory):
I couldn't understand the connection between different causal graphs and covariate balancing-- the authors say they shuffle units from different datasets using training-- how is the differentiation between how different the causal structures are, performed? And why does this not help with increasing the number of covariate balancing problems without acquiring more data?

Author response:
In the dataset augmentation described in Appendix H.3, we added another paragraph to clarify this point. This reshuffling trick is only applied to a set of datasets that share the **same causal graph** in the data-generation process. For example, we do this for simulation A and IHDP-resampled (this shares the same causal graph because different datasets share the same set of covariates and the treatments/outcomes are sampled conditioned on these covariates). 

For a set of datasets where each contains units generated by a causal graph that varies between datasets, e.g. ER-5000 and ACIC, this is disabled. The main intuition is that if we reshuffle units among these datasets, units in a reshuffled dataset could follow different causal graphs, which means there is potentially no underlying causal structure that can explain the data.

Reviewer comment (Novelty):
Overall, I believe the framing of the work as a step towards causal foundation models is misleading, as the method introduced is limited to estimating the ATE under the classic ignorability assumption. The theoretical novelty of the work does not seem strong, as Theorem 1 is a direct implication of results from prior work cited in the paper (Tarr & Imai, 2021; Kallus, 2020; Nguyen et al., 2022). Finally, the strength of the empirical results is tempered by the facts that there is no comparison in performance to various neural ATE estimation methods, and that it is not possible to rule out overfitting as an explanation for outperformance of zero-shot (supervised) learning (see point 7).

Author response:
The **task-universality** of foundation models does not necessarily require adaptability across different **ML formulations** (e.g., prediction, imputation, ATE, counterfactuals); it can also refer to adaptability across different **contexts** of a given task. This viewpoint is shared by many recent studies on foundation models beyond LLMs `[1,2,3,4,5]`, where they all focus on one single type of task but evaluate in-context generalization across datasets. In our case, we focus on treatment effect estimation tasks and evaluate our model in unseen datasets/contexts. We have clarified this implication of foundation models in our revision (Appendix A).

Ignorability and non-interference are classic canonical assumptions that ensure ATE and CATE identifiability.  In practice, our method performs well even with certain levels of violations, as shown in ER-5000 and real-world datasets. Evidence from literature `[6,7]` indicates that optimal balancing can still be a good objective when certain assumptions are relaxed, e.g., when there is latent confounding. We acknowledge the importance of these extensions; however, they fall into orthogonal directions to this work.

As clarified throughout the paper, our work is intended as a stepping stone towards causal foundation models, not a final solution. We have emphasized this in our title, introduction, and the revised abstract. We acknowledge the importance of solving causal problems beyond ATE under relaxed assumptions. Nevertheless, these require much more development that is beyond the scope of this paper and would be best addressed in separate follow-up works.

> **“Are there novel theoretical results in this work beyond direct implications from (Tarr & Imai, 2021; Kallus, 2020; Nguyen et al., 2022)?”**

The novelty of our result is clarified below. We only borrow `(Nguyen et al., 2022)` for the existence of feature maps in Eq. (6); however the SVM optimization problem derived from our formulation is different from `(Nguyen et al., 2022)`. Moreover, our theoretical result focuses on gradient-based optimization in Eq. (7) whereas `(Nguyen et al., 2022)` focusing on deriving new attention architectures. Finally, we would like to highlight that our result gives rise to a novel and important algorithmic implication that attention can be exactly mapped as a treatment effect solver. This enables a theory-grounded method for zero-shot causal inference. To the best of our knowledge, this finding is new. 

## Rebuttal 2 (ID: IZXdppyN8R)
> **“There is no comparison to methods such as (Shi et al., 2019) ... There is also no comparison to (Xia et al., 2021; Xia et al., 2022) ...”** 

Thank you for this suggestion. We added a comparison to two popular neural estimation methods: Dragonnet `(Shi et al., 2019)` and RieszNet `[10]`, which can be found in Appendix I.1. We summarize the results below for simulation A (Multi+OOD+diffsize), ER-5000, and IHDP datasets.

|Method|Simulation-A |ER-5000|IHDP|
| ----------------------------------------- | ------------------------- | ------------------------- | ------------------------- |
|DragonNet|0.386 ± 0.01|11.21 ± 3.17|0.146 ± 0.01|
|RieszNet|**0.045 ± 0.01**|12.90 ± 4.54|**0.110 ± 0.01**|
|Ours|0.126 ± 0.02|N/A|**0.114 ± 0.01**|
|Ours (ZS)|0.147 ± 0.01|11.50 ± 1.85|N/A|
|Ours (ZS-S)|N/A|**2.66 ± 0.33**|N/A|

We observe that our zero-shot method in general is still on-par or better than Dragonnet and RieszNet. This indicates that the proposed method can reach current per-dataset state-of-the-art go-to methods, even under zero-shot inference settings

We did not compare to `(Xia et al., 2021 & 2022)` as their method is centered around G-consistent NCMs, which require access to the true graph of the underlying dataset. However, we included a full discussion of these neural estimation methods in Appendix B.

Reviewer comment (Novelty):
The title "Towards Causal Foundation Model" seems to be misleading, along with the claim "we take a first step towards building causal foundation models, focusing on causal inference tasks". The paper does use self-attention, a key element of most 'foundation models' in the language or vision sense. However, the proposed method is limited to estimation of the ATE under the assumptions of ignorability and sample non-interference; indeed, it relies on the structure of the problem to produce a weighted average of observed outcomes. According to (Bommasani et al., 2021), "A foundation model is any model that is trained on broad data (generally using self-supervision at scale) that can be adapted (e.g., fine-tuned) to a wide range of downstream tasks; current examples include BERT [Devlin et al. 2019], GPT-3 [Brown et al. 2020], and CLIP [Radford et al. 2021]." Clearly, this does not hold for a method that can only estimate the ATE under a specific set of assumptions, independent of any efficiency or performance improvements.

Author response:
The **task-universality** of foundation models does not necessarily require adaptability across different **ML formulations** (e.g., prediction, imputation, ATE, counterfactuals); it can also refer to adaptability across different **contexts** of a given task. This viewpoint is shared by many recent studies on foundation models beyond LLMs `[1,2,3,4,5]`, where they all focus on one single type of task but evaluate in-context generalization across datasets. In our case, we focus on treatment effect estimation tasks and evaluate our model in unseen datasets/contexts. We have clarified this implication of foundation models in our revision (Appendix A).

Ignorability and non-interference are classic canonical assumptions that ensure ATE and CATE identifiability.  In practice, our method performs well even with certain levels of violations, as shown in ER-5000 and real-world datasets. Evidence from literature `[6,7]` indicates that optimal balancing can still be a good objective when certain assumptions are relaxed, e.g., when there is latent confounding. We acknowledge the importance of these extensions; however, they fall into orthogonal directions to this work.

As clarified throughout the paper, our work is intended as a stepping stone towards causal foundation models, not a final solution. We have emphasized this in our title, introduction, and the revised abstract. We acknowledge the importance of solving causal problems beyond ATE under relaxed assumptions. Nevertheless, these require much more development that is beyond the scope of this paper and would be best addressed in separate follow-up works.

Reviewer comment (Theory):
The assumption of sample non-interference, along with the fact that the self-attention is applied on the units of a dataset, seems to remove any connection with the "Causal Inference with Large Language Models" section, as one main reason attention is useful in language models is the fact that it handles sequences well. In theory, arbitrarily re-ordering a dataset should not affect the output of Algorithm 1, due to the non-interference assumption. Re-ordering the words in a document inputted to an LLM should, of course, affect the output of an LLM. Therefore, it is difficult to see any connection between self-attention as used in this work and as used in LLMs - or any connection between this work and LLMs, beyond the fact that it is known that LLMs have difficulty solving the problem of ATE estimation from data and assumptions.

Author response:
> **“What is the connection between this work and LLMs?” / “lacks the classic sequence-learning benefits of self-attention in LLMs”**

We apologize for the confusion caused by the related work section on “Causal Reasoning with Large Language Models”. This section was intended to highlight the limitations of LLMs in quantitative causal tasks such as ATE estimation; we did not imply that there is any deep connection between this work and LLMs, since attention was used for very different purposes. We further clarified this by adding the discussion in the corresponding section.

Reviewer comment (Novelty):
Section 3.1 appears not to contain novel theoretical results; the results showing optimal covariate balancing is equivalent to the dual SVM problem appear to be repeated from (Tarr & Imai, 2021; Kallus, 2020). Section 3.2 appears to rely on (Nguyen et al., 2022) for its result showing the dual SVM problem may be reformulated as self-attention. Therefore, the novelty of Theorem 1 does not seem to be particularly strong, as it is a direct implication of prior work's results.

Author response:
> **“Are there novel theoretical results in this work beyond direct implications from (Tarr & Imai, 2021; Kallus, 2020; Nguyen et al., 2022)?”**

The novelty of our result is clarified below. We only borrow `(Nguyen et al., 2022)` for the existence of feature maps in Eq. (6); however the SVM optimization problem derived from our formulation is different from `(Nguyen et al., 2022)`. Moreover, our theoretical result focuses on gradient-based optimization in Eq. (7) whereas `(Nguyen et al., 2022)` focusing on deriving new attention architectures. Finally, we would like to highlight that our result gives rise to a novel and important algorithmic implication that attention can be exactly mapped as a treatment effect solver. This enables a theory-grounded method for zero-shot causal inference. To the best of our knowledge, this finding is new.

Reviewer comment (Experiments):
There is no comparison to methods such as (Shi et al., 2019), cited in the paper, which is a neural estimation method for the average treatment effect. There is also no comparison to (Xia et al., 2021; Xia et al., 2022; https://causalai.net/r80.pdf, https://causalai.net/r87.pdf), a neural identification and estimation technique for arbitrary interventional and counterfactual queries.

Author response:
> **“There is no comparison to methods such as (Shi et al., 2019) ... There is also no comparison to (Xia et al., 2021; Xia et al., 2022) ...”** 

Thank you for this suggestion. We added a comparison to two popular neural estimation methods: Dragonnet `(Shi et al., 2019)` and RieszNet `[10]`, which can be found in Appendix I.1. We summarize the results below for simulation A (Multi+OOD+diffsize), ER-5000, and IHDP datasets.

|Method|Simulation-A |ER-5000|IHDP|
| ----------------------------------------- | ------------------------- | ------------------------- | ------------------------- |
|DragonNet|0.386 ± 0.01|11.21 ± 3.17|0.146 ± 0.01|
|RieszNet|**0.045 ± 0.01**|12.90 ± 4.54|**0.110 ± 0.01**|
|Ours|0.126 ± 0.02|N/A|**0.114 ± 0.01**|
|Ours (ZS)|0.147 ± 0.01|11.50 ± 1.85|N/A|
|Ours (ZS-S)|N/A|**2.66 ± 0.33**|N/A|

We observe that our zero-shot method in general is still on-par or better than Dragonnet and RieszNet. This indicates that the proposed method can reach current per-dataset state-of-the-art go-to methods, even under zero-shot inference settings

We did not compare to `(Xia et al., 2021 & 2022)` as their method is centered around G-consistent NCMs, which require access to the true graph of the underlying dataset. However, we included a full discussion of these neural estimation methods in Appendix B.

Reviewer comment (Experiments):
It is unclear how the method will scale to datasets of sample size n > 1024, as evaluation time scales quadratically with n. For example, ACIC contains 100k samples, but the experiments subsample datasets of size <=1000, as described in Appendix E. The comparison with ACIC seems unfair, because other methods could learn from the 100k-sample dataset in at most linearly scaling time, while the proposed method's runtime would scale quadratically.

Author response:
> **“It is unclear how the method will scale to datasets of sample size n > 1024…” / “I understand that the method proposed in the paper likely cannot scale well to dataset of ~100k samples, such as ACIC 2018. However, it would be informative to apply DML to each full dataset of 100k and compare this performance to the performance of CInA when given an allowance of the same runtime.”**


Thank you for this comment. We would like to address this using the following points:
- The original ACIC dataset `[11]` includes pre-generated datasets of varying sizes (1k, 2.5k, 5k, 10k, 25k, 50k), each with different data generation mechanisms. We used datasets of size 1k without any further down-sampling in the experiment section. We clarified this in Appendix G.
- During revision, we conducted additional experiments on larger ACIC datasets of size 10k, a common choice of scale considered in the previous benchmarking (`Shi et al., 2019` and `[12]`) This was added to Appendix I.2. We attached the results below. Our method still significantly outperformed other baselines.
- Regarding the current model scale, we would like to highlight that our context length (1024) is already larger than some existing foundational models for tabular and time series datasets, such as `[1,3]` which used 512. 
- Our method, with its transformer-like structure, can benefit from efficient transformer techniques `[13,14,15,16]` to handle datasets of larger scale. We added this point to the discussion section.

|Method|ATE MAE|Elapsed Time (s)|
| ---------------------------------------------------- | ------------------------- | ------------------------- |
|Naive|13.07 ± 8.25|0.005|
|IPW|10.29 ± 5.94|48.927|
|Self-normalized IPW|10.30 ± 5.90|49.322|
|DML|8.572 ± 8.96|7391.743|
|RieszNet| 69.39 ± 31.9|8157.498|
|Ours (ZS)|**1.460 ± 0.48**|78.503|
|Ours (ZS-S)|**1.361 ± 0.42**|77.546|