🤖 machine learning

IV-ICL: Bounding Causal Effects with Instrumental Variables via In-Context Learning

This paper introduces IV-ICL, an amortized Bayesian in-context learning method that efficiently recovers causal effect bounds by directly learning marginal posterior distributions through inclusive KL minimization, outperforming existing semi-parametric and Bayesian baselines in validity, informativeness, and computational speed while also proposing a novel procedure for generating realistic IV benchmarks from randomized controlled trials.

Original authors: Vahid Balazadeh, Hamidreza Kamkari, Medha Barath, Ricardo Silva, Rahul G. Krishnan

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vahid Balazadeh, Hamidreza Kamkari, Medha Barath, Ricardo Silva, Rahul G. Krishnan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Hidden Thief" in Causal Data

Imagine you are a detective trying to figure out if a new medicine (the Treatment) actually cures a disease (the Outcome). You look at patient records, but you notice a problem: some patients got sick because of a hidden factor, like a secret genetic flaw or a bad diet, that you can't see. This is called unobserved confounding.

Because of this "hidden thief," you can't say for sure, "The medicine cures 100% of people." The data is too messy. In the old days, statisticians would say, "We can't solve this," or they would try to guess a single number, which often leads to wrong answers.

Instead, the goal of Partial Identification is to stop guessing a single number and start guessing a range. Instead of saying "It cures 100%," you say, "It cures somewhere between 40% and 80%." This range is your "safety net." If the true answer is inside that net, you are safe.

The Old Way: The Slow, Manual Calculator

Before this paper, finding that safety net was incredibly hard.

The Math Bottleneck: For simple cases, mathematicians wrote down specific formulas (like the Balke-Pearl equations) to calculate the edges of the range. But these formulas are like complex, jagged puzzles. If you change the data even a little, the formula breaks or becomes very hard to solve.
The Manual Labor: Every time you had a new dataset, you had to manually design a new calculator to solve that specific puzzle. It was slow and prone to human error.
The "Guessing" Trap: Some people tried using Bayesian inference (a method that updates beliefs with data). But this method often gets stuck. Imagine trying to find a lost cat in a giant, dark warehouse. The old methods would shine a flashlight on just one spot (a single guess) and ignore the rest of the warehouse, even though the cat could be anywhere else. This leads to ranges that are too narrow and often wrong.

The New Solution: IV-ICL (The "Super-Reader")

The authors introduce IV-ICL. Think of this as training a super-smart AI to be a "Super-Reader" of causal puzzles.

Instead of solving the math puzzle from scratch every time, they teach the AI a massive library of "what-if" scenarios first.

The Training Library: They created millions of fake datasets where they knew the "hidden thief" and the "true cure rate." They taught a Transformer model (a type of AI) to look at the messy data and instantly predict the entire range of possible cure rates.
The "Mass-Covering" Trick: This is the paper's secret sauce.
- Old AI (Exclusive KL): Tries to find the most likely spot for the cat. It ignores the rest of the warehouse. Result: A tiny, dangerous range that often misses the truth.
- IV-ICL (Inclusive KL): Tries to make sure the AI's "flashlight" covers everywhere the cat could possibly be. It spreads its attention wide. If the truth is anywhere in the warehouse, the AI's range will include it. This makes the safety net reliable.

How It Works in Real Life

Pre-Training: The AI spends time learning from millions of synthetic (fake) scenarios. It learns the "shape" of all possible answers.
Inference (The Magic Moment): When you give the AI a real dataset (like a hospital record), it doesn't do any heavy math. It just looks at the data and says, "Based on what I've learned, the cure rate is between X and Y." It does this in a single, instant step.
The Result: The paper shows that IV-ICL is:
- Faster: It is 20 to 500 times faster than existing methods.
- More Accurate: It produces ranges that are actually valid (they contain the truth) and are tighter (more useful) than other methods.

The "RCT-to-IV" Bridge: Testing the AI

One of the hardest parts of testing these methods is that in the real world, we rarely know the "true" answer. How do you know if your safety net is good if you don't know where the target is?

The authors invented a clever way to test their AI using Randomized Controlled Trials (RCTs) (the gold standard of medical testing).

The Analogy: Imagine you have a perfect photo of a target (the RCT result). You want to test if your "foggy lens" (the IV method) can still see the target.
The Trick: They take the perfect RCT data and artificially add "fog" and "noise" to it to make it look like a messy, observational study. They do this in a way that mathematically guarantees the "true target" hasn't moved.
The Benefit: Now they can test their AI on "messy" data and know exactly if it found the right answer, because they started with the perfect answer.

Summary of Claims

No More Manual Formulas: You don't need to solve complex linear programs or write custom math for every new dataset.
No More "Stuck" Guesses: Unlike older Bayesian methods that get stuck on one guess, this method spreads its "net" wide enough to catch the truth.
Speed: It turns a process that takes hours or days into a process that takes seconds.
Reliability: On both fake data and data converted from real medical trials, it consistently produces valid ranges that are tighter (more informative) than the competition.

In short, IV-ICL is a fast, reliable AI that learns to draw the safest possible "safety net" around causal effects, even when the data is messy and the truth is hidden.

Technical Summary: IV-ICL

Problem Statement

Estimating causal effects from observational data is frequently hindered by unobserved confounding, rendering point identification of the treatment effect impossible. In such scenarios, the standard approach is partial identification, which aims to characterize the set of causal effects compatible with the observed data under weaker structural assumptions, typically using the Instrumental Variable (IV) model.

Existing methods for partial identification face five critical limitations (labeled L1–L5 in the paper):

Closed-form dependency (L1): Traditional approaches (e.g., Balke-Pearl bounds) require deriving closed-form bound estimands by solving linear programs. This process must be re-solved for every new setting (discrete vs. continuous, different graph structures).
Estimator design complexity (L2): Even when closed-form bounds exist, constructing accurate, low-variance estimators for high-dimensional covariates requires manual, problem-specific design.
Under-dispersion in Bayesian inference (L3): Standard Bayesian inference (MCMC or Variational Inference) applied directly to causal effects often fails in partially identifiable settings. Because the likelihood is flat along unidentifiable directions, Variational Inference (VI) minimizes the "exclusive" KL divergence, causing the posterior to collapse onto a single mode and fail to cover the true identified set.
Prior sensitivity (L4): In unidentifiable directions, the posterior is dictated entirely by the prior. If the prior's support is too narrow, the resulting bounds are invalid or overly optimistic.
Computational infeasibility (L5): Classical Bayesian inference requires re-running the inference procedure from scratch for every new dataset, making large-scale analysis prohibitively expensive.

Methodology: IV-ICL

The authors propose IV-ICL (Instrumental Variables via In-Context Learning), an amortized Bayesian inference method that learns the marginal posterior distribution of the causal effect directly, deriving bounds as quantiles of this distribution.

Core Mechanism

Instead of solving optimization problems or running MCMC per dataset, IV-ICL trains a Transformer model to map an input dataset directly to the marginal posterior of the Sample Average Treatment Effect (SATE).

Amortized Inference: A single Transformer is pre-trained on a vast library of synthetic datasets. At inference time, a new dataset is provided as "context," and the model outputs the posterior distribution in a single forward pass.
Inclusive KL Objective: The training loss is the negative log-likelihood of the ground-truth SATE given the dataset. Theoretically, minimizing this loss is equivalent to minimizing the expected inclusive Kullback-Leibler (KL) divergence between the true marginal posterior and the model's predicted distribution.
- Unlike the "exclusive" KL minimized by standard VI (which seeks a single mode), the inclusive KL is a "mass-covering" objective. It forces the model to assign positive density to any region where the true posterior has mass. This prevents the posterior from collapsing and ensures the estimated interval covers the entire identified set.

Prior Data Generation

To satisfy the mass-covering requirement and address prior sensitivity (L4), the authors design a specific prior over synthetic Data Generating Processes (DGPs):

Diversity: They generate a rich set of tabular datasets using neural networks to ensure coverage of diverse feature distributions.
IV Constraints: The generation process strictly enforces the IV assumptions (Consistency, Positivity, Exclusion Restriction, and Instrument Unconfoundedness).
Wide Support: A post-processing step shifts the logits of potential outcomes to ensure the induced distribution of SATEs spans a wide range (e.g., $[-1, 1]$ ), preventing the model from inheriting a bias toward zero or narrow intervals.

RCT-to-IV Benchmark Conversion

To evaluate methods on realistic data where ground-truth causal effects are known, the authors propose a procedure to convert Randomized Controlled Trial (RCT) data into IV benchmarks:

Split Covariates: Observed covariates ( $O$ ) and hidden confounders ( $U$ ) are separated from the RCT data.
Synthetic Instrument: An instrument $Z$ is generated based only on $O$ .
Synthetic Treatment: A treatment $T$ is generated based on $O$ , $U$ , and $Z$ , introducing confounding via $U$ and relevance via $Z$ .
Accept-Reject: Samples are retained only if the synthetic treatment matches the original RCT treatment assignment.
This procedure preserves the original RCT outcome distribution and the ground-truth causal effect while creating a valid IV estimation problem.

Key Contributions

Methodological Innovation: A method that eliminates the need for closed-form bound estimands (L1) and amortizes inference into a single, computationally cheap forward pass (L2, L5).
Training Objective: The use of an inclusive KL objective (via amortized learning) to resolve the under-dispersion issue (L3) inherent in standard variational inference for partial identification.
Prior Design: A general-purpose prior over IV-consistent synthetic DGPs with wide support, enabling better transfer and generalization (L4).
Evaluation Framework: A theoretically justified procedure to convert RCTs into IV benchmarks with preserved ground-truth causal effects, allowing for realistic evaluation of partial identification methods.
Empirical Validation: Demonstration that IV-ICL produces intervals that are more reliably valid and informative than semi-parametric, Bayesian, and plug-in baselines, with 20–500× lower inference time.

Experimental Results

The authors evaluated IV-ICL on fully synthetic binary-outcome benchmarks and semi-synthetic benchmarks derived from the National Supported Work (Jobs) and Project STAR (class size) RCTs.

Validity: IV-ICL achieved 100% validity (the estimated interval contained the true bounds or true causal effect) across all benchmarks. In contrast, standard Bayesian inference (using exclusive KL) collapsed to 0% validity on synthetic data and 25% on the Jobs benchmark due to under-dispersion.
Informativeness: Among valid methods, IV-ICL produced the tightest intervals. For example, on the Jobs benchmark, it was significantly tighter than plug-in estimators and nonparametric efficient estimators.
Efficiency: IV-ICL was 20–500 times faster than baselines. While baselines required seconds to minutes per dataset (due to optimization or MCMC), IV-ICL required less than 2 seconds per dataset (often <0.2s) after pre-training.
Instrument Strength: As instrument strength increased, the interval width of IV-ICL contracted monotonically while maintaining perfect validity, demonstrating the model's ability to adapt to the information content of the data.

Significance and Claims

The paper claims that IV-ICL resolves the fundamental trade-offs in partial identification: it provides valid bounds (covering the identified set) without the computational cost of per-dataset optimization or the instability of standard Bayesian inference.

By leveraging amortized in-context learning with a mass-covering objective, the method avoids the pitfalls of mode-seeking variational inference. The authors emphasize that while the method relies on a prior, the inclusive KL objective ensures that if the prior covers the true data-generating process, the resulting bounds are robust. The proposed RCT-to-IV conversion is highlighted as a significant contribution for enabling rigorous, realistic evaluation of partial identification methods on real-world data structures.

The authors remain modest regarding limitations, noting that formal frequentist guarantees under arbitrary prior-data mismatch are not yet established and that the current implementation handles continuous outcomes via discretization, leaving fully continuous settings for future work.