# Estimating Continuous Treatment Effects with Two-Stage Kernel Ridge Regression

arXiv: https://arxiv.org/abs/2604.13410 · HTML: https://mlsanity.com/paper/2604.13410/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 52 references that are in the map.
- First posted on arXiv 15 Apr 2026.

## Builds on (most cited first, 27 of 27)

- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Generalized Random Forests](https://mlsanity.com/paper/1610.01271/) (2016)
- [Just Interpolate: Kernel "Ridgeless" Regression Can Generalize](https://mlsanity.com/paper/1808.00387/) (2018)
- [Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: optimal rate and curse of dimensionality](https://mlsanity.com/paper/1810.08033/) (2018)
- [Linearized two-layers neural networks in high dimension](https://mlsanity.com/paper/1904.12191/) (2019)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Generalization error of random features and kernel methods: hypercontractivity and kernel matrix concentration](https://mlsanity.com/paper/2101.10588/) (2021)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [Learning with invariances in random features and kernel models](https://mlsanity.com/paper/2102.13219/) (2021)
- [Deep Equals Shallow for ReLU Networks in Kernel Regimes](https://mlsanity.com/paper/2009.14397/) (2020)
- [Nonparametric methods for doubly robust estimation of continuous treatment effects](https://mlsanity.com/paper/1507.00747/) (2015)
- [Kernel Methods for Causal Functions: Dose, Heterogeneous, and Incremental Response Curves](https://mlsanity.com/paper/2010.04855/) (2020)
- [Augmented Minimax Linear Estimation](https://mlsanity.com/paper/1712.00038/) (2017)
- [On the Optimality of Misspecified Kernel Ridge Regression](https://mlsanity.com/paper/2305.07241/) (2023)
- [Pseudo-Labeling for Kernel Ridge Regression under Covariate Shift](https://mlsanity.com/paper/2302.10160/) (2023)
- [Entropy Balancing for Continuous Treatments](https://mlsanity.com/paper/2001.06281/) (2020)
- [Augmented Transfer Regression Learning with Semi-non-parametric Nuisance Models](https://mlsanity.com/paper/2010.02521/) (2020)
- [Fast convergence rates for dose-response estimation](https://mlsanity.com/paper/2207.11825/) (2022)
- [Kernel-based off-policy estimation without overlap: Instance optimality beyond semiparametric efficiency](https://mlsanity.com/paper/2301.06240/) (2023)
- [Debiased inference for a covariate-adjusted regression function](https://mlsanity.com/paper/2210.06448/) (2022)
- [Independence weights for causal inference with continuous treatments](https://mlsanity.com/paper/2107.07086/) (2021)
- [Doubly Robust Inference on Causal Derivative Effects for Continuous Treatments](https://mlsanity.com/paper/2501.06969/) (2025)
- [Fast kernel methods: Sobolev, physics-informed, and additive models](https://mlsanity.com/paper/2509.02649/) (2025)
- [Transfer Learning of CATE with Kernel Ridge Regression](https://mlsanity.com/paper/2502.11331/) (2025)
- [Testing Monotonicity of Mean Potential Outcomes in a Continuous Treatment with High-Dimensional Data](https://mlsanity.com/paper/2106.04237/) (2021)

## Built on by (most cited first, 1 of 1)

- [Causal Generalization of Continuous Treatment Effects under Covariate Shift](https://mlsanity.com/paper/2608.19383/) (2026)

## Abstract (arXiv)

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.