Orthogonal Representation Learning for Estimating Causal Quantities

[2 / 5]

Quality&Clarity. We are happy to provide an in-depth mathematical analysis in the revised version of the paper. Let us elaborate on the mentioned issues:

1. To answer your question, we added two new theoretical results,new Remark 2 and Proposition 7. In Remark 2, we formalized the favorable theoretical properties (i.e., double robustness and quasi-oracle efficiency) and showed that, even with slow converging estimators of the nuisance functions, our OR-learners consistently and (quasi-oracle) efficiently estimate the main causal quantities: representation-conditional CAPOs/CATE. In the case of a degenerate representation (e.g., $\Phi(X) = 0$), our OR-learners still efficiently estimate the causal quantities that are now APOs/ATE (= best constant estimators of CAPOs/CATE). We showed this in our new Proposition 7.

   Action: We added two formal statements (new Remark 2 and Proposition 7), where we explained how our OR-learners achieve consistency and efficiency even in the case of the degenerate representations (e.g., $\Phi(X) = 0$).

2. We agree, that the performance of the OR-learners as the predictors of the covariate level CAPOs/CATE heavily depends on the choice of the representation (which is expected), as different representations contain different amounts of information about the heterogeneity of the potential outcomes/treatment effect. Yet, in all the mentioned cases, our OR-learners can consistently and (quasi-oracle) efficiently estimate the representation-level causal quantities:

   • a) See the answer to 1.

   • b-c) In this case, the representation preserves some (but not all) of the heterogeneity of the potential outcomes/treatment effect. In this case, our OR-learners will adjust for the RICB but still would be limited in their heterogeneity (see a new Remark 1). We showcased this with our new Figure 3, where we demonstrated finite-sample vs. asymptotic performance of the representation learning methods and our OR-learners.

     Action: We added our new Remark 1 to prove that our OR-learners adjust for the RICB and a new Figure 3 to describe the performance of our OR-learners with different amounts of the RICB.

   • d) Very interesting question. In this case, our OR-learners would try to “calibrate” the outputs of the propensity / prognostic score so that the resulting target model would be Neyman-orthogonal (unlike propensity/prognostic scores alone) and, thus, double robust and efficient.

     Action: We have extended the discussion on the unconstrained representations in Sec. 4.1 (iv) to include propensity / prognostic scores as possible representations.

   Action: We decided to add a paragraph in Sec. 6, where we discuss a choice of the target model’s inputs. Therein, we explained that the good performance of the OR-leaners depends on an important inductive bias: the high-dimensional covariates should lie on a low-dimensional manifold and the learned representations should capture it well. In this case, the target model with the representations $\Phi(X)$ as the input (as in our OR-learners) would be superior over a target model with $X$ as the input (that might struggle to fit $g(X)$ with DR/R-loss due to high-dimensionality of $X$).

Response to “Questions”

1. l.234. Thank you for the suggestion! Action: We added additional references to the idea of overlap-weighting the MSE objective for CATE estimation [10-11].

2. l.304-306. Sorry for the incompleteness of the reference [13]. We decided to be more precise and formulated a new Remark 3 that states that there exists a hidden layer starting from which the (Hölder) smoothness starts increasing.

   In our proof of the new Remark 3, we refer to the intuition provided in the proof of Theorem 1 in [13] (specifically, the proof of Lemma 3 (d)). Therein, the authors formulated results for fixed-width neural networks with locally quadratic activation functions. Specifically, these networks require an increasing number of hidden layers to be able to approximate increasingly non-smoother functions [e.g., Lemma 3 (d) considered Taylor expansion of the square root function with an increasing number of terms]. Similar in intuition theoretical result (more flexibility requires more layers) holds for other types of fixed-width deep networks [14, 15]. Then, our proof follows by contradiction: There should be a hidden layer with larger smoothness since, otherwise, we wouldn’t be able to approximate the function solely with the remaining layers.

   In our context, Remark 3 holds for the broad class of the representation learning methods. For example, in the experiments, the learned low-dimensional representations, $d_\phi \le d_x$, have fixed dimensionality; and the activation function, ELU, can be considered locally quadratic.

   Action: We added a new Remark 3 and an (informal) proof to our revised PDF, where we explained the reference to [13].

[3 / 5]

3. l. 310-316. In this paragraph, we meant the vector of covariate-conditional expected outcomes, $\Phi(X) = (\mu_0^x(X), \mu_1^x(X) )$, as an example of a valid representation with $d_\phi = 2$. We show its validity in our new Proposition 4 in Appendix C. Such representations can be learned arbitrarily well in the asymptotic regime, given a sufficiently deep representation subnetwork FC$_\phi$ with unconstrained representations (follows from the universal approximation theorem).

   Action: We clarified this statement in the updated version of the PDF (see our new Proposition 4 in Appendix C).

4. l. 319-320. Great question! To formalize the efficiency of estimation of the target model $g$, we used the notions of (i) double robustness and (ii) quasi-oracle efficiency (see Sec. 3.2). Also, we added a new Remark 2 where we precisely stated this efficiency result for our OR-learners. Therein, (i) double robustness implies that, if one of the nuisance functions is estimated consistently, then the projection of the representation-conditional CAPOs/CATE on the working model class, $g^*$, is estimated consistently. Additionally, (ii) quasi-oracle efficiency means that the minimizer of the target loss of Neyman-orthogonal learners with estimated nuisance functions is asymptotically indistinguishable from the minimizer of the target loss with the ground truth nuisance functions even if the nuisance functions have slow rates of estimation (e.g., $o(n^{-1/4})$. Importantly, quasi-oracle efficiency does not hold generally for non-orthogonal learners (e.g. RA- and IPTW-learners), as they always contain a first-order estimation error of the nuisance functions. The theoretical properties (i) and (ii) thus render our OR-learners asymptotically optimal [18].

   Action: We formalized the notation of the target model efficiency in a new Remark 2 (see the updated PDF).

5. l.365-367. To demonstrate why the representation network would scale up the parts of the covariate space, we refer to new Remark 3 and Proposition 5. In Remark 3, we showed that to minimize the MSE, a neural network should increase Hölder smoothness with the depth increase (see our answer to Question 2. l.304-306). Proposition 5 then suggests that the Hölder smoothness increase is only possible when a representation network has a Lipschitz constant larger or equal than one. For the normalizing flows as representation networks, this implies a determinant of the Jacobian larger than one (which, in turn, implies “scaling up”).

   We also confirmed these results with new experimental evidence. In a new Figure 7 in Appendix F.2, we showed how an invertible representation subnetwork scales up the original covariate space when no balancing is applied.

   Action: We added new new Remark 3 and Proposition 5, where we proved the original claim that the representation network would scale up the parts of the covariate space to minimize the MSE. Also, we experimentally confirmed this result based on the synthetic dataset (see a new Figure 7 in Appendix F.2).

6. l.381-384. Yes, you are right: By balancing representations in Sec. 4.2 and 4.3, we meant minimizing the distributional difference between the representation of both groups and not balancing with propensity weighting (thank you for this remark).

   We used the term “re-focusing” in the sense that the parts of the covariate space with low overlap would have large inverse propensities in the DR-learner’s loss, thus forcing the target model to focus on them.

   Regarding the degenerate case with $\Phi(X) = 0$, you raised a very important question (this case can happen only with non-invertible representations). Interestingly, in this case, our OR-learners will lose all the heterogeneity of the potential outcomes/treatment effect but still can consistently estimate average potential outcomes (APOs) or an average treatment effect (ATE). This happens as, in the case of the degenerate representations, the working model class can only fit an intercept and, thus, the objective of the DR-learner coincides with semi-parametric efficient (A-IPTW) estimators of APOs/ATE. We have formalized this interesting result in our new Proposition 7. Note that all the theoretical properties from Remark 2 would still hold (now $V$-conditional CAPOs/CATE became APOs/ATE).

   Action: In the revised version of the manuscript, we clarified that balancing (in Sec. 4.2 and 4.3) refers to the minimization of the distributional distance and not to balancing with propensity weighting. Also, we explained what “re-focusing” means for the DR-learner. Additionally, we added our new Proposition 7, where we explained that, in the case of the degenerate representations $\Phi(X) = 0$, our OR-learners yield semi-parametric efficient (A-IPTW) estimators of APOs/ATE.

[1 / 5]

We are very grateful for your very comprehensive and actionable review! First of all, we are very happy to hear that you found our idea original and that you appreciated the extensiveness of our related work and experiments. In the following, we respond to each comment and how we improved our paper as a result (see our changes in $\color{blue}{\text{blue}}$ in our updated PDF).

Response to “Weaknesses”

Originalty&Signifiance. Thank you for recognizing the novelty of our work: “The first work I see learning representations … to be fed into pseudo-outcome learning”. Indeed, there are some prior works on non-neural representations [1-4] for causal inference. We thus saw the need to expand our literature review and discuss in greater detail how our work is different and novel (see a new paragraph in Appendix A.1).

We are also happy to spell out clearly how our work is different from references [1-4]. Indeed, the papers [1-4] and our work proceed by having the representations as some kind of input for the final stage of the estimation. However, we want to highlight that the papers [1-4] focus on different, much simpler settings. Likewise, other papers [5-6] also focus on different, and rather simple settings. We make a comparison in the table below. Therein, we highlight differences across three dimensions: (i) estimation target, (ii) assumptions, and (iii) linearity of the representation. We give a detailed comparison for each (i)--(iii) below the table.

This is how our work is different from [1-4] and [5-6]:

• Estimation target: The above-mentioned works [1-4, 6] focus on averaged causal quantities, like APO, ATE, or ATT. In contrast, our work studies a general class of heterogeneous causal quantities: representation-conditional CAPOs/CATE. Therefore, our work is more general than [1-4, 6].
• Additional assumptions. Several of the works require additional, strong assumptions. For example, many works on prognostic scores [2-4,6] consider its simplified form ($\Phi(X) = \mathbb{E}(Y[a] \mid X)$), which requires an additional strong assumption that $Y[a] \mid X$ follows a generalized linear model. Also, the work of [4] only considers linear deconfounding scores and additionally makes strong parametric assumptions on the data-generating mechanism (Gaussianity of covariates). In contrast, our work does not make such simplifying assumptions.
• Linear representations. Arguably, the prognostic scores [2] can be also used as representations for CAPOs/CATE estimation. Yet, the prognostic scores are rather a theoretical concept (unless simplifying assumptions are made [3,4,6]), and only simple linear prognostic scores have been studied so far [3, 5]. To the best of our knowledge, the first practical method for non-linear, learnable representations was proposed by [7-9], while our work turns [7-9] into Neyman-orthogonal learners. => The methods in [7-9] (i.e., TARNet, CFR, rCFR) are thus used as baselines in our work.

Action: We added a more extensive overview of the literature to our Extended Related Work (see Appendix A.1 of the updated PDF). Therein, we spell out clearly how our paper is different from [1-6] and thus how our paper is novel.

[1 / 2]

Thank you for the time devoted to reviewing our paper, we really appreciate your follow-up questions! We have addressed them below (and also in our freshly updated PDF where the newest changes are highlighted in $\color{red}{\text{red}}$):

• (1) Your understanding is correct: Our OR-learners are optimal in the sense that they achieve the lowest asymptotic error among the structure-agnostic learners (e.g., RA-/IPTW-learners) of $\Phi(X)$-conditional CAPOs/CATE [1].

  Action: We highlighted this fact in the main part of the paper.

  “do you mean that they actually just try to find the estimators of the CAPOs/CATEs depending on $\phi(X)$”

  Yes, you are right. In Remark 1, we established that $\Phi(X)$-conditional CAPOs/CATE are identifiable (i.e., given the infinite amount of the observational data), and, in Remark 2, we provided the asymptotically optimal meta-learners for them.

  Do you have any guarantees on the loss with $V=X$ and the regressor being $g \circ \Phi$, when $g$ comes from minimizing the loss using $V=\Phi(X)$ and $\Phi(X)$ has a given RICB ?

  This is another very interesting question! If we set $V = X$ and define a second-stage model as $g \circ \Phi$ (with a possibility of additional constraints enforced on $\Phi(X)$), we can get an alternative learner of the representation-level CAPOs/CATE. This new learner would have modified target losses (i.e., Eq. 6/9 + $\alpha$ * balancing constraint) and, still, it would also be Neyman-orthogonal and, thus, asymptotically equivalent to our OR-learners (we show it in our new Remark 8). The main reason is that the (balancing) constraint is itself insensitive to the nuisance functions misspecification (i.e., it doesn’t depend on the propensity score or the covariate-conditional expected outcome).

  The natural question arises of how to choose between different variants of Neyman-orthogonal learners: our OR-learners and the newly proposed alternative. Here, there is no universal answer but only finite-sample considerations. The same issue exists e.g. for the orthogonal learners of CAPOs [2] where three options are available (=two types of DR-learners and i-leaner). Our main finite-sample consideration is that it is that the second-stage model is more optimizationally stable when it has the learned representations $\Phi(X)$ as the inputs in comparison to the suggested model with $g \circ \Phi$, which has to learn the representations $\Phi(X)$ from scratch with a potentially much more unstable loss. The larger instability for the newly proposed learner comes from (a) inverse propensity weights of the DR-loss, (b) additional balancing constraint, and (c) the need to train both $g \circ \Phi$. In our experiments (Setting A, $V = X^*$), we verified that our proposed OR-learners with $V=\Phi(X)$ are indeed more effective than the orthogonal training of the full $g \circ \Phi$.

  Action: We added a new Remark 8, where we described an alternative to our OR-learners, as suggested by the reviewer. Furthermore, we added a discussion of how our OR-learners should be preferred in terms of finite-sample considerations.

  … it is possible to not degrade $X$-conditional CAPO/CATE estimation despite a (not too) bad $\Phi(X)$.

  Both our OR-leaner and the newly proposed alternative variant would degrade the $X$-conditional CAPOs/CATE: our OR-learners would use degraded $\Phi(X)$ as inputs and the new alternative would add the balancing term to the loss (thus creating a degrading “bottleneck” wrt. $\Phi(X)$).

  (b) Is there any way you can find any Φ(X) minimizing the RICB or reaching a guaranteed RICB?

  The RICB can be induced e.g., by balancing constraints. If $\Phi(X)$ aims at minimizing the factual MSE and no balancing is applied (e.g., as in vanilla TARNet, see Sec. 4.1), it doesn’t contain the RICB but only a finite-sample estimation error.

  (c) Alternatively, can you extend the OR-learners to any adjustment set V? Is it really necessary to make it depend on X?

  Yes, in our OR-learners, $V$ can be any measurable function of $X$ and even a constant (e.g., as suggested in Proposition 7). Thus, the models with $V$ independent of $X$ effectively aim at the APOs / ATE.

• (2) To the best of our knowledge, all the related works on balancing representations chose specific IPMs (WM, MMD, and sometimes, total variation (TV)) but still used the general term to refer to them. However, we agree with you and we will tone down the statement so that it is specific to WM and MMD. Indeed, our Proposition 6 does not hold for TV as the TV is also an f-divergence and, thus, stays constant under the invertible transformations.

  Action: We fixed the formulation of Proposition 6 to refer to only WM and MMD.

Our apology, it should be a Remark 8.

We additionally highlighted the newest changes in the PDF with $\color{red}{\text{red color}}$.

[1 / 2]

We are grateful for your review! It is great that you found the problem our paper tackles important. Below, we respond to each comment and how we improved our paper as a result (see our changes in $\color{blue}{\text{blue color}}$ in our updated PDF).

Response to “Weaknesses”

1. Thank you for giving us the opportunity to clarify how our work is different from [1]. Despite some similarities with [1], our work can be rather seen as a non-trivial extension of [1]. Let us elaborate on how our work makes important contributions over [1] and how our work is thus novel:

   1. You are right, the biggest difference between [1] and our work is that we consider the learned representations $\Phi(X)$ as an input to the second-stage model instead of the original covariates $X$. This is non-trivial and thus requires a new derivation. There are two main reasons why such a choice of the target model input might be beneficial for CAPOs/CATE estimation:

      • Differences in method: In Sec. 4.1 (iv), we discuss the unconstrained representations. Therein, the target model with $\Phi(X)$ as inputs provides a middle ground between the full retraining (= with $X$ as inputs as in [1]) and a simple calibration of outputs (= with $\{\hat{\mu}^x_0(X), \hat{\mu}_1^x(X)\}$ as inputs). There, our OR-learners perform a conditional calibration of the learned representations and enable the theoretical properties, namely, double robustness & quasi-oracle efficiency. The choice of $\Phi(X)$ as input turns out to be highly effective in the synthetic and semi-synthetic experiments (see Setting A of Sec. 5).
      • Differences in setting: In the context of the constrained representations, the inputs to the target model might be additionally forced to satisfy constraints, e.g., balancing or fairness (if satisfaction of constraints is a crucial requirement, the target model has to be based on the representations). Importantly, this setting was not studied by [1]. In Sec. 4.2 and 4.3, we specifically focus on balancing representations. We discovered that the orthogonal learners would try to “undo” the effect of balancing: either the benefit of reducing the estimation variance or the damage from too much balancing (=representation-induced confounding bias). These insights constitute our contribution.

      Action: We expanded the literature review to spell out more clearly the differences between [1] and our work in the updated PDF.

   2. Thank you. You are right: In our experiments, the setting with $V = X$ can be seen as an implementation of [1] with a one-layer neural network as the target model. Yet, in the other setting $V = X^*$, we already implemented exactly what you suggested. Specifically, we set the target model with the inputs $X$ to match the depth of the original representation network (i.e., the depth of three hidden layers). As seen in Tables 1 and 2, the results $V = X$ and $V = X^*$ are comparable but inferior in comparison with $V = \Phi(X)$, suggesting the effectiveness of our OR-learners.

      Action: We provided a clarification in the updated PDF that the special case with $V = X^*$ provides a fair comparison of different target models in terms of the network depths. In numerous experiments, we demonstrated that our OR-learners (which use the learned representations and one hidden layer target model) should be preferred over the full retraining of the target model (which coincides in depth with the original representation network).

[1 / 2]

Thank you for your review! We appreciate that you found the topic of our paper interesting. We hope we can address your concerns in the comment below and the revised version of the paper (see our changes in $\color{blue}{\text{blue color}}$ in our updated PDF).

Response to “Summary”

Thank you for the feedback! We took your suggestion to heart and are happy to provide more theoretical analyses, formal statements, and proofs in the revised version of the paper. In particular, we added the following:

• We showed that the representation-conditional CAPOs/CATE are identifiable, given the covariate-conditional nuisance functions, and therefore, our OR-learners can adjust for the representation-induced confounding bias (RICB) (see Remark 1).
• In a new Remark 2, we formalized the favourable theoretical properties guaranteed by our OR-learners due to their Neyman-orthogonality (namely, double robustness and quasi-oracle efficiency).
• We demonstrated, that under mild conditions on the representation network, the minimization of the MSE loss requires an increase of (Hölder) smoothness (see Remark 3).
• We added the proof, that a representation with $d_\phi \le 2$ is sufficient not to induce the RICB (see our new Proposition 4).
• In our new Proposition 5, we proved that, in order for the representation network to increase smoothness, it needs to have a Lipschitz constant of larger or equal than 1. At the same time, in our new Proposition 6, we demonstrated that to achieve balancing, the converse needs to happen (namely the Lipshitz constant needs to be less or equal than 1).
• Finally, in our new Proposition 7, we showed that in the extreme case of a constant representation, our OR-learners yield efficient estimators of APOs/ATE.

Action: We added the above statements as new remarks and propositions to our revised paper. We also added the corresponding proofs to a new Appendix C.

Response to “Weaknesses”

The author may want to present some evidence for some claims in the manuscript.

We are happy to provide the proofs and explanations for the claims in Sec. 3 and 4 (see the response to “Summary”).

Action: We added the above statements as new remarks and propositions to our revised paper. We also added the corresponding proofs to a new Appendix C.

In addition, combining existing ideas may not be interesting enough. The writing can be much improved.

Our paper is located at the intersection of two literature streams (see Appendix A): end-to-end representation learning methods and Neyman-orthogonal learners. Yet, to the best of our knowledge, both streams were developing almost independently and no work discussed a unification of both streams, which is our novelty. As such, the combination of both theories requires an extensive, non-trivial discussion, which we provide in Sec. 4. Importantly, the insights provided by our work might be useful for applied CAPOs/CATE estimation (e.g., discovered connections between balancing representations and meta-learners). That is the researchers should not be burdened with implementing and evaluating dozens of seemingly similar representation learning methods (see Table 5) but rather put their effort into choosing a meta-learner (R-/DR-learner).

Action: We added more explanations about the novelty of our work and why the combination of two main streams of work (representation learning methods and Neyman-orthogonal learners) is non-trivial and requires extensive discussion.

Responses to “Questions”

1. Thank you for the question. We do not require conditional independence (i.e., uncofoundedness of representations) $Y(0), Y(1) \perp A \mid \Phi(X)$ to hold. In contrast, we allow for any, possibly confounded, representations to be input to the second stage model, $g$, analogously to the runtime confounding setting [1]. Yet, the downstream estimation (=second stage of learning) is not necessarily biased: We can still identify both representation-level CAPOs and CATE by employing the original covariate-conditional propensity score ($​​\pi^x_a$) and expected outcomes ($\mu_a^x$). We formalize this in a new Remark 1 on the identifiability of representation-level causal quantities (see Sec. 3.2 and new Appendix C).

   Action: We added a new Remark 1 on the identifiability of representation-level causal quantities to the revised version of our paper (see Sec. 3.2 and new Appendix C of the updated PDF). It shows that CAPOs/CATE based on any subset of the covariates (e.g., based on confounded representations) are identifiable.

2. Thanks for noting this. In Eq. (10), $g$ is implicitly present, as it coincides with one of the nuisance functions ($\hat{\mu}_a^x$).

   Action: We have rewritten Eq. 10 to make the dependency on $g$ explicit.

Dear reviewer v7uW,

We truly appreciate your time and invaluable insights!

To sum it up, we hope that we have addressed all of your concerns in our rebuttal:

• "...the authors do not provide any theoretical analysis. In addition, there is no foundation for some claims..." $\to$ We formalized many claims made in our paper by adding several remarks and propositions (see a new Appendix C).
• "The author may want to present some evidence for some claims in the manuscript." $\to$ We provided the proofs of all the key claims in Sec. 3 and 4 (see a new Appendix C).
• "In addition, combining existing ideas may not be interesting enough. The writing can be much improved." $\to$ In our revised discussion (see "Choice of a target model" & "Implications" on page 10), we elaborated on the novelty of unifying representation learning and Neyman-orthogonal learners, emphasizing its non-triviality and interesting implications.
• Question 1. $\to$ We demonstrated the identifiability of representation-level causal quantities in a new Remark 1.
• Question 2. $\to$ We fixed Eq. 10 to make the dependency on $g$ explicit.
• Question 3. $\to$ We extended Sec. 4.1 (iv) and the discussion (see "Choice of a target model" on page 10) where we better explained the difference between the generic $V$ and the specific $V = \Phi(X)$.
• Question 4. $\to$ We added our new Remark 2 that formalizes how the accuracy in estimating nuisance functions affects the final estimation.

As the discussion period is nearing its end, we wanted to check if you have any remaining concerns or questions we can address. If our revisions meet your expectations, we hope you might consider increasing your score.

Best regards,

The authors