stat.MLSep 27, 2026

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

Authors: Shucheng Liu, Chan Park, Guanhua Chen

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

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

May 10, 2026cs.LG

Doubly Robust Proxy Causal Learning with Neural Mean Embeddings

Unobserved confounding prevents standard covariate adjustment from identifying causal response functions in observational studies. Proxy causal learning addresses this problem through bridge equations involving treatment- and outcome-inducing proxies, avoiding direct recovery of the latent confounder. Existing doubly robust proxy estimators combine outcome and treatment bridges, but typically rely on fixed kernels, sieves, or low-dimensional semiparametric models; existing neural proxy methods are more flexible, but are largely single-bridge estimators. We develop a neural doubly robust framework for proxy causal learning with continuous and structured treatments. Our method introduces a neural mean-embedding estimator for the treatment bridge, combines it with a neural outcome bridge, and estimates the doubly robust correction through a final regression stage. The framework covers population, heterogeneous, and conditional dose-response functions, yielding full response-curve estimators rather than binary-treatment effects. The algorithms use two stages for each bridge and history-aware updates of the final linear layers to stabilize stochastic multi-stage training. We prove consistency of the algorithms showing that the doubly robust error is controlled by the final averaging and regression errors together with the smaller of the outcome- and treatment-side weak-norm bridge errors. Across synthetic and image-valued benchmarks, the proposed estimators outperform existing baselines and single-bridge neural estimators, showing the benefit of combining learned outcome and treatment bridges in a doubly robust construction. Our implementation is available at https://github.com/BariscanBozkurt/DRPCL-Neural-Mean-Embedding.
Jul 23, 2026stat.ML

TLRNet: Estimating Individual Treatment Effect based on Local Information and Single Learner Structure

Causal inference has become a central issue across various fields, including computer science, statistics, economics, education, healthcare, and medicine. The broad applicability of this discipline has garnered increased research funding and attention. In recent years, the estimation of causal effects from observational data has gained traction due to the vast amounts of collected data and the lower costs compared to randomized controlled trials. Advances in causal effect estimation methods have enhanced service personalization tools. For instance, these tools can help identify the most effective type of treatment (considering both cost and success rate) for each patient among different medical service options. This paper proposes an innovative method for estimating the heterogeneity of treatment effects. The structure of the proposed model is based on a deep neural network and a pseudo-single learner. The proposed method has been compared with other state-of-the-art methods on the IHDP benchmark. Acceptable results have been obtained by using one estimator to estimate the potential outcomes of two treatment groups. Accordingly, this paves the way for further development and improvement of the proposed method.
May 7, 2026cs.LG

RepFlow: Representation Enhanced Flow Matching for Causal Effect Estimation

Estimating causal effects from observational data has become increasingly critical in diverse fields including healthcare, economics, and social policy. The fundamental challenge in causal inference arises from the missing counterfactuals and the selection bias. Existing methods are largely limited to point estimates and lack the capacity for distribution modeling. In this work, we propose RepFlow, a novel framework that formulates causal effect estimation as a joint optimization problem integrating representation learning with Conditional Flow Matching (CFM). RepFlow mitigates selection bias by minimizing the entropically regularized Wasserstein distance between treated and control representations. To enhance numerical stability, we further introduce an L2L_2 normalization constraint on latent representations. This balanced representation enables the flow model to accurately capture the distribution of potential outcomes. Extensive experiments across a wide range of benchmarks demonstrate that RepFlow consistently outperforms existing methods in both point and distributional causal effect estimation.