Weighted Spline-Expanded Networks with Distributional Balancing for Continuous Treatment Effects
Organizations: University of North Carolina at Chapel Hill · University of Illinois Urbana-Champaign · University of Wisconsin–Madison
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.
Figures & tables
| Dataset | IHDP | News | TCGA | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Num | 200 | 500 | 1000 | 2000 | 3000 | 6000 | ||||||||
| Covariates | 10 | 25 | 10 | 25 | 100 | 300 | 500 | 100 | 300 | 500 | 1000 | 4000 | 1000 | 4000 |
| GPS | 1.39 | 0.80 | 1.30 | 0.78 | 0.102 | 0.106 | 0.190 | 0.089 | 0.093 | 0.092 | 2.07e-02 | 1.62e-02 | 1.81e-02 | 1.35e-02 |
| CBPS | 1.29 | 0.73 | 1.23 | 0.62 | 0.101 | 0.105 | 0.106 | 0.089 | 0.093 | 0.087 | 2.57e-02 | 2.47e-02 | 2.52e-02 | 2.33e-02 |
| GBM | 1.32 | 0.73 | 1.27 | 0.72 | 0.101 | 0.101 | 0.100 | 0.085 | 0.087 | 0.087 | 2.45e-02 | 2.46e-02 | 2.44e-02 | 2.24e-02 |
| DCOW | 1.30 | 0.73 | 1.23 | 0.63 | 0.095 | 0.105 | 0.084 | 0.111 | 0.110 | 0.083 | 2.54e-02 | 1.54e-02 | 1.65e-02 | 1.21e-02 |
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | IHDP | News | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Num | 200 | 500 | 1000 | 2000 | ||||||
| Covariates | 10 | 25 | 10 | 25 | 100 | 300 | 500 | 100 | 300 | 500 |
| Weighted MLP | 1.91 ±0.04 | 1.92 ±0.04 | 1.92 ±0.04 | 1.90 ±0.05 | 0.517 ±0.04 | 0.492 ±0.01 | 0.492 ±0.01 | 0.516 ±0.02 | 0.505 ±0.02 | 0.494 ±0.01 |
| GPSNet | 0.54 ±0.05 | 0.40 ±0.04 | 0.40 ±0.04 | 0.92 ±0.02 | 0.657 ±0.42 | 0.108 ±0.04 | 0.155 ±0.07 | 0.119 ±0.00 | 0.162 ±0.06 | 0.098 ±0.02 |
| GPSNet-WTR | 0.47 ±0.03 | 0.36 ±0.06 | 0.20 ±0.01 | 0.26 ±0.03 | 0.530 ±0.32 | 0.118 ±0.05 | 0.176 ±0.06 | 0.209 ±0.00 | 0.167 ±0.07 | 0.139 ±0.04 |
| WSENet | 0.49 ±0.06 | 0.96 ±0.12 | 0.31 ±0.03 | 0.88 ±0.06 | 0.160 ±0.03 | 0.150 ±0.01 | 0.090 ±0.02 | 0.112 ±0.02 | 0.135 ±0.04 | 0.144 ±0.02 |
| Dataset | IHDP | News | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Num | 200 | 500 | 1000 | 2000 | ||||||
| Covariates | 10 | 25 | 10 | 25 | 100 | 300 | 500 | 100 | 300 | 500 |
| WSENet (2-layer) | 0.49 ±0.06 | 0.96 ±0.12 | 0.31 ±0.03 | 0.88 ±0.06 | 0.160 ±0.03 | 0.150 ±0.01 | 0.090 ±0.02 | 0.112 ±0.02 | 0.135 ±0.04 | 0.144 ±0.02 |
| WSENet-WTR (2-layer) | 0.34 ±0.05 | 0.29 ±0.03 | 0.17 ±0.00 | 0.18 ±0.02 | 0.083 ±0.02 | 0.074 ±0.01 | 0.078 ±0.02 | 0.065 ±0.02 | 0.049 ±0.03 | 0.058 ±0.00 |
| WSENet (3-layer) | 0.35 ±0.06 | 0.28 ±0.03 | 0.26 ±0.00 | 0.23 ±0.00 | 0.122 ±0.02 | 0.099 ±0.02 | 0.092 ±0.02 | 0.086 ±0.02 | 0.099 ±0.03 | 0.092 ±0.01 |
| WSENet-WTR (3-layer) | 0.27 ±0.04 | 0.26 ±0.04 | 0.14 ±0.00 | 0.21 ±0.00 | 0.118 ±0.02 | 0.092 ±0.01 | 0.111 ±0.01 | 0.080 ±0.02 | 0.086 ±0.00 | 0.098 ±0.02 |