Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Treatment Effect Estimation for Exponential Family Outcomes using Neural Networks with Targeted Regularization

Jiahong Li, Zeqin Yang, Jiayi Dan, Jixing Xu, Zhichao Zou, Peng Zhen, Jiecheng Guo

TL;DR

This work tackles confounding bias in observational treatment effect estimation for exponential-family outcomes by developing an end-to-end neural network estimator for the Average Dose Canonical Function (ADCF). It derives the von Mises expansion to identify the first-order bias and constructs a doubly robust estimator, then generalizes targeted regularization to exponential families, with theoretical convergence guarantees. The proposed framework is instantiated with explicit forms for Bernoulli and Poisson outcomes and implemented via a two-head NN predicting $\mu(\mathbf{x},a)$ and $\pi(a|\mathbf{x})$, augmented by a distribution-scale regularization term that mitigates bias. Empirical results on synthetic and semi-synthetic data (News and TCGA) show state-of-the-art performance across binary and continuous treatment regimes, highlighting the method’s applicability to real-world outcomes beyond Gaussian assumptions.

Abstract

Neural Networks (NNs) have became a natural choice for treatment effect estimation due to their strong approximation capabilities. Nevertheless, how to design NN-based estimators with desirable properties, such as low bias and doubly robustness, still remains a significant challenge. A common approach to address this is targeted regularization, which modifies the objective function of NNs. However, existing works on targeted regularization are limited to Gaussian-distributed outcomes, significantly restricting their applicability in real-world scenarios. In this work, we aim to bridge this blank by extending this framework to the boarder exponential family outcomes. Specifically, we first derive the von-Mises expansion of the Average Dose function of Canonical Functions (ADCF), which inspires us how to construct a doubly robust estimator with good properties. Based on this, we develop a NN-based estimator for ADCF by generalizing functional targeted regularization to exponential families, and provide the corresponding theoretical convergence rate. Extensive experimental results demonstrate the effectiveness of our proposed model.

Treatment Effect Estimation for Exponential Family Outcomes using Neural Networks with Targeted Regularization

TL;DR

This work tackles confounding bias in observational treatment effect estimation for exponential-family outcomes by developing an end-to-end neural network estimator for the Average Dose Canonical Function (ADCF). It derives the von Mises expansion to identify the first-order bias and constructs a doubly robust estimator, then generalizes targeted regularization to exponential families, with theoretical convergence guarantees. The proposed framework is instantiated with explicit forms for Bernoulli and Poisson outcomes and implemented via a two-head NN predicting and , augmented by a distribution-scale regularization term that mitigates bias. Empirical results on synthetic and semi-synthetic data (News and TCGA) show state-of-the-art performance across binary and continuous treatment regimes, highlighting the method’s applicability to real-world outcomes beyond Gaussian assumptions.

Abstract

Neural Networks (NNs) have became a natural choice for treatment effect estimation due to their strong approximation capabilities. Nevertheless, how to design NN-based estimators with desirable properties, such as low bias and doubly robustness, still remains a significant challenge. A common approach to address this is targeted regularization, which modifies the objective function of NNs. However, existing works on targeted regularization are limited to Gaussian-distributed outcomes, significantly restricting their applicability in real-world scenarios. In this work, we aim to bridge this blank by extending this framework to the boarder exponential family outcomes. Specifically, we first derive the von-Mises expansion of the Average Dose function of Canonical Functions (ADCF), which inspires us how to construct a doubly robust estimator with good properties. Based on this, we develop a NN-based estimator for ADCF by generalizing functional targeted regularization to exponential families, and provide the corresponding theoretical convergence rate. Extensive experimental results demonstrate the effectiveness of our proposed model.

Paper Structure

This paper contains 23 sections, 4 theorems, 44 equations, 2 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Problem Statement and Notations
  4. Plug-in Estimator
  5. Targeted Regularization for Exponential Family Outcomes
  6. Bias Analysis and Doubly Robust Estimator
  7. Targeted Regularization
  8. Exponential Family Examples
  9. Bernoulli Distribution
  10. Poisson Distribution
  11. Experiments
  12. Dataset
  13. Synthetic Data Generation
  14. Semi-synthetic Data Generation
  15. Baselines
  16. ...and 8 more sections

Key Result

Lemma 5.1

Let $\psi_a(\boldsymbol{Z};\mathbb{P}) = \mathbb{E}\left\{h(\mathbb{E}[Y \mid \boldsymbol{X}, A = a]) \right\}$ for some twise continuously differentiable link function $h$. For another probability measure $\bar{\mathbb{P}}$, the $\psi$ confirms the von Mises expansion where the influence function is and where $\mu^*(\boldsymbol{x},a)$ lies between $\mu(\boldsymbol{x},a)$ and $\bar{\mu}(\boldsy

Figures (2)

  • Figure 1: Network architecture.
  • Figure 2: Sensitivity analysis on simulation data of binary treatment setting

Theorems & Definitions (8)

  • Lemma 5.1
  • Lemma 5.2
  • Remark 5.3
  • Corollary 5.4
  • Theorem 5.6
  • Remark 5.7
  • proof
  • proof