Doubly Robust Proxy Causal Learning with Neural Mean Embeddings

2026-05-10

AI Summary

This work addresses the challenge of accurately estimating causal response functions in observational studies with unobserved confounding. The authors propose a novel neural doubly robust proxy learning framework that, for the first time, integrates neural networks with doubly robust estimation. The method estimates treatment bridge functions via neural mean embeddings and jointly learns neural outcome bridge functions, enabling estimation of full dose–response curves under both continuous and structured treatment variables. To enhance stability, the approach incorporates a history-aware multi-stage training strategy and a linear-layer updating mechanism. Experiments on synthetic and image-based benchmarks demonstrate that the proposed method significantly outperforms existing baselines and single-bridge neural estimators, confirming the efficacy and superior performance of its doubly robust construction.

Abstract

Unobserved confounding prevents standard covariate adjustment from identifying causal response functions in observational studies. Proxy causal learning addresses this problem through bridge equations involving treatment- and outcome-inducing proxies, avoiding direct recovery of the latent confounder. Existing doubly robust proxy estimators combine outcome and treatment bridges, but typically rely on fixed kernels, sieves, or low-dimensional semiparametric models; existing neural proxy methods are more flexible, but are largely single-bridge estimators. We develop a neural doubly robust framework for proxy causal learning with continuous and structured treatments. Our method introduces a neural mean-embedding estimator for the treatment bridge, combines it with a neural outcome bridge, and estimates the doubly robust correction through a final regression stage. The framework covers population, heterogeneous, and conditional dose-response functions, yielding full response-curve estimators rather than binary-treatment effects. The algorithms use two stages for each bridge and history-aware updates of the final linear layers to stabilize stochastic multi-stage training. We prove consistency of the algorithms showing that the doubly robust error is controlled by the final averaging and regression errors together with the smaller of the outcome- and treatment-side weak-norm bridge errors. Across synthetic and image-valued benchmarks, the proposed estimators outperform existing baselines and single-bridge neural estimators, showing the benefit of combining learned outcome and treatment bridges in a doubly robust construction. Our implementation is available at https://github.com/BariscanBozkurt/DRPCL-Neural-Mean-Embedding.

Problem

• unobserved confounding
• proxy causal learning
• causal response functions
• doubly robust estimation
• neural mean embeddings

Innovation

• doubly robust
• proxy causal learning
• neural mean embeddings
• continuous treatment
• bridge functions