stat.MEApr 15, 2026

Estimating Continuous Treatment Effects with Two-Stage Kernel Ridge Regression

Authors: Seok-Jin Kim, Kaizheng Wang

Organizations: Department of IEOR, Columbia University · Department of IEOR and Data Science Institute, Columbia University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We study the problem of estimating the effect function for a continuous treatment, which maps each treatment value to a population-averaged outcome. A central challenge in this setting is confounding: treatment assignment often depends on covariates, creating selection bias that makes direct regression of the response on treatment unreliable. To address this issue, we propose a two-stage kernel ridge regression method. In the first stage, we learn a model for the response as a function of both treatment and covariates; in the second stage, we use this model to construct pseudo-outcomes that correct for distribution shift, and then fit a second model to estimate the treatment effect. Although the response varies with both treatment and covariates, the induced effect function obtained by averaging over covariates is typically much simpler, and our estimator adapts to this structure. Our optimal learning bounds are achieved without estimating the conditional treatment density, thereby bypassing a major bottleneck in existing methods. Furthermore, we introduce a fully data-driven model selection procedure that achieves provable adaptivity to both the unknown degree of overlap and the spectral decay of the underlying kernel.

Figures & tables

 λ^:=argminλ∈Λ∑k=1M(h^λ(a~k)−h~(a~k))2.\displaystyle\hat{\lambda}:=\mathop{\mathrm{argmin}}_{\lambda\in\Lambda}\sum_{k=1}^{M}\left(\hat{h}_{\lambda}(\tilde{a}_{k})-\tilde{h}(\tilde{a}_{k})\right)^{2}. 
Algorithm 2 Data-Driven Model Selection for TEF
Methodn=1000n=1000n=500n=500
Ours6.22 (0.32)8.40 (0.35)
Plug-in9.87 (0.31)11.81 (0.28)
Direct Regression14.52 (0.31)14.84 (0.36)
Table 1: Comparison of MISE over 100 runs (standard error in parentheses) on synthetic data. Reported values are scaled by 100100 .
MethodMean MISE (SE)
Ours1.2466 (0.1209)
Plug-in LOOCV1.6197 (0.1146)
Direct Regression1.6970 (0.1264)
DML (GRF)2.4230 (0.1837)
DML (NN)2.1065 (0.1454)
DML (LASSO)2.8732 (0.2391)
Table 2: MISE (standard errors in parentheses) for the semi-real benchmark, averaged over 100100 independent runs. Our proposed method yields the lowest error among all compared baselines.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

NotationDescription
Z\mathbf{Z}Design operator of {ψ(xi,ai)}i=1n\{\psi(x_{i},a_{i})\}_{i=1}^{n}
A\mathbf{A}Design operator of {ϕ(aj′)}j=1n\{\phi({a}^{\prime}_{j})\}_{j=1}^{n}
ψˉ(a)\bar{\psi}(a)Feature mean, defined as Ex∼PX[ψ(x,a)]\mathbb{E}_{x\sim\mathcal{P}_{X}}[\psi(x,a)]
ψˉ^(a)\hat{\bar{\psi}}(a)Empirical feature mean, defined as 1n∑i=1nψ(xi,a)\frac{1}{n}\sum_{i=1}^{n}\psi(x_{i},a)
W\mathbf{W}Design operator of {ψˉ^(aj′)}j=1n\{\hat{\bar{\psi}}({a}^{\prime}_{j})\}_{j=1}^{n}
Y\mathbf{Y}Vector of responses y1,…,yny_{1},\dots,y_{n} ( Y∈Rn\mathbf{Y}\in\mathbb{R}^{n} )
Appendix
Table 3: Summary of notation
NotationDefinition / Description
Dtrain\mathcal{D}_{\text{train}}{x1i,a1i,y1i}i=1n1\{x_{1i},a_{1i},y_{1i}\}_{i=1}^{n_{1}}
Dval\mathcal{D}_{\text{val}}{x2i,a2i,y2i}i=1n2\{x_{2i},a_{2i},y_{2i}\}_{i=1}^{n_{2}}
Z1\mathbf{Z}_{1}Design operator of {ψ(x1i,a1i)}i=1n1\{\psi(x_{1i},a_{1i})\}_{i=1}^{n_{1}} in Dtrain\mathcal{D}_{\text{train}}
Z2\mathbf{Z}_{2}Design operator of {ψ(x2i,a2i)}i=1n2\{\psi(x_{2i},a_{2i})\}_{i=1}^{n_{2}} in Dval\mathcal{D}_{\text{val}}
{a1j′}j=1n1\{a_{1j}^{\prime}\}_{j=1}^{n_{1}}Sampled treatments from Psamp⁡\mathcal{P}_{\operatorname{samp}} for estimating h^λ\hat{h}_{\lambda}
{a2j′}j=1n2\{a_{2j}^{\prime}\}_{j=1}^{n_{2}}Sampled treatments from Psamp⁡\mathcal{P}_{\operatorname{samp}} for estimating h~\tilde{h}
Appendix
Table 4: Detailed Summary of Notations and Operators
More in this topic

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 264 papers

Explore similar work

CardsList

  1. Modeling Covariate Transition for Efficient Estimation of Longitudinal Treatment Effects in Randomized Experiments

    May 29, 2026Naoki Chihara, Tatsushi Oka, Yasuko Matsubara +2Causal Effect Estimation

    SANKEN, The University of Osaka, Osaka, Japan · CyberAgent AI Lab · Keio University, Tokyo, Japan +1

  2. Multiply Robust Causal Mediation Analysis with Continuous Treatments

    May 19, 2021Yizhen Xu, AmirEmad Ghassami, Numair Sani +1Causal Effect EstimationCausal Mediation Analysis

    Division of Biostatistics, Department of Population Health Sciences, University of Utah, USA · Department of Mathematics and Statistics, Boston University, USA · Department of Computer Science, Johns Hopkins University, USA

  3. Prediction-Powered Data Fusion for Treatment Effect Estimation

    Oct 8, 2026Yonghan Jung, Shu YangCausal Effect EstimationConditional Average Treatment Effect Estimation

    University of Illinois Urbana-Champaign · North Carolina State University