Controlling for Omitted Variable Bias in Deep Neural Networks

Manuel Pfeuffer; Roshan Prakash Rane; Kerstin Ritter; Sonja Greven

Source-linked AI summary

Controlling for Omitted Variable Bias in Deep Neural Networks

Manuel Pfeuffer, Roshan Prakash Rane, Kerstin Ritter, Sonja Greven

arXiv:2608.25930v1Aug 26, 2026stat.MEcs.CVcs.LG

TL;DR

Deep neural networks can encode outcome-associated covariates from inputs, creating shortcut learning and omitted variable bias, while correlation-based controls do not correct that bias. The paper adds covariate effects to a DNN’s final layer, using refitting, cross-fitting, ridge penalisation, and optional orthogonalisation. The resulting additive decomposition supports unbiased direct-effect estimation and prediction marginalisation over covariate distributions.

Problem
Deep learning has underexplored control variables despite image-encoded covariates creating omitted variable bias, and existing correlation-based controls fail to correct it.
Method
The method refits a pre-trained DNN’s final layer with additive covariate effects, using cross-fitting, ridge penalisation, and optional orthogonalisation.
Results
Theoretical results and simulations show last-layer refitting can recover direct effects at lower sample sizes than existing methods when covariate effects are correctly modelled.
Takeaways & Limitations
Additive separation enables predictions to be marginalised over covariate distributions, including replacement of a training distribution with a target distribution under dataset shift.
Takeaways & Limitations
The method assumes confounding covariates are observed, and the estimated image effect depends on the expressiveness of the DNN backbone.

Abstract

from arXiv · show

Control variables are widely used in statistical modelling to account for omitted variable bias of known confounders. However, they have largely been underexplored in deep learning. This is surprising, given that deep learning models encode image-inferable covariates, such as demographic variables, into their predictions when these covariates are correlated with the outcome---a form of omitted variable bias referred to as 'shortcut learning'. While many existing confound-control or fairness methods try to restrict the correlation of such covariates with model predictions, we show that this fails to correct for omitted variable bias. We therefore propose a control variable approach for deep learning models, based on generalised additive modelling of the effects of model inputs and covariates. As flexible additive models can suffer from concurvity, we introduce an estimation procedure that refits the final layer of a pre-trained network to include covariate effects, using cross-fitting with ridge penalisation. We show how these effects can be orthogonalised with respect to covariates to exclude their mediated effects and that model predictions can be marginalised over the covariate distribution to control for their effect. This yields unbiased, interpretable predictions and offers flexibility to model the desired effects depending on the scientific or fairness objective. We verify our approach using simulated images, and demonstrate consistent estimation of true effects. Existing methods either require more data or fail to recover the true effects. We apply our method to real neuroimaging data with experimentally induced confounding, where it recovers prediction performance to near the level of a model trained on unconfounded data. Code is available at https://github.com/mpff/cocodeel.

1 Introduction

Deep neural networks can exploit confounding or sensitive covariates encoded in inputs, producing shortcut learning and omitted variable bias. The paper extends control-variable methods to DNNs by additively modelling covariate effects in the final layer, with refitting and cross-fitting procedures for estimation.

Motivation: Shortcut learning occurs when DNNs use input-encoded covariates, such as acquisition artefacts, sex, or age, that are associated with outcomes.These associations may improve training-distribution performance but typically fail to generalize out-of-sample and hinder interpretation and clinical applicability.
Problem formulation: Omitted variable bias arises when observed covariates Z correlate with both inputs X and outcomes Y, causing variation associated with Z to be attributed incorrectly to X.The paper assumes such confounding covariates are observed and available for modelling.
Approach: The proposed approach incorporates covariates explicitly and additively into a DNN’s final linear layer, connecting control-variable methods with GAMs, NAMs, and SSNs.The final layer operates as an additive model over a learned feature representation, or neural basis.
Motivation: Separating image and covariate components can prevent demographic information inferable from medical images from being attributed to the image effect.The motivating example concerns alcohol-misuse prediction from structural MRI when sex is correlated with the outcome and inferable from scans.
Theoretical contribution: The paper distinguishes direct and mediated covariate effects, showing that correlation-based confound-control and fairness methods control mediated effects but fail to address omitted variable bias.When both pathways exist, the mediated effect may be absorbed into the X-effect, depending on the application.
Estimation: Last-layer refitting includes control variables and optionally orthogonalises learned features with respect to covariates to estimate the unbiased X-effect without mediated effects.Back-fitting with ridge penalisation stabilizes estimation, while cross-fitting separates backbone training from final-layer refitting to avoid biased estimates.

2 Related Literature

Related work includes additive neural architectures and deep-learning methods that disentangle covariate variation from predictions. These approaches differ in whether they model covariate effects directly or adjust data and predictions to reduce covariate associations.

Additive neural models: Neural Additive Models and Semi-structured Networks additively model multiple input features, including images alongside additional confounders.Some prior work jointly trains deep and covariate effects in architectures similar to NAMs.
Confound control and fairness: Most deep-learning confound-control and fairness methods disentangle covariate variation from predictions through dataset adjustment or modelling-based approaches.Dataset adjustment can become infeasible with many covariates, while continuous covariates are difficult to balance without procedures such as binning.

3 Theoretical Framework

The framework models outcomes additively through network inputs and observed covariates, showing how omitted covariates bias input effects and why correlation-based constraints do not generally recover the desired effects. It decomposes input effects into covariate-mediated and residual components and addresses concurvity through identifiable estimation strategies.

3.1 Biased Estimation of fX due to Omitting Covariates Z: E[Y|X,Z] is modelled as an intercept plus an input effect fX(X) and a covariate effect fZ(Z).Effects are centred so the intercept equals the marginal mean outcome.
3.1 Biased Estimation of fX due to Omitting Covariates Z: ˆfX(X) equals fX(X) plus E_Z|X[fZ(Z)|X], so omitting Z absorbs its conditional direct effect into the estimated input effect.The paper identifies the added term as omitted variable bias and links it to biased DNNs trained only on X.
3.2 Decomposing fX into Z-Mediated and Residual Parts: The Z-mediated component of fX is the conditional expectation E_X|Z[fX(X)|Z], while the residual component is fX minus that expectation.This decomposition extends the classical product-of-coefficients formulation to high-dimensional and nonlinear input effects.
3.3 Why Correlation-based Confound Control Fails: Constraining predictions to be mean-independent of Z does not generally remove omitted variable bias and can yield biased estimates when part of the X-effect is mediated by Z.Post-hoc orthogonalisation satisfies the constraint but retains bias that grows with the strength of fZ.
3.4 The Problem of Concurvity in Estimation: Flexible additive models suffer from concurvity because input and covariate functions can match pointwise, making their separate effects unidentified despite unchanged predictions.The required sample size for identifiability depends on the flexibility of the input-function class.
3.4 The Problem of Concurvity in Estimation: End-to-end NAM-style fitting converges to fX only at larger sample sizes, whereas post-hoc orthogonalisation recovers fX^re; last-layer estimation with cross-fitting converges faster for both effects.The comparison uses a backbone with approximately 21,000 parameters.

4 Methodology

The method refits a pre-trained DNN’s final layer as a partial generalized additive model containing learned input features and covariate effects. Ridge regularisation, sample splitting or cross-fitting, orthogonalisation, and prediction marginalisation support stable effect estimation and covariate-controlled predictions.

4.1 Refitting the Last Layer: A pre-trained DNN backbone maps inputs to learned features, and the refitted final layer models those features linearly alongside an intercept and flexible covariate effect.The resulting model is a partial generalized linear model.
4.2 Estimation: The additive refit uses centred learned features and back-fitting equations for the feature coefficients and covariate function.These equations correspond to a semiparametric partial linear model with one smooth term.
4.3 Orthogonalisation: Post-hoc orthogonalisation applies the centred residual-maker matrix ˜M_Z to the centred feature matrix and estimated coefficients to remove the covariate-mediated component.The paper also provides an iteratively reweighted least-squares version for nonlinear output functions.
4.2 Estimation: A ridge penalty stabilises estimation when the number of last-layer features exceeds the refit sample size, with λ controlling the penalty strength.The procedure selects λ along a regularisation path.
4.2 Estimation: Disjoint pre-training and refit samples restore feature exogeneity and consistency, while K-fold cross-fitting reduces the data cost of sample splitting.Each fold trains the backbone on the other folds and refits the final layer on the held-out fold.
4.4 Prediction Controlled for Covariates: Marginalized prediction averages model predictions over the covariate distribution, approximated in practice by the empirical training distribution or a target distribution under dataset shift.This extends covariate control from effect estimates to predictions.

5 Simulation Study

The simulation study evaluates estimation under controlled image-confounding settings and shows that the proposed controls approach is consistent, while convergence slows with more or strongly correlated confounders.

Simulation setup: The study simulates images from strips whose intensities depend on continuous confounders, with continuous or binary outcomes generated from specified data-generating processes.The default setting uses c1 = 0.5, c2 = 0.5, p = 1, and unit outcome coefficients.
Simulation setup: Figure 3 visualizes simulated images constructed from v1, v2, and v3 strips, with the first two varying by confounder index.The three strip intensities are shown left, center, and right, respectively.
Convergence: As Ntrain increases, DNNs with controls consistently estimate fX, with both bias and variance decaying independently of βZ.This consistency holds with and without orthogonalisation and extends to binary outcomes fitted by IRLS.
Convergence: Uncontrolled DNNs and post-hoc orthogonalisation retain non-vanishing bias that grows with βZ.The observed pattern agrees with the paper’s omitted-variable result.
Failure modes: More confounders slow convergence, while perfect confounder encoding in the image can prevent convergence of the estimated image effect.Increasing image–confounder correlation requires larger samples, and c1 = 1 yields no convergence.

6 Application to Neuroimaging Data with Synthetic Confounding

The neuroimaging experiment introduces synthetic age confounding into UK Biobank MRI prediction and compares control, sample-splitting, and cross-fitting strategies using a pretrained backbone.

Experimental design: The study predicts high alcohol consumption from T1-weighted MRI using UK Biobank data with ntrain = 14,617 and ntest = 4,505.Synthetic confounding is introduced by resampling training data under βsex = 2 and βage ∈ {0, −2}.
Experimental design: Age-control models differ in whether backbone fitting and last-layer refitting share data, use a disjoint split, or use K = 2 or K = 3 cross-fitting.All models use a pretrained 3D ResNet-50 backbone.
Results: Under age-balanced training, all five models perform similarly, but sample splitting reduces performance because each stage uses only half the sample.Cross-fitting with K = 2 recovers this performance drop.
Results: With βage = −2, the uncontrolled DNN’s AUC on the age-balanced test set falls from ≈0.70 to ≈0.62.The model absorbs age into its image features under age-confounded training.
Results: Age control recovers most of the performance gap only when backbone fitting and refitting use disjoint data, while K = 3 cross-fitting further reduces fold-to-fold variance.Additional joint age-and-sex control results are reported elsewhere in the paper.

7 Limitations

The method depends on observed, sufficiently informative, and accurately measured confounding covariates, while its additive decomposition imposes a modelling boundary.

Scope assumptions: The method assumes that confounding covariates Z are observed; unmeasured confounding remains outside its scope.Demographic and acquisition-related variables are often recorded in imaging cohorts.
Representation: The estimated image effect depends on the expressiveness of the DNN backbone and becomes the best approximation within the learned-feature span when representation richness is limited.The omitted-variable correction itself does not depend on this richness.
Measurement error: Measurement noise in Z attenuates its estimated effect, leaving only partial correction of the image effect and asymptotically persistent bias.The bias interpolates between the noise-free and pure-noise cases.

8 Conclusion

The paper extends control-variable modelling to deep neural networks by refitting their final layer, separating direct image effects from covariate effects. Theory and simulations support consistent direct-effect estimation, while marginalisation enables covariate-aware prediction adjustment.

Contribution: The method estimates the direct image effect fX by fine-tuning a pretrained network’s last layer with covariates, cross-fitting, and ridge penalisation.An optional orthogonalisation step recovers the residual effect fX^re.
Findings: Theoretical results and simulations show that last-layer refitting can recover the direct effect at lower sample sizes than existing methods when covariate effects are correctly modelled.The paper also shows that correlation-based confound control does not correct omitted-variable bias.
Implications: Separating fX from fZ allows predictions to be marginalised over p(Z), including replacement of the training covariate distribution with a target distribution.This capability is unavailable to models that merely decorrelate predictions from Z.
Statistical interpretation: The conditional expectation learned by an uncontrolled mean-squared-error model combines fX with the expected covariate effect given X.Formally, ˆfX(X) = fX(X) + EZ|X[fZ(Z) | X].

A.2 Proof of Theorem 3

The proof characterizes the constrained estimator as an orthogonal projection onto functions of X that are orthogonal to covariate functions. It also establishes the corresponding sample estimator through centered smoothing and ridge-penalized residualization.

Minimizing the constrained objective is equivalent to minimizing squared distance from the unconstrained estimator over the constrained function set.The first decomposition term is independent of the candidate function, leaving an orthogonal projection problem.
Functions satisfying the conditional-mean constraint are orthogonal to the covariate-function space H_Z.The proof uses E[g(X)|Z] as an element of H_Z to establish the converse characterization.
The constrained estimator is the limit of alternating projections that remove the Z-conditional mean and project onto functions of X.The two conditional-expectation operators generally do not commute, so the limit has no closed form; under independence of X and Z, the constraint is already satisfied.
When the X-effect contains a mediated component, the constrained estimator differs from the direct X-effect and from the full X-effect.The full effect violates the constraint when its mediated component depends on Z, while the constrained estimate removes that component.
The sample procedure centers outcomes and covariate smoothers, then estimates the X-effect by ridge-penalized least squares after residualizing both outcomes and features with respect to Z.This is equivalent to fitting a ridge regression on the residualized data and mirrors the Speckman estimator apart from the ridge penalty.
The resulting estimator uses the ridge solution for residualized data, with the penalty strength selected through standard validation procedures.The ridge penalty stabilizes estimation when observations are few relative to the number of parameters.

A.5 Proof of Theorem 6

The proof establishes conditional consistency of the ridge-refitted X-effect for a pretrained backbone under model, smoother, and penalty assumptions. It extends the additive formulation to generalized outcomes through IRLS, where control variables reduce bias but can slow convergence.

The consistency result is conditional on the pretrained DNN backbone rather than proving consistency of the backbone itself.The latter would require assumptions about DNN asymptotics, which are outside the paper’s scope.
The linear model assumes additive features, covariate effects, and noise for i.i.d. observations, while allowing redundant DNN features.The target is the unique population minimum-ℓ2-norm coefficient rather than an arbitrary coefficient representation.
Consistency requires smoother-quality conditions for estimating covariate effects and a ridge penalty satisfying the stated asymptotic requirements.Regression splines with slowly increasing basis dimension are given as an example satisfying the smoother conditions.
The generalized model uses a monotone link function so that the conditional mean is determined by an additive predictor containing intercept, X-effect, and Z-effect.This framework includes logistic and Poisson regression as generalized additive models.
In generalized models, IRLS converts the nonlinear estimation problem into a sequence of weighted least-squares problems using pseudo-responses and observation weights.The estimates are repeatedly updated until the X- and Z-effects converge.
The proposed controls remain unbiased in the binary-outcome simulation, whereas the standard DNN’s bias increases with β_Z; convergence is slower for the controlled model.Figure 6 measures convergence using MSPE(f̂_X).
Very small ridge penalties can prevent IRLS convergence when the feature dimension q is large, indicating that larger penalty values should be used.The issue is described as known behavior for regularized IWLS rather than a general failure of the method.

C.1 Implications of Correlation-based Confound Control

Correlation-based confound control targets a residual effect that can discard mediated signal, while flexible additive models face concurvity and non-identifiability. The paper’s refitting approach instead estimates the direct effect and shows more robust convergence than end-to-end NAM fitting.

Implications of Correlation-based Confound Control: Decorrelating predictions from Z estimates the residual effect rather than the direct effect of X.This can discard the component of the X-effect that is predictable from Z and may substantially reduce prediction accuracy.
Implications of Correlation-based Confound Control: The mediated component of the X-effect may be scientifically relevant, so removing it depends on whether the goal is fairness or medical effect estimation.Age-related brain morphology in alcohol-misuse prediction illustrates a potentially meaningful mediated pathway.
Concurvity and Mediated Effects: Flexible additive models can make f_X and f_Z unidentifiable when network features reproduce covariate effects, a problem known as concurvity.The issue is especially acute when the number of network features exceeds the number of observations or when the network can interpolate covariate effects.
Concurvity and Mediated Effects: Under concurvity, arbitrary components can be assigned between X and Z without changing the conditional expectation, undermining interpretation of the direct effect.The concurvity component has no causal or associational interpretation and arises from the complexity of the function classes.
Concurvity and Mediated Effects: Orthogonalization removes concurvity components but also removes mediated components, whereas the proposed approach estimates the direct effect without imposing decorrelation.The paper suggests combining the approaches by orthogonalizing selected variables and merely controlling for others.
Dependence of NAM and DNN with Controls on Backbone Size: The proposed refitting method is robust to backbone size, while NAMs converge more slowly and retain higher MSPE across sample sizes.The backbone feature dimension q ranges from 8 to 1024, and DNN-with-Controls curves collapse onto a single trajectory.
Conditional Expectations and Smoother Matrices: A pretrained DNN can be represented as a linear smoother of outcomes, with ridge regularization preventing interpolation when the feature count exceeds the sample size.Covariate smoothers may be linear or nonlinear, including basis expansions such as cubic splines.
Conditional Expectations and Smoother Matrices: The final regularization parameter is selected from a logarithmically spaced grid by minimizing validation loss.The selected parameter determines the final X-effect and corresponding covariate effect.

D.1 Details of the Simulation Study and Additional Figures

The simulations evaluate effect recovery under linear and nonlinear covariate effects, model misspecification, and noisy controls. Results show consistent nonlinear-effect estimation and graceful degradation when controls are measured with increasing noise.

Simulation setup: The DNN backbone uses q = 32 last-layer features, while fZ and fX are mean-zero linear functions of covariates and learned features.The method uses a 2-fold cross-fit ensemble throughout the simulation study.
Evaluation: Estimators are compared using MSPE over test observations and independently drawn training datasets, with 50 fitted models and equal training-validation splits.The test set contains 800 observations.
Bias-variance behaviour: Bias dominates the MSPE of models without covariates, while variance converges for all models as sample size grows.The corresponding MSPE figures also show that only controlled DNNs estimate fZ consistently and predict y near the noise floor.
Nonlinear covariate effects: With nonlinear fZ, cubic B-spline controls absorb the covariate effect while the bias and variance of f̂X remain essentially independent of βZ.The spline uses nine basis evaluations, five inner knots, and 2-fold cross-fitting.
Nonlinear covariate effects: For nonlinear fZ, Var(f̂Z) scales with βZ^2 and shrinks approximately as N^-1, while squared bias remains below variance across the βZ range.These results support consistency of the spline regression for fZ.
Model misspecification: When only linear controls are used despite nonlinear fZ, the simulation evaluates convergence under model misspecification.The supplied figure caption identifies MSPE and bias-variance decompositions for f̂X and f̂Z.
Noisy controls: With noisy controls, squared bias rises toward the uncontrolled-DNN plateau as noise increases, whereas the controlled estimator’s variance is essentially unaffected.At pure noise, additional data do not remove the remaining bias; orthogonalised residual-effect estimates show the same pattern.

D.4 Additional Details on the UK Biobank Experiment

The UK Biobank experiment uses synthetic age and sex confounding with strict fold-level separation. Cross-fitted age marginalisation nearly recovers balanced-data performance, while coefficient gaps remain attributable to data and logistic-model factors.

Synthetic confounding: Synthetic age and sex covariates are generated before sampling highalc, with confounding strength varied through βage and βsex fixed at 2.Age is sampled from a scaled empirical distribution and sex uniformly from Bernoulli(0.5).
Experimental protocol: A 3D ResNet-50 with 2,048 last-layer features is fine-tuned within five-fold stratified outer cross-validation.Train/refit partitions and synthetic resampling are kept independent within each fold.
Performance comparisons: Cross-fitted age marginalisation nearly recovers the predictive performance of a DNN trained on balanced data, although a small AUC drop remains.Cross-fitting also closes most of the sample-split gap and reduces fold-to-fold variance.
Performance comparisons: Marginalising over sex removes most predictive power because sex is the primary signal in the unconfounded synthetic dataset.The jointly age- and sex-controlled model is therefore treated as a sanity check.
Coefficient recovery: Recovered age and sex coefficients are nearly correct but retain a gap that may reflect pre-existing real-data covariate-outcome correlations, logistic non-collapsibility, and finite-data resampling.More cross-fitting folds reduce per-fold sample sizes and can increase dataset bias.