Weighted neural networks improve continuous treatment effect estimation
Weighted Spline-Expanded Networks with Distributional Balancing for Continuous Treatment Effects
Machine Learning
Summary
Estimating the impact of treatments that vary in dose or intensity is tricky because of confusing factors and complex data. The authors created a new neural network method called WSENet that first balances the data to reduce confusion and then models how the treatment affects outcomes smoothly. WSENet also corrects leftover errors to improve accuracy. Tests on synthetic and real data showed WSENet is more reliable and precise than other methods.
What this means in practice
- •For health data teams: Estimate how different treatment doses affect patient outcomes more accurately using routine clinical data with many variables.
- •For environmental analysts: Analyze how continuous environmental exposures influence health effects with improved modeling of complex, high-dimensional datasets.
Authors
Shucheng Liu, Chan Park, Guanhua Chen
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.