Disentangling Doubt in Deep Causal AI

Estimating predictive uncertainty in deep neural networks has been a major focus across vision, language, and beyond. Gal and Ghahramani [5] showed that Monte Carlo Dropout can be interpreted as approximate Bayesian inference, providing a scalable means to capture epistemic uncertainty. Deep ensembles [6] improve on this by averaging predictions from multiple independently trained models, yielding both strong accuracy and reliable uncertainty estimates. Deep Gaussian Processes [7] stack Gaussian Process priors in multiple layers to model both epistemic and aleatoric uncertainty nonparametrically, though scalability remains a challenge. More recent work has revisited the foundational definitions of aleatoric versus epistemic uncertainty, arguing for richer taxonomies and coherence in their interpretation [8, 9].

Heterogeneous treatment-effect estimation has seen extensive methodological development. Meta-learners such as the T-, S-, and X-learners [10] adapt any base learner to estimate individual treatment effects. Representation-learning approaches like TARNet [1] and DragonNet [2] mitigate selection bias by first learning a shared covariate encoder and then branching into separate treatment and control outcome heads. Bayesian methods—most notably Bayesian Additive Regression Trees (BART) [3, 11] and Bayesian Causal Forests [4]—provide interval estimates for average and individual treatment effects but do not decompose uncertainty into distinct epistemic and aleatoric components. A recent survey of deep causal models highlights the absence of structured uncertainty quantification in this literature [12].

Across uncertainty quantification research, a clear distinction is drawn between aleatoric uncertainty (irreducible outcome noise) and epistemic uncertainty (model uncertainty reducible with more data) [13]. In computer vision and reinforcement learning, methods such as multi-headed networks that jointly predict mean and variance [13] or single-model techniques for simultaneous aleatoric and epistemic estimation [14] have begun to disentangle these sources. However, modular deep architectures for causal inference have lacked an analogous structured decomposition. To our knowledge, no prior work has explicitly separated representation-level uncertainty (due to limited or imbalanced feature encoding) from prediction-level uncertainty (due to outcome noise) within a unified deep treatment-effect model.

Let $x\in\mathcal{X}\subseteq\mathbb{R}^{p}$ be a feature vector, $t\in\{0,1\}$ the binary treatment indicator, and $Y(t)$ the potential outcome under treatment $t$. We observe $Y=Y(T)$ for each unit. A twin network comprises

• •

  a shared encoder $\Phi(x;\theta_{e})\colon\mathbb{R}^{p}\to\mathbb{R}^{d}$, and

• •

  two outcome heads $f_{0}(z;\theta_{0}),f_{1}(z;\theta_{1})\colon\mathbb{R}^{d}\to\mathbb{R}$,

so that the estimated individual treatment effect is

We augment this architecture with Monte Carlo Dropout to obtain uncertainty estimates on $\hat{Y}_{t}(x)=f_{t}(\Phi(x))$.

Following Gal and Ghahramani [5], we insert dropout after every hidden unit in both encoder and heads. Each stochastic forward pass samples a dropout mask, yielding

where $(\theta_{e}^{(n)},\theta_{t}^{(n)})\sim q(\theta)$ approximates the posterior. The overall predictive mean and variance are then estimated by

A key insight is that total predictive variance admits a module-level split by the law of total variance. Denote by $\hat{Y}_{t}(x)$ the random output under all dropout masks; then

where

Intuitively, the first term measures uncertainty in the encoder embedding, while the second term measures uncertainty in the prediction heads. In practice we approximate these via controlled dropout:

• •

  Representation uncertainty: enable dropout only in the encoder (heads deterministic) to estimate

  $$\operatorname{Var}_{\rm rep}(x,t)\approx\frac{1}{N}\sum_{n=1}^{N}\bigl{(}\hat{Y}^{(n)}_{t,\mathrm{rep}}(x)-\overline{Y}_{t,\mathrm{rep}}(x)\bigr{)}^{2}.$$

• •

  Prediction uncertainty: enable dropout only in the heads (encoder deterministic) to estimate

  $$\operatorname{Var}_{\rm pred}(x,t)\approx\frac{1}{N}\sum_{n=1}^{N}\bigl{(}\hat{Y}^{(n)}_{t,\mathrm{pred}}(x)-\overline{Y}_{t,\mathrm{pred}}(x)\bigr{)}^{2}.$$

We empirically verify that

thus confirming our decomposition.

Our implementation (Figure 1) builds on the twin-network paradigm [1, 2]. We use a multi-layer perceptron encoder $\Phi(x;\theta_{e})$ with dropout after each hidden block, and two identically-structured heads $f_{0},f_{1}$ each with their own dropout. We train by minimizing the factual mean-squared error

Optionally, one can add an IPM-based regularizer to balance treated and control embeddings [1].

At test time for a new $x$:

1. 1.

   mode=‘total’:  Run $N$ stochastic passes to compute $\bar{Y}_{t}(x)$ and $\operatorname{Var}_{\rm tot}(x,t)$.

2. 2.

   mode=‘rep_only’:  Enable dropout only in the encoder to compute $\operatorname{Var}_{\rm rep}(x,t)$.

3. 3.

   mode=‘pred_only’:  Enable dropout only in the heads to compute $\operatorname{Var}_{\rm pred}(x,t)$.

4. 4.

   Estimate $\hat{\tau}(x)=\bar{Y}_{1}(x)-\bar{Y}_{0}(x)$ and its variance

   $$\operatorname{Var}(\hat{\tau}(x))\approx\operatorname{Var}_{\rm rep}(x,1)+\operatorname{Var}_{\rm rep}(x,0)+\operatorname{Var}_{\rm pred}(x,1)+\operatorname{Var}_{\rm pred}(x,0).$$

We use three versions of our simulator (Appendix A):

• •

  v1 (Sin×Sin): $Y_{0}=\sin(\pi x_{1}x_{2})+\epsilon$, $\tau=1.5+0.5x_{1}$.

• •

  v2 (Polynomial): $Y_{0}=(x_{1}^{2}+x_{2}^{2})/2+\epsilon$, $\tau=2+x_{1}x_{2}$.

• •

  v3 (Sin+Linear): $Y_{0}=\sin(x_{1})+\cos(x_{2})+0.3x_{1}x_{2}+\epsilon$, $\tau=1+\sin(x_{1}+x_{2})$.

For each version we apply sampling-shift and noise-shift. Domain shift is strong in v1–v2 and mild in v3. We use an 80/20 train/test split and define OOD points via a 10-NN density score (top 20%).

We compare:

• •

  MC-Dropout (total): standard MC Dropout (all masks on).

• •

  Our Method: structured MC Dropout with separate rep_only and pred_only modes.

• •

  Ensemble: 5 deterministic models, variance = Var${}_{\rm ens}[\hat{\tau}]$.

For each test point we measure:

• •

  Spearman’s $\rho$ between each uncertainty ($\sigma_{\rm rep}^{2},\sigma_{\rm pred}^{2},\sigma_{\rm tot}^{2}$) and absolute error $|\hat{\tau}-\tau|$.

• •

  For twins: $\Delta\sigma^{2}=\mathbb{E}[\sigma^{2}\mid\mathrm{OOD}]-\mathbb{E}[\sigma^{2}\mid\mathrm{ID}]$ and ROC-AUC for OOD detection.

All models in PyTorch, Adam (lr=1e-3, weight_decay=1e-4), 50 epochs, dropout=0.2, $N=1000$ MC samples, batch size 128. Code and data at github.com/mercury0100/TwinDrop.

Figure 2 illustrates the decomposition on v1 under sampling- and noise-shift. The three panels show encoder ($\sigma_{\rm rep}^{2}$), control-head and treatment-head ($\sigma_{\rm pred}^{2}$) uncertainty over $(x_{1},x_{2})$.

We verify variance additivity ($\sigma_{\mathrm{rep}}^{2}+\sigma_{\mathrm{pred}}^{2}\approx\sigma_{\mathrm{tot}}^{2}$) and plot error-correlation: $\rho(\sigma^{2},|\hat{\tau}-\tau|)$ for each mode (Figure 3).

Table 1 reports Spearman’s $\rho$ between each uncertainty and $|\hat{\tau}-\tau|$ for v1–v3.

Notably, $\sigma_{\rm pred}^{2}$ becomes informative once epistemic uncertainty is controlled. To show this we sweep a maximum-allowed $\sigma_{\rm rep}^{2}$ threshold and recompute Spearman’s $\rho(\sigma_{\rm pred}^{2},|\hat{\tau}-\tau|)$ on the remaining points. Figure 4 (left) shows this for v1 (strong shift), where $\rho_{\rm pred}$ climbs from near zero to above 0.5 as we remove high-$\sigma_{\rm rep}^{2}$ points; the right panel shows a similar effect in v3 (mild shift).

We assess interval calibration both before and after a simple post-hoc conformal adjustment (using a held-out calibration fold). Raw MC-Dropout intervals on v1 and v3 exhibit nontrivial miscalibration (ECE $=0.104$ and $0.130$, respectively). After conformal adjustment, ECE falls to $0.005$ on v1 and $0.022$ on v3. Figure 5 presents the raw and calibrated reliability curves side by side.

We compare overall uncertainty from a 5-model deterministic ensemble against MC-Dropout. Table 2 reports Spearman’s $\rho(\sigma_{\rm tot}^{2},|\hat{\tau}-\tau|)$.

On the same-sex twins cohort, we induce a multivariate bias along PC1 and split test points into OOD (20%) vs. ID via 10-NN density. Table 3 summarizes $\Delta\sigma^{2}$, ROC-AUC, and $\rho(\sigma^{2},|\hat{\tau}-\tau|)$.

These results confirm that under realistic multivariate shift only $\sigma_{\rm rep}^{2}$ substantially increases on OOD points and becomes the primary error predictor, while $\sigma_{\rm pred}^{2}$ remains flat and $\sigma_{\rm tot}^{2}$ tracks the dominant component.