Controlling for Omitted Variable Bias in Deep Neural Networks

Deep neural networks (DNNs) often exhibit shortcut learning, where models rely on spurious correlations between input data and known covariates (e.g., demographic variables) to improve predictive performance. This phenomenon is mathematically equivalent to omitted variable bias (OVB) in classical statistics. The authors investigate how to incorporate control variables into deep learning architectures to obtain unbiased, interpretable estimates of the direct effect of the input (e.g., an image) on the outcome, rather than just maximizing predictive accuracy.

To address OVB, the authors propose a framework that treats the final layer of a pre-trained DNN as an additive model. They incorporate covariates explicitly into this final layer, using a back-fitting algorithm with ridge penalization to stabilize estimation. To avoid the bias introduced by training the backbone and the final layer on the same data, they employ a cross-fitting procedure. The method allows for the optional orthogonalization of features to isolate residual effects and supports marginalization over the covariate distribution to produce controlled predictions.

The study demonstrates that standard confound-control methods—such as adversarial training or post-hoc decorrelation—fail to remove OVB because they target the correlation between predictions and covariates rather than the underlying structural bias. In contrast, the proposed refitting approach consistently recovers true effects in simulated environments, even when the image and confounder are highly correlated. In real-world neuroimaging experiments, the method successfully recovers predictive performance on unconfounded test data, outperforming baseline models that absorb confounding information into their feature representations.

This work bridges the gap between classical statistical additive modeling and modern deep learning. By providing a rigorous framework for controlling for observed confounders, it enables researchers to use DNNs for scientific hypothesis testing in high-stakes domains like medicine, where interpretability and generalizability are as critical as raw predictive power. It also offers a principled way to handle dataset shift by marginalizing predictions over target covariate distributions.