Learning Treatment Representations for Downstream Instrumental Variable Regression

## Abstract

Abstract Traditional instrumental variable (IV) estimators face a fundamental constraint: they can only accommodate as many endogenous treatment variables as available instruments. This limitation becomes particularly challenging in settings where the treatment is presented in a high-dimensional and unstructured manner (e.g. descriptions of patient treatment pathways in a hospital). In such settings, researchers typically resort to applying unsupervised dimension reduction techniques to learn a low-dimensional treatment representation prior to implementing IV regression analysis. We show that such methods can suffer from substantial omitted variable bias due to implicit regularization in the representation learning step. We propose a novel approach to construct treatment representations by explicitly incorporating instrumental variables during the representation learning process. Our approach provides a framework for handling high-dimensional endogenous variables with limited instruments. We demonstrate both theoretically and empirically that fitting IV models on these instrument-guided representations ensures identification of directions that optimize outcome prediction. Our experiments show that our proposed methodology improves upon the conventional two-stage approaches that perform dimension reduction without incorporating instrument information.

## 1 Introduction

Instrumental-variable (IV) methods are among the most widely used tools for
recovering causal effects in the presence of unmeasured confounding.
Unfortunately, classical IV estimators scale poorly when the treatment
variable $X$ is itself high-dimensional, unstructured, or both. In modern
applications—where the treatment might be provided in the form of clinical treatment pathways encoded as free-text,
purchase histories, or genome-wide expression profiles—the number of potentially endogenous coordinates of $X$ can dwarf the number of available instruments $Z$ (e.g. variables related to capacity constraints in a hospital setting, see, e.g., [7, 6, 27]). A common workaround is to compress $X$ to a low-dimensional
summary $D$ with unsupervised techniques (e.g. PCA, auto-encoders) and then
run a standard two-stage least squares (2SLS) on $D$. Because the dimension
reduction step ignores $Z$, however, the resulting regression can suffer from
severe omitted-variable bias: directions of $X$ that matter for the first-stage
relationship between $Z$ and $X$ may be discarded, violating the exclusion
restriction and invalidating the causal inference step (c.f. Figure [3](https://arxiv.org/html/2506.02200v2#S3.F3) for such a failure).

We propose *Instrument-Guided Representation Learning* (IGRL), a methodology for learning low-dimensional treatment representations that preserve the validity of downstream IV
analysis. IGRL folds the instruments directly into the representation learner so that the learned features $D$ capture the variation in $X$ that is driven by $Z$. The procedure can be viewed as a regularization of the unsupervised learner toward directions that satisfy the
exclusion restriction, thereby eliminating the spurious back-door paths that plague two-step approaches. The resulting representation can then be used in an IV analysis, to learn directions of intervention in the representation space that will improve the target outcome and can be translated back to interventions in the original treatment space.

Prior work on that combines elements of representation learning with elements of instrumental variable analysis is limited and confined to linear methods. Rao and Sabatier et al. described a procedure of performing principal component analysis (PCA) of a response variable with respect to its instruments. Y Takane studied constrained principal component analysis, which takes external information into consideration during dimensional reduction [35]. More recently, Kelly et al. and Wang incorporates instrumental variables in estimating factor models that improves rate of convergence and avoid overfiting for high-dimensional data [19],[34]. The desiderata in all of these works are very different from identifying dimensions of variation that align with the instruments so that causal effects can be identified by downstream IV analysis.

Our work is also related to the literature on learning non-linear disentangled representations and causal representation learning [15, 14, 20, 11, 23, 31, 1, 16, 18, 17, 12]. However, the focus of this line of work has primarily been on discovering causal structure in data [32], rather than constructing representations for downstream causal tasks. Our work is closely related to the identifiable VAE (iVAE) [20]. The instrument can be viewed as the auxiliary information that can guide non-linear latent factor analysis. However, a crucial difference of our work is that we view the instrument $Z$ as only privileged information that is available only when estimating the causal effects and not when performing interventions. Hence, crucially we want our encoder to only take as input the treatment $X$ and not the instrument $Z$. Moreover, our desiderata is not the discovery of the true latent factors, but solely the discovery of valid decompositions of the treatment for downstream IV analysis. This allows us to relax many of the assumptions that are prevalent in this line of work.

Similar to our work, Saengkyongam et al.’s Rep4Ex approach addresses interventional outcome prediction under a similar structural equation model, but intervenes on $Z$ rather than our latent treatment space ($D$). Other dimensionality reduction studies for high-dimensional treatments [25, 2] operated without unobserved confounders and used outcome-guided factor selection. Additional discussion appears in the appendix.

Our work aligns closer to the recent contributions by Vafa et al. and Du et al., which also highlights the omitted variable bias problem in learned representations in the context where representation learning is used for a set of high-dimensional observed confounders of a treatment and designs representation learning techniques to alleviate it. In that setting, the learned representation can implicitly omit important parts of the observed confounders, causing bias in the final causal estimate due to implicit unobserved confounding. Our goal is inherently different as we want to learn a latent representation of a highly confounded, high-dimensional treatment, as opposed to learning a latent representation of a high-dimensional confounder.

## 2 Problem Statement: Learning Interventions via Representations

We consider a setting where we are given data that contain samples of variables $(Z,X,Y)$, where $X$ is a high-dimensional “treatment” variable, $Y$ is a scalar outcome of interest and $Z$ is a low-dimensional vector of instruments. The treatment $X$ is heavily confounded via unobserved confounding variables $U$ that have a causal influence on the value of $X$ and also on $Y$, as depicted in Figure [1(a)](https://arxiv.org/html/2506.02200v2#S2.F1.sf1).

Our goal is to learn a latent representation of the highly confounded, high-dimensional treatment, so as to perform instrumental variable analysis on this learned representation and identify an outcome-improving direction of intervention in representation space and hence subsequently also in the original treatment space. Naive representation learning approaches for the treatment run the risk of an omitted variable problem that can invalidate the downstream causal analysis based on instrumental variables. Causal analysis using instrumental variables crucially assumes that the instrument $Z$, the treatment $X$, and the outcome $Y$ respect the causal graph depicted in Figure [1(a)](https://arxiv.org/html/2506.02200v2#S2.F1.sf1). In particular, the instrument $Z$ is assumed to only affect outcome $Y$ through its effect on treatment $X$. When the high-dimensional treatment $X$ is replaced by a learned representation $D$, we run the risk that the part of $X$ that is not represented in $D$ contains elements that are correlated with both the instrument $Z$ and the outcome $Y$. As a result, $D$ no longer absorbs the entire effect of the instrumental variable Z on the outcome $Y$. This creates causal pathways from the instrument $Z$ to the outcome $Y$ that do not flow through the representation $D$, as shown in Figure [1(b)](https://arxiv.org/html/2506.02200v2#S2.F1.sf2). Therefore, we need to regularize the representation learning process to ensure that the causal influence through these omitted paths is minimal.

Figure: (a) Instrumental variable causal graph, with instruments $Z$, high-dimensional treatment $X$, outcome $Y$, unobserved confounders $U$.

An ideal latent representation $D$ should satisfy the causal graph depicted in Figure [1(c)](https://arxiv.org/html/2506.02200v2#S2.F1.sf3). In particular, the instrument $Z$ should not have a causal effect on $X$ that is not absorbed by the latent representation $D$. If the representation encodes all outcome-relevant information, then a direct edge from $X$ to $Y$ should not exist. However, the existence of such an edge does not invalidate the downstream instrumental variable analysis, and hence, it is not essential to exclude it.

#### Structural Equation Model.

To formalize our problem we will consider the following data generating process (structural causal model) for our observed random variables:

$$ $\displaystyle D={}$ $\displaystyle A\cdot Z+U,~{}$ $\displaystyle U\perp\!\!\!\perp Z$ (1) $\displaystyle X={}$ $\displaystyle f(D,V),~{}$ $\displaystyle V\perp\!\!\!\perp Z$ $\displaystyle Y={}$ $\displaystyle h(D)+\eta(U,V,\epsilon),~{}$ $\displaystyle\epsilon\perp\!\!\!\perp Z$ $$

where the random variables $U,V,D,\epsilon$ are latent and $A$ is an $r\times k$ matrix that captures the effect of the instruments $Z\in\mathbb{R}^{k}$ on a vector of latent decisions $D\in\mathbb{R}^{r}$. For convenience of notation, we will assume that $\mathbb{E}[U]=\mathbb{E}[V]=\mathbb{E}[\eta(U,V,\epsilon)]=0$.(^1^11Appropriate intercept constants need to be added to the equations in the absence of this convention.) $U$ represents the unobserved confounder that drives the elements of the treatment that are also driven by the instrument. $\epsilon$ represents an outcome noise variable and is allowed to be correlated with $U,V$. $D$ represents the aspects of the treatment $X$ that are affected by the instrument and $V$ represents the remaining aspects that describe the treatment $X$, but are independent of the instrument. In particular, we assume that the encoding/decoding between the latent representations and the observed treatment is invertible:

###### Assumption 2.1 (Invertible Encoding) .

The function $f$ is invertible, and write the encoding function $e(X)=f^{-1}(X)=(D,V)$, i.e. there is a one-to-one correspondence between the high-dimensional treatment $X$ and the characteristics $(D,V)$ that describe the treatment.

From this perspective, $(D,V)$ can be thought as a non-linear decomposition of the treatment into the instrument-dependent and the instrument-independent components. We will further denote with $e_{D}(X)=D$ and $e_{V}(X)=V$ for the encodings of the treatment that return the corresponding components. Moreover, we assume that the transformation between the instrument to the latent representation $D$ is full rank.

###### Assumption 2.2 (Full-Rank Latents) .

Assume that the matrix $A$ has full row-rank and $\mathbb{E}[ZZ^{\top}]\succ 0$.

Note that the full rank assumption on $\mathbb{E}[ZZ^{\top}]$ can always be satisfied by a preprocessing step that applies a PCA transformation to the instruments and removes co-linear or almost co-linear instruments.

###### Assumption 2.1 (Invertible Encoding) .

###### Assumption 2.2 (Full-Rank Latents) .

#### Learning Good Interventions via Representations.

Given data containing observations $(Z,X,Y)$ stemming from such a structural equation model, our goal is to learn a soft intervention mapping $t(X)$, such that the average intervened outcome is larger than the original outcome. We will denote with $Y^{(X\leftarrow x)}$ the random outcome from the intervention where we fix the value of $X$ to be $x$. Thus we are searching for a soft intervention $t(X)$ such that:

$$ $\displaystyle\mathbb{E}\left[Y^{(X\leftarrow t(X))}\right]>\mathbb{E}[Y]$ (2) $$

Note that due to the one-to-one correspondence of $X$ with its decomposition, any such interventional outcome can equivalently be thought as an intervention on the latent components of the treatment, i.e. $Y^{(D\leftarrow e_{D}(x),V\leftarrow e_{V}(x))}$. Given the structural Equation ([1](https://arxiv.org/html/2506.02200v2#S2.E1)), the expected outcome under a soft intervention $t(X)$ can be written as:

$$ $\displaystyle\mathbb{E}\left[Y^{(X\leftarrow t(X))}\right]=\mathbb{E}[h(e_{D}( t(X)))+\eta(U,e_{V}(t(X)),\epsilon)]$ (3) $$

We will identify such an intervention via the means of intervention on a learned representation. In particular, given observations, we will learn an encoding $\tilde{e}_{D}(X)=\tilde{D}$ that respects the properties in Equation ([1](https://arxiv.org/html/2506.02200v2#S2.E1)) (potentially together with a learned encoding $\tilde{e}_{V}(X)=\tilde{V}$) and a corresponding decoder $\tilde{f}(\tilde{D})$ (potentially also taking as input $\tilde{V}$) that maps the learned encoding back into a high-dimensional treatment. Subsequently, we will estimate an outcome improving direction $u$ in the learned representation space via instrumental variable analysis, viewing $\tilde{D}$ as the “treatment” and $Z$ as the instrument. We will apply the direction $u$ to the learned representations, i.e. $\tilde{D}+\alpha u$, for some scalar intervention amount $\alpha$. For ease of notation, we denote with $(\cdot)_{\alpha u}$ to be the corresponding random variable $(\cdot)$ after this intervention. Then decode back to the high-dimensional treatment space $X_{\alpha u}=\tilde{f}(\tilde{D}+\alpha u)$ (potentially $X_{\alpha u}=\tilde{f}(\tilde{D}+\alpha u,\tilde{V})$ if an encoding of $V$ was also learned). This process (depicted also visually in Figure [2](https://arxiv.org/html/2506.02200v2#S2.F2) and described algorithmically in Algorithm [1](https://arxiv.org/html/2506.02200v2#alg1)) defines our soft-intervention mapping, formally defined as:

$$ $\displaystyle t(X)=\tilde{f}(\tilde{e}_{D}(X)+\alpha u,\tilde{e}_{V}(X)),$ (4) $$

with the second input of $\tilde{f}$ omitted if an encoding $\tilde{e}_{V}$ is not learned.

Figure: Figure 2: Intervention on learned representation.

Figure: Algorithm 1 Intervention in Latent Representation Space and evaluation

## 3 Instrument Guided Representation Learning: The Linear Setting

To make matters more concrete, we will start this analysis with the case where the structural equation model that is associated with the causal graph in Figure [1(c)](https://arxiv.org/html/2506.02200v2#S2.F1.sf3) contains only linear relationships:

$$ $\displaystyle D={}$ $\displaystyle A\cdot Z+U,$ $\displaystyle U\perp\!\!\!\perp Z$ (5) $\displaystyle X={}$ $\displaystyle B\cdot D+B_{\perp}\cdot V,$ $\displaystyle V\perp\!\!\!\perp Z$ $\displaystyle Y={}$ $\displaystyle\theta^{\top}D+\eta(U,V,\epsilon),$ $\displaystyle\epsilon\perp\!\!\!\perp Z$ $$

where $B$ is an $m\times r$ dimensional matrix that maps the $k$ instrument-driven latent decisions $D$ to the observed high-dimensional treatments $X\in\mathbb{R}^{m}$ *and is assumed to be full column rank*. $B_{\perp}$ is a matrix whose column space is orthogonal to the column space of $B$ and is also *assumed to be full column rank*. $U$ corresponds to a random vector of latent unobserved confounders that also affect decisions and outcomes. $\theta$ is an $r$ dimension vector capturing the direct effects of the latent decisions on the outcome. We will assume that the matrix $A$ is of full row rank, i.e., we have more instruments $Z$ than latent decisions $D$, and the instruments vary these latent dimensions in a full-rank manner.

Note that this setting falls under our general model since the function $f(D,V)=BD+B_{\perp}V$ is invertible. In particular, by the orthogonality of the column space of the two matrices and the fact that they are both full column rank, we have that:

$$ $\displaystyle e_{D}(X)\triangleq~{}$ $\displaystyle B^{+}X=D$ $\displaystyle e_{V}(X)\triangleq~{}$ $\displaystyle B_{\perp}^{+}X=V$ (6) $$

where $B^{+}$ denotes the Moore-Penrose pseudo-inverse of a matrix and which is a left inverse for full column rank matrices, i.e. $B^{+}=(B^{\top}B)^{-1}B^{\top}$.
Moreover, note that we could have equivalently defined the structural equation for $X$ as:

$$ $\displaystyle X=~{}$ $\displaystyle B\cdot D+V,$ $\displaystyle V\perp\!\!\!\perp Z$ (7) $$

We could always split the second part into $B\cdot V+B_{\perp}V$ and redefine $D\to D+V$, or equivalently redefine $U\to U+V$. The formulation in Equation [5](https://arxiv.org/html/2506.02200v2#S3.E5) is chosen for notational convenience.

Our target quantity of interest is the overall effect $\theta$ of the latent factors $D$ on the outcome $Y$. If we could identify the latent factors $D,V$ from the observed variables, then we could simply use the improving intervention direction $u=\theta/\|\theta\|$. In this case, our improving intervention corresponds to $t(X)=B(e_{D}(X)+\alpha u)+B_{\perp}V=X+\alpha Bu$, with $D_{\alpha u}=D+\alpha u$ and $V_{\alpha u}=V$, which would lead to an internventional outcome of $Y_{\alpha u}=\theta^{\top}(D+\alpha u)+\eta(V,U,\epsilon)$, hence:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}]=\mathbb{E}[Y]+\alpha\theta^{\top}u= \mathbb{E}[Y]+\alpha\|\theta\|$ (8) $$

Note that in this linear setting, to perform the intervention, it suffices that solely learn a linear encoder $e_{D}(X)=B^{+}X$, since the intervention can be performed implicitly as $t(X)=X+\alpha Bu$, which would not require learning an encoding for $V$. Hence we will take this approach in the remainder of this section.
We will show that in this setting it is feasible to identify improving interventions, even though the natural latent decomposition $D$ might not be necessarily identifiable. We will show that we can always identify a representation $\tilde{D}$, such that $\tilde{D}$ is an invertible linear transformation of $D$.

Note that in this setting, a linear regression of $X$ on $Z$ uncovers the matrix $C=B\cdot A$ since our structural equation model implies the regression equation:

$$ $\displaystyle X=~{}$ $\displaystyle B\cdot A\cdot Z+B_{\perp}\cdot V+B\cdot U,$ $\displaystyle\mathbb{E}[B_{\perp}\cdot V+B\cdot U\mid Z]=B_{\perp}\mathbb{E}[V ]+B\cdot\mathbb{E}[U]=~{}$ $\displaystyle 0$ (9) $$

Moreover, since $A$ is full row rank, the column space of $C$ can be proven to be the same as the column space of $B$. Thus, if we perform a *thin* singular value decomposition of $C={\cal U}\Sigma{\cal V}^{\top}$, then the $m\times k$ matrix of left eigenvectors ${\cal U}$ can be used as matrix $\hat{B}$, as they correspond to an orthonormal basis of the column space of $C$ and, therefore, also of the column space of $B$. Consequently, $\Sigma{\cal V}^{\top}$ can be used as $\hat{A}$. Subsequently, we can take $\tilde{D}=\hat{B}^{\top}X=\hat{B}^{\top}BD$. Since the column space of $\hat{B}$ is the same as the column space of $B$, the square matrix $P=\hat{B}^{\top}B$ is invertible.
An intervention in the direction of $u$ in the learned representation can be thought of as an intervention in the direction of $P^{-1}u$ in the natural representation. An instrumental variable regression estimate, using $Z$ as the instrument, $\tilde{D}$ as the treatment, and $Y$ as the outcome, is characterized as the solution to the moment restriction: $\mathbb{E}[Z(Y-\tilde{\theta}^{\top}\tilde{D})]=0$. It can be shown that as long as matrix $A$ is full row rank and the instruments are not co-linear, i.e. $\mathbb{E}[ZZ^{\top}]\succ 0$, then the above system has a unique solution, $\tilde{\theta}=(P^{-1})^{\top}\theta$, which is the correct causal effect of interventions on $\tilde{D}$. We will then learned representation space in the direction $u=\tilde{\theta}/\|\tilde{\theta}\|$. The implied intervention in the $X$-space is $t(X)=X+\alpha\hat{B}u$.
Algorithm [2](https://arxiv.org/html/2506.02200v2#alg2) formalizes this procedure and the following theorem formalizes these arguments and provides the outcome improvement guarantee for this intervention.

###### Theorem 3.1 .

Under the linear structural equation model in Equation ([5](https://arxiv.org/html/2506.02200v2#S3.E5)) and assuming $B,B_{\perp}$ have full column rank and Assumption [2.2](https://arxiv.org/html/2506.02200v2#S2.Thmtheorem2) holds, then the representation and intervention produced by the LIRR algorithm satisfy: $\tilde{D}=PD$, for the invertible matrix $P\triangleq\hat{B}^{\top}B$. Moreover, $\tilde{\theta}=(P^{-1})^{\top}\theta$ and the interventional outcome satisfies the guaranteed improvement property:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}]=\mathbb{E}[Y]+\alpha\|(P^{-1})^{\top}\theta\|$ $$

Figure: Algorithm 2 Linear Instrument Regularized Representation (LIRR) and Intervention

The LIRR algorithm offers substantial improvements over typical approaches to dimensionality reduction when one is faced with high-dimensional treatments and low dimensional instruments. For instance, if instead one performed a typical dimensionality reduction approach of taking the top-$k$ principal components of the treatment $X$ and using that as a learned representation (with $k$ being the dimension of the instrument), then this top-$k$ components can miss many of the dimensions that the instrument is varying and hence lead to an erroneous downstream intervention via the IV analysis. Such a stark comparison is presented in Figure [3](https://arxiv.org/html/2506.02200v2#S3.F3), where we depict the outcome pre and post intervention in a synthetic example where we know the ground truth. While our LIRR approach consistently produces intervened outcomes with larger values, the PCA followed by IV approach produces worse outcomes.

Figure: Figure 3: Comparing LIRR with a PCA followed by IV approach to constructing improving interventions. Data is generated following the variant of the linear SEM presented in Equations [5](https://arxiv.org/html/2506.02200v2#S3.E5)&[7](https://arxiv.org/html/2506.02200v2#S3.E7), where $U$ and $V$ are mixtures of independently sampled uniform random variables and $\eta(U,V,\epsilon)=U+\epsilon$, where $\epsilon$ is Gaussian.
Refer to caption: extracted/6565704/plots/example.png

###### Theorem 3.1 .

## 4 Instrument Guided Representation Learning: The Non-Linear Setting

We will now investigate the general setting introduced in Equation ([1](https://arxiv.org/html/2506.02200v2#S2.E1)). In this non-linear setting, we will require some further assumptions on the latent factors. In particular, we will be assuming that the latent components $D$ are independent of the orthogonal components $V$ that constitute $X$ that are not driven by the instrument. In particular, we will assume the slightly stronger property of joint independence of $Z,U,V$, which implies that $D\perp\!\!\!\perp V$.

###### Assumption 4.1 (Joint Independence) .

Assume that $Z\perp\!\!\!\perp U\perp\!\!\!\perp V$ (jointly independent).

Moreover, we will assume the regularity conditions that the mixing function $f$ is differentiable and that the instrument is supported on an open subset of $\mathbb{R}^{k}$. For convenience of notation, we will also assume that $\mathbb{E}[Z]=0$, which can be achieved by a pre-processing step that demeans the instruments.

###### Assumption 4.2 (Differentiable Decoding Function) .

$f$ is a differentiable function with uniformly bounded derivatives.

###### Assumption 4.3 .

$\mathbb{E}[Z]=0$ and the support of $Z$, ${\cal Z}$, is an open subset of $\mathbb{R}^{k}$.

Moreover, we will make a completeness assumption on the strength of the instrument. Completeness is a standard assumption for non-parametric identification using instrumental variable analysis [5].

###### Assumption 4.4 (Bounded Completeness) .

$D$ is bounded complete for Z, that is, for all bounded real functions $h$, we have that:

$$ $\displaystyle\mathbb{E}[h(D)|Z]=0\quad\text{a.s.}\quad\Rightarrow\quad h(D)=0 \quad\text{a.s.}$ $$

We discuss sufficient conditions for the bounded completeness assumption in the Appendix (Lemma [C.3](https://arxiv.org/html/2506.02200v2#A3.Thmtheorem3)). In particular, it involves characteristic function assumptions that have also been typical in the identifiable latent factor literature [22].

###### Theorem 4.5 .

Suppose that the data generating process follows the SEM described in Equation [1](https://arxiv.org/html/2506.02200v2#S2.E1), and satisfies Assumptions [2.1](https://arxiv.org/html/2506.02200v2#S2.Thmtheorem1) & [2.2](https://arxiv.org/html/2506.02200v2#S2.Thmtheorem2) & [4.1](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem1) & [4.2](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem2) & [4.3](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem3) & [4.4](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem4). Let $(\tilde{D},\tilde{V}):=(\tilde{e}_{D}(X),\tilde{e}_{V}(X))=\tilde{e}(X)$ denote the learned representations. Consider encoder-decoder pairs with perfect reconstruction, i.e. $X=\tilde{f}\circ\tilde{e}(X)$. Then, for the solution $\tilde{e},\tilde{f}$, and full row rank matrix $\tilde{A}$ that minimizes the objective function

$$ $\displaystyle\mathbb{E}[\|\tilde{e}_{D}(X)-\tilde{A}Z\|^{2}]$ (10) $$

subject to the following constraints:

- •
$\tilde{e}$ is a differentiable function with uniformly bounded derivatives.
- •
$\tilde{A}$ has full row rank.
- •
$\tilde{D}=\tilde{A}Z+\tilde{U}$ with $\tilde{U}\perp\!\!\!\perp Z$ and $\mathbb{E}[\tilde{U}]=0$.

we have that, with probability $1$, $\tilde{D}=PD$ and $\tilde{U}=PU$, for $P=\tilde{A}A^{+}$. Moreover, the matrix $P$ is invertible.

###### Lemma 4.6 .

Suppose the assumptions of Theorem [4.5](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem5) hold, and additionally impose the following constraints on the learned functions $\tilde{e},\tilde{f},\tilde{A}$ that minimize the objective in Equation  [10](https://arxiv.org/html/2506.02200v2#S4.E10):

- •
$\tilde{D}\perp\!\!\!\perp\tilde{V}$
- •
$\tilde{e}$ is an invertible function when restricted to inputs in the image of $f$ and $\tilde{f}\circ\tilde{e}(x)=x$ for all $x\in\text{Im}(f)$.

Let $q\triangleq\tilde{e}\circ f$ and $q^{-1}\triangleq e\circ\tilde{f}$. Let $q_{2}$ denote the $\tilde{V}$-component of the output of $q$ and $q_{2}^{-1}$ the $V$-component of the output of $q^{-1}$.
Then we have that $(\tilde{D},\tilde{V})=q(D,V)$ and $(D,V)=q^{-1}(\tilde{D},\tilde{V})$, almost surely. Moreover, it must also hold with probability 1 that:

- •
$\tilde{V}=q_{2}(D,V)=(\tilde{e}\circ f)_{2}(D,V)$(^2^22With $(\tilde{e}\circ f)_{2}$ we denote the $V$ component of the output of the function $\tilde{e}\circ f_{0}$.) with the property that for all $d,d^{\prime}\in{\cal D}$:
$\text{Law}(q_{2}(d,V))=\text{Law}(q_{2}(d^{\prime},V)).$
- •
$V=q_{2}^{-1}(D,V)=(e\circ\tilde{f})_{2}(\tilde{D},\tilde{V})$ with the property that for all $d,d^{\prime}\in{\cal\tilde{D}}\triangleq\{Pd:d\in{\cal D}\}$:
$\text{Law}(q_{2}^{-1}(d,\tilde{V}))=\text{Law}(q_{2}^{-1}(d^{\prime},\tilde{V}
)).$

###### Remark 4.7 .

Note that the assumptions that $\mathbb{E}[Z]=0$ and $\mathbb{E}[ZZ^{\prime}]\succ 0$ are without loss of generality as we can always pre-process $Z$ by centering it and removing co-linear instruments. Moreover, in practice the assumption that $\tilde{D}=\tilde{A}Z+\tilde{U}$, with $\tilde{U}\perp\!\!\!\perp Z$ and $\mathbb{E}[\tilde{U}]=0$ can be achieved by minimizing a square loss with an intercept, i.e.

$$ $\displaystyle\min_{e,f,A,c:e,f\text{invertible},e\circ f=\text{identity}} \mathbb{E}[\|e_{D}(X)-AZ-c\|^{2}]$ $$

and then defining $\tilde{D}=\tilde{e}_{D}(X)\triangleq e_{D}(X)-c$, $\tilde{f}=f+c$.

Subsequently, we identify an intervention as described in Algorithm [1](https://arxiv.org/html/2506.02200v2#alg1). In particular, we will run an IV analysis, with $Z$ as the instrument, $\tilde{D}$ as the treatment, and $Y$ as the outcome, to estimate a causal model in representation space by finding a solution to the conditional moment restrictions:

$$ $\displaystyle\mathbb{E}[Y-\tilde{h}(\tilde{D})\mid Z]=0$ (11) $$

Note that since $\tilde{D}=PD$ and since $\mathbb{E}[Y\mid Z]=\mathbb{E}[h(D)\mid Z]$, we have by the completeness assumption that:

$$ $\displaystyle\mathbb{E}[h(D)-\tilde{h}(PD)\mid Z]=0\Rightarrow h(D)=\tilde{h}( PD)\text{ a.s. }\implies h(P^{-1}\tilde{D})=\tilde{h}(\tilde{D})\text{ a.s. }$ $$

If for instance, $h$ is assumed to be linear, then $\tilde{h}$ is also a linear function and it suffices to run a linear instrumental variable analysis (e.g. two-stage-least-squares). If $h$ is non-linear, then we calculate the average derivative of $\tilde{h}$, i.e.

$$ $\tilde{\theta}=\mathbb{E}[\nabla_{\tilde{D}}\tilde{h}(\tilde{D})]=(P^{-1})^{ \top}\mathbb{E}[\nabla_{D}h(D)]$ $$

and perform the intervention

$$ $u=\tilde{\theta}/\|\tilde{\theta}\|$ $$

as described in Algorithm [1](https://arxiv.org/html/2506.02200v2#alg1). In finite samples, recently introduced doubly robust methods for estimation of average derivatives of solutions to non-parametric IV problems can be used [3, 4].

###### Theorem 4.8 .

Assume that:

$$ $\displaystyle Y=h(D)+\eta(U,V,\epsilon),\quad\epsilon\perp\!\!\!\perp\{Z,U,V\}$ $$

and that $h$ is twice differentiable with a bounded second derivative. Let $\tilde{e},\tilde{f},\tilde{A}$ be an optimal solution as prescribed in Lemma [4.6](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem6) with the extra constraint that:

$$ $\displaystyle\tilde{U}\perp\!\!\!\perp\tilde{V}\perp\!\!\!\perp Z$ (joint independence) $$

and assume that the assumptions of Lemma [4.6](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem6) are satisfied. Furthermore, assume that the variable $D$ has full support in $\mathbb{R}^{r}$, i.e. ${\cal D}=\mathbb{R}^{r}$. Then setting $u=\tilde{\theta}/\|\tilde{\theta}\|$, in Algorithm [1](https://arxiv.org/html/2506.02200v2#alg1), with $\tilde{\theta}=\mathbb{E}[\nabla_{\tilde{D}}\tilde{h}(\tilde{D})]$ and $\tilde{h}$ the solution to the conditional moment restriction problem in Equation ([11](https://arxiv.org/html/2506.02200v2#S4.E11)), we have that:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}-Y]=\alpha\|(P^{-1})^{\top}\mathbb{E}[ \nabla_{D}h(D)]\|+O(\alpha^{2})$ $$

Hence, for small enough step size $\alpha$, the identified intervention will achieve a positive improvement on the outcome (assuming that $\mathbb{E}[\nabla_{D}h(D)]\neq 0$).

###### Assumption 4.1 (Joint Independence) .

###### Assumption 4.2 (Differentiable Decoding Function) .

###### Assumption 4.3 .

###### Assumption 4.4 (Bounded Completeness) .

###### Theorem 4.5 .

###### Lemma 4.6 .

###### Remark 4.7 .

###### Theorem 4.8 .

#### Instrument Regularized Auto-Encoder

To achieve the positive improvement as described in Theorem [4.8](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem8), then we need to incorporate loss components that are minimized only when i) $e,f$ reconstruct the input $X$, ii) $e_{D}(X)$ is predicted linearly by $Z$ with a full rank matrix $A$, iii) the residual of this regression $e_{D}(X)-AZ-c$, which approximates $U$, needs to be independent of $Z$, iv) $Z$ needs to be independent of $e_{V}(X)$ and v) $e_{D}(X)$ needs to be independent of $e_{V}(X)$,$e_{D}(X)-AZ-c$ needs to be independent of $e_{V}(X)$, and the variables $(e_{D}(X)-AZ-c,e_{V}(X),Z)$ are jointly independent. While we do not explicitly enforce $\tilde{A}$ to be full row rank, we expect this to be satisfied due to the reconstruction loss and the condition that $\tilde{D}\perp\!\!\!\perp\tilde{V}$. Moreover, note that instead of only joint independence of $Z,\tilde{U},\tilde{V}$ we also explicitly enforce pairwise independencies for computational reasons.

We introduce the instrument-regularized auto-encoder loss, which incorporates all these elements:

$$ $\displaystyle\min_{e,f,A,c}{}$ $\displaystyle\mathbb{E}\left[\|X-f\circ e(X)\|^{2}\right]+\lambda\mathbb{E} \left[\|e_{D}(X)-AZ-c\|^{2}\right]$ (IRAE) $\displaystyle+\mu_{1}{\cal R}(e_{D}(X)-AZ-c,Z)+\mu_{2}{\cal R}(Z,e_{V}(X))$ $\displaystyle+\mu_{3}\big{(}c_{1}{\cal R}(e_{D}(X),e_{V}(X))+c_{2}{\cal R}(e_{ D}(X)-AZ-c,e_{V}(X))$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}+c_{3}{\cal R}(e_{D}(X)-AZ-c,Z,e_{V }(X))\big{)}$ $$

${\cal R}(A,B)$ or ${\cal R}(A,B,C)$, denotes any regularizer that can be evaluated on a set of $n$ samples and which takes small values when the random variables $A,B$ or $A,B,C$ are jointly independent. Many examples of such independence-regularizers have been introduced in the literature. Our methodology is agnostic to the exact regularizer used. In our experiments, we used a kernel-based independence test statistic (HSIC) [10] and its generalization to joint independence (d-HSIC) [26].

In experiments, for the purposes of ablation analysis, we will denote with IRAE[0] the variant that contains only the regularization parts that are multiplied by $\lambda$, with IRAE[1] the variant that contains the parts that are multiplied by $\lambda,\mu_{1}$, with IRAE[2] the variant that contains the parts multiplied by $\lambda,\mu_{1},\mu_{2}$ and IRAE the variant that contains all regularizers.

## 5 Experimental Evaluation

#### Linear setting

We benchmark LIRR ([Algorithm 2](https://arxiv.org/html/2506.02200v2#alg2)) against PCA under the setting of linear data generating process. As a baseline, we consider using PCA to extract the top $k=4$ principal components of X as the learned latent representation. After the representation is generated, we run 2SLS with representation $\tilde{D}$ as “treatment”, outcome $Y$, and instruments $Z$ to identify the direction of perturbation. We apply steps 4-6 in [Algorithm 1](https://arxiv.org/html/2506.02200v2#alg1) with $\alpha=1$ to compute the improvement $\mathbb{E}[Y_{\alpha u}-Y]$.

For each experiment, we randomly generate elements of $A,B,\theta$ in [Equation 5](https://arxiv.org/html/2506.02200v2#S3.E5) from normal distributions and test our method across three distinct noise cases: 1) independent Gaussian distributions for both U and V, 2) correlated Uniform distribution for U, independent Gaussian distribution for V, 3) correlated Uniform distribution for U, correlated Gaussian distribution for V. Each experiment was repeated 100 times with different random seeds, each containing a sample size of 10000 with 80-20 train-test split. We also varied the dimensionality of X, $m$, to examine the dimension effects while holding the dimension of Z constant ($k=4$). The distribution of average improvements across seeds is presented in [Table 1](https://arxiv.org/html/2506.02200v2#S5.T1). Detailed data generation procedures are provided in the appendix.

We note that when noise follows independent Gaussian distributions across coordinates of $U$ and $V$, PCA method performs comparably to LIRR. However, PCA fails to generalize effectively under non-independent noise conditions. The average improvement of our proposed method exceeds that of SVD in case 2 and 3, and being more than 1 standard deviation from zero except for the case of DGP 3 and $m=500$. Curse of dimensionality still exists, as improvement decreases as the $m$ increases.

**Table 1: Average Test Improvement Comparison on Linear Data: LIRR vs. PCA (Mean $\pm$ Std). DGP 1 corresponds to independent U and V, DGP 2 corresponds to correlated U and independent V, and lastly DGP 3 corresponds to correlated U and V.**
| Size $m$ | Method |  | DGP 1 | DGP 2 | DGP 3 |
| --- | --- | --- | --- | --- | --- |
| 50 | LIRR |  | $\mathbf{3.7283\pm 2.7360}$ | $\mathbf{5.4706\pm 4.1242}$ | $\mathbf{5.4944\pm 4.0596}$ |
| PCA |  | $3.1035\pm 3.6229$ | $3.1717\pm 4.0468$ | $2.5171\pm 4.8519$ |  |
| 100 | LIRR |  | $\mathbf{2.4189\pm 2.0164}$ | $\mathbf{4.0806\pm 3.5969}$ | $\mathbf{3.8931\pm 3.3116}$ |
| PCA |  | $2.1249\pm 2.7203$ | $2.4044\pm 3.5741$ | $2.5713\pm 3.7491$ |  |
| 500 | LIRR |  | $\mathbf{1.0355\pm 0.9698}$ | $\mathbf{1.6996\pm 1.5957}$ | $\mathbf{1.5934\pm 1.7305}$ |
| PCA |  | $1.0098\pm 1.0786$ | $0.9005\pm 1.3995$ | $1.1716\pm 1.6904$ |  |

#### Non-linear setting

Next we consider a non-linear data generating process, where the data is generated by [Equation 1](https://arxiv.org/html/2506.02200v2#S2.E1) where $f$ is quadratic and $h$ is linear. We benchmark LIRR and IRAE against PCA and vanilla Autoencoder (vanilla AE), variational autoencoder (VAE), and iVAE under the setting of a quadratic data generating process. Here, vanilla AE refers to autoencoder with only reconstruction loss. VAE refers variational autoencoder that maximizes the likelihood $p_{f}(X)$ with Gaussian latent representation. iVAE [20] utilizes both $Z$ and $X$ in encoding and decoding, maximizing the conditional likelihood of $p_{f,A}(X|Z)$ as information of $Z$ is available in simulations. For LIRR, PCA, IRAE[1], vanilla AE, VAE, iVAE the bottleneck is of the same dimension as the instrument, i.e. $k=4$, so that downstream 2SLS will not be ill-posed, whereas the bottleneck size of IRAE[2] and IRAE was $10$. [Algorithm 1](https://arxiv.org/html/2506.02200v2#alg1) is then applied to evaluate the average improvement in outcome, when each of the aforementioned representation learning methods is used. For probabilistic autoencoder VAE and iVAE, we sampled 10 representations for each observation X and compared them to the original outcome.

We repeated the experiment 30 times across different random seeds, each containing a sample size of 10000 with 70-10-20 train-val-test split. Results on average improvement are depicted in Table [2](https://arxiv.org/html/2506.02200v2#S5.T2). Our findings reveal that dimension reduction methods which operate without Z information (PCA, vanilla AE, vanilla VAE) yield minimal outcome improvement. In contrast, methods that incorporate Z consistently demonstrate positive mean improvements. The most substantial improvement is achieved by our IRAE[1] and IRAE method, with IRAE having performance gains at more than one standard deviation above zero.

**Table 2: Average Test Improvement Comparison of 9 Methods on Quadratic Data (Mean $\pm$ Std). DGP 1 corresponds to independent U and V, DGP 2 corresponds to correlated U and independent V, and lastly DGP 3 corresponds to correlated U and V.**
| Method | Case 1 | Case 2 | Case 3 |
| --- | --- | --- | --- |
| PCA | $0.1322\pm 0.3216$ | $0.0545\pm 0.2994$ | $0.0848\pm 0.2382$ |
| LIRR | $3.5086\pm 2.0455$ | $3.4711\pm 1.9683$ | $3.5682\pm 2.1296$ |
| Vanilla AE | $0.4138\pm 2.2000$ | $0.8418\pm 1.1560$ | $0.7801\pm 1.7335$ |
| IRAE[0] | $6.1055\pm 7.1634$ | $2.2898\pm 6.9957$ | $4.8993\pm 6.3310$ |
| IRAE[1] | $\mathbf{6.4174\pm 5.2602}$ | $4.6175\pm 5.0479$ | $\mathbf{5.8023\pm 7.1041}$ |
| IRAE[2] | $5.5471\pm 4.6573$ | $4.5554\pm 4.0707$ | $5.2145\pm 4.4358$ |
| IRAE | $5.7740\pm 4.7664$ | $\mathbf{6.5253\pm 6.0132}$ | $4.9113\pm 4.0009$ |
| Vanilla VAE | $0.3651\pm 0.4629$ | $0.2725\pm 0.5071$ | $0.2055\pm 0.3394$ |
| iVAE | $0.2709\pm 0.3672$ | $0.1192\pm 0.2503$ | $0.1652\pm 0.2929$ |

#### MNIST experiment 1

We examine a case where the outcome is determined by the color of MNIST digits. In this experiment, we independently generated 2-dimensional instrumental variables $Z$ and 2-dimensional confounders $U$. The color features $D$ are represented as 3-dimensional RGB values determined by both $Z$ and $U$. The outcome variable is calculated as the sum of R, G, and B values. The observed data X consists of MNIST digit pixels. If our methods successfully identify the correct causal direction, we expect intervened images to display increased brightness.
All except IRAE[2] and IRAE has bottleneck size same as dimension $Z$ and the IRAE[2] and IRAE methods had a bottleneck of size $10$. Performance improvement results across $40$ seeds, with each experiment being run on a subsample of the MNIST dataset, are reported in [Table 3](https://arxiv.org/html/2506.02200v2#S5.T3) for the leading methods (AE and IRAE) as well as for ablation variants of IRAE. Additional visualizations are available in the appendix. As a remark, we note that having multiple dependence penalty terms may be difficult to train. For this reason, we trained IRAE[1] first and transferred the knowledge into the larger bottleneck models in IRAE[2] and IRAE.

Our experiments reveal important insights about latent space representation and instrumental variables. The vanilla AE, with no specialized latent regularization, produces reconstructed digits that closely resemble the originals, indicating the latent space primarily focuses on digit reconstruction. This can be seen in [Figure 5](https://arxiv.org/html/2506.02200v2#S5.F5) where the latent distribution is best explained by digit (right most plot). When IV regression are applied to this representation, no meaningful directional information can be extracted since digit morphology has no relation to the target variable Y, resulting in no improvement. When we introduce instrument regularization while maintaining the same dimensionality as $Z$ in IRAE[1], the representation is forced to capture more color information at the expense of digit reconstruction. Ideally, we could increase the prediction error weight to infinity to enforce full capture of $Z$ information, but in practice, some digit information remains in the representation, leading to better improvement than Vanilla AE but worse than the following two methods. By expanding the latent dimension, we achieve both better digit reconstruction and color information preservation. The larger dimensional space accommodates more digit morphology without needing to compete space with color information, bringing reconstruction error closer to zero while enabling instrumental variables to recover the target direction. This can be seen in [Figure 4](https://arxiv.org/html/2506.02200v2#S5.F4) where the latent distribution is well explained by instruments, RGB values, and $Y$ (three plots on the left). The improvement in IRAE[2] is less pronounced than in IRAE due to information leakage between components $D$ and $V$, resulting in acceptable reconstruction and prediction error but less identifiable direction when IVs are applied solely to the $D$ component. By adding a dependence penalty between $D$ and $V$, IRAE achieves marginally better improvement.

**Table 3: Average Test Improvement Comparison of 5 Methods on MNIST Data (Mean $\pm$ Std)**
|  |  | Vanilla AE | IRAE[0] | IRAE[1] | IRAE[2] | IRAE |
| --- | --- | --- | --- | --- | --- | --- |
| sample size | image |  |  |  |  |  |
| 1000 | reconstructed | -0.57 (0.03) | -0.64 (0.17) | -0.6 (0.15) | -0.63 (0.09) | -0.62 (0.09) |
| intervened(0.2) | -0.56 (0.04) | -0.44 (0.14) | -0.5 (0.14) | -0.52 (0.16) | -0.46 (0.19) |  |
| intervened(1.0) | -0.51 (0.05) | -0.39 (0.15) | -0.42 (0.15) | -0.34 (0.2) | -0.31 (0.24) |  |
| 10000 | reconstructed | -0.51 (0.03) | -0.73 (0.04) | -0.73 (0.04) | -0.33 (0.05) | -0.33 (0.04) |
| intervened(0.2) | -0.5 (0.03) | -0.07 (0.15) | 0.07 (0.26) | 0.72 (0.48) | 0.87 (0.43) |  |
| intervened(1.0) | -0.47 (0.04) | 0.04 (0.2) | 0.15 (0.29) | 0.92 (0.47) | 1.1 (0.46) |  |
| 30000 | reconstructed | -0.51 (0.03) | -0.74 (0.04) | -0.73 (0.04) | -0.33 (0.05) | -0.33 (0.06) |
| intervened(0.2) | -0.49 (0.03) | -0.11 (0.08) | -0.17 (0.11) | 0.38 (0.4) | 0.33 (0.44) |  |
| intervened(1.0) | -0.44 (0.05) | -0.06 (0.12) | -0.13 (0.2) | 0.78 (0.43) | 0.69 (0.56) |  |
| -1 | reconstructed | -0.47 (0.02) | -0.71 (0.05) | -0.71 (0.04) | -0.25 (0.05) | -0.25 (0.06) |
| intervened(0.2) | -0.46 (0.03) | -0.06 (0.26) | 0.14 (0.3) | 0.98 (0.42) | 0.9 (0.52) |  |
| intervened(1.0) | -0.41 (0.06) | 0.05 (0.29) | 0.26 (0.36) | 1.13 (0.44) | 1.05 (0.51) |  |

Figure: Figure 4: Alignment of recovered latent variables with instrument, true representation [R,G,B], reward (Y) and digit for the IRAE model (Case 1 DGP). Data points with similar instrument, color, and reward are grouped together in the latent space.
Refer to caption: extracted/6565704/paper_images/irae3_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_True_50_scatterplot.png

Figure: Figure 5: Alignment of recovered latent variables with instrument, true representation [R,G,B], reward and digit for the Vanilla AE model (Case 1 DGP). Data points with same digits are grouped together in the latent space.
Refer to caption: extracted/6565704/paper_images/vanilla_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_scatterplot.png

Figure: Figure 6: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the IRAE trained model (Case 1 DGP).
Refer to caption: extracted/6565704/paper_images/irae3_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_True_50_examples.png

Figure: Figure 7: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the Vanilla AE trained model (Case 1 DGP).
Refer to caption: extracted/6565704/paper_images/vanilla_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_examples.png

Figure: Figure 8: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the IRAE trained model with latent dimension 32 (Case 2 DGP). We note that when we allow the latent dimension to be larger, we obtained better digit preservation. See appendix E.4 for hyperparameter tuning details.
Refer to caption: extracted/6565704/paper_images/irae3_22_config_dgp2_dense_32_1.0_linear_hsic_50_256_5_False_50_examples.png

## References

- Ahuja et al. [2022]
Kartik Ahuja, Divyat Mahajan, Vasilis Syrgkanis, and Ioannis Mitliagkas.
Towards efficient representation identification in supervised learning.
In *Conference on Causal Learning and Reasoning*, pages 19–43. PMLR, 2022.
- Andreu et al. [2024]
Oriol Corcoll Andreu, Athanasios Vlontzos, Michael O’Riordan, and Ciaran M Gilligan-Lee.
Contrastive representations of high-dimensional, structured treatments.
*arXiv preprint arXiv:2411.19245*, 2024.
- Bennett et al. [2022]
Andrew Bennett, Nathan Kallus, Xiaojie Mao, Whitney Newey, Vasilis Syrgkanis, and Masatoshi Uehara.
Inference on strongly identified functionals of weakly identified functions.
*arXiv preprint arXiv:2208.08291*, 2022.
- Bennett et al. [2023]
Andrew Bennett, Nathan Kallus, Xiaojie Mao, Whitney Newey, Vasilis Syrgkanis, and Masatoshi Uehara.
Source condition double robust inference on functionals of inverse problems.
*arXiv preprint arXiv:2307.13793*, 2023.
- Cui et al. [2024]
Yifan Cui, Hongming Pu, Xu Shi, Wang Miao, and Eric Tchetgen Tchetgen.
Semiparametric proximal causal inference.
*Journal of the American Statistical Association*, 119(546):1348–1359, 2024.
- [6]
Jing Dong, Pengyi Shi, Fanyin Zheng, and Xin Jin.
Capacity management in networks: A structural estimation approach for hospital inpatient wards.
- Dong et al. [2019]
Jing Dong, Pengyi Shi, Fanyin Zheng, and Xin Jin.
Off-service placement in inpatient ward network: Resource pooling versus service slowdown.
*Columbia Business School Research Paper Forthcoming*, 2019.
- Du et al. [2024]
Tianyu Du, Ayush Kanodia, Herman Brunborg, Keyon Vafa, and Susan Athey.
Labor-llm: Language-based occupational representations with large language models.
*arXiv preprint arXiv:2406.17972*, 2024.
- D’Haultfoeuille [2011]
Xavier D’Haultfoeuille.
On the completeness condition in nonparametric instrumental problems.
*Econometric Theory*, 27(3):460–471, 2011.
- Gretton et al. [2007]
Arthur Gretton, Kenji Fukumizu, Choon Teo, Le Song, Bernhard Schölkopf, and Alex Smola.
A kernel statistical test of independence.
*Advances in neural information processing systems*, 20, 2007.
- Hälvä and Hyvarinen [2020]
Hermanni Hälvä and Aapo Hyvarinen.
Hidden markov nonlinear ica: Unsupervised learning from nonstationary time series.
In *Conference on Uncertainty in Artificial Intelligence*, pages 939–948. PMLR, 2020.
- Hälvä et al. [2024]
Hermanni Hälvä, Jonathan So, Richard E Turner, and Aapo Hyvärinen.
Identifiable feature learning for spatial data with nonlinear ica.
In *International Conference on Artificial Intelligence and Statistics*, pages 3331–3339. PMLR, 2024.
- Harada and Kashima [2021]
Shonosuke Harada and Hisashi Kashima.
Graphite: Estimating individual effects of graph-structured treatments.
In *Proceedings of the 30th ACM international conference on information & knowledge management*, pages 659–668, 2021.
- Hyvärinen [2013]
Aapo Hyvärinen.
Independent component analysis: recent advances.
*Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences*, 371(1984):20110534, 2013.
- Hyvärinen and Oja [2000]
Aapo Hyvärinen and Erkki Oja.
Independent component analysis: algorithms and applications.
*Neural networks*, 13(4-5):411–430, 2000.
- Hyvärinen et al. [2023]
Aapo Hyvärinen, Ilyes Khemakhem, and Hiroshi Morioka.
Nonlinear independent component analysis for principled disentanglement in unsupervised deep learning.
*Patterns*, 4(10), 2023.
- Hyvärinen et al. [2024]
Aapo Hyvärinen, Ilyes Khemakhem, and Ricardo Monti.
Identifiability of latent-variable and structural-equation models: from linear to nonlinear.
*Annals of the Institute of Statistical Mathematics*, 76(1):1–33, 2024.
- Jin and Syrgkanis [2023]
Jikai Jin and Vasilis Syrgkanis.
Learning causal representations from general environments: Identifiability and intrinsic ambiguity.
*arXiv preprint arXiv:2311.12267, (to appear at NeurIPS24)*, 2023.
- Kelly et al. [2020]
Bryan T Kelly, Seth Pruitt, and Yinan Su.
Instrumented principal component analysis.
*Available at SSRN 2983919*, 2020.
- Khemakhem et al. [2020]
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen.
Variational autoencoders and nonlinear ica: A unifying framework.
In *International conference on artificial intelligence and statistics*, pages 2207–2217. PMLR, 2020.
- Lopez et al. [2020]
Romain Lopez, Chenchen Li, Xiang Yan, Junwu Xiong, Michael Jordan, Yuan Qi, and Le Song.
Cost-effective incentive allocation via structured counterfactual inference.
In *Proceedings of the AAAI conference on artificial intelligence*, volume 34, pages 4997–5004, 2020.
- Lu et al. [2021]
Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobato, and Bernhard Schölkopf.
Invariant causal representation learning for out-of-distribution generalization.
In *International Conference on Learning Representations*, 2021.
- Monti et al. [2020]
Ricardo Pio Monti, Kun Zhang, and Aapo Hyvärinen.
Causal discovery with general non-linear relationships using non-linear ica.
In *Uncertainty in artificial intelligence*, pages 186–195. PMLR, 2020.
- Muandet et al. [2020]
Krikamol Muandet, Wittawat Jitkrittum, and Jonas Kübler.
Kernel conditional moment test via maximum moment restriction.
In *Conference on Uncertainty in Artificial Intelligence*, pages 41–50. PMLR, 2020.
- Nabi et al. [2022]
Razieh Nabi, Todd McNutt, and Ilya Shpitser.
Semiparametric causal sufficient dimension reduction of multidimensional treatments.
In *Uncertainty in Artificial Intelligence*, pages 1445–1455. PMLR, 2022.
- Pfister et al. [2018]
Niklas Pfister, Peter Bühlmann, Bernhard Schölkopf, and Jonas Peters.
Kernel-based tests for joint independence.
*Journal of the Royal Statistical Society Series B: Statistical Methodology*, 80(1):5–31, 2018.
- Qin et al. [2023]
Jimmy Qin, Carri W Chan, Jing Dong, Shunichi Homma, and Siqin Ye.
Waiting online versus in-person: An empirical study on outpatient clinic visit incompletion.
2023.
- Rao [1964]
C Radhakrishna Rao.
The use and interpretation of principal component analysis in applied research.
*Sankhyā: The Indian Journal of Statistics, Series A*, pages 329–358, 1964.
- Sabatier et al. [1989]
Robert Sabatier, Jean-Dominique Lebreton, and D Chessel.
Principal component analysis with instrumental variables as a tool for modelling composition data.
*Multiway data analysis*, pages 341–352, 1989.
- Saengkyongam et al. [2023]
Sorawit Saengkyongam, Elan Rosenfeld, Pradeep Ravikumar, Niklas Pfister, and Jonas Peters.
Identifying representations for intervention extrapolation.
*arXiv preprint arXiv:2310.04295*, 2023.
- Schölkopf et al. [2021a]
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio.
Toward causal representation learning.
*Proceedings of the IEEE*, 109(5):612–634, 2021a.
- Schölkopf et al. [2021b]
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio.
Toward causal representation learning.
*Proceedings of the IEEE*, 109(5):612–634, 2021b.
- Vafa et al. [2024]
Keyon Vafa, Susan Athey, and David M Blei.
Estimating wage disparities using foundation models.
*arXiv preprint arXiv:2409.09894*, 2024.
- Wang [2024]
Cong Wang.
Counterfactual and synthetic control method: Causal inference with instrumented principal component analysis.
*arXiv preprint arXiv:2408.09271*, 2024.
- Y Takane [2001]
M Hunter Y Takane.
Constrained principal component analysis: A comprehensive theory.
2001.
URL [https://doi.org/10.1007/s002000100081](https://doi.org/10.1007/s002000100081).

## Appendix A Further Related Work

In this section we provide a more discussion on related work that is not covered in the main text.

Identifying Representations for Intervention
Extrapolation Similar to our work, Saengkyongam et al. proposed the Rep4Ex approach which tries to solve the task of interventional outcome prediction by identifying the SCM. Importantly, although they work with a similar SCM as we do (Equation [1](https://arxiv.org/html/2506.02200v2#S2.E1)), the level of intervention differs - our work considers interventions on the latent treatment space ($D$), while Saengkyongam et al. considers intervening on $Z$ (using notations in Equation [1](https://arxiv.org/html/2506.02200v2#S2.E1)). Moreover, our work is motivated by the presence of unobserved confounding between the latent representation of the treatment and the outcome, whereas their work is motivated by the need to extrapolate to unseen interventions, while the treatment that they consider is fully exogenous. Like our approach, they employ autoencoders to learn latent representations from potentially high-dimensional observed features, but use maximum moment restriction (MMR) regularization [24] to enforce the constraint $E[e_{D}(X)-AZ|Z]=0$. This can be achieved when $E[e_{D}(X)-AZ]=0$ and $e_{D}(X)-AZ\perp\!\!\!\perp Z$, corresponding to our $\lambda$ and $\mu_{1}$ term in [Equation IRAE](https://arxiv.org/html/2506.02200v2#S4.Ex12). Additionally, while Rep4Ex assumes a deterministic mixing function from the latent representation to the observables $X$, our method explicitly handles noisy observations of X through $e_{V}(X)$, which allows for broader generalization.

Dimensionality Reduction for High Dimensional Treatments
When learning a representation for the treatment, it is important for the learned representation to capture all causal factors so that the causal relationship is preserved for downstream estimation tasks like treatment effect estimation. Nabi et al.utilize semi-parametric inference theory for structural models to provide a generalized the sufficient dimension reduction approach for learning lower-dimensional representation for treatment, while capturing the relationship between the treatment and the mean counterfactual outcome. Andreu et al. employed a contrastive approach to learn a representation of the high-dimensional treatments. These works studied settings that did not involve the presence of unobserved confounders of the treatment, while we focus on heavily confounded high dimensional structured treatments. Moreover, in these works, the selection of causally relevant factors are guided by the outcome, where as we take an inherently different approach that learns the latent representations using auxiliary information from instrumental variables instead of the treatment.

Independence Conditions In our work, we show that independence between certain variables (for more details, see Theorem [4.5](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem5)) is desirable for identification. We enforce the independence condition by incorporating a Hilbert-Schmidt Independence Criterion (HSIC) [10] regularizer. This approach has also been adopted in prior research: for instance, Lopez et al. employed HSIC regularization to mitigate bias in observational datasets for applications in counterfactual policy optimization, while Harada and Kashima use it to learn a representations of the treatment that is independent with the target individual in order to mitigate selection bias.

## Appendix B Proof of Linear Identification

Before proving the main theorem, we first present some useful lemma.

###### Lemma B.1 .

Suppose $A$ is a $n\times k$ matrix with full row rank ($k>n$), and $B$ is a $m\times n$ matrix, with full column rank ($m>n$). Then the columns of $C=BA$ spans the same space as the columns of $B$.

###### Proof of Lemma B.1 .

Let $\mathcal{R}(\cdot)$ denote the column space of a matrix.

For any $x\in\mathcal{R}(B)$, there exist vector $y$ such that $x=By$. Since $A$ is full row rank, we know that $AA^{+}=I_{n}$, and $x=By=BAA^{+}y=C(A^{+}y)$. Therefore $x\in\mathcal{R}(C)$, so $\mathcal{R}(B)\subseteq\mathcal{R}(C)$.

Similarly, for any $x\in\mathcal{R}(C)$, there exist vector $y$ such that $x=BAy=B(Ay)$. So $x\in\mathcal{R}(B)$, and we have $\mathcal{R}(C)\subseteq\mathcal{R}(B)$.

Together, we have $\mathcal{R}(C)=\mathcal{R}(B)$.

∎

Now we proceed to prove Theorem [3.1](https://arxiv.org/html/2506.02200v2#S3.Thmtheorem1).

###### Proof of Theorem 3.1 .

From Equation [9](https://arxiv.org/html/2506.02200v2#S3.E9), we have that:

$$ $\displaystyle X=BAZ+B_{\perp}V+BU$ $$

Then taking the conditional expectation over $Z$, we have:

$$ $\displaystyle\mathbb{E}[X|Z]=~{}$ $\displaystyle BAZ+\mathbb{E}[B_{\perp}V+BU]$ $\displaystyle=~{}$ $\displaystyle BAZ+\mathbb{E}[B_{\perp}\mathbb{E}[V|Z]]+\mathbb{E}[B\mathbb{E}[ U|Z]]$ $\displaystyle=~{}$ $\displaystyle BAZ+B_{\perp}\mathbb{E}[V]+B\mathbb{E}[U]$ (Since $V\perp\!\!\!\perp Z$ and $U\perp\!\!\!\perp Z$ ) $\displaystyle=~{}$ $\displaystyle BAZ$ $$

Thus $C:=BA$ can be uniquely identified as the solution to the linear regression problem, regressing $X$ on $Z$.
Consider the SVD decomposition of $C={\cal U}\Sigma{\cal V}^{\top}$. Let $\hat{B}={\cal U}$, and $\hat{A}=\Sigma{\cal V}^{\top}$. Then by Lemma [B.1](https://arxiv.org/html/2506.02200v2#A2.Thmtheorem1), we have that the columns of $\hat{B}$ spans the same space as the columns of $B$. In other words, there exist an invertible change of basis matrix $P$ such that $B=\hat{B}P$. Since $\hat{B}$ is orthonormal (by construction of SVD), we have that $\hat{B}^{T}\hat{B}=I_{r}$, and $P=\hat{B}^{T}B$. As a result, we also have:

$$ $\displaystyle D=~{}$ $\displaystyle B^{+}X=(B^{T}B)^{-1}B^{T}X$ $\displaystyle=~{}$ $\displaystyle(P^{T}\hat{B}^{T}\hat{B}P)^{-1}P^{T}\hat{B}^{T}X$ $\displaystyle=~{}$ $\displaystyle(P^{T}P)^{-1}P^{T}\hat{B}^{T}X$ $\displaystyle=~{}$ $\displaystyle P^{-1}\hat{B}^{T}X=P^{-1}\tilde{D}$ $$

Next, we show that $\tilde{\theta}=(P^{-1})^{T}\theta$. The LIRR algorithm solves for $\tilde{\theta}$ from the following moment equation:

$$ $\displaystyle 0=~{}$ $\displaystyle\mathbb{E}[Z(Y-\tilde{\theta}^{T}\hat{D})]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[Z(\theta^{T}D+\eta(V,U,\epsilon)-\hat{\theta}^{T}PD)]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[Z(\theta^{T}D-\tilde{\theta}^{T}PD)]$ (Since $U,V,\epsilon\perp\!\!\!\perp Z$ and $\mathbb{E}[\eta(U,v,\epsilon)]=0$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}[ZD^{T}](\theta-P^{T}\tilde{\theta})$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[Z(Z^{T}A^{T}+U^{T})](\theta-P^{T}\tilde{\theta})$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[ZZ^{T}]A^{T}(\theta-P^{T}\tilde{\theta})$ $$

Since the instruments are not co-linear, we have that $\mathbb{E}[ZZ^{T}]\succ 0$, i.e. $\mathbb{E}[ZZ^{T}]$ is invertible. Thus $\mathbb{E}[ZZ^{T}]A^{T}(\theta-P^{T}\tilde{\theta})=0$ if and only if $A^{T}(\theta-P^{T}\tilde{\theta})=0$. Since $A^{T}$ has full column rank, then by the Rank-Nullity theorem, the null space of $A^{T}={0}$. Together, this shows that $\tilde{\theta}=(P^{-1})^{T}\theta$ is the unique solution to the moment condition.

Lastly, we show that the intervened outcome is guaranteed improvement in expectation. Consider an intervention in the direction of $u=\tilde{\theta}/\|\tilde{\theta}\|$ in the $\tilde{D}$ space, this maps to an intervention in the $D$ space as:

$$ $\displaystyle e_{D}(t(X))=~{}$ $\displaystyle B^{+}t(X)=D+\alpha B^{+}\hat{B}\tilde{\theta}$ $\displaystyle=~{}$ $\displaystyle D+\alpha P^{-1}\frac{\tilde{\theta}}{\|\tilde{\theta}\|}=D+ \alpha P^{-1}\frac{(P^{-1})^{\top}\theta}{\|(P^{-1})^{\top}\theta\|}$ $$

Since, we intervene only in $D$, $e_{V}(t(X))=V$.
Then, we can compute the intervened outcome:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}]=~{}$ $\displaystyle\mathbb{E}[\theta^{T}e_{D}(t(X))+\eta(e_{V}(t(X)),U,\epsilon)]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\theta^{T}e_{D}(t(X))]$ ( $e_{V}(t(X))=V$ , and $\mathbb{E}[\eta(U,v,\epsilon)]=0$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}\left[\theta^{T}\left(D+\alpha P^{-1}\frac{(P^{-1})^{ \top}\theta}{\|(P^{-1})^{\top}\theta\|}\right)\right]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\theta^{T}D+\alpha\|(P^{-1})^{\top}\theta\|]=\mathbb{E }[Y]+\alpha\|(P^{-1})^{\top}\theta\|$ $$

∎

###### Lemma B.1 .

###### Proof of Lemma B.1 .

###### Proof of Theorem 3.1 .

## Appendix C Proof of Non-linear Identification

###### Proof of Theorem 4.5 .

By definition of $(\tilde{D},\tilde{V})$, we have:

$$ $\displaystyle(\tilde{D},\tilde{V})=\tilde{e}(X)=\tilde{e}\circ f(D,V)=:q(D,V)$ $$

Denote with $q_{1}(D,V)$ the $\tilde{D}$ component of the output of $q$ and $q_{2}$ the $\tilde{V}$ component.

Since we have that $\tilde{D}=\tilde{A}Z+\tilde{U}$, with $\tilde{U}\perp\!\!\!\perp Z$ and $\mathbb{E}[\tilde{U}]=0$, we can write:

$$ $\displaystyle\mathbb{E}[\tilde{D}\mid Z=z]=\mathbb{E}[\tilde{A}Z+\tilde{U}\mid Z =z]=\tilde{A}z$ $$

Moreover:

$$ $\displaystyle\mathbb{E}[\tilde{D}\mid Z=z]=~{}$ $\displaystyle\mathbb{E}[q_{1}(D,V)\mid Z=z]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[q_{1}(Az+U,V)\mid Z=z]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[q_{1}(Az+U,V)]$ ( $Z\perp\!\!\!\perp\{U,V\}$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}_{U}[\mathbb{E}_{V}[q_{1}(Az+U,V)]]$ ( $U\perp\!\!\!\perp V$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}_{U}[\tilde{q}_{1}(Az+U)]$ ( $\tilde{q}_{1}(d)\triangleq\mathbb{E}_{V}[q_{1}(d,V)]$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}(Az+U)]$ $$

Thus we can conclude that:

$$ $\displaystyle\tilde{A}z=\mathbb{E}[\tilde{D}\mid Z=z]=\mathbb{E}[\tilde{q}_{1} (Az+U)]$ $$

Since this holds for all $z\in{\cal Z}$ and since ${\cal Z}$ is an open set, we can take the derivative with respect to $z$, to derive:

$$ $\displaystyle\forall z\in{\cal Z}:\tilde{A}=\partial_{z}\mathbb{E}[\tilde{q}_{ 1}(Az+U)]$ $$

Since $q_{1}$ is continuously differentiable with bounded derivatives, the same holds for $\tilde{q}_{1}$ and therefore we can exchange the order of differentiation and expectation:

$$ $\displaystyle\tilde{A}=\mathbb{E}[\partial_{z}\tilde{q}_{1}(Az+U)]$ $$

Letting $\tilde{q}_{1}^{(1)}$ denote the Jacobian of the function $\tilde{q}_{1}(d)$, we can write by the chain rule:

$$ $\displaystyle\tilde{A}=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}^{(1)}(Az+U)A]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}^{(1)}(Az+U)]A$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}^{(1)}(Az+U)\mid Z=z]A$ ( $Z\perp\!\!\!\perp U$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}^{(1)}(AZ+U)\mid Z=z]A$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}^{(1)}(D)\mid Z=z]A$ $$

Since $A$ is full row rank, we have that $AA^{+}$ is invertible. Thus we can write:

$$ $\displaystyle\tilde{A}A^{+}=~{}$ $\displaystyle\mathbb{E}[\tilde{q}^{(1)}(D)\mid Z=z]$ $$

or equivalently:

$$ $\displaystyle\forall z\in{\cal Z}:\mathbb{E}[\tilde{q}_{1}^{(1)}(D)-\tilde{A}A ^{+}\mid Z=z]=0$ $$

By the bounded completeness assumption and since both $\tilde{A}A^{+}$ and $\tilde{q}_{1}^{(1)}$ are bounded, the latter implies that:

$$ $\displaystyle\forall d\in{\cal D}:\tilde{q}_{1}^{(1)}(d)=\tilde{A}A^{+}$ $$

or equivalently that:

$$ $\displaystyle\tilde{q}_{1}(d)=\tilde{A}A^{+}d+\tilde{\nu}$ $$

for some constant vector $\nu$. Moreover,

$$ $\displaystyle\mathbb{E}[\tilde{D}]=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}(D)]$ $\displaystyle=~{}$ $\displaystyle\tilde{A}A^{+}\mathbb{E}[D]+\tilde{\nu}$ $\displaystyle=~{}$ $\displaystyle\tilde{A}A^{+}A\mathbb{E}[Z]+\tilde{A}A^{+}\mathbb{E}[U]+\tilde{\nu}$ $\displaystyle=~{}$ $\displaystyle\tilde{\nu}$ $$

But we also have $\mathbb{E}[\tilde{D}]=A\mathbb{E}[Z]+\mathbb{E}[\tilde{U}]=0$.
Hence, we have that $\tilde{\nu}=0$. Thus:

$$ $\displaystyle\forall d\in{\cal D}:\tilde{q}_{1}(d)=\tilde{A}A^{+}d$ $$

Next, we argue that $\tilde{A}A^{+}$ is an invertible matrix. Note that:

$$ $\displaystyle\mathbb{E}[\tilde{D}Z^{\top}]=~{}$ $\displaystyle\mathbb{E}[(\tilde{A}Z+\tilde{U})Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\tilde{A}\mathbb{E}[ZZ^{\top}]$ ( $\tilde{U}\perp\!\!\!\perp Z$ , $\mathbb{E}[Z]=0$ , $\mathbb{E}[\tilde{U}=0]$ ) $$

Moreover:

$$ $\displaystyle\mathbb{E}[\tilde{D}Z^{\top}]=~{}$ $\displaystyle\mathbb{E}[q_{1}(D,V)Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[q_{1}(AZ+U,V)Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\mathbb{E}[q_{1}(AZ+U,V)\mid Z,U]Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}(AZ+U)Z^{\top}]$ ( $Z\perp\!\!\!\perp U\perp\!\!\!\perp V$ ) $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\tilde{q}_{1}(D)Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[(\tilde{A}A^{+}D)Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\tilde{A}A^{+}\mathbb{E}[DZ^{\top}]$ $\displaystyle=~{}$ $\displaystyle\tilde{A}A^{+}\mathbb{E}[(AZ+U)Z^{\top}]$ $\displaystyle=~{}$ $\displaystyle\tilde{A}A^{+}A\mathbb{E}[ZZ^{\top}]$ ( $Z\perp\!\!\!\perp U$ , $\mathbb{E}[U]=0$ ) $$

Thus we have concluded that:

$$ $\displaystyle\tilde{A}\mathbb{E}[ZZ^{\top}]=\mathbb{E}[\tilde{D}Z^{\top}]= \tilde{A}A^{+}A\mathbb{E}[ZZ^{\top}]$ $$

Since $\mathbb{E}[ZZ^{\top}]$ is assumed to be invertible, the latter implies that:

$$ $\displaystyle\tilde{A}=\tilde{A}A^{+}A$ $$

By Lemma [C.1](https://arxiv.org/html/2506.02200v2#A3.Thmtheorem1) in Appendix [C.2](https://arxiv.org/html/2506.02200v2#A3.SS2), since $\tilde{A}$ and $A$ have full row rank, the row span of $\tilde{A}$ is equal to the row span of $A$ and the matrix $\tilde{A}A^{+}$ is invertible.

We have thus concluded that:

$$ $\displaystyle\forall d\in{\cal D}:\tilde{q}_{1}(d)=\tilde{A}A^{+}d$ $$

and $\tilde{A}A^{+}$ is invertible.

Consider any solution with perfect encoder-decoder pair $(\tilde{e},\tilde{f})$, and $\tilde{A}$ that satisfies the conditions of the theorem and minimizes the objective function:

$$ $\displaystyle\mathbb{E}[\|\tilde{e}_{D}(X)-\tilde{A}Z\|^{2}]=\mathbb{E}[\| \tilde{D}-\tilde{A}Z\|^{2}]$ $$

For any feasible solution, we can decompose this objective into two components by centering around

$$ $\mu_{\tilde{A}}(d)\triangleq\tilde{A}A^{+}d$ $$

i.e.:

$$ $\displaystyle\mathbb{E}[\|\tilde{D}-\tilde{A}Z\|^{2}]=~{}$ $\displaystyle\mathbb{E}[\|\tilde{D}-\mu_{\tilde{A}}(D)+\mu_{\tilde{A}}(D)- \tilde{A}Z\|^{2}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\|\tilde{D}-\mu_{\tilde{A}}(D)\|^{2}+\|\mu_{\tilde{A}} (D)-\tilde{A}Z\|^{2}]+2\mathbb{E}[(\tilde{D}-\mu_{\tilde{A}}(D))^{\top}(\mu_{ \tilde{A}}(D)-\tilde{A}Z)]$ $$

Consider the inner product term. Since we have that:

$$ $\displaystyle\mathbb{E}[\tilde{D}-\mu_{\tilde{A}}(D)\mid D,Z]=~{}$ $\displaystyle\mathbb{E}[q_{1}(D,V)-\tilde{q}_{1}(D)\mid D,Z]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[q_{1}(D,V)\mid D,Z]-\tilde{q}_{1}(D)$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[q_{1}(D,V)\mid D,Z]-\mathbb{E}[q_{1}(D,V)\mid D]$ $$

Since $Z\perp\!\!\!\perp U\perp\!\!\!\perp V$, we have by Lemma [C.2](https://arxiv.org/html/2506.02200v2#A3.Thmtheorem2) in Appendix [C.2](https://arxiv.org/html/2506.02200v2#A3.SS2) that $Z\perp\!\!\!\perp V\mid\mathbbm{1}\{AZ+U=d\}$:

$$ $\displaystyle\mathbb{E}[q_{1}(D,V)\mid D=d,Z]=\mathbb{E}[q_{1}(d,V)\mid D=d,Z] =\mathbb{E}[q_{1}(d,V)\mid D=d]=\mathbb{E}[q_{1}(D,V)\mid D=d]$ $$

Thus:

$$ $\displaystyle\mathbb{E}[\tilde{D}-\mu_{\tilde{A}}(D)\mid D,Z]=0$ $$

From this we conclude that for any feasible solution $\tilde{e},\tilde{f},\tilde{A}$, we have that the objective can be decomposed as:

$$ $\displaystyle\mathbb{E}[\|\tilde{e}_{D}(X)-\tilde{A}Z\|^{2}]=~{}$ $\displaystyle\mathbb{E}[\|\tilde{D}-\mu_{\tilde{A}}(D)\|^{2}]+\mathbb{E}[\|\mu _{\tilde{A}}(D)-\tilde{A}Z\|^{2}]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[\|q_{1}(D,V)-\mu_{\tilde{A}}(D)\|^{2}]+\mathbb{E}[\| \mu_{\tilde{A}}(D)-\tilde{A}Z\|^{2}]$ $$

Suppose that with positive probability, we have that $q_{1}(D,V)\neq\mu_{\tilde{A}}(D)=\tilde{A}A^{+}D$. Then we have that:

$$ $\displaystyle\mathbb{E}[\|q_{1}(D,V)-\mu_{\tilde{A}}(D)\|^{2}]>0$ $$

In this case, we will provide an alternative feasible solution, which achieves smaller objective than $\tilde{e},\tilde{f},\tilde{A}$. Consider the solution:

$$ $\displaystyle\tilde{e}^{\prime}(x)=~{}$ $\displaystyle(\tilde{A}A^{+}e_{D}(x),e_{V}(x))$ $\displaystyle\tilde{f}^{\prime}(d,v)=~{}$ $\displaystyle f((\tilde{A}A^{+})^{-1}d,v)$ $$

Note that we used the fact that for any feasible solution, we have already shown that $\tilde{A}A^{+}$ is invertible. Moreover, note that for this solution we have that, $\tilde{e}^{\prime},\tilde{f}^{\prime}$ is invertible, since $e,f$ is invertible and $\tilde{A}A^{+}$ is invertible. Finally,

$$ $\displaystyle\tilde{f}^{\prime}\circ\tilde{e}^{\prime}(x)=~{}$ $\displaystyle x$ $\displaystyle\tilde{e}^{\prime}(f(d,v))=~{}$ $\displaystyle(\tilde{A}A^{+}e_{D}(f(d,v)),e_{V}(f(d,v)))=(\tilde{A}A^{+}d,v)$ $$

Thus:

$$ $\displaystyle\tilde{D}=\tilde{e}_{D}^{\prime}(X)=\tilde{e}_{D}^{\prime}(f(D,V) )=\tilde{A}A^{+}D=\tilde{A}A^{+}AZ+\tilde{A}A^{+}U=\tilde{A}Z+\tilde{A}A^{+}U$ $$

Where we used the fact that we have already shown (in the proof of Theorem [4.5](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem5)) that $\tilde{A}A^{+}A=\tilde{A}$. Thus, we also have that:

$$ $\displaystyle\tilde{D}=\tilde{A}Z+\tilde{U}$ $$

where $\tilde{U}=\tilde{A}A^{+}U$ and satisfies $\tilde{U}\perp\!\!\!\perp Z$ and $\mathbb{E}[\tilde{U}]=0$.

Therefore, this new solution is a feasible solution. Moreover, since under this solution we have that $\tilde{D}=\tilde{A}A^{+}D=\mu_{\tilde{A}}(D)$, the first part of the objective vanishes and the objective takes the value:

$$ $\displaystyle\mathbb{E}[\|\mu_{\tilde{A}}(D)-\tilde{A}Z\|^{2}]<\mathbb{E}[\|q_ {1}(D,V)-\mu_{\tilde{A}}(D)\|^{2}]+\mathbb{E}[\|\mu_{\tilde{A}}(D)-\tilde{A}Z \|^{2}]$ $$

contradicting the optimality of the original solution.

Thus we have derived that for any optimal feasible solution, it must hold that with probability $1$:

$$ $\displaystyle\tilde{D}=q_{1}(D,V)=\mu_{\tilde{A}}(D)=\tilde{A}A^{+}D$ (12) $$

with $\tilde{A}A^{+}$ an invertible matrix. Moreover, this implies that $\tilde{U}=\tilde{D}-\tilde{A}Z=\tilde{A}A^{+}U$.
∎

###### Proof of Lemma 4.6 .

In this proof, we argue about the properties of the second part of the function $q$. Note that since $\tilde{D}\perp\!\!\!\perp\tilde{V}$ and since $\tilde{D}=PD$, for some invertible $P$, with probability $1$, we have that $\tilde{V}\perp\!\!\!\perp D$. Thus:

$$ $\displaystyle q_{2}(D,V)\equiv\tilde{V}\perp\!\!\!\perp D$ $$

Since, $D\perp\!\!\!\perp V$, this implies that $\text{Law}(q_{2}(d,V))=\text{Law}(q_{2}(d^{\prime},V))$ for all $d,d^{\prime}\in{\cal D}$.

Since $\tilde{e}$ is a bijection when restricted to inputs that are outputs of $f$ and since $f$ is an injection, we have that $\tilde{e}\circ f$ is an injection. Thus there exists a well-defined inverse function $q^{-1}=e\circ\tilde{f}$, such that $D,V=q^{-1}(D,V)$. Let $q_{2}^{-1}$ be the $V$ component of its output. Since $D\perp\!\!\!\perp V$ and $\tilde{D}=PD$, we have that:

$$ $\displaystyle V\equiv q_{2}^{-1}(\tilde{D},\tilde{V})\perp\!\!\!\perp\tilde{D}$ $$

Since $\tilde{D}\perp\!\!\!\perp\tilde{V}$, this implies that $\text{Law}(q_{2}^{-1}(d,\tilde{V}))=\text{Law}(q_{2}^{-1}(d^{\prime},\tilde{V}))$ for all $d,d^{\prime}\in{\cal D}$
∎

###### Proof of Theorem 4.5 .

###### Proof of Lemma 4.6 .

### C.1 Proof of Positive Improvement

###### Proof of Theorem 4.8 .

In this proof, we show that intervention in the direction of average derivatives of $\tilde{h}$ guarantees positive improvement for sufficiently small $\alpha$, assuming that $h$ is twice differentiable. If we perform the intervention $\tilde{D}+\alpha u$, then we have by Lemma [4.6](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem6) that:

$$ $\displaystyle D_{\alpha u},V_{\alpha u}=(D+\alpha P^{-1}u,q_{2}^{-1}(\tilde{D} +\alpha u,\tilde{V}))$ $$

Since ${\cal D}=\mathbb{R}^{r}$ and since $P$ is an invertible matrix, we have that ${\cal\tilde{D}}=\mathbb{R}^{r}$. Thus, for all $d\in{\cal\tilde{D}}$, we also have that $d+au\in{\cal\tilde{D}}$. By Lemma [4.6](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem6), have that for all $d\in{\cal\tilde{D}}$:

$$ $\displaystyle\text{Law}(q_{2}^{-1}(d,\tilde{V}))=\text{Law}(q_{2}^{-1}(d+ \alpha u,\tilde{V}))$ $$

By Theorem [4.5](https://arxiv.org/html/2506.02200v2#S4.Thmtheorem5), we have that, with probability $1$, $\tilde{D}=PD$ and $\tilde{U}=PU$. Moreover, by assumption, we have that $\tilde{V}\perp\!\!\!\perp\tilde{U}\perp\!\!\!\perp Z$, which implies

$$ $\tilde{V}\perp\!\!\!\perp\{\tilde{A}Z+\tilde{U},P^{-1}\tilde{U}\}\implies \tilde{V}\perp\!\!\!\perp\{\tilde{D},U\},$ $$

we also have that:

$$ $\text{Law}(q_{2}^{-1}(d,\tilde{V})\mid\tilde{D}=d,U)=\text{Law}(q_{2}^{-1}(d+ \alpha u,\tilde{V})\mid\tilde{D}=d,U)\\ \implies\text{Law}(q_{2}^{-1}(\tilde{D},\tilde{V})\mid\tilde{D},U)=\text{Law}( q_{2}^{-1}(\tilde{D}+\alpha u,\tilde{V})\mid\tilde{D},U)$ $$

which by the definition of $V$ and $V_{\alpha}$ is equivalent to:;

$$ $\displaystyle\text{Law}(V\mid\tilde{D},U)=\text{Law}(V_{\alpha u}\mid\tilde{D} ,U)\implies\text{Law}(V\mid U)=\text{Law}(V_{\alpha u}\mid U)$ $$

By the outcome structural equation

$$ $\displaystyle Y=h(D)+\eta(U,V,\epsilon)$ $$

we have that:

$$ $\displaystyle Y_{\alpha u}=h(D+\alpha P^{-1}u)+\eta(U,V_{\alpha u},\epsilon_{Y})$ $$

and that:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}-Y]=~{}$ $\displaystyle\mathbb{E}[h(D+\alpha P^{-1}u)-h(D)]+\mathbb{E}[\eta(U,V_{\alpha u },\epsilon)-\eta(U,V,\epsilon)]$ $$

Since $\epsilon\perp\!\!\!\perp\{Z,U,V\}$ and since $V_{\alpha u}$ is a measurable function of these random variables, we have that $\epsilon\perp\!\!\!\perp\{V_{\alpha u},V,U\}$. Letting $\tilde{\eta}(u,v)=\mathbb{E}_{\epsilon}[\eta(u,v,\epsilon)]$, we can write:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}-Y]=~{}$ $\displaystyle\mathbb{E}[h(D+\alpha P^{-1}u)-h(D)]+\mathbb{E}[\tilde{\eta}(U,V_ {\alpha u})-\tilde{\eta}(U,V)]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[h(D+\alpha P^{-1}u)-h(D)]+\mathbb{E}[\mathbb{E}[\tilde {\eta}(U,V_{\alpha u})-\tilde{\eta}(U,V)\mid U]]$ $\displaystyle=~{}$ $\displaystyle\mathbb{E}[h(D+\alpha P^{-1}u)-h(D)]$ ( $\text{Law}(V\mid U)=\text{Law}(V_{\alpha u}\mid U)$ ) $$

By a first-order Taylor expansion and since $h$ is twice differentiable with bounded first and second derivatives:

$$ $\displaystyle\mathbb{E}[Y_{\alpha u}-Y]=\mathbb{E}[\alpha\nabla_{D}h(D)^{\top} P^{-1}u]+O(\alpha^{2})=\alpha\|(P^{-1})^{\top}\mathbb{E}[\nabla_{D}h(D)]\|+O( \alpha^{2})$ $$

∎

###### Proof of Theorem 4.8 .

### C.2 Auxiliary Lemmas

###### Lemma C.1 .

Suppose ${A}$ and $B$ are $r\times k$ matrices with full row rank. If $A=AB^{+}B$, then $\text{rowspan}({A})=\text{rowspan}(B)$ and ${A}B^{+}$ is invertible.

###### Proof.

Consider the thin SVDs of $B=U_{B}\Sigma_{B}V_{B}^{\top}$ and $A=U_{A}\Sigma_{A}V_{A}^{\top}$. Then

$$ $\displaystyle B^{+}B$ $\displaystyle=V_{B}\Sigma_{B}^{-1}U_{B}^{\top}U_{B}\Sigma_{B}V_{B}^{\top}=V_{B }V_{B}^{\top}$ $$

is the projection onto the row space of $B$. Then, we have:

$$ $\displaystyle A=AB^{+}B\quad\Leftrightarrow\quad AV_{B}V_{B}^{\top}=A$ (13) $$

First, we prove by contradiction that $\text{rowspan}(A)=\text{rowspan}(B)$. Suppose $x\in\text{rowspan}(A)=\text{span}(V_{A})$, but $x\notin\text{rowspan}(B)=\text{span}(V_{B})$. Let $V_{B}^{\perp}$ denote an orthogonal completion of $V_{B}$, then

$$ $\displaystyle x\notin\text{span}(V_{B})$ $\displaystyle\Rightarrow x=V_{B}V_{B}^{\top}x+V_{B}^{\perp}(V_{B}^{\perp})^{ \top}x$ $$

where $V_{B}^{\perp}(V_{B}^{\perp})^{\top}x\neq 0$, which implies $(V_{B}^{\perp})^{\top}x\neq 0$ as $V_{B}^{\perp}$ is orthogonal. Hence, we have the following:

$$ $\displaystyle\|x\|^{2}=x^{\top}x=x^{\top}V_{B}V_{B}^{\top}x+x^{\top}V_{B}^{ \perp}(V_{B}^{\perp})^{\top}x=\|V_{B}^{\top}x\|^{2}+\|(V_{B}^{\perp})^{\top}x \|^{2}>\|V_{B}^{\top}x\|^{2}$ $$

However, we also have that $AV_{B}V_{B}^{\top}x-Ax=0$, which implies $u\triangleq V_{B}V_{B}^{\top}x-x\in\text{null-space}(A)=\text{span}(V_{A}^{
\perp})$. Thus, it should be orthogonal with $x\in\text{span}(V_{A})$.

$$ $\displaystyle 0=x^{\top}u=x^{\top}(V_{B}V_{B}^{\top}x-x)=\|V_{B}^{\top}x\|^{2} -\|x\|^{2}\neq 0$ $$

This yields a contradiction! Thus, $\text{rowspan}(A)\subseteq\text{rowspan}(B)$. Since, both matrices have full row rank, then $A$ and $B$ have the same row space.

Now we show that $AB^{+}$ is invertible:

$$ $\displaystyle AB^{+}$ $\displaystyle=U_{A}\Sigma_{A}V_{A}^{\top}V_{B}\Sigma_{B}^{-1}U_{B}^{\top}$ $$

Since $A,B$ are full row rank, $U_{A},U_{B},\Sigma_{A},\Sigma_{B}$ are $r\times r$ invertible matrices. So it suffices to show that $V_{A}^{\top}V_{B}$ is invertible.
Since $\text{span}(V_{A})=\text{span}(V_{B})$, there exists an invertible change-of-basis matrix $P$ such that

$$ $\displaystyle V_{B}=V_{A}P\Rightarrow V_{A}^{\top}V_{B}=V_{A}^{\top}V_{A}P=I_{ r}P=P\Rightarrow AB^{+}\text{ is invertible.}$ $$

∎

###### Lemma C.2 .

If $U\perp\!\!\!\perp V\perp\!\!\!\perp Z$ (jointly independent), then $V\perp\!\!\!\perp Z\mid f(Z,U)$, for any measurable function $U$.

Let $W=f(Z,U)$. Then:

$$ $\displaystyle p(v\mid w,z)=~{}$ $\displaystyle\int_{u}p(v\mid w,z,u)p(u\mid w,z)du$ $\displaystyle=~{}$ $\displaystyle\int_{u}p(v\mid z,u)p(u\mid w,z)du$ $\displaystyle=~{}$ $\displaystyle\int_{u}p(v)p(u\mid w,z)du$ ( $V\perp\!\!\!\perp Z\perp\!\!\!\perp U$ ) $\displaystyle=~{}$ $\displaystyle p(v)\int_{u}p(u\mid w,z)du$ $\displaystyle=~{}$ $\displaystyle p(v)$ $$

###### Lemma C.3 (Sufficient Conditions for Bounded Completeness) .

Consider $D=A\cdot Z+U,\quad U\perp\!\!\!\perp Z$. $D$ is bounded complete for $Z$ if the following holds:

- •
The measure of $AZ$ is continuous and is supported on $\mathbb{R}^{r}$.
- •
The density of $U$ is continuous.
- •
The characteristic function of the distribution of $U$ is infinitely often differentiable and does not vanish on the real line.

###### Proof of Lemma C.3 .

This result follows as a Corollary of Theorem 2.1 in D’Haultfoeuille, where we consider the special case of linear mappings from $Z$ to $D$.
∎

###### Lemma C.1 .

###### Proof.

###### Lemma C.2 .

###### Lemma C.3 (Sufficient Conditions for Bounded Completeness) .

###### Proof of Lemma C.3 .

## Appendix D Further Details on Experimental Evaluation

### D.1 Linear

This section provides details of the linear experiments briefly described in Section 5 of the main paper. While the main paper presents summary statistics of average improvements and key findings, here we included detailed data generating equations and histograms of the improvements across runs.

The data are generated using the three following cases.

To determine the true outcome after perturbation, We used the formula

$$ $Y_{\alpha u}=\theta^{T}(B^{\dagger}X_{\alpha u}).$ $$

In addition to the summary statistics included in the main paper, we also plotted the distribution of average test improvements across seeds in [Figure 9](https://arxiv.org/html/2506.02200v2#A4.F9). We can observe that the test improvements of LIRR are shifted more to the right compared to the baseline PCA method.

Figure: (a) Linear DGP 1: Independent U and V
Refer to caption: x1.png

### D.2 Quadratic

This section provides details of the nonlinear experiments briefly described in Section 5 of the main paper. While the main paper presents summary statistics of average improvements and key findings, here we included detailed data generating equations, model hyperparameter, and histograms of the improvements across runs.

The data are generated using the following 3 cases.

All encoder architectures incorporate a Random Fourier Feature layer, followed by three feedforward layers and a final linear projection. Decoders consist of three feedforward layers and a final linear projection layer. For our IRAE[2] and IRAE models, we set the bottleneck dimension to 10, larger than the instrumental variable dimension $r=k=4$. By construction, Vanilla and IRAE[1] has bottleneck equal to $k=4$. To determine the true outcome after perturbation, we used the formula

$$ $Y_{\alpha u}=\theta^{T}((B^{\dagger}X_{\alpha u})[:r]),$ $$

where $[:r]$ index into the first order terms (excluding the quadratic and cross terms) of $D$.

The hyperparameters used in the training procedure are described in [Table 4](https://arxiv.org/html/2506.02200v2#A4.T4).

**Table 4: Training Parameters for Quadratic Simulations**
|  | Vanilla AE | IRAE[0] | IRAE[1] | IRAE[2] | IRAE | VAE | iVAE |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Architecture |  |  |  |  |  |  |  |
| Encoder dimensions | $100\rightarrow 50\rightarrow 20$ |  |  |  |  |  |  |
| Decoder dimensions | $20\rightarrow 50\rightarrow 100$ |  |  |  |  |  |  |
| RFF bandwidth $\sigma$ | $20$ |  |  |  |  |  |  |
| Bottleneck dimension | 4 | 4 | 4 | 10 | 10 | 4 | 4 |
| Optimization |  |  |  |  |  |  |  |
| Optimizer | RMSprop |  |  |  |  |  |  |
| Learning rate | $5\times 10^{-4}$ | $1\times 10^{-4}$ | $5\times 10^{-4}$ |  |  |  |  |
| Alpha | $0.9$ |  |  |  |  |  |  |
| Epsilon | $1\times 10^{-8}$ |  |  |  |  |  |  |
| Weight decay | $1\times 10^{-6}$ |  |  |  |  |  |  |
| Momentum | None |  |  |  |  |  |  |
| Regularization Parameters |  |  |  |  |  |  |  |
| $\lambda$ | 0 0 | $1$ | $1$ | $1$ | $1$ | NA | NA |
| $\mu_{1}$ | 0 0 | 0 0 | $1$ | $1$ | $1$ | NA | NA |
| $\mu_{2}$ | 0 0 | 0 0 | 0 0 | $1$ | $1$ | NA | NA |
| $\mu_{3}$ | 0 0 | 0 0 | 0 0 | 0 0 | $1$ | NA | NA |
| weight for kl term | NA | NA | NA | NA | NA | 3 | 3 |
|  | $c_{1}=1.0$, $c_{2}=c_{3}=0.0$ |  |  |  |  |  |  |
| Training Protocol (with early stopping of patience 20) |  |  |  |  |  |  |  |
|  | 1000 epochs |  |  |  |  |  |  |

Additional plots corresponding to [Table 4](https://arxiv.org/html/2506.02200v2#A4.T4) are included in [Figure 10](https://arxiv.org/html/2506.02200v2#A4.F10).

Figure: (a) Quadratic DGP 1: Independent U and V
Refer to caption: x4.png

### D.3 MNIST Experiment 1

This section provides details of the MNIST experiments briefly described in Section 5 of the main paper. Here we included detailed data generating equations, model hyperparameter, and plots for IRAE[0], IRAE[1], IRAE[2] that were not included in the main paper.

The data for MNIST experiment is generated using *Case 1 DGP*.

All encoders consist of three Conv2D layers, followed by additional feedforward layers, and conclude with a linear projection. Decoders mirror this architecture in reverse order. For our IRAE[2] and IRAE models, we set the bottleneck dimension to 10 which is larger than $k=2$. For vanilla and IRAE[0], IRAE[1], the bottle neck is 2. The autoencoder with multiple HSIC regularization terms presents greater training challenges due to the complexity of term. To address this, we initialized IRAE[2] and IRAE with weights from the simpler IRAE[1] model. All of models are trained with 60k training samples and evaluated on 10k test set. More training details can be found in [Table 5](https://arxiv.org/html/2506.02200v2#A4.T5).

**Table 5: Training Parameters for MNIST Simulations**
| Parameter | Vanilla AE | IRAE[0] | IRAE[1] | IRAE[2] | IRAE |
| --- | --- | --- | --- | --- | --- |
| Architecture |  |  |  |  |  |
| Kernel Size | 3 |  |  |  |  |
| Encoder channels | 16 $\rightarrow$ 32 $\rightarrow$ 64 |  |  |  |  |
| Decoder channels | 64 $\rightarrow$ 32 $\rightarrow$ 16 |  |  |  |  |
| Bottleneck dimension | 2 | 2 | 2 | 10 | 10 |
| Optimization |  |  |  |  |  |
| Optimizer | Adam (default parameters in torch) |  |  |  |  |
| Learning rate | $1\times 10^{-3}$ |  |  |  |  |
| Weight initialization | None | None | None | From IRAE[1] | From IRAE[1] |
| Loss Weights |  |  |  |  |  |
| $\lambda$ | 0 | 10 | 10 | 10 | 10 |
| $\mu_{1}$ | 0 | 0 | 10 | 10 | 10 |
| $\mu_{2}$ | 0 | 0 | 0 | 10 | 10 |
| $\mu_{3}$ | 0 | 0 | 0 | 0 | 10 |
|  | $c_{1}=0.8$, $c_{2}=0.15$, $c_{3}=0.05$ |  |  |  |  |
| Weights $\lambda,\mu_{2},\mu_{3}$ took value $1$ whenever non-zero, instead of $10$ for experiments in the Case 2-4 DGPs. |  |  |  |  |  |
| Training Epochs (with early stopping of patience 5) |  |  |  |  |  |
|  | 50 | 50 | 50 | 50* | 50* |

###### Remark D.1 (Calculation of Outcome from Image) .

To calculate expected $Y_{\alpha u}$, we first perform 2-mean clustering on the image pixels and extract the red, green, blue values from the center of the colored cluster. Then, we take the sum of these values as $Y$. Note that is this similar to taking the average colors over the gray scale mask so the colors would be slightly smaller than the original colors. We tested the methods on the original image and the result is  0.2 smaller on average.

###### Remark D.2 (Calculation of Outcome Improvement) .

When calculating the outcome improvement of the intervention, take the difference between the kmeans calculation described in the previous paragraph applied to the image produced by the intervention and we subtract the outcome of the kmeans calculation when applied to the original image.

###### Remark D.3 .

We use a linear kernel for HSIC in order to perform benchmarking at a large scale in fast speed, which may not capture all nonlinear dependencies in this complex image representation setting. More complex independence statistics based on domain knowledge, could perhaps lead to more disentanglement, albeit they might also be harder to train. In subsequent section experiments we also examine a pairwise RBF Kernel based HSIC and we find that it does not lead to improved performance as compared to the linear kernel.

###### Remark D.4 .

We observe that this example does not perfectly align with the formulation in [Equation 1](https://arxiv.org/html/2506.02200v2#S2.E1). Here, the number of instruments is 2, which is fewer than the natural representation of D of 3 colors. We may be able to interpret the learned representation as a 2-dimensional subspace of the 3-dimensional color representation, but the mapping from Z to D is still not immediately invertible as assumed in the theory. Additionally, while our theoretical analysis assumes a mapping from color $D$ to outcome directly, our calculation employs k-means clustering on $X$ instead. Nevertheless, this example demonstrates that our method performs robustly even in settings beyond those covered by our theoretical guarantees, and offers potential future directions of theoretical investigation.

Figure: Figure 11: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the IRAE[2] trained model (Case 1 DGP).
Refer to caption: extracted/6565704/paper_images/irae2_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_True_50_scatterplot.png

Figure: Figure 12: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the IRAE[1] trained model (Case 1 DGP).
Refer to caption: extracted/6565704/paper_images/irae1_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_scatterplot.png

Figure: Figure 13: Original gray, original color, reconstructed, treated($\alpha=0.2$) and treated($\alpha=1.0$) for the IRAE[0] trained model (Case 1 DGP).
Refer to caption: extracted/6565704/paper_images/justZ_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_scatterplot.png

Figure: Figure 14: IRAE[2] on Case 1 DGP for one random seed (random seed 22), with a Conv AutoEncoder, linear HSIC as independence criterion, latent dimension 10, regularization weights $\lambda=\mu_{1}=\mu_{2}=10$ and training for 50 epochs with early stopping (patience 5 epochs) warm start from IRAE1.
Refer to caption: extracted/6565704/paper_images/irae2_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_True_50_examples.png

Figure: Figure 15: IRAE[1] on Case 1 DGP for one random seed (random seed 22), with a Conv AutoEncoder, linear HSIC as independence criterion, latent dimension 2, regularization weights $\lambda=\mu_{1}=10$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch
Refer to caption: extracted/6565704/paper_images/irae1_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_examples.png

Figure: Figure 16: IRAE[0] on Case 1 DGP for one random seed (random seed 22), with a Conv AutoEncoder, linear HSIC as independence criterion, latent dimension 2, regularization weights $\lambda=\mu_{1}=10$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch
Refer to caption: extracted/6565704/paper_images/justZ_22_config_dgp1_conv_10_10.0_linear_hsic_50_256_5_examples.png

###### Remark D.1 (Calculation of Outcome from Image) .

###### Remark D.2 (Calculation of Outcome Improvement) .

###### Remark D.3 .

###### Remark D.4 .

### D.4 MNIST Experiment 2

Building on the results from our MNIST experiments in Section 5, we conducted a more comprehensive evaluation by exploring additional hyperparameter configurations and data generating processes. Given that independence test statistics are often complex and challenging to train, we systematically investigated various model architectures, independence test statistics calculation, and initialization strategies to identify optimal configurations. To align with our theoretical requirements outlined in [assumption 2.2](https://arxiv.org/html/2506.02200v2#S2.Thmtheorem2), we evaluated our approach on a supplementary dataset with three instruments, denoted as *Case 2 DGP*.

Our findings reveal that simpler dense architectures perform at least as well as, and often better than, more complex convolutional neural networks for this task. Furthermore, we observed that larger bottleneck dimensions in IRAE[2] and IRAE models better preserve the original digit morphology in treated images — a potentially valuable property when morphological features is confounded the outcome variable.

The full set of hyperparameters explored are included in [Table 6](https://arxiv.org/html/2506.02200v2#A4.T6). All of models are trained with 60k training samples and evaluated on 10k test set, for 40 random seeds. Regularization weights are 0 or 1. All models are trained with 50 epochs after initialization with early stopping of patience 5.

**Table 6: Summary of parameters explored in MNIST Experiment 2**
| Setting<br>Category | Options | Description |
| --- | --- | --- |
| Data<br>Generating Process | DGP2 | Three Instruments |
| Autoencoder<br>Architecture | Dense | Encoder: Dense layer $3\times 28\times 28\rightarrow 512$, followed by linear projection to latent dimension |
| Decoder: Linear layer from latent dimension to 512, followed by dense layer $512\rightarrow 3\times 28\times 28$ |  |  |
| Convolution | Encoder: Three Conv2D layers with channel $16\rightarrow 32\rightarrow 64$ of kernel size 3, followed by a dense layer of size 256 and linear projection to latent dimension |  |
| Decoder: Linear layer from latent dimension to size 256, followed by dense layer and three Conv2D layers with channel $64\rightarrow 32\rightarrow 16$ of kernel size 3 |  |  |
| Latent<br>Dimension<br>IRAE[2] and IRAE | 10 | Used for IRAE[2] and IRAE models |
| 32 | Used for IRAE[2] and IRAE models |  |
| Regularization<br>Type | Linear HSIC | Applied as independence measure on the entire vector |
| Pairwise HSIC | Applied between pairwise coordinates |  |
| Weight<br>Initialization<br>IRAE[2] and IRAE | Without warmstart | Training from randomly initialized weights for 50 epochs |
| With warmstart | Initializing with weights transferred from a pre-trained IRAE[1] model, and training for additional 50 epochs |  |

We highlight some findings from our exploration of the performance of our proposed methods across various hyperparameter dimensions:

Architecture: We found that simple dense layers can achieve better performance than convolutional architectures for this task, suggesting that Conv2D layers may be unnecessarily complex for this particular example.

Data Generating Process: Our experimental results demonstrate that the relative performance of our methods remains consistent across both DGP1 and DGP2.

Latent Dimension: When using larger latent dimensions (32), both the reconstructed and treated images preserved more of the original digit morphology although the improvement is smaller (c.f. [Figures 18](https://arxiv.org/html/2506.02200v2#A4.F18), [19](https://arxiv.org/html/2506.02200v2#A4.F19), [20](https://arxiv.org/html/2506.02200v2#A4.F20), [21](https://arxiv.org/html/2506.02200v2#A4.F21), [22](https://arxiv.org/html/2506.02200v2#A4.F22) and [23](https://arxiv.org/html/2506.02200v2#A4.F23)). This may be a desired property in some cases, especially in the case that the digit morphology is a confounder (not tested in our experiment) and has a direct effect on the outcome.

Regularization Type: While pairwise HSIC may theoretically capture more nonlinear dependencies, we found that it was often more difficult to train in practice. Linear HSIC consistently yielded better performance with greater training stability.

Weight Initialization: Dense architectures performed well without warm start initialization, while convolutional architectures benefited significantly from weight transfer. This difference likely stems from the higher complexity and larger parameter space of convolutional networks.

Overall, the best improvement model stems from the IRAE method with all regularizers, a Dense architecture, latent = 10, linear HSIC with no warm start.

Figure: Figure 17: Experimental results for the Case 2 data generating process. Mean improvement and standard deviation of improvement is reported. *reconstructed* refers to the mean outcome improvement of the reconstructed image from the autoencoder with no intervention in the latents, as compared to the original image. *intervened($\alpha$)* refers to the mean outcome improvement of the image produced by intervening on the latents in direction $\alpha\cdot u$, where $u=\theta/\|\theta\|$ and $\theta$ is estimated by 2SLS in latent space.

Figure: Figure 18: IRAE on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 32, regularization weights $\lambda=\mu_{1}=\mu_{2}=\mu_{3}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE1).
Refer to caption: extracted/6565704/paper_images/irae3_22_config_dgp2_dense_32_1.0_linear_hsic_50_256_5_False_50_examples.png

Figure: Figure 19: IRAE on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 10, regularization weights $\lambda=\mu_{1}=\mu_{2}=\mu_{3}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE1).
Refer to caption: extracted/6565704/paper_images/irae3_22_config_dgp2_dense_10_1.0_linear_hsic_50_256_5_False_50_examples.png

Figure: Figure 20: IRAE[2] on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 32, regularization weights $\lambda=\mu_{1}=\mu_{2}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE[1]).
Refer to caption: extracted/6565704/paper_images/irae2_22_config_dgp2_dense_32_1.0_linear_hsic_50_256_5_False_50_examples.png

Figure: Figure 21: IRAE[2] on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 10, regularization weights $\lambda=\mu_{1}=\mu_{2}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE[1]).
Refer to caption: extracted/6565704/paper_images/irae2_22_config_dgp2_dense_10_1.0_linear_hsic_50_256_5_False_50_examples.png

Figure: Figure 22: IRAE[1] on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 3 = number of instruments, regularization weights $\lambda=\mu_{1}=1$ and $\mu_{2}=\mu_{3}=0$ and training for 50 epochs with early stopping (patience 5 epochs).
Refer to caption: extracted/6565704/paper_images/irae1_22_config_dgp2_dense_10_1.0_linear_hsic_50_256_5_examples.png

Figure: Figure 23: IRAE[0] on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 3 = number of instruments, regularization weights $\lambda=\mu_{1}=1$ and $\mu_{2}=\mu_{3}=0$ and training for 50 epochs with early stopping (patience 5 epochs).
Refer to caption: extracted/6565704/paper_images/justZ_22_config_dgp2_dense_10_1.0_linear_hsic_50_256_5_examples.png

Figure: Figure 24: Vanilla AE on Case 2 DGP for one random seed (random seed 22), with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 3 = number of instruments, regularization weights $\lambda=\mu_{1}=\mu_{2}=\mu_{3}=0$ and training for 50 epochs with early stopping (patience 5 epochs).
Refer to caption: extracted/6565704/paper_images/vanilla_22_config_dgp2_dense_10_1.0_linear_hsic_50_256_5_examples.png

### D.5 Case 3: Confounded Outcome

We examine the following confounded outcome generating process, where the instruments now affect the colors in a more convoluted intertwined manner. We denote this as *Case 3 DGP*.

All of models are trained with 60k training samples and evaluated on 10k test set, for 40 random seeds. Regularization weights are 0 or 1. All models are trained with 50 epochs after initialization with early stopping of patience 5.

In this confounding setting, we found that IRAE[0], IRAE[1], IRAE[2], IRAE still led to improved outcome, whereas Vanilla AE did not.

Figure: Figure 25: Experimental results for the Case 3 data generating process. Mean improvement and standard deviation of improvement is reported.

### D.6 Case 4: Confounded DGP with One Outcome Relevant Dimension

We examine the following confounded outcome generating process, where the instruments now affect the colors in a more convoluted intertwined manner. Moreover, only the red channel is relevant for the outcome and the outcome is confounded. We denote this as *Case 4 DGP*.

All of models are trained with 60k training samples and evaluated on 10k test set, for 40 random seeds. Regularization weights are 0 or 1. All models are trained with 50 epochs after initialization with early stopping of patience 5.

We demonstrate in this data generating process the importance of running an instrumental variable regression in the latent space. We see below that if instead we had run OLS regressing the outcome on the identified latent factors, then the direction would be erroneous and the interventional images will not be moving the image towards more red colors.

Figure: Figure 26: Experimental results for the Case 4 data generating process. Mean improvement and standard deviation of improvement is reported.

Figure: Figure 27: IRAE on Case 4 DGP for one random seed, with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 32, regularization weights $\lambda=\mu_{1}=\mu_{2}=\mu_{3}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE1). Interventional images are intervened in the direction identified by 2SLS in the latent space with instrument $Z$, treatment $D$ and outcome $Y$. The outcome is larger when the color of the image is changed to red.
Refer to caption: extracted/6565704/irae3examples_dpg4_dense_32_1.0_linear_hsic_50_False_50.png

Figure: Figure 28: IRAE on Case 4 DGP for one random seed, with a Dense AutoEncoder, linear HSIC as independence criterion, latent dimension 32, regularization weights $\lambda=\mu_{1}=\mu_{2}=\mu_{3}=1$ and training for 50 epochs with early stopping (patience 5 epochs) from scratch (no warm start from IRAE[1]). Interventional images are intervened in the direction identified by OLS($Y\sim D$) in the latent space. The outcome is larger when the color of the image is changed to red.
Refer to caption: extracted/6565704/irae3examples_ols_dpg4_dense_32_1.0_linear_hsic_50_False_50.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 4
presentation: 1
questions: - How realistic are Assumptions 2.1 and 2.2 in practice and how would minor violations of these assumptions affect the framework (theoretically/some intuition and maybe even provide some ablation study on simulated data)?
- Could the authors clarify why in ll. 1.66 they consider soft interventions and not deterministic ones and how they define them?
- Could the authors elaborate more on when it is useful/necessary to learn representations of V and when learning only D is sufficient?
soundness: 3
strengths: - The IV setting with potentially high dimensional confounded treatments is an important and underexplored research direction in causal inference research.
- The motivation showing that omitted variable bias can also come from dimensionality reduction/representation learning of treatments (and not only of cofounders like most previous work showed) is an interesting and important finding
- Their method, ensuring that the IV information is maintained in the treatment representation, is a nice and intuitive idea and the theoretical analysis in linear and non-linear setting helps to motivate this choice.
summary: This paper considers an instrument variable setting with a potentially higher dimensional confounded treatment. The paper argues that naively adapted methods for such a setting, adopting dimensionality reduction on the treatment first and then performing IV regression, may introduce omitted variable bias. To address this challenge, the authors propose a framework to incorporate the instrument information during the dimensionality reduction/representation learning of the treatment to ensure improvement guarantees on soft interventions in this setting. The authors derive a theory and apply their method to both linear and non-linear setups, providing experimental results on (semi-)synthetic datasets.
weaknesses: The clarity of the paper could be improved and implications of different parts of the method could be mentioned.
- Part of the main motivation is that IV estimators “can only accommodate as many endogenous treatment variables as available instruments”. I think while in general it makes sense that effect estimation with high dimensional treatments is challenging, this statement should be explained more and shown in more detail.
- I think the related work is a bit limited to some specific works and could be extended to better show the contribution of their method with respect to existing works, in particular around representation learning for high dimensional confounders and OMV bias (e.g., https://www.pnas.org/doi/abs/10.1073/pnas.2427298122, https://arxiv.org/abs/2311.11321, https://arxiv.org/abs/2502.04274 ) , highdimensional treatments (e.g., https://arxiv.org/abs/2009.14061, https://arxiv.org/abs/2106.01939, https://arxiv.org/abs/2301.12292 ) , and representation learning for IVs (e.g., https://arxiv.org/abs/1612.09596 ).
- Assumptions 2.1 and 2.2, i.e. invertibility of encodings and existence and full rank of A, are specific to their novel setting and method. Thus, the implications of these assumptions should be discussed in more detail. How realistic are these in practice and how would minor violations of these assumptions affect the framework (theoretically/some intuition and maybe even provide some ablation study on simulated data)?
- The motivation of choosing soft interventions only and not deterministic ones is not fully clear to me. From a practical perspective, I think this would assume that during inference time we first can sample from the observational policy which might limit the practicality in real-world applications?

Presentation could be improved, especially in the experiments section
- The baselines should be introduced and described properly, also there is currently a mix between linear and non-linear setting (baselines are already referenced first and then described later).
- In the tables, it would be better to display their method in the last rows and not in the middle to better show their method.
- Table 3 over sample size could also be a plot to better compare the methods with sample size as x axis
- Intuition on why the different IRAE[x] variants perform differently and and some discussion on in which general settings the respective regularization terms are helpful and maybe some guideline when to use which variant could be helpful.
- real world applicability not clear, no motivation for real world application or any experiments to give some intuition provided (validation not possible but some insights could be shown in real world data)
- No discussion/conclusion section is provided. I think a short summary and mentioning of potential limitations would strengthen the paper.
- Presentation and fontsize of the plots needs to be improved clearly. Also, Figure 3 is hard to interpret and understand in general.

Overall, I think when my main concerns regarding unclarities are answered, I think this is a nice and interesting paper from the content side and I am open to increase my score. However, I think especially regarding the presentation (Figures, providing discussion, etc.), substantive changes are required for publication.

---

confidence: 2
contribution: 2
initial_rating: 4
presentation: 2
questions: Please see my concerns in weakness.
soundness: 2
strengths: The authors pinpoint how standard unsupervised dimensionality reduction of high dimensional treatments can violate the exclusion restriction by discarding instrument driven variation, a problem they term omitted treatment bias. The proposed solution of instrument guided representation learning represents a creative fusion of causal inference and representation learning, directly addressing this limitation in prior two stage approaches.

The authors develop a complete framework with specialized methods for both linear and nonlinear settings, supported by strong theoretical guarantees on the identifiability of valid intervention directions.  The writing is remarkably clear, with logical progression from linear to nonlinear cases.
summary: This paper addresses a fundamental limitation in instrumental variable (IV) regression when dealing with high-dimensional treatments: the violation of the exclusion restriction due to unsupervised dimensionality reduction. The authors propose Instrument-Guided Representation Learning (IGRL), a novel framework that integrates instruments directly into the representation learning process to ensure the latent treatment representation captures all instrument-driven variation. Through both linear (LIRR) and non-linear (IRAE) implementations, supported by theoretical guarantees and extensive experiments, the method demonstrates superior performance over conventional approaches in identifying outcome-improving interventions, effectively mitigating the omitted variable bias that plagues standard two-stage procedures.
weaknesses: The theoretical foundation of this paper relies on a set of strong structural assumptions that may be difficult to satisfy in real world applications. A key example is the core assumption of joint independence. This assumption requires the instrument, the confounder representation, and the orthogonal components to be fully independent. This is particularly challenging to guarantee with high dimensional and complex data. Furthermore, prerequisites such as the invertibility of the encoding and decoding functions and specific linear relationships among latent variables significantly limit the direct applicability of the theoretical model to practical scenarios. While these assumptions provide a necessary foundation for theoretical derivation, the paper does not sufficiently discuss the robustness of the estimates when these conditions are violated. This omission creates uncertainty for the method's practical deployment.

The paper exhibits notable shortcomings in reproducibility. This is particularly problematic for the non linear IRAE model. This model involves tuning regularization weights for up to six hyperparameters. The strategy for setting these critical parameters and their sensitivity to model performance are not elaborated. I encourage author provide anonymous github link.

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: - Line 95: The data does not include any observed confounders. Why is this reasonable? If we assume all confounders are unobserved, we introduce a lot of  uncertainty in the modelling
- Why exactly do naïve representation learning approaches introduce omitted variable bias? This needs to be explained further.
- Assumption 2.1/ structural equation for X: This assumption is very strong. Why would such an encoding exist? How can we validate it? 
- Experiments: Why is the proposed method not compared against other treatment representation learning methods? Even if they are not designed for IV 
analysis, the comparison is of high value to assess the usefulness of the proposed method. A comparison with PCA is insufficient.
soundness: 2
strengths: - The problem setting is interesting and novel to the best of my knowledge
summary: The paper proposes a method for representation learning of high-dimensional treatments that allows for a valid downstream instrumental variable treatment effect estimation setting. To do so, the paper proposes to incorporate IVs in the representation learning step. The procedure can be interpreted as a regularization of the unsupervised representation learning step to eliminate spurious backdoor paths. The paper compares the presented method with common statistical dimension reduction techniques, such as PCA, and ML methods, such as VAEs.
weaknesses: - Line 67: The paper states that causal representation learning is uncommon in causal inference, but common in causal discovery. This is not true.
The paper thus lacks a sufficient discussion of related work on representation learning in causal inference tasks. 
- The method requires very strong untestable assumptions, limiting its applicability in practice.
- The paper neither contains a discussion on the method, its limitations, or the results, nor a conclusion, leaving it an unfinished work. Potentially, these important parts of a research paper are neglected because of the page limit. However, this cannot be seen as an excuse.
- Mathematical statements are neither proven in the main paper nor are proofs referenced.
- The paper is lacking a proper motivation with precise use case examples. The problem statement seems a bit far-fetched and unrealistic. For example, what exactly are the available instruments in line 39? In my opinion, it is rather rare in practice to have suitable instruments
- Many references are incomplete or not stated correctly. This is not in line with proper scientific work.
- Many language/grammar errors. Sentences are incomplete, making it difficult for the reader to follow the line of thought.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - I have one general question regarding the representation learning approach: Usually, we lose identifiability of the causal target quantity when using a treatment representation due to violation of the consistency assumption (no hidden variability within the treatment). Could you elaborate on how exactly your method addresses this problem? The results seem reasonable as they only provide improvement guarantees and no identifiable treatment effect.
soundness: 3
strengths: - Learning representations of high-dimensional treatments is a difficult but relevant problem in many disciplines, e.g., when the treatment is a text
- The proposed method comes with theoretical improvement guarantees and promising empirical results
- The idea seems novel and rigorously executed
summary: This paper addresses the inability of IV models to handle high-dimensional treatments when instruments are scarce. The authors identify that existing two-step approaches that apply unsupervised dimension reduction (e.g., PCA) before IV regression violate the exclusion restriction and introduce omitted variable bias. To resolve this, they propose Instrument-Guided Representation Learning (IGRL), which incorporates instrumental variables directly into the representation learning process to ensure learned features capture only Z-driven variation in X. The approach provides both theoretical guarantees that the learned representations identify outcome-prediction-optimizing directions and empirical evidence of improvements over conventional two-stage methods.
weaknesses: - The current version of the method does not allow for incorporating observed covariates, which could strengthen improvements in practice. Is this a straightforward extension or would this require more elaborate theory?
- As is common for IV methods, the proposed method depends on assumptions that may be difficult to verify in practice
- The specific setting (instruments + high-dimensional treatment) seems somewhat niche. A few clear motivating examples might strengthen the story of the paper.


Decision: Reject