# Learning Treatment Representations for Downstream Instrumental Variable Regression

arXiv: https://arxiv.org/abs/2506.02200 · HTML: https://mlsanity.com/paper/2506.02200/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 35 references that are in the map.
- First posted on arXiv 2 Jun 2025.

## Builds on (most cited first, 20 of 20)

- [Variational Autoencoders and Nonlinear ICA: A Unifying Framework](https://mlsanity.com/paper/1907.04809/) (2019)
- [Hidden Markov Nonlinear ICA: Unsupervised Learning from Nonstationary Time Series](https://mlsanity.com/paper/2006.12107/) (2020)
- [Semiparametric proximal causal inference](https://mlsanity.com/paper/2011.08411/) (2020)
- [Causal Discovery with General Non-Linear Relationships Using Non-Linear ICA](https://mlsanity.com/paper/1904.09096/) (2019)
- [Identifiability of latent-variable and structural-equation models: from linear to nonlinear](https://mlsanity.com/paper/2302.02672/) (2023)
- [Nonlinear Independent Component Analysis for Principled Disentanglement in Unsupervised Deep Learning](https://mlsanity.com/paper/2303.16535/) (2023)
- [Kernel-based Tests for Joint Independence](https://mlsanity.com/paper/1603.00285/) (2016)
- [Identifying Representations for Intervention Extrapolation](https://mlsanity.com/paper/2310.04295/) (2023)
- [Kernel Conditional Moment Test via Maximum Moment Restriction](https://mlsanity.com/paper/2002.09225/) (2020)
- [GraphITE: Estimating Individual Effects of Graph-structured Treatments](https://mlsanity.com/paper/2009.14061/) (2020)
- [Learning Causal Representations from General Environments: Identifiability and Intrinsic Ambiguity](https://mlsanity.com/paper/2311.12267/) (2023)
- [Towards efficient representation identification in supervised learning](https://mlsanity.com/paper/2204.04606/) (2022)
- [Cost-Effective Incentive Allocation via Structured Counterfactual Inference](https://mlsanity.com/paper/1902.02495/) (2019)
- [Inference on Strongly Identified Functionals of Weakly Identified Functions](https://mlsanity.com/paper/2208.08291/) (2022)
- [Source Condition Double Robust Inference on Functionals of Inverse Problems](https://mlsanity.com/paper/2307.13793/) (2023)
- [Contrastive representations of high-dimensional, structured treatments](https://mlsanity.com/paper/2411.19245/) (2024)
- [LABOR-LLM: Language-Based Occupational Representations with Large Language Models](https://mlsanity.com/paper/2406.17972/) (2024)
- [Semiparametric Causal Sufficient Dimension Reduction Of Multidimensional Treatments](https://mlsanity.com/paper/1710.06727/) (2017)
- [Identifiable Feature Learning for Spatial Data with Nonlinear ICA](https://mlsanity.com/paper/2311.16849/) (2023)
- [Estimating Wage Disparities Using Foundation Models](https://mlsanity.com/paper/2409.09894/) (2024)

## Abstract (arXiv)

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-informed 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.