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_DTL2025()).

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 propogated 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 confidence intervals which do not include the true effect.

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.

Hypothesis

By sampling from the joint probability distribution for each prediction from the nuisance models, the hypothesis is that we can maintain Neyman orthogonality and avoid over-fitting bias, and make better inference decisions by improving uncertainty quantification.

begin
    using DoubleML
    using ConformalPrediction  # This triggers loading of DoubleMLConformalExt
    using MLJ
    using StableRNGs
    using Random
    using DataFrames
end

Load MLJ models

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

begin # loading MLJ models
    EvoTreeRegressor = @load EvoTreeRegressor pkg = EvoTrees verbosity = 0
    RandomForestRegressor = @load RandomForestRegressor pkg = BetaML verbosity = 0
    SRRegressor = @load SRRegressor pkg = SymbolicRegression verbosity = 0
end;

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.

const Ext = Base.get_extension(DoubleML, :DoubleMLConformalExt)

DoubleMLConformalExt

Data generation

Note

The below is specifically a counter-example to show where FDML may fail in terms of coverage. A more comprehensive assessment across multiple random seeds would be needed for more a proper evaluation.

Here, we use the data generation process from Section 6 of DiTraglia and Liu (2025), and use their default paramater choices of:

\(\alpha =\) 2.0 (the true causal effect)
\(n =\) 200 (number of observations)
\(p =\) 100, (number of covariates), and
\(\sigma_{\varepsilon} = 2\)

Let's compare the results of fitting a Conformal Double Machine Learning model, vs the a standard DML model with 5-fold cross-fitting.

Note

Note that the illustrative example below is specifically a counter-example to illustrate where standard DML may fail, where conformal model may succeed. In reality, on this specific problem, both methods (conformal and non-conformal FDML) often fail to capture the true causal effect.

Larger-scale simulations across a wide range of random seeds are needed for more comprehensive evaluation of the empirical performance of the different methods to truly evaluate the performance of each method.

As set out in the paper, Bayesian Double Machine Learning (BDML) performs very well on this problem. Keep an eye out for a forthcoming BayesianDoubleML.jl package which implements the models set out in DiTraglia and Liu (2025)!

begin
    seed = 60
    rng = StableRNG(seed)

    true_alpha = 2.0

    n = 200
    p = 100
    sigma_epsilon = 2.0

    data = make_plr_DTL2025(n, p, sigma_epsilon; alpha = true_alpha, rng = rng)
end

DoubleMLData{Float32, Vector{Float32}}(Float32[-8.547153, 2.537241, 2.5340595, -0.4896508, -4.6916947, 7.371019, 2.0968182, 0.22039314, -3.3462834, 4.2125087  …  -4.048314, 1.3922781, 0.025425058, 2.4426374, -1.2662034, -4.7019324, -0.8984496, 2.6580913, 1.6113335, -0.037568532], Float32[-3.8231306, 1.3155899, 0.91706735, 1.2340404, -1.9633923, 0.707324, 0.7664218, -1.1376064, -1.4266428, 1.698602  …  -0.90348893, -0.3089625, -0.57444024, -0.040365264, -0.7284262, -0.42811778, -2.2260504, 0.31699485, 0.84488803, -0.39746693], Float32[-0.64719266 0.3935131 … 0.43429512 -0.1457195; 0.06077252 -0.046264276 … 0.42417604 0.0051582777; … ; 0.68033195 -3.2544427 … 0.22436728 1.3658123; 1.3604089 1.5732734 … -0.3286756 0.60283846], 200, 100, :y, :d, [:X1, :X2, :X3, :X4, :X5, :X6, :X7, :X8, :X9, :X10  …  :X91, :X92, :X93, :X94, :X95, :X96, :X97, :X98, :X99, :X100])

Estimating a Conformal Double Machine Learning (CDML) model

begin
    Random.seed!(seed)

    # Set the coverage for the nuisance models.
    coverage = 0.99

    ml_l = conformal_model(
        EvoTreeRegressor(seed = seed);
        method = :simple_inductive,
        coverage = coverage
    )
    ml_m = conformal_model(
        EvoTreeRegressor(seed = seed);
        method = :simple_inductive,
        coverage = coverage
    )
    # Create and fit conformal model
    model_conformal = Ext.DoubleMLPLRConformal(
        data,
        ml_l,
        ml_m;
        n_mc_samples = 1_000
    )
    @time Ext.fit!(model_conformal, rng = rng, verbose = 0)

    coeftable(model_conformal)
end

────────────────────────────────────────────────────────────────────
   Estimate  Std. Error  z value  Pr(>|z|)  Lower 95.0%  Upper 95.0%
────────────────────────────────────────────────────────────────────
d   1.78495    0.133532    13.37    <1e-40      1.51971      2.04702
────────────────────────────────────────────────────────────────────

# Test against standard PLR model
begin
    Random.seed!(seed)

    model = DoubleMLPLR(
        data,
        EvoTreeRegressor(seed = seed),
        EvoTreeRegressor(seed = seed);
        n_folds = 5,
        n_rep = 1
    )

    @time DoubleML.fit!(model; verbose = 0)

    coeftable(model)
end

────────────────────────────────────────────────────────────────────
   Estimate  Std. Error  z value  Pr(>|z|)  Lower 95.0%  Upper 95.0%
────────────────────────────────────────────────────────────────────
d   1.57854    0.110481    14.29    <1e-45        1.362      1.79508
────────────────────────────────────────────────────────────────────

Results

The above example shows that the conformal model ran roughly 4-5x faster than the standard model that uses 5-fold cross-fitting. In addition, the conformal model includes the true causal effect, whereas the standard model does not.

How does Conformal Double Machine Learning work?

Warning

This package, and the implementation of Conformal Double Machine Learning, remain experimental.

Currently, the implementation of CDML 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)\).