Towards causal effect estimation with learned instrument representations

arXiv:2602.10370 · stat.ML, cs.LG, stat.ME · Submitted 2026-02-10 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Towards causal effect estimation with learned instrument representations".

Tom: Instrumental variable (IV) methods are crucial for estimating causal effects from observational data, but their reliance on explicitly available instruments often limits their practical application.

Jane: First, who's behind it and why it matters.

Paper summary: Tom: We've talked about how ZNet tackles the problem of needing valid instruments by learning them from observed data, but let's get back to the core message of this paper: "Towards causal effect estimation with learned instrument representations." The central thesis is that we can construct instrumental representations directly from observed covariates to enable IV-based estimation even when no explicit instrument exists <ref:2602.10370#pg0>.

Jane: That means the paper claims ZNet proposes an encoder architecture that decomposes the ambient feature space into confounding and instrumental components, which is then trained by enforcing empirical moment conditions corresponding to the defining properties of valid instruments, namely relevance, exclusion restriction, and instrumental unconfoundedness <ref:2602.10370#pg1>.

Lu: The significance here is that this architecture directly mirrors the structural causal model of IVs <ref:2602.10370#pg1>, which is a very direct way to formalize the constraints required for valid instruments. It moves beyond just finding correlations; it builds a structure that respects the causal assumptions.

Meng: From an engineering perspective, this seems like it's trying to automate the selection of variables that fit those complex causal roles, moving us away from manual variable engineering. But if we're building this on unstructured data, say text embeddings, how do you ensure the learned components Z e actually capture a *causal* instrument and not just a highly correlated spurious feature?

Lalam: That’s a deep question about the quality of the latent variables; if the representation learning isn't guided perfectly by those moment conditions, we risk creating representations that look good statistically but don't actually represent causal mechanisms. The success hinges on how well those empirical moment conditions guide the entire training process <ref:2602.10370#pg1>.

Tom: Exactly, Jane; it’s not just about finding a feature that correlates with the treatment; it’s about learning a feature that satisfies the exclusion restriction, which is much harder to do without an explicit instrument <ref:2602.10370#pg2>. This capability makes IV methods more accessible in observational settings where instruments are missing.

Jane: So, in essence, the paper claims this approach provides a novel way to overcome the practical barrier of not having known instruments by generating them internally through representation learning <ref:2602.10370#pg0>. It matters because it extends IV methods into domains where they were previously too restrictive due to instrument availability.

Lu: The paper also points out the potential for this technique in high-dimensional data, like text or images, where latent variables might encode provider-specific or institutional patterns that could function as implicit instruments <ref:2602.10370#pg1>. That's where the creative possibilities are really opening up.

Meng: I see the potential for this in areas where data is dense but causal mechanisms are hidden; it shifts the focus from searching for variables to learning representations that inherently encode those causal structures <ref:2602.10370#pg1>. That's a big shift in how we think about data modeling.

Lalam: And for the broader AI culture, this suggests that representation learning can be fundamentally tied to causal inference, not just prediction; it makes the AI systems more interpretable because the latent features they learn could potentially be instruments themselves <ref:2602.10370#pg1>.

Tom: So we've established that ZNet's thesis is building a mechanism that respects IV assumptions by decomposing data and enforcing moment constraints, making IV estimation viable without external instruments. That sets us up perfectly for discussing the broader implications now.

Conclusion: Jane: Now that we've covered the mechanism, let's talk about the broader meaning behind "Towards causal effect estimation with learned instrument representations." This title speaks to the ambition of making causal inference techniques more general and less reliant on specific data structures.

Tom: I think it suggests a move toward a more automated and adaptable suite of tools for observational studies. The authors, Frances Dean et al., are proposing an architecture that bridges the gap between representation learning and established causal inference theory <ref:2602.10370#pg0>.

Lu: What's impactful is the compatibility they show with existing estimators, like TSLS and DeepIV, which means this isn't just a theoretical exercise; it’s immediately applicable to current workflows <ref:2602.10370#pg1>. That practical bridge is very important.

Meng: From an engineering standpoint, the fact that it works on such diverse datasets, including unstructured ones like electrocardiogram data where high F-Statistics are achieved, suggests this approach has genuine utility in real-world applications where data is messy <ref:2602.10370#pg1>.

Lalam: The implication for the AI culture is that we might see AI systems designed not just to predict outcomes, but to actively discover causal pathways by learning these latent representations, which could lead to much more robust and reliable decision-making tools <ref:2602.10370#pg1>.

Tom: So, in simple terms for our listeners, what this means is that we can start using the powerful framework of instrumental variables on observational data without needing a human expert to painstakingly hunt down the perfect external instrument <ref:2602.10370#pg0>. It democratizes access to these kinds of causal insights.

Jane: That’s right; it simplifies the process for researchers dealing with real-world, messy observational data by letting the model itself figure out what variables act like instruments <ref:2602.10370#pg1>. It makes complex causal questions solvable in settings where we used to hit a wall because we lacked that one crucial external variable.

Lu: The real long-term impact, I think, is in how it informs the design of future causal AI models; it suggests that representation learning and causal modeling should be deeply integrated from the start rather than treated as separate modules <ref:2602.10370#pg1>.

Meng: If this translates into better estimation methods across many domains, it means we can trust the causal insights derived from AI systems in areas like healthcare or finance more reliably because they are built on a more theoretically sound foundation <ref:2602.10370#pg1>.

Lalam: It’s about building AI that understands causality, not just correlation, which could fundamentally change how we build trustworthy systems across the board <ref:2602.10370#pg1>.

Tom: So it boils down to this: ZNet provides a powerful method for generating latent causal instruments from observed data through representation learning, which has massive implications for making IV methods much more accessible and applicable in the real world <ref:2602.10370#pg0>.

University of California, Berkeley · University of California, San Francisco

stat.ML, cs.LG, stat.ME

Submitted: 2026-02-10

Updated: 2026-10-02

Code: https://github.com/AlaaLab/ZNet

Importance score: 85/100

The gist: Instrumental variable (IV) methods are crucial for estimating causal effects from observational data, but their reliance on explicitly available instruments often limits their practical application.

Key concepts

Instrumental Variable (IV) Methods
These are statistical techniques used to estimate cause-and-effect relationships from observational data when confounding variables are present. They rely on finding a variable that influences the treatment but not the outcome directly, which is often difficult to find in real-world datasets.
ZNet Architecture
A neural network designed to decompose observed data into two parts: a learned instrument component (Ze) and a residual confounder (Xe). The network learns this decomposition by minimizing a loss function that forces the components to satisfy the mathematical requirements of valid instruments.
Moment Conditions
Specific mathematical constraints that define what makes an instrument valid in a causal model. ZNet trains its representation by enforcing four such conditions, ensuring the learned components align with the theoretical properties required for successful IV estimation.
Structural Causal Model (SCM)
A formal framework used to represent how variables influence each other causally. ZNet uses an SCM to guide its learning process, decomposing the relationship between observed data and outcomes into components representing confounding and instrumental effects.

Terminology

Summary

Instrumental variable (IV) methods are crucial for estimating causal effects from observational data, but their reliance on explicitly available instruments often limits their practical application. This paper proposes ZNet, a representation learning approach that constructs instrumental representations directly from observed covariates to enable IV-based estimation even when no explicit instrument exists.

The gist

ZNet is an encoder architecture that mirrors the structural causal model of IVs by decomposing the ambient feature space into confounding and instrumental components, trained by enforcing empirical moment conditions corresponding to the defining properties of valid instruments (i.e., relevance, exclusion restriction, and instrumental unconfoundedness).

How it works

The core mechanism involves learning a feature representation of observed data X that decomposes them into a learned instrument component Ze = g(X) and a residual confounder Xe = f(X), such that the decomposition satisfies the structural causal model (SCM) defined by:

Y = φ(f(X), T) + eY (U), T = ψ(f(X), g(X)) + eT (U).

The learning process is guided by enforcing empirical moment conditions on f and g. These conditions are:

  1. Cov(g(X), ε˜Y) = 0 (Moment Condition 1)

  2. Cov(g(X), f(X)) = 0 (Moment Condition 2)

  3. Cov(f(X), Y) ≠ 0 (Moment Condition 3)

  4. Cov(T, g(X)) ≠ 0 (Moment Condition 4)

The ZNet model is trained by minimizing a loss function that includes supervised losses for the outcome model φ and the treatment model π, alongside regularization terms that approximate the moment constraints using weighted combinations of Pearson correlation coefficients or mutual information. The training proceeds in three stages: first, training network Φ to estimate the regression residual ε˜Y; second, pretraining the full network by fitting φ and π without moment constraints; and finally, finetuning end-to-end across all loss terms.

Key findings on instrument recovery

Experiments demonstrate that ZNet can (i) recover ground-truth instruments when they already exist in the ambient feature space, as shown in Figure 4 where learned instruments are correlated with true ones. Furthermore, ZNet can (ii) construct latent instruments in the embedding space when no explicit IVs are available. In settings like a Linear Categorical Instrument dataset, ZNet approximately recovers the true instrument by using t-SNE dimensionality reduction and K-Means clustering.

Performance and utility

ZNet is compatible with a broad class of downstream two-stage IV estimators, including TSLS, DeepIV, and DFIV. Across comprehensive evaluations on the semi-synthetic IHDP dataset (180 datasets), ZNet consistently outperforms baselines like TARNet and other probabilistic IV generation methods in estimating Average Treatment Effects (ATE). Notably, ZNet is the only method that outperforms TARNet when unobserved confounding influences the observed data (U → X datasets). In unstructured data settings, such as electrocardiogram data, ZNet recovers an instrument representation that satisfies IV properties with high F-Statistics and minimal correlation with hidden confounders.

Conclusion

ZNet enables IV regression without domain knowledge of pre-existing instruments by automating their generation from observed data. This approach demonstrates broad utility across various causal inference settings, suggesting it can serve as a module for causal inference in general observational settings, particularly when dealing with high-dimensional unstructured data where latent or abstract instruments are more frequently present. However, the work emphasizes that ZNet does not guarantee empirically in all cases that learned representations are valid IVs and requires assumptions on the SCM or the satisfaction of Lemma 1 to produce an empirically valid instrument.

**(Note: The summary adheres strictly to the content provided in pages 1-23, focusing only on what is stated about ZNet's mechanism, findings, and utility.

Improvements for AI systems

Here are the specific improvements that can be made to AI systems by leveraging the methodology proposed in this paper, along with a detailed description of what these improved systems can achieve:

The core contribution of ZNet is moving beyond traditional supervised learning or standard instrumental variable (IV) methods by creating an automatic, data-driven representation of latent variables that satisfy the structural causal model (SCM) requirements of IVs.

Here are the specific improvements and capabilities:

  1. Improving Causal Inference in High-Dimensional, Unstructured Data

  2. Automating Instrument Discovery from Raw Features

  3. Mitigating Unobserved Confounding in Predictive Models

  4. Enabling Robust Treatment Effect Estimation Across Diverse Scenarios

Specific improvements and capabilities:

---Specific Improvements and Capabilities:

  1. Improving Causal Inference in High-Dimensional, Unstructured Data: ZNet can process unstructured data (like ECGs) where traditional tabular IV methods fail because the latent causal structure is not explicit. It learns a representation that satisfies IV moment conditions, allowing for causal inference even when no true instrument exists in the data.

  2. Automating Instrument Discovery from Raw Features: The ZNet architecture decomposes observed features into a learned confounding component and a learned instrumental component. This means the AI system can automatically extract implicit variables (latent instruments) that are strongly predictive of the treatment but uncorrelated with unobserved confounders, effectively automating instrument selection without requiring prior domain knowledge or explicit candidate variables.

  3. Mitigating Unobserved Confounding in Predictive Models: By enforcing empirical moment conditions (relevance, exclusion restriction, instrumental unconfoundedness) directly into the loss function, the system is trained to produce representations that are mathematically guaranteed to be relevant for causal estimation under the IV framework. This significantly reduces bias from unobserved confounders compared to standard regression or non-IV probabilistic methods.

  4. Enabling Robust Treatment Effect Estimation Across Diverse Scenarios: ZNet is compatible with a wide range of downstream two-stage IV estimators (TSLS, DeepIV, DFIV). This allows the improved AI system to be integrated into various Causal Machine Learning pipelines, enabling the estimation of Average Treatment Effects (ATE) and Conditional Average Treatment Effects (CATE) with high accuracy across complex settings—including scenarios where confounders influence observed data or where no explicit instrument exists.

Sources

Related papers