Learning Representations of Instruments for Partial Identification of Treatment Effects

## Abstract

Abstract Reliable estimation of treatment effects from observational data is important in many disciplines such as medicine. However, estimation is challenging when unconfoundedness as a standard assumption in the causal inference literature is violated. In this work, we leverage arbitrary (potentially high-dimensional) instruments to estimate bounds on the conditional average treatment effect (CATE). Our contributions are three-fold: (1) We propose a novel approach for partial identification through a mapping of instruments to a discrete representation space so that we yield valid bounds on the CATE. This is crucial for reliable decision-making in real-world applications. (2) We derive a two-step procedure that learns tight bounds using a tailored neural partitioning of the latent instrument space. As a result, we avoid instability issues due to numerical approximations or adversarial training. Furthermore, our procedure aims to reduce the estimation variance in finite-sample settings to yield more reliable estimates. (3) We show theoretically that our procedure obtains valid bounds while reducing estimation variance. We further perform extensive experiments to demonstrate the effectiveness across various settings. Overall, our procedure offers a novel path for practitioners to make use of potentially high-dimensional instruments (e.g., as in Mendelian randomization).

## 1 Introduction

Estimating the conditional average treatment effect (CATE) from observational is an important task for personalized decision-making in medicine (Feuerriegel et al., ). For example, a common question in medicine is to estimate the effect of alcohol consumption on the onset of cardiovascular diseases (Holmes et al., ). There are many reasons, including costs and ethical concerns, why CATE estimation is often based on observational data (such as, e.g., electronic health records, clinical registries).

Figure: Figure 1: Overview of the IV setting. We consider complex instruments $Z$ (e.g., gene data, text, images), observed confounders $X$, unobserved confounders $U$, a binary treatment $A$, and an outcome $Y$.
Refer to caption: x1.png

However, identifying the CATE from observational data is challenging as it typically requires *strong* assumptions in the form of *unconfoundedness* (Rubin, ). Unconfoundedness assumes there exist no additional unobserved confounders $U$ between treatment $A$ and outcome $Y$. If the unconfoundedness assumption is violated, a common strategy is to leverage instrumental variables (IVs) $Z$. IVs affect only the treatment $A$ but exclude unobserved confounding between $Z$ and $Y$, which often can be ensured by design such as for randomized studies with non-compliance (Imbens & Angrist, ). The causal graph for the IV setting is shown in Fig. 1.

Motivational example: *Mendelian randomization.* Mendelian randomization (Pierce et al., ) refers to the use of genetic information as instruments $Z$ to estimate the effect of a treatment or exposure $A$ (e.g., alcohol consumption) on some medical outcome $Y$ (e.g., cardiovascular diseases). In this setting, there are further patient characteristics that are observed ($X$) but also unobserved ($U$), which one accounts for through the instrument. Yet, common challenges are that (i) instruments with genetic information are often *high-dimensional* and (ii) involve *complex, non-linear relationships between instruments and treatment intake or exposure*.

However, existing IV methods using machine learning for point estimation of the CATE rely on *strong simplifying assumptions* ($\rightarrow$ violating (ii) from above). For example, some methods assume linearity in some feature space in the CATE and make other, strict parametric assumptions on the unobserved confounders such as additivity or homogeneity (Hartford et al., ; Singh et al., ; Xu et al., ). Yet, such simplifying assumptions are often *not* realistic and can even lead to unreliable and false conclusions by the mis-specification of the CATE.

A potential remedy is to use IVs for partial identification of the CATE where one circumvents any hard parametric assumptions by estimating upper and lower bounds of the CATE (Manski, ). This is usually sufficient in medical practice when one is merely interested in whether a treatment variable (e.g., exposure as in Mendelian randomization) has a positive or a negative effect. So far, methods for partial identification of the CATE in IV settings are rare. There exist closed-form bounds (i.e., via a fixed target estimand that can be learned), yet only for the setting with both *discrete* instruments and *discrete* treatments (Balke & Pearl, ).(^6^66Originally, the paper (Balke & Pearl, ) derives the bounds for the average treatment effect (ATE) but which can be straightforwardly extended to the CATE.)

Existing machine learning methods for partial identification are typically designed for *simple* instruments that are binary or discrete ($\rightarrow$ violating (i) from above).(^7^77In Mendelian randomization, this then requires further simplifications. For example, many applications of Mendelian randomization in medicine treat the genetic variants as *discrete* instruments (sometimes also by grouping) or use a predefined risk or polygenic score as a *one-dimensional* continuous instrument (Pierce et al., ). Hence, rich information from complex gene data is lost.) Alternatively, methods that extend partial identification for continuous instruments require *unstable* training paradigms such as adversarial learning (Kilbertus et al., ; Padh et al., ) which becomes even more unstable for more complex instruments. In contrast, there is a scarcity of methods that can deal robustly with continuous, as well as *complex* and potentially high-dimensional instruments such as, e.g., gene expressions as in Mendelian randomization but also text, images, or graphs.(^8^88In Appendix B, we provide an extended discussion about the real-world relevance and applicability of our method.)

Our paper: In this work, we leverage complex instruments for partial identification of the CATE. Specifically, we allow for instruments that can be continuous and potentially high-dimensional (such as gene information) and, on top of that, we explicitly allow for complex, non-linear relationships between instruments and treatment intake or exposure. In the rest of this paper, we refer to this setting as “complex” instruments.

To this end, we proceed as follows. (1) We propose a novel approach for partial identification through a mapping of complex instruments to a discrete representation space so that we yield valid bounds on the CATE. We motivate our approach in Fig. 2. (2) We derive a two-step procedure that learns tight bounds using a neural partitioning of the latent instrument space. As a result, we avoid instability issues due to numerical approximations or adversarial training, which is a key limitation of prior works. We further improve the performance of our procedure by explicitly reducing the estimation variance in finite-sample settings to yield more reliable estimates. (3) We provide a theoretical analysis of our procedure and perform extensive experiments to demonstrate the effectiveness across various settings.

Contributions:(^9^99Code is available at [https://github.com/JSchweisthal/ComplexPartialIdentif](https://github.com/JSchweisthal/ComplexPartialIdentif).) (1) To the best of our knowledge, this is the first IV method for partial identification of the CATE based on complex instruments. (2) We derive a two-step procedure to learn tight bounds. (3) We demonstrate the effectiveness of our method both theoretically and numerically.

## 2 Related Work

Machine learning for CATE estimation with IV: Existing works have different objectives. One literature stream leverages IVs for CATE estimation but focuses on settings where the treatment effect can be identified from the data. This includes work that extends the classical two-stage least-squares estimation to non-linear settings by learning non-linear feature spaces (Singh et al., ; Xu et al., ), deep conditional density estimation in the first stage (Hartford et al., ), or using moment conditions (Bennett et al., ). Another literature stream aims at new machine learning methods with favorable properties such as being doubly robust (Kennedy et al., ; Ogburn et al., ; Semenova & Chernozhukov, ; Syrgkanis et al., ) or multiply robust (Frauen & Feuerriegel, ).

Figure: Figure 2: Leveraging complex instruments for partial identification of the CATE through discrete representations of $Z$. Naive discretization on the IV input space leads to wide, and thus non-informative, bounds. Our method learns a latent representation $\phi(Z)$ to yield tight bounds.
Refer to caption: x2.png

Recently, researchers started applying machine learning methods to IVs from Mendelian randomization (Legault et al., ; Malina et al., ), which is our motivational example from above. However, these works aim at point-identified CATE estimation with IVs. As a result, these rely on *hard* and generally untestable assumptions on some effects in the causal graph, such as linearity, monotonicity, additivity, or homogeneity (Wang & Tchetgen Tchetgen, ). This is unlike our method for partial identification that does *not* require such hard assumptions and that is non-parametric.

Partial identification: Partial identification aims to identify and learn upper and lower bounds of some causal quantity (e.g., the CATE) when the causal quantity itself cannot be point identified from the data and assumptions. In a general setting with binary treatments, Robins () and Manski () derived closed-form bounds on the ATE for bounded outcomes $Y$. Further work extended these ideas to settings with binary instrumental variables, binary treatments, and binary outcomes (Balke & Pearl, ; ) to derive tighter bounds. Newer approaches for discrete variables include the works of Duarte et al. () and Guo et al. (). Swanson et al. () provide an extensive overview of partial identification in this setting. Other works focus on how to leverage additional observed confounders to further tighten bounds on the ATE (see, e.g., Levis et al., ). However, these works do not focus on efficiently leveraging continuous or even high-dimensional instruments for learning tight bounds, unlike our work that is tailored to such complex instruments.

Another literature stream focuses on partial identification under general causal graphs (Balazadeh et al., ), including IV settings with continuous variables such as continuous treatments (Gunsilius, ; Hu et al., ; Kilbertus et al., ; Padh et al., ). However, these methods either make strong assumptions about the treatment response functions or require unstable optimization via adversarial training and/or generative modeling such as through using GANs. This can easily result in *unreliable* estimates of bounds for finite data, especially with high-dimensional instruments. Further, these methods are *not* directly tailored for binary treatments, unlike our method.

Research gap: To the best of our knowledge, reliable machine learning methods for partial identification of the CATE with complex instruments are missing. To draw conclusions about CATEs (as in, e.g., Mendelian randomization), our method is the first to: (i) make use of the complex instrument information (e.g., continuous or high-dimensional), (ii) avoid making strong parametric assumptions by focusing on partial identification, and (iii) avoid unstable training procedures such as adversarial learning.

## 3 Problem Setup

Setting: We focus on the standard IV setting (Angrist et al., ; Wooldridge, ). Hence, we consider instruments (e.g., gene data, text, images) given by $Z\in\mathcal{Z}\subseteq\mathbb{R}^{d}$ but, unlike previous research, allow the instruments to be complex. As such, we allow the instruments to be continuous and potentially high-dimensional. We further have access to an observational dataset $\mathcal{D}=\left\{z_{i},x_{i},a_{i},y_{i}\right\}_{i=1}^{n}$ of size $n$. The data is sampled i.i.d. from a population $(Z,X,A,Y)\sim\mathbb{P}$, with observed confounders $X\in\mathcal{X}\subseteq\mathbb{R}^{p}$, binary treatments $A\in\mathcal{A}\subseteq\{0,1\}$, and bounded outcomes $Y\in\mathcal{Y}\subseteq[s_{1},s_{2}]\subseteq\mathbb{R}$. Additionally, we allow for unobserved confounders $U$ of arbitrary form between $A$ and $Y$.

We further assume a causal structure as shown in Fig. 1. In particular, we assume that $Z$ is an instrumental variable that has an effect on the treatment $A$ but no direct effect on the outcome $Y$ except through $A$. Further, we assume that $Z$ is independent of $X$, e.g., by randomization. In Appendix B, we provide an extended discussion to show the real-world relevance and validity of our assumptions in different settings.

Notation: Throughout our work, we denote the *response function* by $\mu^{a}(x,z):=\mathbb{E}[Y|X=x,A=a,Z=z]$ and the *propensity score* by $\pi(x,z):=\mathbb{P}(A=1|X=x,Z=z)$.

CATE: We use the potential outcomes framework (Rubin, ) to formalize our causal inference problem. Let $Y(a)\in\mathcal{Y}$ denote the potential outcome under treatment $A=a$. We are thus interested in the CATE $\tau(x)=\mathbb{E}[Y(1)-Y(0)|X=x]$.

Identifiability: We make the following standard assumptions from the literature in partial identification with IVs (Angrist et al., ).

###### Assumption 1 (Consistency) .

$Y(A)=Y$.

###### Assumption 2 (Exclusion) .

$Z\perp\!\!\!\perp Y(A)\mid(X,A,U)$, where $U$ is an unobserved confounder between $A$ and $Y$. This implies that $Z$ affects $Y$ only through $A$ and there is no additional unobserved confounding between $Z$ and $Y$.

###### Assumption 3 (Independence) .

$Z\perp\!\!\!\perp(U,X)$.

Note that, however, Assumptions 1–3 from the standard IV setting are *not* sufficient to ensure identifiability of the CATE (Gunsilius, ). To ensure identifiability, one would require additional assumptions, such as linearity or, more generally, additive noise assumptions (Hartford et al., ; Wang & Tchetgen Tchetgen, ). Yet, such assumptions are highly restrictive and are neither testable nor typically ensured in real-world scenarios. Hence, this motivates our objective to perform partial identification instead.

Objective: We frame our objective as a *partial identification* problem and thus focus on estimating *valid* bounds $(b^{-}(x),b^{+}(x))$ for the CATE $\tau(x)$ such that $b^{-}(x)\leq\tau(x)\leq b^{+}(x)$ holds for all possible $x\in\mathcal{X}$. Furthermore, the bounds should be *informative*, i.e., we would like to minimize the expected bound width $\mathbb{E}_{X}[b^{+}(X)-b^{-}(X)]$, while still ensuring validity. Formally, we aim to solve

$$ $b_{\ast}^{-},b_{\ast}^{+}\in\;\underset{b^{-},b^{+}}{\arg\min}\,\mathbb{E}_{X} [b^{+}(X)-b^{-}(X)]\quad\text{s.t.}\quad b^{-}(x)\leq\tau(x)\leq b^{+}(x)\quad \text{for all}\quad x\in\mathcal{X}.$ (1) $$

###### Assumption 1 (Consistency) .

###### Assumption 2 (Exclusion) .

###### Assumption 3 (Independence) .

## 4 Partial identification of CATE with complex instruments

### 4.1 Overview

We now present our proposed method to solve the partial identification problem from Eq. (1). Solving Eq. (1) directly is *infeasible* because it involves the unknown CATE $\tau(x)$. Hence, we propose the following approach:

| Outline: 1 We learn a discretized representation (also called partitioning) $\phi(Z)$ of the instrumental variable $Z$. 2 We then derive closed-form bounds given the discrete representation $\phi$.<br>3 We transform the closed-form bounds back to our original bounding problem and, in particular, express all quantities involved as quantities that can be estimated from observational data. |
| --- |

Below, we first explain why existing closed-form bounds are *not* directly applicable and why deriving such bounds is non-trivial. We then proceed by providing the corresponding theory for the above method. Specifically, we first take a population view to show theoretically that our bounds are valid (Sec. 4.2). Then, we take a finite-sample view and present an estimator (Sec. 4.3).

*Limitations of existing bounds:* There exist different approaches for bounding treatment effects (see Sec. 2) using continuous instruments, yet these either require additional assumptions or can easily become unstable, especially for high-dimensional $Z$. Furthermore, these bounds consider continuous treatments but are *not tailored* for binary treatments (e.g., whether a drug is administered). Hence, we derive custom bounds for our setting.

*Why is the derivation non-trivial?* For binary treatments, it turns out that there exist closed-form solutions for bounds whenever the instrument $Z$ is discrete. That is, the existing bounds for the average treatment effect (ATE) with continuous bounded outcome proposed in (Manski, ) can be extended to non-parametric closed-form bounds for the CATE (Schweisthal et al., ). While these bounds are useful in a setting with discrete instruments $Z$, they are not directly applicable to continuous or even high-dimensional $Z$ due to two main reasons: (1) The bounds need to be evaluated for *all* combinations $l,m\in\mathcal{Z}^{2}\subseteq\mathbb{R}^{d}\times\mathbb{R}^{d}$, which is *intractable*. (2) Evaluating the bounds only on a random subset of combinations $l,m$ can result in *arbitrary high* estimation variance for regions with a low joint density of $p(X=x,Z=l)$ or $p(X=x,Z=m)$.
Hence, we must derive a novel method for estimating bounds based on complex instruments (that are, e.g., continuous or high-dimensional), yet this is a highly *non-trivial* task.

### 4.2 Population view

In the following theorem, we provide a novel theoretical result of how to obtain valid bounds based on discrete representations $\Phi(Z)$ of the instrument $Z$.

###### Theorem 1 (Bounds for arbitrary instrument discretizations) .

Let
$\phi:\mathcal{Z}\xrightarrow{}\{0,1,\ldots,k\}$ be an arbitrary mapping from the high-dimensional instrument $Z$ to a discrete representation. We define

$$ $\mu_{\phi}^{a}(x,\ell)=\int_{Z}\frac{\mu^{a}(x,z)\mathbb{P}(\phi(Z)=\ell|Z=z)} {\mathbb{P}(A=a,\phi(Z)=\ell)}\mathbb{P}(A=a|Z=z)\mathbb{P}(Z=z)\mathop{}\! \mathrm{d}z\quad\text{and}$ (2) $$

$$ $\pi_{\phi}(x,\ell)=\int_{Z}\frac{\pi(x,z)\mathbb{P}(\phi(Z)=\ell|Z=z)}{\mathbb {P}(\phi(Z)=\ell)}\mathbb{P}(Z=z)\mathop{}\!\mathrm{d}z.$ (3) $$

Then, under Assumptions 1, 2, and 3, the CATE $\tau(x)$ is bounded by

$$ $b_{\phi}^{-}(x)\leq\tau(x)\leq b_{\phi}^{+}(x),$ (4) $$

with

$$ $b_{\phi}^{+}(x)=\min_{l,m}b^{+}_{\phi;l,m}(x)\quad\text{and}\quad b_{\phi}^{-} (x)=\max_{l,m}b^{-}_{\phi;l,m}(x)$ (5) $$

where

$$ $b^{+}_{\phi;l,m}(x)=\pi_{\phi}(x,l)\mu_{\phi}^{1}(x,l)+(1-\pi_{\phi}(x,l))s_{2 }-(1-{\pi}_{\phi}(x,m)){\mu}_{\phi}^{0}(x,m)-{\pi}_{\phi}(x,m)s_{1},$ (6) $$

$$ $b^{-}_{\phi;l,m}(x)={\pi}_{\phi}(x,l){\mu}_{\phi}^{1}(x,l)+(1-{\pi}_{\phi}(x,l ))s_{1}-(1-\pi_{\phi}(x,m)){\mu}_{\phi}^{0}(x,m)-{\pi}_{\phi}(x,m)s_{2}.$ (7) $$

###### Proof.

See Appendix A.
∎

Theorem 1 states that, in population, we yield valid closed-form bounds for $\tau(x)$ for arbitrary representations $\phi$. In particular, we can relax the optimization problem from Eq. (1) and obtain valid bounds $b_{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi^{\ast
}}^{+}(X)\geq b_{\ast}^{+}(X)$ and $b_{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi^{\ast
}}^{-}(X)\leq b_{\ast}^{-}(X)$ by solving

$$ $\phi^{\ast}\in\underset{\phi\in\Phi}{\arg\min}\;\mathbb{E}_{X}[b_{\color[rgb]{ 0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi}^{+}(X)-b_{\color[ rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi}^{-}(X)].$ (8) $$

Here, we highlight the dependence of variables on the representation ${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi}$ in blue to show the differences to Eq. (1). Note the following differences: In contrast to Eq. (1), we do not impose any validity constraints in Eq. (8) because Theorem 1 automatically ensures the validity of our bounds. Furthermore, in contrast to Eq. (1), the objective from Eq. (8) only depends on identifiable quantities that can be estimated from observational data.

Implications of Theorem 1: Our derivation of closed-forms bounds for arbitrary discrete representations of complex $Z$ comes with an important additional benefit: The bounds only depend on (i) discrete probabilities, (ii) quantities which are independent of $\phi$ and thus do not change for different $\phi$, and (iii) the discrete representation mapping to be learned itself. As a result, this allows us to *directly* learn $\phi$ wrt. Eq. (8). As such, we circumvent the need for adversarial or alternating training, which results in more robust estimation.

###### Theorem 1 (Bounds for arbitrary instrument discretizations) .

###### Proof.

### 4.3 Finite-sample view

In practice, we have to estimate the bounds from Theorem 1 from finite observational data. For this purpose, we start with arbitrary initial estimators: $\hat{\pi}(x,z)$ is the estimator of the propensity score ${\pi}(x,z)$, $\hat{\mu}^{a}(x,z)$ of the response function ${\mu}^{a}(x,z)$, and $\hat{\eta}(z)$ of ${\eta}(z)=\mathbb{P}(A=1\mid Z=z)$.

Once the initial estimators are obtained, we can estimate our second-stage nuisance functions defined in Eq. (22) and (23) via

$$ $\hat{\mu}_{\phi}^{a}(x,\ell)=\frac{1}{\sum_{j=1}^{n}\mathds{1}\{{\phi}(z_{j})= \ell,a_{j}=a\}}\sum_{j=1}^{n}\hat{\mu}^{a}(x,z_{j})\mathds{1}\{{\phi}(z_{j})= \ell\}(a\hat{\eta}(z_{j})+(1-a)(1-\hat{\eta}(z_{j})))$ (9) $$

and via

$$ $\hat{\pi}_{\phi}(x,\ell)=\frac{1}{\sum_{j=1}^{n}\mathds{1}\{\phi(z_{j})=\ell\} }\sum_{j}^{n}\hat{\pi}(x,z_{j})\mathds{1}\{\phi(z_{j})=\ell\}.$ (10) $$

Finally, we can directly ‘plug in’ these estimators into Eq. (5) to compute estimates of the upper and lower bound $\hat{b}_{\phi}^{-}(x),\hat{b}_{\phi}^{+}(x)$.

A naïve approach would now directly use $(\hat{b}_{\phi}^{-}(x),\hat{b}_{\phi}^{+}(x))$ to solve the optimization in Eq. (8). However, for finite samples, it turns out this is infeasible without restricting the complexity of the representation function. The reason is outlined in the following theoretical results.

###### Lemma 1 (Tightness-bias-variance tradeoff) .

Let $\mathbb{E}_{n}$ and $\text{{Var}}_{n}$ denote the expectation and variance with respect to the observational data (of size $n$). Then, it holds

$$ $\mathbb{E}_{n}\Big{[}\left(b_{\ast}^{+}(x)-\hat{b}_{\phi^{\ast}}^{+}(x)\right) ^{2}\Big{]}\leq 2\Big{(}\underbrace{\left(b_{\ast}^{+}(x)-{b}_{\phi^{\ast}}^{+ }(x)\right)^{2}}_{\text{(i) Population tightness}}+\underbrace{\mathbb{E}_{n} \left[b_{\phi^{\ast}}^{+}(x)-\hat{b}_{\phi^{\ast}}^{+}(x)\right]^{2}}_{\text{( ii) Estimation bias}}+\underbrace{\text{{Var}}_{n}(\hat{b}_{\phi^{\ast}}^{+}(x ))}_{\text{(iii) Estimation variance}}\Big{)}.$ (11) $$

###### Proof.

See Appendix A.
∎

Interpretation of Lemma 1: Lemma 1 shows that the mean squared error (MSE) between the estimated representation-based bound $\hat{b}_{\phi^{\ast}}^{+}(x)$ and the ground-truth optimal bound $b_{\ast}^{+}(x)$ can be decomposed into the following three components: (i) *population tightness*, (ii) *estimation bias*, and (iii) *estimation variance*. $\bullet$ Term (i) describes the discrepancy between the tightest achievable representation-based bound ${b}_{\phi^{\ast}}^{+}$ and the ground-truth bound $b_{\ast}^{+}(x)$. It will *decrease* if we allow for more complex representations $\Phi$, for example by increasing the number of partitions $k$. $\bullet$ Term (ii) describes the estimation bias due to using finite-sample estimators for estimating the bounds. It will generally depend on the type of estimators we employ for $\hat{\pi}(x,z)$, $\hat{\mu}^{a}(x,z)$, and $\hat{\eta}(z)$. $\bullet$ Finally, term (iii) characterizes the variance due to using finite-sample estimators. In contrast to term (i), it will *increase* when we allow the representation to be more complex.

To make point (iii) more explicit, we derive the asymptotic distributions of the estimators from Eq. (9) and Eq. (10) that are used to estimate the final bounds.

###### Theorem 2 (Asymptotic distributions of estimators) .

It holds that

$$ $\displaystyle\sqrt{n}\hat{\mu}_{\phi}^{a}(x,\ell)$ $\displaystyle\xrightarrow{d}\mathcal{N}\left(\mu_{\phi}^{a}(x,\ell),\frac{1}{p _{\ell,\Phi}}\left(\frac{\operatorname{Var}(g(Z)\mid\Phi(Z)=\ell)}{c}+d\right)\right)$ (12) $\displaystyle\sqrt{n}\hat{\pi}_{\phi}(x,\ell)$ $\displaystyle\xrightarrow{d}\mathcal{N}\left({\pi}_{\phi}(x,\ell),\frac{1}{p_{ \ell,\Phi}}\operatorname{Var}(h(Z)\mid\Phi(Z)=\ell)\right)$ (13) $$

for $c,d>0$ and where $p_{\ell,\Phi}=\mathbb{P}(\Phi(Z)=\ell)$, $g(Z)=\hat{\mu}^{a}(x,Z)(a\hat{\eta}(Z)+(1-a)(1-\hat{\eta}(Z))$, and $h(Z)=\hat{\pi}(x,Z)$.

###### Proof.

See Appendix A.
∎

We observe that the variance of the estimators (and, thus, of the estimated bounds) explodes for small values of $p_{\ell,\Phi}=\mathbb{P}(\Phi(Z)=\ell)$. Hence, to reduce the estimation variance, we aim to learn a representation $\phi$ that avoids low $p_{\ell,\Phi}$ for some $\ell$, e.g., by limiting the number of partitions $k$.
$\mathbf{\Rightarrow}$ Altogether, as a consequence of Lemma 1 and Theorem 2, we obtain an *inherent trade-off between tightness of the bounds in populations and estimation variance in finite-samples*.

Learning objective for the representation $\phi$: Due to the inherent trade-off between tightness of the bounds and estimation variance, the aim for learning the representation $\phi$ is two-fold. On the one hand, we (a) aim to learn tight bounds, which is given in the objective in Eq. (8). On the other hand, we (b) also have to account for controlling the variance in finite-sample settings, especially for high-dimensional $Z$. Motivated by Theorem 2, we ensure $\hat{p}_{\ell,\Phi}>\varepsilon$ for some $\varepsilon>0$, where $\hat{p}_{\ell,\Phi}$ is an estimator of $p_{\ell,\Phi}=\mathbb{P}(\Phi(Z)=\ell)$. Combining both (a) and (b) yields the following objective:

$$ $\phi^{*}\in\underset{\phi\in\Phi}{\arg\min}\;\mathbb{E}_{X}[\hat{b}_{\phi}^{+} (X)-\hat{b}_{\phi}^{-}(X)]\quad\text{s.t.}\quad\hat{p}_{\ell,\Phi}>\varepsilon,$ (14) $$

for some $\varepsilon>0$ and all $\ell\in\{1,\dots,k\}$.
We next present a neural method to learn tight bounds using the above objective.

###### Lemma 1 (Tightness-bias-variance tradeoff) .

###### Proof.

###### Theorem 2 (Asymptotic distributions of estimators) .

###### Proof.

## 5 Neural method for learning CATE bounds with complex instruments

In this section, we propose a neural method for our objective to learn tight and valid bounds. Our method consists of two separate stages (see Algorithm 1): 1 we learn initial estimators of the three nuisance functions, and 2 we learn an optimal representation $\phi^{\ast}$, so that the width of the bounds is minimized. Note that our method is completely model-agnostic. Hence, arbitrary machine learning models can be used in the first and second stages in order to account for the properties of the data. For example, for instruments with gene data, one could use pre-trained encoders to further optimize the downstream performance. We give an overview of the workflow of our method in Fig. 3 (see Algorithm 1 for pseudocode).

Figure: Figure 3: Workflow of the second stage of our method for calculating bounds on the CATE: The representation network $\phi_{\theta}$ learns discrete latent representations of the complex $Z$ (e.g., continuous or high-dimensional). By employing the pre-trained $\hat{\mu}$, $\hat{\pi}$, and $\hat{\eta}$, we can directly calculate the nuisance estimates conditional on the latent representation $\phi(z)$ by using Eq. (9) and Eq. (10) to yield the bounds.
Refer to caption: x3.png

Initial nuisance estimation: In the first stage, we can use arbitrary machine learning models (e.g., feed-forward neural network) to learn the first-stage nuisance functions $\hat{\mu}^{a}(x,z)=\hat{\mathbb{E}}[Y\mid X=x,A=a,Z=z]$ ,
$\hat{\pi}(x,z)=\hat{\mathbb{P}}(A=1\mid X=x,Z=z)$, and
$\hat{\eta}(z)=\hat{\mathbb{P}}(A=1\mid Z=z)$.

Recall that we consider $Z$ and $X$, which are both potentially high-dimensional. Hence, for $\hat{\mu}^{a}(x,z)$ and $\hat{\pi}(x,z)$, we use network architectures that have (i) different encoding layers for $X$ and $Z$, so that we capture structured information within the variables and (ii) shared layers on top of the encoding to learn common structures. Further, for $\hat{\mu}^{a}(x,z)$, we use two outcome heads for both treatment options $A\in\{0,1\}$ to ensure that the influence of the treatment on the outcome prediction does not ‘get lost’ in the high-dimensional space of $X$ and $Z$ (Shalit et al., ).

Figure: Algorithm 1 Two-stage learner for estimating bounds with complex instruments

Representation learning: In the second stage, we train a neural network to learn discrete representations of the instruments with the objective of obtaining tight bounds but with constraints on the estimation variance. To learn the function $\phi(z)$, we use a neural network $\phi_{\theta}$ with trainable parameters $\theta$. Then, on top of the final layer of the encoder, we leverage the Gumbel-softmax trick (Jang et al., ), which allows us to learn $k$ *discrete* representations of the latent space of the instruments, where $k$ can be flexibly chosen as a hyperparameter.

Custom loss function: We further transform our objective into a loss function to train the network $\phi_{\theta}$. For that, we design a compositional loss consisting of three terms:

A *bound-width minimization loss* that aims at our objective in Eq. (14), defined via

$$ $\mathcal{L}_{\text{{b}}}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\hat{b}_{\phi_{ \theta}}^{+}(x_{i})-\hat{b}_{\phi_{\theta}}^{-}(x_{i})$ (15) $$

A *regularization loss* to enforce the constraints in Eq. (14), i.e., enforcing that $\hat{p}_{\ell,\Phi}=\hat{\mathbb{P}}(\phi_{\theta}(Z)=\ell)>\varepsilon$, $\forall\ell\in{1,\ldots,k},\text{ for some }\varepsilon>0$. For that, we aim to penalize the negative log-likelihood $-\sum_{j=1}^{k}\log(\mathbb{P}({\phi_{\theta}}(Z)=j))$, which we can estimate via

$$ $\mathcal{L}_{\text{{reg}}}(\theta){=}-\sum_{j=1}^{k}\log\Big{(}\frac{1}{n}\sum _{i=1}^{n}\mathds{1}\{{\phi_{\theta}}(z_{i})=j\}\Big{)}.$ (16) $$

An *auxiliary guidance loss* $\mathcal{L}_{\text{{aux}}}(\theta)$, which enforces more heterogeneity between $\mathbb{P}(Z\mid\phi_{\theta}(Z)=l)$ and $\mathbb{P}(Z\mid\phi_{\theta}(Z)=m)$, for all $l,m$. We implement this by adding an additional linear classification head on top of the last hidden layer of $\phi_{\theta}$ before the discretization and then calculating the cross entropy loss between the predictions of the head and the discrete representation $\phi_{\theta}(z)$. In principle, the term $\mathcal{L}_{\text{{aux}}}(\theta)$ is not needed for our objective, but we empirically observed that it helps to stabilize the training further by avoiding to ‘get stuck’ in non-informative local minima. Hence, we yield our final training loss

$$ $\mathcal{L}(\theta)=\mathcal{L}_{\text{{b}}}(\theta)+\lambda\mathcal{L}_{\text {{reg}}}(\theta)+\gamma\mathcal{L}_{\text{{aux}}}(\theta),$ (17) $$

with hyperparameters $\lambda$ and $\gamma$. Here, $\lambda$ controls the trade-off between bound tightness and estimation variance, and can thus be tailored depending on the application. The hyperparameter $\gamma$ can be simply tuned as usual.

A benefit of our custom loss is that it is particularly efficient and robust compared to other learning procedures (such as alternating learning procedure or adversarial training). In the second stage, we solely update the parameters $\theta$ of the discretization network $\phi_{\theta}$ to minimize $\mathcal{L_{\theta}}$. In contrast, the networks of the first-stage nuisance estimators have frozen weights. In the second stage, networks of the first-stage nuisance estimators are only evaluated but are *not* updated.

## 6 Experiments

Baselines: Existing methods (see Sec. 2) focus either on (a) point identification with strong assumptions, (b) partial identification with continuous treatment variables, or (c) discrete instruments. We instead focus on a setting with complex instruments and binary treatments. Hence, existing methods are not tailored to our setting, because of which a fair comparison is precluded. Instead, we thus demonstrate the validity and tightness of our bounds. Further, for comparison, we propose an additional Naïve baseline, which first learns a discretization of the instruments (via $k$-means clustering) and then learns the nuisance functions wrt. to the discretized instruments to apply the existing bounds for discrete instruments from Lemma 2 on top.

Figure: Figure 4: Datasets 1 and 2: Estimated bounds on the CATE. Shown: mean $\pm$ sd over 5 runs for different number of discretizations $k$. *Left*: Dataset 1 with a simple $\pi(x,z)$. *Right*: Dataset 2 with a complex $\pi(x,z)$.
Refer to caption: x4.png

Data: We perform experiments mimicking Mendelian Randomization but where we simulate the data to have access to the ground-truth CATE for performance evaluations, so that we can check for coverage and validity of the bounds. We consider three different realistic settings. For Datasets 1 and 2, we consider a one-dimensional continuous instrument representing a polygenic risk score (Pierce et al., ). Further, in Dataset 1, we model the true $\pi(x,z)$ as a rather simple function to check if our method is already competitive in such settings. In Dataset 2, we model $\pi(x,z)$ as a complex function to evaluate the performance in more challenging settings. We use the same CATE for Dataset 1 and Dataset 2 to allow for comparisons between both. In Dataset 3, we model high-dimensional instruments with single nucleotide polymorphisms (SNPs, i.e., genetic variants; Burgess et al., ) to test our method in an additional realistic and even more complex setting. In all datasets, we model the CATE to be heterogeneously conditioned on $X$ to check whether the bounds adapt to different subpopulations. Details are in Appendix D.

Performance metrics: We report the following metrics to assess the validity and robustness of the estimated bounds: (i) The *coverage*, i.e., how often the true CATE lies within the estimated bounds. (ii) The average *width* of bounds, where lower values indicate more informative bounds. (iii) The *mean squared difference* MSD($k$) of the predicted bounds over different values of $k$, indicating the robustness wrt. to the selection of the hyperparameter. Further, for Dataset 3, we model $\pi(x,z)$ to be dependent on some latent discrete representation of the observed $Z$, such that we can approximate oracle bounds. Thus, we can evaluate the (iv) MSE and (v) the coverage wrt. to the oracle bounds.

**Table 1: Datasets 1 and 2: Comparison of the Naïve baseline and our method. Reported are the coverage, width, and mean squared distance MSD($k$) between the estimated bounds over 5 runs. Best results in bold.**
| Dataset | Method | $k$ | Coverage[$\uparrow$] | Width[$\downarrow$] | MSD($k$)[$\downarrow$] |
| --- | --- | --- | --- | --- | --- |
| Dataset 1 | Naïve | 2 | $\textbf{1.00}\pm 0.00$ | $1.62\pm 0.06$ | $0.28\pm 0.06$ |
| Naïve | 3 | $\textbf{1.00}\pm 0.00$ | $\textbf{0.83}\pm 0.16$ |  |  |
| Ours | 2 | $\textbf{1.00}\pm 0.00$ | $\textbf{1.01}\pm 0.05$ | $\textbf{0.03}\pm 0.03$ |  |
| Ours | 3 | $\textbf{1.00}\pm 0.00$ | $1.09\pm 0.04$ |  |  |
| Dataset 2 | Naïve | 2 | $\textbf{1.00}\pm 0.00$ | $1.34\pm 0.19$ | $0.09\pm 0.06$ |
| Naïve | 3 | $\textbf{1.00}\pm 0.00$ | $1.28\pm 0.20$ |  |  |
| Ours | 2 | $\textbf{1.00}\pm 0.00$ | $\textbf{1.13}\pm 0.19$ | $\textbf{0.06}\pm 0.06$ |  |
| Ours | 3 | $\textbf{1.00}\pm 0.00$ | $\textbf{1.15}\pm 0.31$ |  |  |

Implementation details: For our method, we use multi-layer-perceptrons (MLPs) for the first-stage nuisance estimation and an MLP with Gumbel-softmax (Jang et al., ) discretization on the last layer for learning $\phi_{\theta}$. For the Naïve baseline, we use $k$-means clustering in the first step to learn discretized instruments and then use MLPs with identical architecture for the nuisance estimation to ensure a fair comparison. We provide further details in Appendix C.

**Table 2: Dataset 3 (high-dimensional): Comparison of the Naïve baseline and our method. Reported are the coverage, width, the coverage and MSE wrt. to the oracle bounds, and the mean squared distance MSD($k$) between the estimated bounds over $5$ runs.**
| Method | $k$ | Coverage[$\uparrow$] | Width[$\downarrow$] | Coverage (oracle bounds)[$\uparrow$] | MSE (oracle bounds)[$\downarrow$] | MSD($k$)[$\downarrow$] |
| --- | --- | --- | --- | --- | --- | --- |
| Naive | 2 | $1.00\pm 0.00$ | $1.96\pm 0.05$ | $1.00\pm 0.00$ | $0.15\pm 0.02$ | $0.10\pm 0.10$ |
|  | 4 | $1.00\pm 0.00$ | $1.91\pm 0.03$ | $1.00\pm 0.00$ | $0.13\pm 0.02$ |  |
|  | 6 | $1.00\pm 0.00$ | $1.74\pm 0.26$ | $0.75\pm 0.50$ | $0.09\pm 0.05$ |  |
|  | 8 | $1.00\pm 0.00$ | $1.89\pm 0.09$ | $1.00\pm 0.00$ | $0.12\pm 0.04$ |  |
| Ours | 2 | $1.00\pm 0.00$ | $1.87\pm 0.05$ | $1.00\pm 0.00$ | $0.12\pm 0.02$ | $\textbf{0.03}\pm 0.02$ |
|  | 4 | $1.00\pm 0.00$ | $1.87\pm 0.08$ | $1.00\pm 0.00$ | $0.12\pm 0.03$ |  |
|  | 6 | $1.00\pm 0.00$ | $1.85\pm 0.06$ | $1.00\pm 0.00$ | $0.11\pm 0.02$ |  |
|  | 8 | $1.00\pm 0.00$ | $1.83\pm 0.07$ | $0.99\pm 0.01$ | $0.11\pm 0.03$ |  |

Figure: Figure 5: Dataset 3 (high-dimensional): Average bound width. Sensitivity analysis wrt. to the number of instruments $k$ where we show the average bound width and estimation variance over 5 runs.
Refer to caption: x6.png

Results: We present the results of our experiments for Datasets 1 and 2 in Table 1 and Fig. 4, and for Dataset 3 in Table 2 and Fig. 5. Here, we compare our method against the Naïve baseline and vary the number of clusters $k$, with higher $k$ for the more complex Dataset 3.

Overall, we observe the following patterns: (i) Both methods (i.e., ours and the Naïve baseline) almost always reach a perfect coverage of 100%, which shows the validity of the bounds. (ii) As expected, our method is able to learn *tighter bounds* on the more complex Datasets 2 and 3 over the different runs and for different $k$. This demonstrates that our method aimed at learning representations with the objective of tightening bounds can clearly improve over a discretization that uses solely information of $Z$ in the first step (Naïve). (iii) Our method is robust, especially also over different values of $k$. This is demonstrated by a low MSD($k$), in particular for Datasets 1 and 3. Yet, this is unlike the Naïve baseline, which performs *not* robust wrt. $k$. The latter is also clearly visible by the large differences in the learned bounds in Fig. 4 on the left. Further, Fig. 5 indicates unstable behavior of the Naïve baseline with higher variation, even resulting in learning falsely overconfident bounds for $k=6$. In contrast, our method yields bounds that are robust for given $k$ as well as over varying values of $k$, which is naturally encouraged by our objective of flexibly learning representations while ensuring robustness in bound estimation.

Takeaways: Our method can successfully learn bounds that have a high coverage and a low width. Further, our method outperforms the Naïve baseline clearly while ensuring robustness. Here, our results show that the source of the performance gain is the way how we learn the representation $\phi$ and that the performance gain from our method becomes larger for more complex datasets.

Limitations: Our method for partial identification allows us to relax multiple assumptions that are inherent to methods for point identification. Nevertheless, we still rely on the standard assumptions of IV settings. However, such assumptions often hold by design or can be ensured by expert knowledge such as in Mendelian randomization. We provide an extended discussion in Appendix B.

Conclusion: We propose a novel method for learning tight bounds on treatment effects by making use of complex instruments (e.g., instruments that are continuous, potentially high-dimensional, and that have non-trivial relationships with the treatment intake or exposure).

## Appendix A Proofs

### A.1 Proof of Theorem 1

We begin by stating a result from the literature that obtains valid bounds for discrete instruments.

###### Lemma 2 ( (Swanson et al., 2018 ; Schweisthal et al., 2024 ) ) .

Under Assumptions 1 and 2, the CATE is bounded via

$$ $b^{-}(x)\leq\tau(x)\leq b^{+}(x),$ (18) $$

with

$$ $b^{+}(x)=\min_{l,m}b^{+}_{l,m}(x)\quad\text{and}\quad b^{-}(x)=\max_{l,m}b^{-} _{l,m}(x)$ (19) $$

where

$$ $b^{+}_{l,m}(x)=\pi(x,l)\mu^{1}(x,l)+(1-\pi(x,l))s_{2}-(1-\pi(x,m))\mu^{0}(x,m) -\pi(x,m)s_{1},$ (20) $$

$$ $b^{-}_{l,m}(x)=\pi(x,l)\mu^{1}(x,l)+(1-\pi(x,l))s_{1}-(1-\pi(x,m))\mu^{0}(x,m) -\pi(x,m)s_{2}.$ (21) $$

###### Proof of Theorem 1 .

First, note that, for a given representation $\phi$, the representation $\phi(Z)$ is still a valid (discrete) instrument that satisfies Assumptions 1 and 2. Hence, we can apply Lemma 2 using $\phi(Z)$ as an instrument and immediately obtain the bounds from Theorem 1, but with *representation-induced nuisance functions*
$\mu_{\phi}^{a}(x,\ell)=\mathbb{E}[Y|X=x,A=a,\phi(Z)=\ell]$ and $\pi_{\phi}(x,\ell)=\mathbb{P}(A=1|X=x,\phi(Z)=\ell)$ for $\ell\in\{0,\dots,k\}$.

We can write the representation-induced response function as

$$ $\displaystyle\mathbb{E}[Y|X=x,A=a,\phi(Z)=\ell]\stackrel{{\scriptstyle Z\perp \!\!\!\perp X}}{{=}}\int_{Z}\mathbb{E}[Y|X=x,A=a,Z=z]\mathbb{P}(Z=z|A=a,\phi(Z )=\ell)\mathop{}\!\mathrm{d}z$ (22) $\displaystyle=\int_{Z}\mathbb{E}[Y|X=x,A=a,Z=z]\frac{\mathbb{P}(\phi(Z)=\ell|A =a,Z=z)\mathbb{P}(A=a|Z=z)\mathbb{P}(Z=z)}{\mathbb{P}(A=a|\phi(Z)=\ell)\mathbb {P}(\phi(Z)=\ell)}\mathop{}\!\mathrm{d}z$ $\displaystyle=\frac{1}{\mathbb{P}(A=a|\phi(Z)=\ell)\mathbb{P}(\phi(Z)=\ell)}$ $\displaystyle\qquad\int_{Z}\mathbb{E}[Y|X=x,A=a,Z=z]\mathbb{P}(\phi(Z)=\ell|A= a,Z=z)\mathbb{P}(A=a|Z=z)\mathbb{P}(Z=z)\mathop{}\!\mathrm{d}z$ $\displaystyle=\frac{1}{\mathbb{P}(A=a|\phi(Z)=\ell)\mathbb{P}(\phi(Z)=\ell)}$ $\displaystyle\qquad\qquad\int_{Z}\mathbb{E}[Y|X=x,A=a,Z=z]\mathbb{P}(\phi(Z)= \ell|Z=z)\mathbb{P}(A=a|Z=z)\mathbb{P}(Z=z)\mathop{}\!\mathrm{d}z$ $$

and the representation-induced propensity score as

$$ $\displaystyle\mathbb{P}(A=1|X=x,\phi(Z)=\ell)\stackrel{{\scriptstyle Z\perp\! \!\!\perp X}}{{=}}\int_{Z}\mathbb{P}(A=1|X=x,Z=z)\mathbb{P}(Z=z|\phi(Z)=\ell) \mathop{}\!\mathrm{d}z$ (23) $\displaystyle=$ $\displaystyle\int_{Z}\mathbb{P}(A=1|X=x,Z=z)\mathbb{P}(\phi(Z)=\ell|Z=z)\frac{ \mathbb{P}(Z=z)}{\mathbb{P}(\phi(Z)=\ell)}\mathop{}\!\mathrm{d}z$ $\displaystyle=$ $\displaystyle\frac{1}{\mathbb{P}(\phi(Z)=\ell)}\int_{Z}\mathbb{P}(A=1|X=x,Z=z) \mathbb{P}(\phi(Z)=\ell|Z=z)\mathbb{P}(Z=z)\mathop{}\!\mathrm{d}z,$ $$

which completes the proof.
∎

###### Lemma 2 ( (Swanson et al., 2018 ; Schweisthal et al., 2024 ) ) .

###### Proof of Theorem 1 .

### A.2 Proof of Lemma 1

###### Proof.

The result follows from

$$ $\displaystyle\mathbb{E}_{n}\left[\left(b_{\ast}^{+}(x)-\hat{b}_{\phi^{\ast}}^{ +}(x)\right)^{2}\right]$ $\displaystyle=\mathbb{E}_{n}\left[\left(b_{\ast}^{+}(x)-{b}_{\phi^{\ast}}^{+}( x)+{b}_{\phi^{\ast}}^{+}(x)-\hat{b}_{\phi^{\ast}}^{+}(x)\right)^{2}\right]$ (24) $\displaystyle\leq 2\left(\left(b_{\ast}^{+}(x)-{b}_{\phi^{\ast}}^{+}(x)\right) ^{2}+\mathbb{E}_{n}\left[\left({b}_{\phi^{\ast}}^{+}(x)-\hat{b}_{\phi^{\ast}}^ {+}(x)\right)^{2}\right]\right)$ (25) $\displaystyle\overset{(\ast)}{(=)}2\left(\left(b_{\ast}^{+}(x)-{b}_{\phi^{\ast }}^{+}(x)\right)^{2}+\mathbb{E}_{n}\left[b_{\phi^{\ast}}^{+}(x)-\hat{b}_{\phi^ {\ast}}^{+}(x)\right]^{2}+\text{{Var}}_{n}(\hat{b}_{\phi^{\ast}}^{+}(x))\right),$ (26) $$

where we used the bias-variance decomposition for the MSE for $(\ast)$.
∎

###### Proof.

### A.3 Proof of Theorem 2

###### Proof.

We derive the asymptotic distributions of the estimators $\hat{\mu}_{\phi}^{a}(x,\ell)$ from Eq. (9) and $\hat{\pi}_{\phi}(x,\ell)$ from Eq. (10). We proceed by analyzing the numerator and denominator of each estimator. First, we show that both are asymptotically normal and then we apply the delta method to obtain the asymptotic distribution of the ratios.

Distribution of $\hat{\mu}_{\phi}^{a}(x,\ell)$: Recall from Equation (9) that we can write $\hat{\mu}_{\phi}^{a}(x,\ell)$ as

$$ $\hat{\mu}_{\phi}^{a}(x,\ell)=\frac{S_{n}}{N_{n}},$ (27) $$

where

$$ $\displaystyle S_{n}$ $\displaystyle=\frac{1}{n}\sum_{j=1}^{n}W_{j},\quad\text{with}\quad W_{j}=\hat{ \mu}^{a}(x,z_{j})\mathds{1}\{\phi(z_{j})=\ell\}[a\hat{\eta}(z_{j})+(1-a)(1- \hat{\eta}(z_{j}))],$ (28) $\displaystyle N_{n}$ $\displaystyle=\frac{1}{n}\sum_{j=1}^{n}D_{j},\quad\text{with}\quad D_{j}= \mathds{1}\{\phi(z_{j})=\ell,a_{j}=a\}.$ (29) $$

We define the moments

$$ $\displaystyle\mu_{W}$ $\displaystyle=\mathbb{E}[W]=p_{\ell}\theta_{\ell}$ (30) $\displaystyle\sigma_{W}^{2}$ $\displaystyle=\operatorname{Var}(W)=p_{\ell}(\gamma_{\ell}-p_{\ell}\theta^{2}_ {\ell})$ (31) $\displaystyle\mu_{D}$ $\displaystyle=\mathbb{E}[D]=p_{\ell}q_{\ell}$ (32) $\displaystyle\sigma_{D}^{2}$ $\displaystyle=\operatorname{Var}(D)=p_{\ell}q_{\ell}(1-p_{\ell}q_{\ell})$ (33) $\displaystyle\mathrm{c}_{WD}$ $\displaystyle=\operatorname{Cov}(W,D)=p_{\ell}q_{\ell}\theta_{\ell}(1-p_{\ell}),$ (34) $$

where $p_{\ell}=\mathbb{P}(\Phi(Z)=\ell)$, $q_{\ell}=\mathbb{P}(A=a\mid\Phi(Z)=\ell)$, $\theta_{\ell}=\mathbb{E}[g(Z)\mid\Phi(Z)=\ell]$, and $\gamma_{\ell}=\mathbb{E}[g(Z)^{2}\mid\Phi(Z)=\ell]$, with $g(Z)=\hat{\mu}^{a}(x,Z)(a\hat{\eta}(Z)+(1-a)(1-\hat{\eta}(Z))$.

By the central limit theorem, we know that

$$ $\sqrt{n}\begin{pmatrix}S_{n}\\ N_{n}\end{pmatrix}\xrightarrow{d}\mathcal{N}_{2}\left(\mu=\begin{pmatrix}\mu_{ W}\\ \mu_{D}\end{pmatrix},\Sigma=\begin{pmatrix}\sigma_{W}^{2}&\mathrm{c}_{WD}\\ \mathrm{c}_{WD}&\sigma_{D}^{2}\end{pmatrix}\right).$ (35) $$

Let $f(s,n)=\frac{s}{n}$. We are interested in the asymptotic distribution of the ratio $\hat{\mu}_{\phi}^{a}(x,\ell)=f(S_{n},{N_{n}})$. The delta method states that

$$ $\sqrt{n}f(S_{n},N_{n})\xrightarrow{d}\mathcal{N}_{2}\left(f(\mu_{W},\mu_{D}), \nabla f^{\top}(\mu_{W},\mu_{D})\Sigma\nabla f(\mu_{W},\mu_{D})\right)$ (36) $$

Using that the gradient is $\nabla f^{\top}(\mu_{W},\mu_{D})=\left(\dfrac{1}{\mu_{D}},-\dfrac{\mu_{W}}{\mu
_{D}^{2}}\right)$, we can obtain the asymptotic variance via

$$ $\displaystyle\nabla f^{\top}(\mu_{W},\mu_{D})\Sigma\nabla f(\mu_{W},\mu_{D})$ $\displaystyle=\dfrac{\sigma_{W}^{2}}{\mu_{D}^{2}}-2\dfrac{\mu_{W}\mathrm{c}_{ WD}}{\mu_{D}^{3}}+\dfrac{\mu_{W}^{2}\sigma_{D}^{2}}{\mu_{D}^{4}}$ (37) $\displaystyle=\frac{1}{p_{\ell}}\left(\frac{(\gamma_{\ell}-\theta_{\ell}^{2})} {q_{\ell}^{2}}+\frac{\theta_{\ell}^{2}(1-p_{\ell}q_{\ell})}{q_{\ell}^{3}}\right)$ (38) $\displaystyle=\frac{1}{p_{\ell}}\left(\frac{\operatorname{Var}(g(Z)\mid\Phi(Z) =\ell)}{q_{\ell}^{2}}+\frac{\theta_{\ell}^{2}(1-p_{\ell}q_{\ell})}{q_{\ell}^{3 }}\right).$ (39) $$

Distribution of $\hat{\pi}_{\phi}(x,\ell)$: Recall from Equation (10) that we can write $\hat{\pi}_{\phi}(x,\ell)$ as

$$ $\hat{\pi}_{\phi}(x,\ell)=\frac{S_{n}}{N_{n}},$ (40) $$

where

$$ $\displaystyle S_{n}$ $\displaystyle=\frac{1}{n}\sum_{j=1}^{n}W_{j},\quad\text{with}\quad W_{j}=\hat{ \pi}(x,z_{j})\mathds{1}\{\phi(z_{j})=l\},$ (41) $\displaystyle N_{n}$ $\displaystyle=\frac{1}{n}\sum_{j=1}^{n}D_{j},\quad\text{with}\quad D_{j}= \mathds{1}\{\phi(z_{j})=l\}.$ (42) $$

We define the moments

$$ $\displaystyle\mu_{W}$ $\displaystyle=\mathbb{E}[W]=p_{\ell}\theta_{\ell}$ (43) $\displaystyle\sigma_{W}^{2}$ $\displaystyle=\operatorname{Var}(W)=p_{\ell}(\gamma_{\ell}-p_{\ell}\theta^{2}_ {\ell})$ (44) $\displaystyle\mu_{D}$ $\displaystyle=\mathbb{E}[D]=p_{\ell}$ (45) $\displaystyle\sigma_{D}^{2}$ $\displaystyle=\operatorname{Var}(D)=p_{\ell}(1-p_{\ell})$ (46) $\displaystyle\mathrm{c}_{WD}$ $\displaystyle=\operatorname{Cov}(W,D)=p_{\ell}\theta_{\ell}(1-p_{\ell}),$ (47) $$

where $p_{\ell}=\mathbb{P}(\Phi(Z)=\ell)$, $\theta_{\ell}=\mathbb{E}[h(Z)\mid\Phi(Z)=\ell]$, and $\gamma_{\ell}=\mathbb{E}[h(Z)^{2}\mid\Phi(Z)=\ell]$, with $h(Z)=\hat{\pi}(x,Z)$.

By the central limit theorem, we know that

$$ $\sqrt{n}\begin{pmatrix}S_{n}\\ N_{n}\end{pmatrix}\xrightarrow{d}\mathcal{N}_{2}\left(\mu=\begin{pmatrix}\mu_{ W}\\ \mu_{D}\end{pmatrix},\Sigma=\begin{pmatrix}\sigma_{W}^{2}&\mathrm{c}_{WD}\\ \mathrm{c}_{WD}&\sigma_{D}^{2}\end{pmatrix}\right).$ (48) $$

We can then calculate the asymptotic variance using the delta method as above and obtain

$$ $\displaystyle\nabla f^{\top}(\mu_{W},\mu_{D})\Sigma\nabla f(\mu_{W},\mu_{D})$ $\displaystyle=\dfrac{\sigma_{W}^{2}}{\mu_{D}^{2}}-2\dfrac{\mu_{W}\mathrm{c}_{ WD}}{\mu_{D}^{3}}+\dfrac{\mu_{W}^{2}\sigma_{D}^{2}}{\mu_{D}^{4}}$ (49) $\displaystyle=\frac{1}{p_{\ell}}(\gamma_{\ell}-\theta_{\ell}^{2})$ (50) $\displaystyle=\frac{1}{p_{\ell}}\operatorname{Var}(h(Z)\mid\Phi(Z)=\ell).$ (51) $$

∎

###### Proof.

## Appendix B Real-world relevance and validity of assumptions

In this section, we elaborate on the real-world relevance of our considered setting and show that our assumptions often hold and are even weaker than the ones of existing approaches. For that, we draw upon two real-world settings.

### B.1 Mendelian randomization

Mendelian randomization (MR; the main motivational example from our paper) is a widely used method from biostatistics to estimate the causal effect of some treatment or exposure (such as alcohol consumption) on some outcome (such as cardiovascular diseases). We refer to Pierce et al. () for an introduction to MR, which also shows that MR is widely used in medicine. For that, genetic variants (such as different single nucleotide polymorphisms, SNPs) are used as instruments where it is known that they only influence the exposure but not directly the outcome. Our method for partial identification with complex instruments is perfectly suited for this common real-world application. Depending on the use case, either a predefined genetic risk score (Burgess et al., ) as a continuous variable, or up to hundreds of SNPs are used simultaneously as IVs to strengthen the power of the analysis, resulting in high-dimensional instruments (Pierce et al., ).

Validity of assumptions: The IV assumptions used in our paper such as the exclusion and independence assumptions can be ensured by expert knowledge (e.g., given some observed confounder age ($X)$, genetic variations (Z) do not affect age) or, in some cases, they can be even directly tested for (Glymour et al., ). In contrast, as explained in Sec. 2, existing methods for MR rely on additional hard assumptions on top such as the knowledge about the parametric form of the underlying data-generating process. Especially with such high-dimensional IVs, misspecification of these models may result in significantly biased effect estimates. In contrast, our method does not rely on any parametric assumption and also no additional assumptions compared to previous methods, thus enabling more reliable causal inferences in the real-world application of MR by using *strictly weaker* assumptions than existing work.

### B.2 Indirect experiments

With indirect experiments (IEs), we show that, in principle, our method is not constrained to medical applications but is also highly useful in various other domains. IEs are widely applied in various areas such as social sciences or public health to estimate causal effects in settings with non-adherence, i.e., where people cannot be forced to take treatments but rather be encouraged by some nudge (Pearl, ). For instance, researchers might be interested in estimating the effect of some treatment such as participating in a healthcare program ($T$) on some health outcome $Y$ by randomly assigning nudges $Z$ (IVs) in the form of different text messages on social media promoting participation. Here, common nudges (IVs) are in the form of, for instance, text or even image data and thus high-dimensional, showing the necessity of a method capable of handling complex IVs such as ours.

Validity of assumptions: As a major benefit of IEs, the IV assumptions are *ensured per design* as the IVs are randomly assigned, and, thus they always hold. Hence, our method provides a promising tool for evaluating the effects of IEs.

## Appendix C Implementation and training details

Model architecture: For all our models, we use MLPs with ReLU activation function. For $\hat{\mu}_{\phi}^{a}$, we use 2 layers to encode $X$ and 3 layers to encode $Z$. Then, we concatenate the outputs and add 2 additional shared layers. Finally, we calculate the outputs by a separate treatment head for $A=0$ and $A=1$ to ensure the expressiveness of $A$ for predicting $Y$. For $\hat{\pi}$, we use the same architecture. For $\hat{\eta}$, we use 3 layers. For $\phi_{\theta}$, we also use 3 layers and apply discretization on top of the $K$ outputs (Jang et al., ). For the nuisance parameters of the $k$-means baseline, we use the same models as for $\hat{\mu}_{\phi}^{a}$ and $\hat{\pi}$ for a fair comparison. We use a neuron size of 10 for all hidden layers.

Training details: For training our nuisance functions, we use an MSE loss for the functions learning the continuous outcome $Y$ and a cross-entropy loss for functions learning the binary treatment $A$. For all models, we use the Adam optimizer with a learning rate of $0.03$. We train our models for a maximum of 100 epochs and apply early stopping. For our method, we fixed $\lambda=1$ and performed random search to tune for $[0,1]$ for $\gamma$. We use PyTorch Lightning for implementation. Each training run of the experiments could be performed on a CPU with 8 cores in under 15 minutes.

## Appendix D Data description

Dataset 1: We simulate an observed confounder $X\sim\mathrm{Uniform}[-1,1]$ and an unobserved confounder $U\sim\mathrm{Uniform}[-1,1]$.

The instrument $Z$ is defined as

$$ $Z\sim\text{Mixture}\left(\frac{1}{2}\mathrm{Uniform}[-1,1]+\frac{1}{4}\mathrm{ Beta}(2,2)+\frac{1}{4}(-\mathrm{Beta}(2,2))\right).$ (52) $$

We define $\rho$ as

$$ $\rho=\frac{1}{1+\exp\left(-\left(\left(2|Z|-\max(Z)\right)+X+0.5\cdot U\right) \right)}.$ (53) $$

Then, the propensity score is given by

$$ $\pi=(\rho-0.5)\cdot 0.9+0.5.$ (54) $$

We then sample our treatment assignments from the propensity scores as

$$ $A\sim\textrm{Bernoulli}({\pi}).$ (55) $$

The conditional average treatment effect (CATE) is defined as

$$ $\tau(X)=-\frac{(2.5X)^{4}+12\sin(6X)+0.5\cos(X)}{80}+0.5.$ (56) $$

The outcome $Y$ is then generated by

$$ $Y=\left(X+0.5U+0.1\cdot\text{Laplace}(0,1)\right)\cdot 0.25+\tau(X)\cdot A.$ (57) $$

Dataset 2: We keep the other properties but change the propensity score to be more complex, which results in harder-to-learn optimal representations of $Z$ for tightening the bounds.
The propensity score is given by

$$ $\pi=\sin(2.5Z+X+U)\cdot 0.48+0.48+\frac{0.04}{1+\exp(-3|Z|)}.$ (58) $$

Dataset 3: We simulate $X$ and $U$ as above. Then, we sample a $d$-dimensional $Z\in\{0,1\}^{d}$ with $d=20$ as

$$ $Z\sim\text{Binomial}(d,0.5).$ (59) $$

Thus, our modeling is here inspired by using multiple SNPs (appearances of genetic variations) as instruments (Burgess et al., ), where we simulate potential variations for 20 genes.

Then, we define

$$ $\rho=\sum_{j=1}^{d}[\mathds{1}\{j\leq 5\}Z_{j}]$ (60) $$

and the propensity score, inspired by the more complex setting of Dataset 2, as

$$ $\pi=0.48\;\text{sin}(10\rho+X+U)+0.48+\frac{0.04}{1+\exp(-3|5\rho|)}.$ (61) $$

Then, we define the CATE as

$$ $\tau(X)=-\frac{-(1.6X+0.5)^{4}+12\sin(4X+1.5)+\cos(X)}{80}+0.5.$ (62) $$

and the outcome dependent on $\tau$, $X$ and $U$ analogously as for Datasets 1 and 2.

By Eq. (60), we ensure that some of the modeled SNPs are irrelevant for $\pi$ and thus do not affect the treatment or exposure $A$. Thereby, we focus on realistic settings in practice, where the relevance of instruments cannot always be ensured which imposes challenges especially for existing methods for point identification, but not for our approach. Further, we ensure that the latent score $\rho$ can only take 5 discrete levels. This allows us to approximate oracle bounds using the discrete bounds on top of $\rho$ by leveraging Lemma 2 such that we can evaluate our method and the baseline in comparison to oracle bounds.

To create the simulated data used in Sec. 6, we sample $n=2000$ from the data-generating process above. We then split the data into train (40%), val (20%), and test (40%) sets such that the bounds and deviation can be calculated on the same amount of data for training and testing.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: I have no further questions.
soundness: 3
strengths: This paper, to my knowledge, provides a novel solution to a well motivated problem. It provides thorough theoretical results and adequate empirical evaluation. The writing quality and organization is excellent. The availability of high dimensional instruments is common in many settings beyond the medical domain, so I believe the potential for impact is high.
summary: This paper provides a method for the partial identification of treatment effects when potentially high-dimensional instruments belonging to a broad class of domains are available. To overcome the potential of overly conservative bounds, it proposes a discrete embedding scheme for the instruments. It provides theoretical results for the validity of the proposed bounds and for estimator efficiency. It provides an empirical evaluation of the method using a well motivated simulated data generating process.
weaknesses: A discussion on potential applications outside of medicine may help to increase the visibility of this work. Perhaps by extending Appendix B.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - Even the experiements in the high dimensional settings had only 20 dimensional instruments. I am not familar with Mendelian randomization with SNPs, so I wonder if this falls into the normal range of such applications. 
- Is the bound from Theorem 1 sharp?
soundness: 3
strengths: - I appreciate the clear comparison with existing works. 
- The paper provided clear description of a neural implementation of the proposed approach, and there is some innovation in the custom loss function.
summary: This paper proposes a partial idetification approach of the conditional average treatment effect (CATE) that allows for arbitrary instruments. In particular, the proposed appraoch first maps the intstruments to a discrete representation space, and then leverages partial identification results for discrete intruments to bound the CATE. The representation mapping is learned by minimizing the length of the bounding interval while stabilizing finite sample performance. The authors also provided a neural implementation of the proposed method and performed simulation experiments.
weaknesses: Overall, the theoretical contribution of the paper appears limited (some of them are very straight forward applications of the existing results for discrete instruments) and may benefit from further exploration.
- The paper only provided finite samle guarantees for the second stage nuisance functions, but not for the final bounds, which might be more useful (for e.g. guaging the quality of the estimation) and easier to interpret. For instance it would be nice to have some theoretical results on the coverage (or scaled coverage - coverage/length of interval) of the estimated bounds.
- The number of partitions ($k$) is an important hyperparameter, but the paper did not provide anytheoretical results on how the estimation error or the tightness of the bounds would scale with $k$.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Marks will be adjusted according to the replies. Below are some feedback/questions:
1. What is the point of building bounding estimators? As mentioned in the paper, existing literature focuses on point estimators. Indeed, we can build the confidence interval based on the point estimator. It is well-known that these are two concepts, but it is odds to build bounding estimators present in the paper.
2. Theorem 2 gives the asymptotic properties of the nuisance estimators $\hat{\mu}$ and $\hat{pi}$ when $\phi$ is arbitrary. It would be better to study the asymptotic properties of the nuisance estimators due to optimal $\phi^*$ after Eqn. (8).
3. It is strange to have $N_b$ in Algorithm 1 (see lines 387, 388).
4. The auxiliary guidance loss is not clear enough. It would be better if the authors could provide an explicit formula for the loss similar to bandwidth minimization loss and regularization loss.
5. The size of $k$ should affect the bounded estimators. Please demonstrate a scientific way to choose $k$. In the existing experiments, authors study results for different $k$. In any case, it seems that increasing $k$ would give better results. It should not be correct in general.
6. Could you provide the results of the following experiment:
assuming linear generative model, nonlinear ML methods are used to estimate the nuisance functions and then construct the corresponding bandwidth;
assuming a nonlinear generative model, linear ML methods are used to estimate the nuisance functions and then construct the corresponding bandwidth.
7. Would the constructed bounding estimators be doubly robust?
soundness: 3
strengths: This paper investigates the conditional average treatment effect (CATE), assuming the presence of unobserved confounders. Further, researchers encounter complex instrument information (e.g., continuous/high-dimensional), which makes the computation hard in many circumstances. The authors aim to develop estimators that can provide reliable analysis of the above assumptions. Unlike the existing literature, the authors establish bounding estimators of the target quantity and provide theoretical studies of the constructed estimators. A merit of the paper is that it is well-written, so readers can easily follow its notions.
summary: Reliable estimation of treatment effects from observational data is crucial in many disciplines and is a challenge when unconfoundedness assumption in the causal inference literature is violated. In this paper, the authors leverage arbitrary instruments to estimate bounds on the conditional average treatment effect (CATE). First, the authors propose a novel approach for partial identification by mapping instruments to a discrete representation space such that valid bounds on the CATE are yielded. Second, the authors derive a two-step procedure that learns tight bounds using a tailored neural partitioning of the latent instrument space. Thus, instability issues due to numerical approximations or adversarial training are avoided. Third, the authors provide theoretical studies on the proposed approach. Finally, the authors perform extensive experiments to demonstrate the effectiveness across various settings.
weaknesses: There are some notational problems in the paper. For example:
1. Notational problems in Eqn. (10).
2. According to Eqn. (8) (or Eqn. (14)), $\phi$ should be a representation map, and $\Phi$ should be a set of all representations. It is quite confusing to write $\Phi(Z)=l$ (e.g., see line 323). 
3. Please also state the explicit form of the constants $c$, $d$ in Theorem 2.
Please recheck the notations to avoid any confusions.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: See Weaknesses part.
soundness: 2
strengths: 1.	The study addresses an intriguing and valuable problem in causal inference.

2.	The writing is clear and accessible, making complex concepts easier to understand.
summary: This paper introduces a method for estimating treatment effects from observational data despite violations of unconfoundedness. It presents a two-step procedure using high-dimensional latent instruments to derive valid bounds on the Conditional Average Treatment Effect.
weaknesses: 1. The method appears to be a straightforward combination of existing techniques, lacking clear theoretical innovation.  

2. Although the paper highlights the advantage of utilizing high-dimensional instruments, there are limited experiments involving such instruments in the evaluation phase.  

3. The experimental results show no substantial improvements, casting doubt on the method's effectiveness.  

4. While the authors discuss the lack of comparison implementations, it would strengthen the argument to adapt baseline methods to the current setting and compare their performance with that of the proposed method, making the results more persuasive.

5. The final loss function includes two hyperparameters, making it crucial to conduct a sensitivity analysis to evaluate the method's robustness. However, the absence of such experiments raises concerns about the reliability and generalizability of the proposed approach.

6.  The assumption that all observed variables are confounders may be stringent and difficult to satisfy in practical scenarios.

7. The two-stage method presents several limitations, such as potential inefficiencies in estimation and sensitivity to the parameters selected at each stage, which may compromise the overall robustness of the results.

Meta-review: This paper tackles the problem of bounding treatment effects (i.e. partial identification) when there is unobserved confounding in observational data. Here, the paper leverages instruments to estimate bounds on the conditional average treatment effect (CATE). This is done by first transforming instruments into a discrete representation space via the Gumbel Softmax trick followed by a two-step procedure to learn bounds on the CATE by partitioning the representation space. The work proposes using a loss function to learn the mapping by simultaneously minimizing the length of the bounding interval while stabilizing finite sample performance. The methodology was tested on simulated data. 

Both the problem setting and the idea are interesting; but at the end of the discussion period this came down to a borderline decision. I closely followed the discussions and rebuttal. I think the authors did an excellent job of responding to the reviewers' questions but my reading of the discussion and responses was that the reviewers felt the manuscript in its current form read more like a proof of concept rather than a strong standalone piece of work either theoretically or empirically. Such a paper is reasonable but the reviewers did not find common ground on whether the idea alone, with simple experiments, sufficed for publication.

The authors clarify that their their major goal is "*not* to derive new identification bounds or new theoretical results for finite sample guarantees. Instead, the main contribution is to "provide a novel stable learning procedure in an important widely overlooked setting and to show its strong performance". Section 6 in the paper starts off by saying that no fair comparison is possible because there is not other work tailored to this setting. To that end, one suggestion to strengthen the manuscript is with further incorporation of baselines and empirical results in similar, even if not identical settings. For example, Appendix D1 in https://arxiv.org/abs/2202.10806 contains a semi-synthetic experiment based on Mendelian randomization that could form a point of comparison to the work. Similarly, the line of work on partial identification with neural causal models (https://arxiv.org/abs/2405.15673v2) also contains experiments on binary data (e.g. Table 1 in Appendix E https://arxiv.org/abs/2210.08139); this too could form yet another point of comparison. These would strengthen the claims around the efficiency of the estimator. Theoretically, an understanding of how the hyperparameter k influences the bounds would also improve the work. I encourage the authors to resubmit with improvements along the aforementioned axes.

Decision: Reject