🤖 machine learning

Targeted Regularization for Causal Effect Estimation with Exponential Dispersion Family Outcomes

This paper proposes a unified targeted regularization framework for neural networks that extends causal effect estimation to Exponential Dispersion Family outcomes (including binary, count, and skewed data) by deriving von Mises expansions to correct first-order bias and jointly estimating outcome, propensity, and fluctuation models end-to-end.

Original authors: Jiahong Li, Zeqin Yang, Jixing Xu, Enzheng Hua, Zhichao Zou, Peng Zhen, Jiecheng Guo

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiahong Li, Zeqin Yang, Jixing Xu, Enzheng Hua, Zhichao Zou, Peng Zhen, Jiecheng Guo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to figure out if a specific action (like showing an ad) actually causes a result (like a user buying a product). In the real world, you can't just run a perfect experiment where you show the ad to half the people and hide it from the other half in a controlled lab. Instead, you have to look at messy, real-world data where people chose to see the ads themselves. This is called causal effect estimation.

For a long time, computer scientists have used powerful tools called Neural Networks (AI brains) to solve this. But there's a catch: these AI brains are great at predicting what happened, but they often struggle to tell you why it happened without making mistakes. Specifically, they tend to be biased (skewed) and slow to get the right answer.

To fix this, researchers developed a technique called Targeted Regularization. Think of this as a "correction knob" you turn on the AI while it's learning. It forces the AI to adjust its internal calculations so that the final answer is mathematically guaranteed to be fair and accurate.

The Problem with the Old Way
Until now, this "correction knob" only worked well for one specific type of data: continuous numbers (like temperature or height) that follow a nice, smooth bell curve (Gaussian distribution).

But the real world is messy.

Did a user buy something? That's binary (Yes/No).
How many times did they click? That's a count (0, 1, 2, 3...).
These don't follow a smooth bell curve; they are "skewed" or "jagged."

The old methods were like trying to use a ruler designed for measuring smooth water to measure jagged rocks. It just didn't fit.

The New Solution: A Universal Adapter
This paper introduces a new, unified framework that acts like a universal adapter. It allows the "correction knob" (Targeted Regularization) to work on any type of data that follows a common statistical family called the Exponential Dispersion Family (EDF). This covers everything from simple Yes/No answers to complex counts.

Here is how they did it, explained through three simple metaphors:

1. The "Blueprint" (Von Mises Expansion)

Before you can fix a broken machine, you need a blueprint that shows exactly where the error is.

What they did: The authors derived a mathematical "blueprint" (called a von Mises expansion) for these messy, skewed data types.
The Metaphor: Imagine the AI's initial guess is a slightly crooked picture frame. The blueprint shows them exactly how crooked it is and why. This blueprint reveals the "bias term"—the specific mathematical error that needs to be removed.

2. The "Smoothie" vs. The "Chunky Soup" (Continuous Treatment)

Sometimes, the "treatment" isn't just "Yes/No" (like an ad) but a continuous amount (like the volume of an ad shown).

The Problem: When the treatment is a continuous amount, the math gets very "chunky" and breaks down. You can't measure the effect at a single, exact point because it's too sensitive.
The Solution: They used a technique called Sieve Projection.
The Metaphor: Imagine trying to draw a perfect curve through a bunch of scattered dots. If you try to hit every single dot exactly, the line gets jagged and breaks. Instead, they put a "sieve" (like a colander) over the dots. The sieve smooths out the jagged edges, creating a clean, manageable curve that the math can handle without breaking. This allows them to apply the correction knob even to continuous data.

3. The "Self-Correcting Chef" (Unified Targeted Regularization)

Now that they have the blueprint and the smooth curve, they built a new cooking method.

The Old Way: The AI would guess a recipe, taste it, and then someone would shout, "Add salt!" (adding a correction term). But sometimes, adding salt made the soup too salty or ruined the flavor (pushing the answer outside the possible range).
The New Way: The AI is now trained with a Unified Targeted Regularization loss function.
The Metaphor: Instead of shouting corrections after the meal is cooked, the chef (the Neural Network) is given a special ingredient (the fluctuation parameter) while cooking. The chef learns to adjust the heat and seasoning simultaneously so that the final dish comes out perfect without needing to be "fixed" afterward. The AI learns the outcome, the probability of the treatment, and the correction factor all at the same time, end-to-end.

The Results

The authors tested this new "Universal Adapter" on synthetic data (simulated worlds) and real-world datasets (like news articles and cancer data).

The Outcome: Their method consistently outperformed existing methods. It was more accurate at predicting the true cause-and-effect relationship, whether the data was a simple Yes/No or a complex count.
The Guarantee: They proved mathematically that this method is "doubly robust." This means even if the AI makes a mistake in one part of its calculation (like guessing the wrong probability), it can still get the right answer as long as the other part is correct. It's like having two safety nets; if you fall through one, the other catches you.

In Summary
This paper takes a powerful AI technique that was previously limited to smooth, simple data and upgrades it with a universal adapter. By creating a new mathematical blueprint and a "self-correcting" training method, they allow AI to accurately measure cause-and-effect in the messy, real-world scenarios we actually encounter every day.

Technical Summary: Targeted Regularization for Causal Effect Estimation with Exponential Dispersion Family Outcomes

1. Problem Statement

The estimation of causal effects from observational data using Neural Networks (NNs) has shown strong empirical performance, yet achieving desirable semiparametric properties—specifically doubly robustness and fast convergence rates—remains challenging. Existing NN-based estimators, such as those utilizing Targeted Maximum Likelihood Estimation (TMLE) or Targeted Regularization (e.g., Shi et al., 2019; Nie et al., 2021), are largely limited to continuous outcomes assuming a Gaussian distribution.

This limitation restricts applicability to real-world scenarios involving binary outcomes (e.g., purchase decisions), count outcomes (e.g., number of clicks), or other skewed distributions common in practice. While previous work (Gao and Hastie, 2022) addressed exponential family outcomes using the Difference in Natural Parameters (DINA) within a partially linear R-learner framework, it did not extend the targeted regularization paradigm to NNs for these distributions. Consequently, there is a lack of NN-based estimators that can correct first-order bias at the distributional level for Exponential Dispersion Family (EDF) outcomes while maintaining theoretical guarantees.

The paper aims to estimate the Average Dose Canonical Function (ADCF), defined as $\psi_a = \theta[\text{do}(A=a)]$ , where $\theta$ is the canonical (natural) parameter of the EDF. This choice aligns with standard practices for comparing effects on the canonical scale and avoids heterogeneity issues often found in conditional means.

2. Methodology

The proposed framework unifies targeted regularization for EDF outcomes under both discrete (binary) and continuous treatment settings. The methodology proceeds through three main theoretical and algorithmic steps:

A. Semiparametric Analysis and Bias Characterization

The authors first derive the von Mises expansion of the plug-in estimator for the ADCF to identify the structure of the first-order bias.

Discrete Treatments: For binary treatments, they derive the expansion showing that the bias consists of a second-order remainder term involving the product of estimation errors for the outcome model ( $\mu$ ) and the propensity score ( $\pi$ ). This product structure implies that if either model is consistently estimated, the bias vanishes (the doubly robust property).
Continuous Treatments: For continuous treatments, the pointwise ADCF is not pathwise differentiable, rendering standard influence functions undefined. To overcome this, the authors introduce a sieve-projected ADCF. They project the infinite-dimensional dose-response curve onto a finite-dimensional B-spline sieve space. This projection restores pathwise differentiability, allowing for a well-defined influence function and a valid von Mises expansion for the projected curve.

B. Unified Targeted Regularization

Leveraging the bias structure revealed by the von Mises expansion, the authors construct a unified targeted regularization objective.

Fluctuation Model: Instead of adding a correction term post-hoc (which may push estimates outside the parameter space), the method introduces a scalar perturbation function $\epsilon(a)$ to the canonical parameter. The fluctuated parameter is defined as:
$\hat{\theta}^*(x, a) = h(\hat{\mu}(x, a)) + \frac{\epsilon(a)}{\hat{\pi}(a|x)} h'(\hat{\mu}(x, a))$
where $h$ is the link function and $h'$ is its derivative.
Loss Function: The method integrates this fluctuation into the NN training via a unified loss function:
$L_{TR}(\mu, \pi, \epsilon) = L(\mu, \pi) + \beta R(\mu, \pi, \epsilon)$
Here, $L(\mu, \pi)$ is the standard negative log-likelihood for the EDF, and $R$ is the targeted regularization term derived from the negative log-likelihood of the fluctuated model. The term $R$ is constructed such that its derivative with respect to $\epsilon$ equals the score equation, effectively zeroing out the bias term at convergence.
Implementation: The fluctuation parameter $\epsilon(a)$ is parameterized using the same B-spline basis used for the sieve projection in continuous treatments, ensuring consistency between the theoretical target and the optimization.

C. Theoretical Guarantees

The paper proves that the resulting NN-based estimator achieves a doubly robust convergence rate in the $L_2$ norm. Specifically, the error is bounded by:
$O_p\left(n^{-1/3}\sqrt{\log n} + r_1(n)r_2(n) + r_2(n)^2\right)$
where $r_1(n)$ and $r_2(n)$ are the convergence rates of the propensity score and outcome estimators, respectively. This extends the theoretical guarantees of Nie et al. (2021) from Gaussian outcomes to the full Exponential Dispersion Family.

3. Key Contributions

Semiparametric Analysis of ADCF: The authors derive the von Mises expansion for the ADCF of EDF outcomes. For continuous treatments, they introduce the sieve-projected ADCF to handle the non-pathwise differentiability of the pointwise functional, establishing a well-defined influence function.
Unified Targeted Regularization: They propose a unified regularization objective that corrects bias at the distributional level for any EDF member. This allows the NN to jointly learn the outcome model, propensity score, and fluctuation parameter end-to-end while remaining within the parameter space.
Theoretical Extension: The work provides a doubly robust convergence rate proof, generalizing previous results from Gaussian distributions to the broader class of EDFs (including Bernoulli and Poisson).
Empirical Validation: The method is validated on synthetic and semi-synthetic datasets (News, TCGA) with binary and continuous treatments, demonstrating state-of-the-art performance in both Bernoulli and Poisson outcome settings.

4. Experimental Results

The authors evaluated their method against several baselines, including Causal Forest, Dragonnet (and its adapted version for non-Gaussian outcomes), VCNet, DINA-learner, and a standard Doubly Robust (DR) estimator.

Binary Treatment: On datasets with Bernoulli and Poisson outcomes, the proposed method ("Ours") achieved the lowest Mean Absolute Error (MAE) for the Average Treatment Effect (ATE), outperforming Dragonnet, VCNet, and the standard DR estimator. Notably, the DR estimator was found to be less stable in finite samples due to division by the propensity score in the correction term.
Continuous Treatment: For continuous treatments, the method achieved the lowest Average Mean Squared Error (AMSE) for the ADCF. The results confirmed that simply adapting the loss function of existing models (e.g., replacing MSE with negative log-likelihood) is insufficient; the specific targeted regularization mechanism is crucial for performance.
Sensitivity Analysis: The hyperparameter $\beta$ , controlling the strength of regularization, was shown to have a moderate impact; values that are too small fail to correct bias, while values that are too large interfere with nuisance function estimation.

5. Significance and Claims

The paper claims to address a critical gap in causal inference with deep learning: the inability of existing targeted regularization frameworks to handle non-Gaussian, skewed outcomes common in real-world applications (such as binary conversion rates or count data).

By deriving the necessary semiparametric expansions and constructing a unified loss function, the authors demonstrate that targeted regularization can be successfully generalized to the Exponential Dispersion Family. The significance of this work lies in:

Theoretical Rigor: Providing the first derivation of the von Mises expansion and influence function for EDF outcomes under both discrete and continuous treatments, enabling the construction of doubly robust estimators.
Practical Applicability: Offering an end-to-end NN architecture that handles binary and count outcomes without relying on the restrictive partially linear assumptions of previous methods (like DINA-learner).
Robustness: Achieving convergence rates that are faster than individual nuisance estimators when both are consistent, while maintaining stability within the parameter space, unlike additive bias-correction methods.

The authors conclude that their framework extends the theoretical guarantees of functional targeted regularization to a much broader class of outcome distributions, verified by state-of-the-art empirical results.