Kernel methods for long term dose response curves

TL;DR

This work tackles the problem of estimating long-term causal effects of continuous actions by fusing short-term randomized experiments with long-term observational data. It introduces a nonlinear, nonparametric estimator based on reproducing kernel Hilbert space (RKHS) methods and kernel ridge regression, leveraging kernel mean embeddings to represent conditional distributions and construct a closed-form estimator for the long-term dose response $\theta_0(d)$. The authors prove uniform consistency with finite-sample rates that depend on the effective dimension and smoothness of the data, and demonstrate the method on Project STAR, showing nonlinear long-term effects and performance close to an oracle with access to long-term experimental data. They further discuss extensions to long-term counterfactual distributions and outline future directions, including neural or Nyström-based scalability. The practical impact lies in enabling robust extrapolation of long-term outcomes for continuous actions across domains with distribution shifts and heterogeneous links between short-term and long-term effects.

Abstract

A core challenge in causal inference is how to extrapolate long term effects, of possibly continuous actions, from short term experimental data. It arises in artificial intelligence: the long term consequences of continuous actions may be of interest, yet only short term rewards may be collected in exploration. For this estimand, called the long term dose response curve, we propose a simple nonparametric estimator based on kernel ridge regression. By embedding the distribution of the short term experimental data with kernels, we derive interpretable weights for extrapolating long term effects. Our method allows actions, short term rewards, and long term rewards to be continuous in general spaces. It also allows for nonlinearity and heterogeneity in the link between short term effects and long term effects. We prove uniform consistency, with nonasymptotic error bounds reflecting the effective dimension of the data. As an application, we estimate the long term dose response curve of Project STAR, a social program which randomly assigned students to various class sizes. We extend our results to long term counterfactual distributions, proving weak convergence.

Figures (6)

• Figure 1: Real data are reasonably modeled as continuous actions. Much empirical work focuses on "small" (light blue) versus "large" (dark blue) class sizes, yet the supports are rich and overlapping.
• Figure 2: Real data have a low effective dimension in the sense of \ref{['eq:b']}. Figures \ref{['fig:eigen1']} and \ref{['fig:eigen2']} visualize the initial $25$ eigenvalues of $K_{S^{\textsc{obs}}S^{\textsc{obs}}}\odot K_{X^{\textsc{obs}}X^{\textsc{obs}}}$ and $K_{D^{\textsc{exp}}D^{\textsc{exp}}}\odot K_{X^{\textsc{exp}}X^{\textsc{exp}}}$, defined in Algorithm \ref{['algo:dose']}, using Project STAR data. These are empirical analogues of Assumptions \ref{['assumption:smooth']} and \ref{['assumption:smooth_op']}.
• Figure 3: The source condition in Sobolev space means the number of square integrable derivatives. For illustration, we take $H$ to be the Sobolev space over $\mathbb{R}$ with one square integrable derivative, with respect to Lebesgue measure, so $b=2$ in the sense of \ref{['eq:b']}. Figures \ref{['fig:source1']} and \ref{['fig:source2']} visualize $f_0\in H^1$ versus $f_0\in H^2$, i.e. $c=1$ versus $c=2$ in the sense of \ref{['eq:c']}.
• Figure 4:
• Figure 5: Our method is robust to the choice of sample. We compare our proposal for the missing-at-random model, either with the full sample (dark blue, solid) or a subsetted sample (dark blue, dashed). In the background, we indicate ranges of class sizes in the Project STAR protocol (gray).

Theorems & Definitions (37)

• Definition 2.1: Long term dose response curves
• Lemma 2.1: Surrogate formula; e.g. Theorem 1 of athey2020estimating
• Remark 2.1: The surrogate formula is generally unbounded when actions are continuous
• Remark 2.2: Alternative long term models
• Remark 4.1: Subscripts
• Example 4.1: Quadratic kernels
• Example 4.2: Gaussian kernels
• Theorem 4.1: Representation via short term kernel embedding
• proof : Proof sketch
• Remark 4.2: Embeddings that fuse experimental and observational data