Learning Treatment Representations for Downstream Instrumental Variable Regression

Traditional instrumental variable (IV) estimators cannot accommodate more treatments than instruments, a limitation that is critical for high-dimensional, unstructured data like clinical treatment pathways. Current practice—applying unsupervised dimension reduction before IV estimation—suffers from substantial omitted treatment bias because the representation learning step ignores the instrument. We propose a novel framework that constructs treatment representations by explicitly incorporating instrumental variables. We prove that this instrument-guided approach ensures the identification of optimal outcome-prediction directions even with limited instruments. Validation on large-scale, semi-synthetic clinical data derived from a major hospital, along with other simulations, shows that our approach significantly outperforms conventional two-stage methods.

In healthcare research, calculating the true impact of a treatment is challenging because unobserved factors—like a patient's lifestyle or socioeconomic status—can skew the results. To solve this, researchers use Instrumental Variables (IVs): natural, random "nudges" (such as a doctor's personal prescribing preference) that help isolate a treatment's true cause-and-effect relationship. Traditional IV methods break down when dealing with modern, complex medical care—such as unstructured clinical pathways involving hundreds of sequential decisions. Current solutions use AI to compress this high-dimensional data into a simplified summary before applying the IV. However, because this compression step ignores the random nudge entirely, critical information is lost, leading to biased and inaccurate conclusions. This study introduces a new framework that fixes this blind spot by explicitly incorporating the random nudges (IVs) directly into the AI's data-compression process. The authors mathematically prove that this "instrument-guided" approach successfully extracts the exact treatment patterns that matter for patient outcomes, even when the number of available nudges is limited. Tested on large-scale, semi-synthetic clinical data from a major hospital, this new method significantly outperformed conventional approaches, providing a more accurate tool for understanding which complex treatments actually improve patient health.