Weighted Spline-Expanded Networks with Distributional Balancing for Continuous Treatment Effects

2026-09-27

AI Summary

This study addresses the challenges of confounding bias, model misspecification, and high-dimensional covariates in estimating continuous treatment effects from observational data by proposing WSENet, an end-to-end deep learning framework. The method introduces a novel optimal weighting strategy based on distance covariance to achieve distributional independence. By integrating structured treatment embeddings with weighted objective regularization, it constructs a doubly robust estimator, further incorporating spline expansions and efficient influence functions for bias correction. Experiments on semi-synthetic and real-world datasets demonstrate that WSENet significantly outperforms existing baseline methods in both estimation accuracy and stability.

Abstract

Estimating causal effects with continuous treatments in observational studies is challenging due to confounding, model misspecification, and high-dimensional covariates. We propose the Weighted Spline-Expanded Network (WSENet), an end-to-end neural framework that addresses these challenges by combining covariate balancing, structured treatment embedding, and bias-corrected outcome estimation. WSENet first applies Distance Covariate Optimal Weights to induce distributional independence between covariates and treatment without relying on parametric models. It then learns the conditional outcome via a structured network that fuses outcome-relevant representations of covariates with a spline-expanded treatment input, enabling smooth and flexible modeling of the dose-response relationship. To mitigate residual bias, we introduce Weighted Targeted Regularization, a correction technique based on efficient influence functions that yields a doubly robust estimator. Extensive evaluations on semi-synthetic and real-world datasets, including high-dimensional genomic and environmental health data, demonstrate that WSENet consistently outperforms existing baselines in both accuracy and stability.

Problem

• Continuous Treatment Effects
• Causal Inference
• Observational Studies
• Confounding
• High-dimensional Covariates

Innovation

• Continuous Treatment Effects
• Spline-Expanded Network
• Covariate Balancing
• Doubly Robust Estimator
• Targeted Regularization