🤖 AI

Conformal Prediction for Causal Effects of Continuous Treatments

This paper introduces a novel conformal prediction method that provides valid finite-sample prediction intervals for the potential outcomes of continuous treatments, even when the propensity score is unknown and must be estimated from data.

Original authors: Maresa Schröder, Dennis Frauen, Jonas Schweisthal, Konstantin Heß, Valentyn Melnychuk, Stefan Feuerriegel

Published 2026-02-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Maresa Schröder, Dennis Frauen, Jonas Schweisthal, Konstantin Heß, Valentyn Melnychuk, Stefan Feuerriegel

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 Picture: Why We Need a "Safety Net"

Imagine you are a doctor trying to decide how much chemotherapy to give a cancer patient. You have a computer program (an AI) that predicts how the tumor will shrink based on the dosage.

The Problem: Most AI programs only give you a single number: "If you give 50mg, the tumor will shrink by 2mm." This is called a point estimate. It's like a weather forecast saying, "It will be exactly 72°F." But what if it's actually 60°F or 80°F? In medicine, being wrong can be dangerous. You need to know the range of possibilities: "The tumor will likely shrink between 1mm and 3mm."

The Solution: This paper introduces a new mathematical "safety net" called Conformal Prediction. Instead of just giving a single guess, it draws a box around the answer to say, "We are 95% sure the real answer is inside this box."

The Specific Challenge: Continuous Doses

Most existing safety nets only work for simple choices, like "Give the drug" or "Don't give the drug" (Binary). But in real life, doctors deal with continuous treatments—like adjusting a dial from 1mg to 100mg.

The authors had to solve two tricky problems to make their safety net work for these continuous dials:

1. The "New Policy" Problem (The Shift)

Imagine you are training a driver in a simulation where they always drive on the right side of the road (Observational Data). Now, you want to test them in a country where they must drive on the left (Intervention).

The Issue: The driver's habits (the data distribution) are different in the new country. If you use the old safety rules, they might fail.
The Paper's Fix: The authors created a way to mathematically "re-calibrate" the safety net. They account for the fact that the rules of the road have changed, ensuring the safety net still holds even when the patient's treatment is different from what the AI has seen before.

2. The "Hidden Map" Problem (Unknown Propensity)

In the real world, we often don't know why a doctor prescribed a specific dose in the past. Maybe Dr. Smith always gives high doses, while Dr. Jones gives low doses. We call this the propensity score (the hidden map of how treatments are chosen).

The Issue: Usually, to build a perfect safety net, you need to know this hidden map exactly. But in real hospitals, we don't have that map; we have to guess it based on past records. Guessing introduces new errors.
The Paper's Fix: This is the paper's biggest breakthrough. They built a safety net that still works even if you have to guess the hidden map. They mathematically proved that even if your guess about the doctor's habits is slightly off, the safety net remains valid and won't break.

How They Did It (The Analogy of the "Stretched Rubber Band")

Think of the AI's prediction error as a rubber band.

Standard Method: You stretch the rubber band based on past data to see how far it might snap.
The Authors' Method: They realized that when you change the treatment (the intervention), the rubber band stretches differently because the "wind" (the data distribution) has changed.
The Innovation: They developed a new way to measure that stretch.
- If you know the wind pattern (Known Propensity), they use a precise formula to adjust the rubber band.
- If you don't know the wind pattern (Unknown Propensity), they use a "worst-case scenario" approach. They assume the wind might be slightly stronger or weaker than your guess, and they stretch the rubber band just enough to cover those possibilities.

What They Proved (The Results)

The authors didn't just make up a theory; they tested it.

Synthetic Tests: They created fake medical data where they knew the "true" answer. They compared their method against other popular AI uncertainty methods (like "MC Dropout," which is a common but often unreliable trick).
- Result: The other methods often gave "false confidence"—their boxes were too small and missed the true answer. The authors' method consistently kept the true answer inside the box, exactly as promised.
Real-World Test: They tested it on a real dataset of ICU patients (MIMIC-III), predicting blood pressure based on how long a patient was on a ventilator.
- Result: Their method showed wider, more honest intervals for patients who received unusual treatment durations (where data was scarce), whereas other methods gave deceptively narrow, risky intervals.

The Bottom Line

This paper provides a new tool for doctors and data scientists. It allows them to use AI to predict the effects of continuous treatments (like drug dosages) with a guaranteed safety net.

Even if the AI has to guess how doctors usually prescribe medicine, this new method ensures that the "safety box" around the prediction is wide enough to be trustworthy. This is crucial for safety-critical fields like medicine, where being wrong can have serious consequences.

Key Takeaway: They are the first to create a mathematically proven "safety box" for continuous drug dosages that works even when we don't fully understand the history of how those dosages were chosen.

Technical Summary: Conformal Prediction for Causal Effects of Continuous Treatments

1. Problem Formulation

The paper addresses the critical need for Uncertainty Quantification (UQ) in causal machine learning, specifically for continuous treatments. While causal ML is widely used in safety-critical domains like personalized medicine (e.g., chemotherapy dosage), existing methods typically provide only point estimates, neglecting prediction uncertainty.

The authors identify three specific challenges in applying Conformal Prediction (CP) to continuous causal effects:

Distribution Shift (Challenge a): Intervening on a treatment variable shifts the covariate distribution (specifically the propensity score). This violates the exchangeability assumption inherent to standard CP, rendering standard intervals invalid between observational and interventional distributions.
Unknown Propensity Scores (Challenge b): In observational data, the propensity score is unknown and must be estimated. This estimation introduces additional uncertainty that standard CP methods do not account for, potentially leading to under-coverage.
Data Sparsity for Continuous Values (Challenge c): Unlike binary treatments, specific continuous treatment values are rarely observed in the data, making direct application of discrete CP methods impossible.

The objective is to derive finite-sample valid prediction intervals $C(X_{n+1}, \diamond)$ for potential outcomes $Y_{n+1}(\diamond)$ under either hard interventions ( $a^*$ ) or soft interventions ( $A^*(X_{n+1})$ ), even when the propensity score is unknown and estimated.

2. Methodology

The proposed method extends Split Conformal Prediction to handle distribution shifts induced by interventions and uncertainty from propensity estimation.

Core Framework

The method utilizes a non-conformity score, typically the residual $S_i = |Y_i - \phi(X_i, A_i)|$ , where $\phi$ is an arbitrary machine learning model. The goal is to construct an interval such that $P(Y_{n+1}(\diamond) \in C) \geq 1 - \alpha$ .

The authors distinguish two scenarios based on the availability of the propensity score $\pi(a|x)$ :

Scenario 1: Known Propensity Score

When the observational policy is known (e.g., in a randomized trial or known policy setting), the shift is modeled as a change from $\pi(a|x)$ to $\pi(a+\Delta A|x)$ .

Approach: The authors reformulate the problem as an augmented quantile regression. They define a function class $\mathcal{F}$ representing the propensity shift.
Optimization: They derive a dual optimization problem (Theorem 4.2) to find the $(1-\alpha)$ -quantile of the non-conformity scores under the shifted distribution. This avoids the infeasibility of solving for all possible imputed scores by leveraging strong duality properties.
Result: Theorem 4.2 provides finite-sample coverage guarantees for potential outcomes under known propensity shifts.

Scenario 2: Unknown Propensity Score

When the policy is unknown, the method must handle hard interventions ( $a^*$ ) where the propensity becomes a Dirac delta function, and the score $\pi$ must be estimated ( $\hat{\pi}$ ).

Assumption: The estimation error is bounded such that the ratio $c_a = \hat{\pi}/\pi$ lies within $[1/M, M]$ for some constant $M$ .
Smoothing: To handle the discontinuity of the Dirac delta function, the authors approximate it using a Gaussian kernel with variance $\sigma^2 \to 0$ .
Optimization: The problem is formulated as a non-convex optimization (Eq. 18) involving the smoothing parameter $\sigma$ and the error bound $c_a$ .
Theoretical Guarantee: The authors prove (Lemma 4.4) that the problem is Type-I invex and satisfies the Linear Independence Constraint Qualification (LICQ). This allows the use of Karush-Kuhn-Tucker (KKT) conditions to derive global optimality.
Result: Theorem 4.5 establishes finite-sample coverage guarantees for unknown propensity scores by optimizing over the uncertainty bounds and smoothing parameters.

3. Key Contributions

Novel Method for Continuous Treatments: The first proposal of conformal prediction intervals for potential outcomes of continuous treatments, addressing the gap left by existing methods limited to binary/discrete treatments.
Finite-Sample Guarantees under Unknown Propensity: Derivation of valid prediction intervals that account for the uncertainty introduced by estimating propensity scores, without requiring asymptotic assumptions.
Algorithmic Implementation: An efficient algorithm (Alg. 1) for calculating these intervals using dual formulations and KKT conditions, avoiding brute-force search over imputed scores.
Empirical Validation: Demonstration of effectiveness on synthetic and real-world datasets.

4. Experimental Results

The authors evaluated their method against baselines including Monte Carlo (MC) Dropout, Deep Ensembles, and Gaussian Process (GP) regression.

Synthetic Datasets:
- Faithfulness: The proposed CP intervals achieved empirical coverage consistent with the target significance level ( $1-\alpha$ ) across various soft and hard interventions.
- Baselines: MC Dropout and Deep Ensembles showed significantly lower coverage (unfaithful intervals), often failing to capture the true outcome. GP regression provided valid coverage only for very small distribution shifts but failed under larger shifts due to narrow intervals.
- Stability: The proposed method exhibited low variability in empirical coverage compared to the high variance observed in MC Dropout.
Real-World Dataset (MIMIC-III):
- Task: Predicting blood pressure outcomes under varying durations of mechanical ventilation.
- Observation: The CP intervals correctly reflected higher uncertainty in treatment regions with sparse training data (high ventilation durations). In contrast, MC Dropout intervals were narrower, suggesting under-coverage in these critical regions.

5. Significance and Claims

The paper claims to be the first to propose conformal prediction for continuous treatments where the propensity score is unknown and must be estimated from data.

Safety-Critical Applicability: The method provides finite-sample validity guarantees, making it suitable for safety-critical applications like personalized medicine where asymptotic guarantees (which may fail in small samples) are insufficient.
Model Agnosticism: The approach works with any underlying machine learning model $\phi$ for outcome prediction.
Robustness to Distribution Shift: Unlike standard CP, the method explicitly accounts for the distribution shift caused by the intervention, ensuring validity even when the test distribution differs from the training distribution.
Practical Utility: The authors note that while the method is conservative (potentially wider intervals) when propensity estimation is poor or data is sparse, this conservatism ensures reliability, preventing the dangerous under-coverage seen in other UQ methods.

The work fills a critical gap in causal ML, enabling reliable decision-making under uncertainty for continuous treatment variables in both experimental and observational settings.