An example of Double Machine Learning using Conformal Prediction

This notebook illustrates an early-stage demonstration of the potential for the use of conformal predictions in the double machine learning framework.

The main motivation for bringing conformal predictions into double machine learning framework is to:

• Propagate uncertainty in nuisance model predictions to causal inference

• Reduce computational burden by avoiding cross-fitting

The below introduces some of this motivation in more detail, and shows a particular example using simulated data (from make_plr_CCDDHNR2018()).

What is Conformal Prediction?

Conformal prediction (a.k.a. conformal inference) is a user-friendly paradigm for creating statistically rigorous uncertainty sets/intervals for the predictions of such models. Critically, the sets are valid in a distribution-free sense: they possess explicit, non-asymptotic guarantees even without distributional assumptions or model assumptions.

— Angelopoulos and Bates (2022)

Why Conformal Double Machine Learning?

Sample splitting, typically in the form of cross-fitting, is one of the key features of standard Frequentist Double Machine Learning (FDML) which aims to solve the issue of over-fitting bias.

Cross-fitting alleviates the "potential dependence between nuisance estimates and parts of the data used for estimating the target parameter" (Ahrens et al (2025)).

As stated in Ahrens et al (2025), p 4:

"...Because \(\hat{\eta}\) is an estimator, it is itself a random function of the data. \(\hat{\eta}\) is thus generally correlated with the observations \(\{W_i\}_{i=1}^n\) also used in the estimating equation \(\frac{1}{n}\sum_{i=1}^n m(W_i; \theta, \hat{\eta})\). When this dependence is strong, for example due to "overfitting", it may generate large differences between \(\frac{1}{n}\sum_{i=1}^n m(W_i; \theta, \hat{\eta})\) and \(\frac{1}{n}\sum_{i=1}^n m(W_i; \theta, \eta_0)\), which results in poor performance of \(\hat{\theta}\)."

In practice, however, there are a few issues that cross-fitting does not resolve:

• First, in the presence of large data, a practical issue is that cross-fitting can be computationally costly as it requires fitting a model at least 1 time for each fold of the cross-validation set.

• Second, and more fundamentally, a theoretical issue is that cross-fitting does not account for uncertainty in the predictions \(\hat{\eta}\), but rather treats them as point estimates. Any uncertainty in these point estimates is not propagated into the causal inference for \(\theta\). Thus, FDML estimates of the causal parameter, \(\hat{\theta}\), often do not have good coverage - e.g., using simulated data where the true causal effect is known, FDML often leads to 95% confidence intervals which include the true effect less than 95% of the time.

As shown in this notebook, however, over-fitting bias can be alleviated without cross-fitting! If we instead think of our estimates for \(\hat{\eta}\) as following some joint probability distribution, we can simply fit one time, using one holdout set for the conformal prediction calibration; we can then directly account for the uncertainty in our nuisance estimates, and propagate this uncertainty through to the final inference stage.

Load MLJ models

Let's experiment with EvoTrees, RandomForest, and Symbolic Regression.

Ensure the DoubleMLConformalExt is accessible

Below, we get the extension for estimating Conformal Double Machine Learning models.

This is implemented as a package extension as it remains experimental.

Data generation

Estimating a Conformal Double Machine Learning (CDML) model

As noted above, this is a bit of a contrived example. In anecdotal testing, the conformal model appears to give better empirical coverage. However, more extensive testing is needed to properly evaluate the performance of the different methods.

How does Conformal Double Machine Learning work?

Currently, the main implementation of CDML (DoubleMLPLRConformal) works by:

• Training the conformal models without cross-fitting. Users may specify train_ratio for some conformal prediction methods, but predictions are made on the full dataset

• Obtaining conformal predictions (i.e, a tuple of a lower and upper bound for each prediction). These conformal predictions guaranteed a user-defined coverage level (e.g., 95%).

• Use Monte Carlo sampling from conformal prediction intervals to propagate uncertainty, using Beta(2,2) marginals with Gaussian copula to account for correlation between the uncertainties in predictions for the outcome \(\hat{l}(x)\) and treatment \(\hat{m}(x)\).

Alternative CDML variant using the Unscented Transform (UT)

Monte Carlo sampling can be computationally intensive and stochastic. Instead of MC, this variant propagates the uncertainty in the nuisance parameters deterministically through the DML2 score function using the unscented transform at second-order accuracy.

The UT variant of the CDML model (DoubleMLPLRConformalUT) currently works as follows:

Workflow

1. Conformal model fitting: The conformal-wrapped learners ml_l and ml_m are trained to obtain prediction intervals. Users may specify n_folds and n_rep for K-fold cross-fitting (default n_folds=1 fits on the full dataset).

2. Conformal predictions: A lower and upper bound are obtained for each prediction, guaranteeing the user-defined coverage level (e.g., 95%).

3. Uncertainty modeling: Prediction errors within each interval are modeled with Beta(2,2) marginals. The correlation between the errors in \(\hat{l}(x)\) and \(\hat{m}(x)\) is estimated from residuals and its uncertainty is represented via the Fisher z-transform\(z \sim N(\hat{z}, 1/(n-3))\).

4. Closed-form moment propagation: For any fixed correlation, the means and covariances of the aggregated score statistics \((\bar{A}, \bar{B})\) are computed in closed form (O(n)) using a Gaussian copula (Isserlis cross-moments).

5. 2D Unscented Transform: The score function is propagated via the standard 2D UT (5 sigma points), yielding a mean that corrects for Jensen bias and a variance that captures second-order effects.

6. Correlation uncertainty integration: The UT is evaluated at Gauss-Hermite quadrature points (default 3, optional 5) over the Fisher z-transform distribution. The final result mixes these components by the law of total variance.

7. Combined inference: The reported coefficient is the UT mean. The reported standard error combines the standard DML sampling SE with the conformal-only UT variance: \(SE_{total} = \sqrt{SE_{DML}^2 + Var_{UT}}\).