Localising Dropout Variance in Twin Networks

TL;DR

This paper introduces a layer-wise variance decomposition for twin-network models, enabling practitioners to locate sources of model failure and improve uncertainty estimation in treatment-effect predictions.

Contribution

It proposes a novel variance decomposition method that distinguishes between encoder and head uncertainties, enhancing diagnostic capabilities in out-of-distribution scenarios.

Findings

Encoder variance dominates under distributional shift.

Encoder uncertainty effectively predicts out-of-distribution errors.

Decomposition adds minimal computational cost.

Abstract

Accurate individual treatment-effect estimation demands not only reliable point predictions but also uncertainty measures that help practitioners \emph{locate} the source of model failure. We introduce a layer-wise variance decomposition for deep twin-network models: by toggling Monte Carlo Dropout independently in the shared encoder and the outcome heads, we split total predictive variance into an \emph{encoder component} ($\sigma_{\mathrm{enc}}^2$) and a \emph{head component} ($\sigma_{\mathrm{head}}^2$), with $\sigma_{\mathrm{enc}}^2 + \sigma_{\mathrm{head}}^2 \approx \sigma_{\mathrm{tot}}^2$ by the law of total variance. Across three synthetic covariate-shift regimes, the encoder component dominates under distributional shift ($\rho_{\mathrm{enc}}=0.53$) while the head component becomes informative only once encoder uncertainty is controlled. On a real-world twins cohort with induced multivariate shift, only $\sigma_{\mathrm{enc}}^2$ spikes on out-of-distribution samples and becomes the primary error predictor ($\rho_{\mathrm{enc}}\!\approx\!0.89$), while $\sigma_{\mathrm{head}}^2$ remains flat. The decomposition adds negligible cost over standard MC Dropout and provides a practical diagnostic for deciding whether to collect more diverse covariates or more outcome data.