Prediction-Powered Causal Inferences

We thank the reviewer @N6pk for the feedback. Overall, the general objective and purpose of the paper seem delivered and appreciated (“novel idea”, “promising results”), even if the raised concerns suggest a potential misunderstanding of the prediction-powered setting and related literature. We clarify here.

Causal Model Assumption (Weakness 2)

• The PPCI formulation does not require introducing a causal model, whose measurement-annotation relationship could also be controversial, as suggested by reviewer @7mz6. It simply relies on extending a well-defined treatment effect estimation problem (with its local causal assumptions), replacing the ground-truth (factual) outcomes with model predictions from high-dimensional predictive measurements.

Latent confounding (Weakness 7)

• Note that we are assuming that the causal estimand is statistically identifiable. If that is not the case, then even with ground truth $Y$, the causal estimation is ill-posed and so it doesn’t make much sense to investigate the Prediction-Powered extension of it. In our experiment, this is ruled out by the fact that we consider randomized trials, which are one of the many possible ways to guarantee identification (one could have instruments, or feasible backdoor adjustments).

The reviewer raised a number of valuable minor concerns and clarification requests that we now address, and we plan to adapt in the final version of the manuscript, hoping to improve its clarity accordingly.

• Introduction (Weakness 1)
  • The first paragraph in the introduction doesn’t refer to explainability, if not for the expression “black-box” referring to a model with generally no guarantees as to how it reaches its predictions, especially required in scientific applications. We are happy to make this clearer.
  • Given the 2-stage nature of the problem, i.e., factual outcome estimation and causal effect estimation, we need to specify the object of prediction.
  • A real-world problem (zero-shot generalization to ISTAnt) is used as a motivating example from the introduction to the conclusion to support problem interpretation, additionally illustrated in Fig. 1-3 and Example 1-2. As suggested, we plan to extend the dataset presentation in the introduction.
• Typos (Weaknesses 3a, 4 , 5).
  • We will replace $\mathcal{I}$ with $i$ ranging from $1$ to $n$. We thank the reviewer for spotting it, together with the other typos.
• Task Clarification (Weakness 3b, 3c)
  • ISTAnt task is explicitly presented in L89-90, “ATE estimation on the ISTAnt dataset zero-shot, i.e., without behaviour annotations. Instead, we leverage our own experiment variant with different treatment and recording platform” and additionally illustrated in Fig. 1. After extending the example presentation in the introduction, we hope this will be clearer.
  • AIPW [Robins et al. 1994] is the gold standard for ATE estimation and is assumed to be known according to common practices in the field.
• Clarification by Real-world Example (Weakness 6)
  • The Casual Lifting constraint is already explained by our motivating example (see L182-189), while the Conditional Calibration interpretation is self-explicatory, still interpreted several times, e.g., in Example 1 (see L149-152).

[1] J.M. Robins et al, Estimation of regression coefficients when some regressors are not always observed, JASA 1994.

We thank the reviewer @4Lsh for the curated feedback, which summary and strengths/weaknesses show the in-depth understanding and appreciation of the paper. We answer here to the raised questions and plan to integrate such clarifications in the final version of the manuscript, together with the suggested limitations emphasis.

Causal Lifting constraint: novelty and related works (Weakness 1, Questions 2-3)

• In this paper, we characterize which properties a (factual outcome) predictor requires to be incorporated in a causal inference pipeline without training on the target population. Alongside, we also propose a methodology for achieving them, relying on similar data. Particularly, we show that enforcing conditional independence between the measurement representation and the experimental settings, given the outcome, transfers Causal Validity among similar experiments (Thm. 3.1). The novelty lies in such formulation, which implementation happens to be a well-studied problem in Representation Learning for other reasons. We reviewed this in the Invariant Representation paragraphs in Section 4.
• Methodologically, we chose a specific implementation of the invariance constraint, which is far simpler than adversarial [Louizos et al. 2016] or information bottleneck [Alemi et al. 2017] approaches which are known to be unstable and inefficient respectively. We are happy to expand in Section 4 on the advantages of the different approaches and why we chose DERM for our particular application, see next.

DERM Limitations: full support and finite sample setting (Weaknesses 3-4, Question 1)

• We practically designed our ad-hoc implementation, motivated by our application of interest (high-dimensional representation, with discrete and limited covariates), but other approaches could and should be considered depending on the problem characteristics. As recognised in the limitations, DERM has no guarantees beyond the full support assumption (as shown with Hard shifts on CausalMNIST experiments), and new methods, tailored for such a regime, or just a strongly unbalanced distribution, surely deserve to be investigated.

Latent confounding (Weakness 2, Question 1)

• Note that we are assuming that the causal estimand is statistically identifiable. If that is not the case, then even with ground truth $Y$, the causal estimation is ill-posed. In our real-world experiment, this is ruled out by the fact that we consider randomized trials, which are one of the many possible ways to guarantee identification.

[1] Louizos, Christos et al. "The Variational Fair Autoencoder" ICLR 2016

[2] Alemi, Alexander A. et al. "Deep variational information bottleneck" ICLR 2017

We thank the reviewer @7mz6 for the careful and hands-on feedback. Overall, the general message to “tackle a timely, scientifically motivated problem via its formalization” and “offering a straightforward, model-agnostic” solution seems delivered. Two concerns are raised, which we respond to here and plan to adapt the final version of the manuscript accordingly.

Model selection (Weakness 1)

• We considered the best-performing model hyperparameters (which encoder tokens to use, head architecture, batch-size, etc.) on ISTAnt according to the extensive ablation study in Cadei et al. (2024). These were the only arbitrary choices and were performed using ERM, so they may at most favour the ERM baseline, which still drastically fails. For each method (ERM, vREX, IRM, ours), we tuned the learning rate and invariance regularizer of IRM and vREX for enough epochs (15) to allow convergence. We performed a grid search, discarding hyperparameters that led to factual outcome prediction accuracy smaller than 0.8 on a validation set from the training experiment (1000 samples). Next, we selected for each method the hyperparameters leading to the smallest Treatment Effect Bias [Cadei et al., 2024] estimated on the full prediction-powered training experiment (i.e., no leaking from the target distribution). We plan to clarify these steps in the final version of the paper, in Appendix B.2.
• Adam's second moment coefficient is a typo, and we confirm we used instead default values ($\beta_1=0.9$, $\beta_2=0.999$), as shared in the code.
• It is true that some encoders are more performant/confounded than others, also in agreement with the original paper from Cadei et al. (2024), yet our approach always succeeded in enforcing valid estimates. New benchmarks are required for a more exhaustive comparison among pre-trained models and inner causal entanglement evaluation, i.e., $I(\phi(X), Z|Y)$, but this is out of the scope of this paper.

Annotation Mechanism Assumption (Weakness 2)

• Annotation Mechanism invariance is assumed anyway by domain (human) experts, even without relying on artificial predictions. Taking our application as an example, experimental ecologists practically learn by experience how to recognize behaviours, and then assume it is transferable when they run new experiments and collect the same labels. Let’s further observe that in practice, we are never relying on their causal relationship, i.e., between the high-dimensional recordings (X) and the outcome behaviours (Y).

[1] Cadei, Riccardo et al. “Smoke and Mirrors in Causal Downstream Tasks” NeurIPS 2024

We thank the reviewer @xo43 for the detailed feedback. Overall, the general message to “connect predictive modeling to causal identifiability”, detecting “rigorous conditions” and proposing a “simple and effective method” seems delivered. The concerns relate to (i) assumptions test and justification, and (ii) causal model baselines. We comment on the former, while we disagree on the latter for factual points, and plan to adapt the final version of the manuscript with corresponding clarifications.

(i) Assumptions clarifications (Weaknesses 1 and 2)

• Conditional Calibration

  • Conditional Calibration on the target distribution is not an assumption, but a sufficient condition for valid causal inference. It can be enforced and tested on the training distribution, and we show in Thm 3.1 under which assumptions it generalizes to a similar experiment even if out-of-distribution.

• Causal Lifting Constraint

  • The key assumption enabling Thm. 3.1 is that the Causal Lifting constraint holds on the target distribution. We can train and test whether this constraint holds on the training, but the reviewer is correct that, in general, it may not transfer.
  • Between the two evils, we propose to enforce the Causal Lifting constraint and not directly Conditional Calibration. Indeed, a model optimized for standalone Conditional Calibration may still fail to transfer when dealing with new effect modifiers with similar appearance but different effects from the ones in the training experiment (see Example 2). This approach is similar in spirit to the considered Domain Generalization baselines, i.e., vREx and IRM, enforcing invariant prediction performances (ideally unbiased) varying the experiment settings. Instead, the Causal Lifting constraint robustly prevents such shortcuts directly in the representation space by ignoring any source of possible spurious correlations.
  • Ultimately, it is true that the Causal Lifting constraint is not a testable assumption on the target distribution (at least zero-shot). However, we believe that the approach has merit as we demonstrated empirically just enforcing the Causal Lifting constraint on the training experiment and validly estimating the treatment effect on a similar target experiment with different treatments and recording platforms.

• Full Support Assumption

  • DERM is our practical implementation of the Causal Lifting constraint, by sample reweighting, relying on the full support of outcomes across covariates, i.e., Eq. 9. This is easy to test, e.g., check that the behaviour of interest is happening at least once for each experimental setting; still, we agree that there could be scenarios where such a support assumption is challenged/violated, so losing its guarantees (we acknowledge it in the Conclusion). In practice, DERM still outperforms ERM, IRM,vREx, even in such a regime (see Hard Shifts experiments on CausalMNIST), and the assumption is not violated in our real-world application.

(ii) Comparison with other Causal Estimators (Weakness 3)

• The reviewer wonders if additionally benchmarking our pipeline for PPCI, i.e., causal lifting of foundational model and AIPW estimator, against standalone established causal effect estimators. However, it is not possible by design, since all the proposed causal estimators rely on the factual outcomes of the target experiment, which are not observed directly in PPCI, and only observed as raw measurements, e.g., pixels. More precisely:
  • In XLearner, E[Y|T=t, W=w] cannot be estimated on the target experiment because Y is not observed. On the training population, the corresponding conditional expectation can be completely different, as the effect can vary.
  • TARNet requires unavailable (target) factual outcome supervision for the network training. Possible variants, additionally relying on a pre-trained encoder, may be considered for fine-tuning on the training distribution (e.g., concatenating X and W) but without generalization guarantees (similar in spirit to our ERM baseline).
  • Similarly, BART requires unavailable (target) factual outcome supervision for each regression. Even if trained on the training distribution, it has no guarantee to generalize to a target experiment with different causal effects, and anyway, it wouldn’t scale on complex data structures, e.g., images.
  • CEVAE still requires unavailable (target) factual outcome supervision for training. Even if trained on the training distribution, it has no guarantee to generalize, and anyway, it is an overkill to assume unobserved confounders when they are actually observed (W) and potentially controlled too, e.g., RCT.

We thank reviewer @pzYf for the feedback. The methodological contribution of the paper seems properly delivered, while our problem formulation and theoretical analysis seems ignored as a contribution (no mention in reviewer’s summary and questioned Causal Lifting contribution) and partially misinterpreted (Zero-Shot meaning and implementation, Prediction-Powered Estimation Theorem) despite appreciated (“excellent and clearly written”, “neat, appealing, and well-presented”, “smooth and pleasant”, “highly relevant”). No major concern is raised, and weaknesses mostly revolve around unclear phrasing that we commit to improving. We address the raised questions in order, hoping to clarify potential misinterpretation:

• Q1: “Is ‘retrievable zero-shot from a pre-trained model’ a predicate for the factual outcomes? What exactly is retrievable, and how?”

  • Retrievable refers to the factual outcomes estimated by a predictive model. We agree it is not the best wording, and we will replace ‘‘retrievable” with “predicted”.

• Q2: “What does ‘this problem’ refer to in line 30? Why is it labeled a zero-shot problem before the pretrained model is introduced?”

  • Zero-shot refers to predicting the factual outcomes “without training on the data from the target experiment” (L35), transferring knowledge from a pre-trained model and similar training data (see Def. 2.3), while preserving the properties in the downstream causal estimation.

• Q3: “Can the authors elaborate on the source of bias in an RCT setting mentioned in lines 33–34?”

  • Biases in prediction-powered average treatment effect estimation can arise from different reasons, as discussed by Cadei et al. (2024). In L33-34, we refer to Lemma 3.1 from the same paper, describing how small biases in prediction (i.e., $1-\epsilon$ accuracy with systematic errors in a subgroup) can invalidate the downstream causal estimation even in the simplest causal setting, i.e., a Randomized Controlled Trial.

• Q4: “Would Lemma 2.2 hold for any standard ATE estimator under the given assumptions, not just under (g)?”

  • We prove Thm. 2.2 only for AIPW, gold standard for ATE estimation on RCT (e.g., in our experiments on ISTAnt), and observational studies with suitable backdoor adjustment. It shows how the properties of a causal estimator, i.e., AIPW, are preserved by a conditional calibrated factual outcome model $g$. Similar results may be shown for other estimators.

• Q5: “In the causal setting, would (X) be a parent of (Y) in a causal graph, or vice versa? Clarification would improve causal interpretation.”

  • For ease of interpretation, we assumed the outcome ($Y$) as a cause of the high-dimensional observation ($X$), see L139. Intuitively, $Y$ is something that happens in the real world, and X is its measurement, e.g., video recording (changing $Y$ affects X, but changing the pixels in $X$ does not affect what happens in the experiment). However, we also never explicitly rely on such an assumption, but rather on $P(Y|X)$ invariance.

• Q6: “How is DERM defined given that no labels for (Y) are available? How are the weights (w_i) constructed in this case?” and weakness “lines 33–35 create confusion in my opinion, if no labels are available, how is training done at all?”

  • DERM is the method we propose to enforce the Causal Lifting constraint by training (finetuning) a model on annotated data from similar experiments (see Def. 2.3), generally out-of-distribution for the non-annotated target experiment (by design). It reflects a realistic and common setup for the type of applications we target, e.g., ecologists study ants' behaviors across sequential similar experiments. In practice, it estimates the weights $w_i$ and operates on such annotated similar experiments, attempting to transfer causally valid predictions to the target experiment.

• Q7: “Is the definition of causal lifting introduced by the authors, or does it extend prior work?”

  • Yes, the definition is new, and it answers the challenge raised by Cadei et al. (2024) to estimate the treatment effect on the predictions of pre-trained models. Indeed, without further assumptions, a pre-trained model is generally biased, e.g., under-performing in a specific subgroup, and we introduce the expression Causal Lifting to refer to the procedure of correcting such biases.

• Q8: “Could the authors include Conditional Average Treatment Effect (CATE) results, given their relevance to ML applications?”

  • To the best of our knowledge, ISTAnt by Cadei et al. (2024) is still the unique benchmark for PPCI on high-dimensional measurements, and only evaluates ATE estimation. Through this paper, we have already contributed to it by introducing a similar real experiment for training (simulating a realistic PPCI pipeline). We agree on the relevance of extending PPCI benchmarks to heterogeneous effects, despite being outside the scope of this paper.

• Q9: “Would a footnote on valid adjustment sets and identifiability improve clarity and completeness?”

  • Fair point, we are happy to stress this further. We are indeed assuming that the causal estimand is statistically identifiable. If that is not the case, then even with ground truth Y, the causal estimation is ill-posed. In our experiment, this is ruled out by the fact that we consider randomized trials, which are one of the many possible ways to guarantee identification (one could have instruments, or feasible backdoor adjustments).
    ​​

The reviewer suggested a few other minor wording improvements, which are listed as the paper's weaknesses. Here is what we suggest:

• We will replace “Retrievable zero-shot” with “zero-shot predictions”.
• We will replace “Powering causal inferences” with “prediction powered treatment effect estimation”, which more precisely follows established Prediction Powered Inference terminology in statistics [Anastasios N. et al. 2023].
• We will avoid mentioning “Conditional Calibration” before its definition in Definition 2.2, unless to explicitly claim it as a novel contribution and point to the definition.
• In Definition 2.1 we assume that the causal estimand is statistically identifiable in $Y$ and define that a predictor for $Y$ is valid only if the causal estimand is also identifiable in $g(X)$. This definition ties the identifiability of $g(X)$ to that of $Y$, not viceversa. We’ll make sure this is clearer in the draft using the footnote on identifiability suggested in Q9.
• We thank the reviewer for the other specific suggestions where clarity can be improved. We will take this into account in the final revision and improve the presentation.

[1] Cadei, Riccardo et al. “Smoke and Mirrors in Causal Downstream Tasks” NeurIPS 2024

[2] Angelopoulos, Anastasios N. et al. “Prediction-Powered Inference” Science 2023