Sequential Kernel Embedding for Mediated and Time-Varying Dose Response Curves

Sequential kernel embedding for mediated and time-varying dose response curves

\def\spacingset#1{ {#1}} \spacingset{1}

{\it Keywords:} Continuous treatment, reproducing kernel Hilbert space, treatment-confounder feedback

\spacingset{1.65}

Introduction

We study mediation analysis and time-varying treatment effects with possibly continuous treatments. Mediation analysis asks, how much of the total effect of the treatment $D$ on the outcome $Y$ is mediated by a particular mechanism $M$ that takes place between the treatment and outcome? Time-varying analysis asks, what would be the effect of a sequence of treatments $D_{1:T}$ on the outcome $Y$, even when that sequence may not have been implemented? We consider nonparametric causal functions of continuous treatments. For example, the time-varying dose response curve of two continuous treatments is the function $\theta_0^{GF}(d_1,d_2):=E\{Y^{(d_1,d_2)}\}$, which may refer to medical dosages, lifestyle habits, occupational exposures, or training durations.

The time-varying dose response curve arises when evaluating social programs from several rounds of surveys. For example, the National Job Corps Study randomized access to a large scale job training program in the US and collected several rounds of surveys schochet2008does. Individuals could decide whether to participate and for how many hours, possibly over multiple years. A natural question is: how much would the average individual benefit from a certain number of class hours in year one and a possibly different number of class hours in year two? This quantity is an example of a time-varying dose response curve, where the number of class hours is the time-varying dose, and the expected benefit is the response.

The difficulty in estimating time-varying dose response curves is the complex feedback loop resulting from the initial dose. Formally, we model class hours in different years as a sequence of continuous treatments subject to treatment-confounder feedback. In other words, we consider the possibility that class hours in one year may affect health behaviors such as drug use in a subsequent year, which may then affect subsequent class hours. Though several rounds of Job Corps surveys were collected, economists typically study only the initial survey due to a concern for treatment-confounder feedback, and a lack of simple yet flexible estimators that can adjust for it, while allowing treatments to be continuous. We propose such an estimator using kernels, establish its properties, share a cleaned data set that includes the additional surveys, and carry out empirical analysis on the cleaned data as well as on simulated data. Previous methods using kernels for continuous treatments singh2020kernel do not handle treatment-confounder feedback, and therefore cannot analyze the additional surveys. This paper's technical innovation is a sequential kernel embedding to do so.

While mediated and time-varying causal functions are identified in theory, they are challenging to estimate in practice. For example, under standard assumptions on time-varying covariates $X_{1:T}$, the time-varying dose response is identified as the g-formula, i.e. the sequential integral $ \theta_0^{GF}(d_1,d_2)=\int \gamma_0(d_1,d_2,x_1,x_2) \mathrm{d}P(x_2|d_1,x_1)\mathrm{d}P(x_1) $, where $\gamma_0(d_1,d_2,x_1,x_2)=E(Y|D_1=d_1,D_2=d_2,X_1=x_1,X_2=x_2)$ robins1986new. When treatments are continuous, the functional $ \gamma $ $ \mapsto $ $ \int $ $ \gamma(d_1,d_2,x_1,x_2) $ $\mathrm{d}P(x_2|d_1,x_1)$ $\mathrm{d}P(x_1) $ is generally not bounded van1991differentiable,newey1994asymptotic or pathwise differentiable bickel1993efficient. Popular estimators restrict attention to a binary treatment, parametric models, Markov simplifications, or constrained effect modification for tractability, and may even redefine the estimand. See vansteelandt2014structural for a review. Each of these restrictions simplifies the sequential integral in order to simplify estimation. Our research question is: can we devise simple machine learning estimators for causal functions that preserve the richness of the sequential integral, and therefore the generality of treatment-confounder feedback in classic identification, while also achieving nonasymptotic uniform rates?

In this paper, we match the generality of mediated and time-varying identification with the flexibility and simplicity of kernel ridge regression estimation. We propose a new family of nonparametric estimators for causal inference over short horizons. Our algorithms combine kernel ridge regressions, so they inherit the practical and theoretical virtues that make kernel ridge regression widely used. Crucially, we preserve the nonlinearity, dependence, and effect modification of identification theory with time-varying confounders. Our contribution has three aspects.

First, we introduce an algorithmic technique that appears to be an innovation in the reproducing kernel Hilbert space (RKHS) literature: sequential kernel embeddings, i.e. RKHS representations of mediator and covariate conditional distributions given a hypothetical treatment sequence, which account for treatment-confounder feedback. For example, we introduce the sequential embedding $\mu_{x_1,x_2}(d_1)$ such that the inner product $\langle f,\mu_{x_1,x_2}(d_1) \rangle$ in an appropriately defined Hilbert space equals the g-formula's sequential integral $\int f(x_1,x_2) \mathrm{d}P(x_2|d_1,x_1)\mathrm{d}P(x_1)$. We prove that the sequential kernel embedding exists because the RKHS restores boundedness of the g-formula's functional $\gamma\mapsto \int \gamma(d_1,d_2,x_1,x_2)$ $ \mathrm{d}P(x_2|d_1,x_1) \mathrm{d}P(x_1) $, even when the treatments are continuous.

Second, we use our new technique to derive estimators with simple closed forms that combine kernel ridge regressions, extending the regression product baron1986moderator and recursive regression bang2005doubly insights to machine learning. We use sequential embeddings to propose uniformly consistent machine learning estimators of time-varying dose response curves without restrictive linearity, Markov, or no-effect-modification assumptions, which to our knowledge is new. As extensions, we propose what may be the first unrestricted incremental response curves and counterfactual distributions for time-varying treatments, relaxing the restrictions of the structural nested distribution model robins1992estimation. In Section 9 of rahul2024supplement, for discrete treatments, we use sequential embeddings to propose simpler nuisance parameter estimators for known inferential procedures. In particular, we avoid multiple levels of sample splitting and iterative fitting.

Third, we prove that our simple estimators based on sequential embedding achieve nonasymptotic uniform rates for causal functions. Specifically, for the continuous treatment case, we prove uniform consistency with finite sample rates that combine minimax rates for kernel ridge regression caponnetto2007optimal,fischer2017sobolev,li2022optimal. The rates do not directly depend on the data dimension, but rather smoothness and effective dimension, generalizing Sobolev rates. We extend these results to incremental response curves and counterfactual outcome distributions. We relate our results to semiparametric analysis in Section 9 of rahul2024supplement. In particular, for the discrete treatment case, we verify $n^{-1/2}$ Gaussian approximation and semiparametric efficiency, articulating a double spectral robustness whereby some kernels may have higher effective dimensions as long as others have sufficiently low effective dimensions.

We illustrate the practicality of our approach by conducting comparative simulations and estimating the mediated and time-varying response curves of the Jobs Corps. In nonlinear simulations over short horizons, the algorithms reliably outperform some state of the art alternatives. Under standard identifying assumptions, our direct and indirect dose response curve estimates suggest that job training reduces arrests via social mechanisms besides employment. By allowing for continuous treatments and treatment-confounder feedback, our time-varying dose response curve estimates suggest that relatively few class hours in the first and second year confer most of the benefit to counterfactual employment. Of independent interest, we clean and share a version of the Job Corps data that may serve as a benchmark for new approaches to time-varying estimation.

Related work

In seminal works, robins1992identifiability,pearl2001direct,imai2010identification and robins1986new rigorously develop identification theory for mediation analysis and time-varying treatment effects, respectively. A rich class of mediated and time-varying causal functions are estimable in principle if the analyst has a sufficiently rich set of covariates over time $X_{1:T}$. Treatments may be continuous; relationships among the outcome, treatments, and covariates may be nonlinear; and dependences may include treatment-confounder feedback and effect modification by time-varying confounders gill2001causal,vanderweele2009conceptual.

For continuous treatments, nonparametric estimators for mediated response curves of huber2018direct,ghassami2021multiply use density estimation, which can be challenging as dimension increases. Machine learning estimators for time-varying dose response curves of lewis2020double rely on restrictive linearity, Markov, and no-effect-modification assumptions, which imply additive effects of time-varying treatments.

For binary treatments, a rich literature provides abstract conditions for $n^{-1/2}$ semiparametric estimation, to which we relate our results in Section 9 of rahul2024supplement; see e.g. scharfstein1999adjusting,van2006targeted,zheng2011cross,van2012targeted,tchetgen2012semiparametric,petersen2014targeted,molina2017multiple,luedtke2017sequential,rotnitzky2017multiply,chernozhukov2018original,farbmacher2020causal,bodory2021evaluating,singh2021finite and references therein. Still, estimators that preserve the full generality of identification for binary treatments are not widely used in empirical research vansteelandt2014structural, perhaps due to the complexity of nuisance parameter estimation.

Unlike previous work that incorporates the RKHS into causal inference, we provide a framework for mediated and time-varying estimands. Previous work incorporates the RKHS into time-fixed causal inference. nie2017quasi,foster2019orthogonal,kennedy2020optimal propose methods based on orthogonal loss minimization for heterogeneous treatment effects, and wong2018kernel,zhao2019covariate,kallus2020generalized,hirshberg2019minimax,singh2021debiased propose methods based on balancing weights for average treatment effects. muandet2020counterfactual propose counterfactual distributions for a binary treatment, while singh2020kernel propose dose responses and counterfactual distributions for a continuous treatment. Whereas previous work studies the time-fixed setting, we study sequential settings and prove strong results despite the additional challenges of treatment-confounder feedback and effect modification by time-varying confounders.

This paper subsumes our earlier draft singh2021kernel, which subsumed singh2020workshop.

RKHS assumptions

We summarize RKHS notation, interpretation, and assumptions that we use in this paper. Let $k:\mathcal{W}\times\mathcal{W}\rightarrow \mathbb{R}$ be a function that is continuous, symmetric, and positive definite. We call $k$ the kernel, and we call $\phi:w\mapsto k(w,\cdot)$ the feature map. The kernel is the inner product of features $k(w,w')=\langle \phi(w),\phi(w')\rangle_{\mathcal{H}}$, and we formally define the inner product below. The RKHS is the closure of the span of the features $\{\phi(w)\}_{w\in\mathcal{W}}$. As such, the features are interpretable as the dictionary of basis functions for the RKHS: for $f\in\mathcal{H}$, we have that $f(w)=\langle f,\phi(w) \rangle_{\mathcal{H}}$.

Kernel ridge regression uses the RKHS $\mathcal{H}$ as the hypothesis space in an infinite dimensional optimization problem with a ridge penalty, and it has a well known closed form solution:

where $K_{WW}\in\mathbb{R}^{n\times n}$ is the kernel matrix with $(i,j)$th entry $k(W_i,W_j)$ and $K_{Ww}\in\mathbb{R}^n$ is the kernel vector with $i$th entry $k(W_i,w)$. To tune the ridge penalty hyperparameter $\lambda$, both generalized cross validation and leave one out cross validation have closed form solutions, and the former is asymptotically optimal li1986asymptotic. To analyze the bias and variance of kernel ridge regression, the statistical learning literature places assumptions on the smoothness of $f_0$ and effective dimension of $\mathcal{H}$. Both assumptions describe the kernel's spectrum, which we now define.

Denote by $\mathbb{L}^2_{\nu}(\mathcal{W})$ the space of square integrable functions with respect to the measure $\nu$. Consider the convolution operator $ L:\mathbb{L}_{\nu}^2(\mathcal{W})\rightarrow \mathbb{L}_{\nu}^2(\mathcal{W}),\; f\mapsto \int k(\cdot,w)f(w)\mathrm{d}\nu(w) $. By the spectral theorem, we define its spectrum, where $(\eta_j)$ are weakly decreasing eigenvalues and $(\varphi_j)$ are orthonormal eigenfunctions that form a basis of $\mathbb{L}_{\nu}^2(\mathcal{W})$. As such, $ Lf=\sum_{j=1}^{\infty} \eta_j\langle \varphi_j,f \rangle_{\mathbb{L}^2_{\nu}(\mathcal{W})}\cdot \varphi_j $.

The space $\mathcal{H}$ is the subset of $\mathbb{L}^2_{\nu}(\mathcal{W})$ for which higher order terms in $(\varphi_j)$ have a smaller contribution, subject to $\nu$ satisfying the conditions of Mercer's theorem steinwart2012mercer. Under those conditions, $k(w,w')=\sum_{j=1}^{\infty}\eta_j \varphi_j(w)\varphi_j(w')$; $(\eta_j)$ and $(\varphi_j)$ describe the kernel's spectrum.

An analyst can place smoothness and effective dimension assumptions on the spectral properties of a statistical target $f_0$ estimated in the RKHS $\mathcal{H}$ caponnetto2007optimal,fischer2017sobolev. These assumptions are formalized by parameters $(b,c)$:

The value $c$ quantifies how well the leading terms in $(\varphi_j)$ approximate $f_0$; a larger value of $c$ corresponds to a smoother target $f_0$. A larger value of $b$ corresponds to a faster rate of spectral decay and therefore a lower effective dimension. Both $(b,c)$ are joint assumptions on the kernel and data distribution. Correct specification implies $c\geq 1$ and a bounded kernel implies $b\geq 1$ fischer2017sobolev. Minimax optimal rates for regression are governed by $(b,c)$ caponnetto2007optimal,fischer2017sobolev, with faster rates corresponding to higher values. In our analysis of causal estimands, we obtain nonparametric rates and semiparametric rate conditions that combine regression rates in terms of $(b,c)$.

Spectral assumptions are easy to interpret in Sobolev spaces fischer2017sobolev. Let $\mathcal{W}\subset \mathbb{R}^p$. Denote by $\mathbb{H}_2^s$ the Sobolev space with $s>p/2$ square integrable derivatives, which can be generated by the Mat\`ern kernel. Suppose $\mathcal{H}=\mathbb{H}_2^s$ is chosen as the RKHS for estimation. If $f_0\in \mathbb{H}_2^{s_0}$, then $c=s_0/s$; $c$ quantifies the additional smoothness of $f_0$ relative to $\mathcal{H}$. In this Sobolev space, $b=2s/p>1$. The effective dimension is increasing in the original dimension $p$ and decreasing in the degree of smoothness $s$. The minimax optimal regression rates are

and both are achieved by kernel ridge regression with $\lambda=n^{-1/(c+1/b)}=n^{-2s/(2s_0+p)}$.

We place five types of assumptions in this paper, generalizing the standard RKHS learning theory assumptions from kernel ridge regression to mediated and time-varying causal inference: identification, RKHS regularity, original space regularity, smoothness, and effective dimension. We formally instantiate these assumptions for mediated responses in Section (ref), time-varying responses in Section (ref), and counterfactual distributions in Section 10 of rahul2024supplement. We uncover a double spectral robustness in related semiparametric inferential theory in Section 9 of rahul2024supplement: some kernels may have higher effective dimensions, as long as other kernels have lower effective dimensions.

Mediated response curves

Pearl's mediation formula

Mediation analysis decomposes the total effect of a treatment $D$ on an outcome $Y$ into the direct effect versus the indirect effect mediated via the mechanism $M$. The problem is sequential since $D$ causes $M$ and $Y$, then $M$ also causes $Y$. We denote the counterfactual mediator $M^{(d)}$ given a hypothetical intervention on the treatment $D=d$. We denote the counterfactual outcome $Y^{(d,m)}$ given a hypothetical intervention on the treatment $D=d$ and the mediator $M=m$.

$\theta_0^{TE}(d,d')$ generalizes average total effects. The average total effect of a binary treatment is $E[Y^{\{1,M^{(1)}\}}-Y^{\{0,M^{(0)}\}}]$. For a continuous treatment, the function $d\mapsto E[Y^{\{d,M^{(d)}\}}]$ may be infinite dimensional, which makes this problem fully nonparametric.

An analyst may wish to measure how much of the total effect is indirect: how much of the total effect would be achieved by simply intervening on the distribution of the mediator $M$? For example, in Section (ref), we investigate the extent to which employment mediates the effect of job training on arrests. With a binary treatment, the indirect effect is $E[Y^{\{1,M^{(1)}\}}-Y^{\{1,M^{(0)}\}}]$. In the former term, the mediator follows the counterfactual distribution under the intervention $D=1$, and in the latter, it follows the counterfactual distribution under the intervention $D=0$.

The remaining component of the total effect is the direct effect: if the mediator were held at the original distribution corresponding to $D=d$, what would be the impact of the treatment $D=d'$? For example, in Section (ref), we investigate the effect of job training on arrests holding employment at the original distribution. With a binary treatment, the direct effect is $E[Y^{\{1,M^{(0)}\}}-Y^{\{0,M^{(0)}\}}]$.

The final target parameter is $\theta_0^{ME}(d,d')$. It is useful because $\theta_0^{TE}(d,d')$, $\theta_0^{IE}(d,d')$, and $\theta_0^{DE}(d,d')$ can be expressed in terms of $\theta_0^{ME}(d,d')$. With a binary treatment, this quantity is a matrix in $\mathbb{R}^{2\times 2}$. With a continuous treatment, it is a surface over $\mathcal{D}\times\mathcal{D}$.

In seminal works, robins1992identifiability,pearl2001direct,imai2010identification state sufficient conditions under which the mediated response curves can be measured from the outcome $Y$, treatment $D$, mediator $M$, and covariates $X$, which we call selection on observables for mediation. We modestly extend the classic identification result from $\theta_0^{ME}(d,d')$ to its incremental version $\theta_0^{ME,\nabla}(d,d')$. Let $ \gamma_0(d,m,x)=E(Y|D=d,M=m,X=x) $.

See Section 12 of rahul2024supplement for the identifying assumptions and proof of our extension. Proposition (ref) identifies the other quantities in Definition (ref). For subsequent analysis, it helps to define $\omega_0(d,d';x)=\int \gamma_0(d',m,x) \mathrm{d}P(m|d,x)$, so that $\theta_0^{ME}(d,d')=\int \omega_0(d,d';x) \mathrm{d}P(x)$.

Sequential kernel embedding

Lemma (ref) makes precise how each mediated response curve is identified as a sequential integral of the form $\int \gamma_0(d',m,x)\mathrm{d}Q$ for the distribution $Q=P(m|d,x)P(x)$. Since $x$ appears in $\gamma_0(d',m,x)$, $P(m|d,x)$, and $P(x)$, the sequential integral is coupled and therefore challenging to estimate. We prove that, with the appropriate RKHS construction, the components $\gamma_0(d',m,x)$, $P(m|d,x)$, and $P(x)$ can be decoupled. Moreover, the sequential distribution $Q$ can be encoded by a sequential kernel embedding, which is our key innovation. We use these techniques to reduce sequential causal inference into the combination of kernel ridge regressions, which then allows us to propose simple estimators with closed form solutions.

To begin, we construct the appropriate RKHS for $\gamma_0$. In our construction, we define RKHSs for the treatment $D$, mediator $M$, and covariates $X$, then assume that the regression is an element of a certain composite space. To lighten notation, we will suppress subscripts when arguments are provided. We assume the regression $\gamma_0$ is an element of the RKHS $\mathcal{H}$ with the kernel $k(d,m,x;d',m',x') = k_{\mathcal{D}}(d,d') k_{\mathcal{M}}(m,m') k_{\mathcal{X}}(x,x')$. Formally, this choice of kernel corresponds to the tensor product: $\gamma_0\in\mathcal{H}=\mathcal{H}_{\mathcal{D}}\otimes \mathcal{H}_{\mathcal{M}} \otimes \mathcal{H}_{\mathcal{X}}$, with the tensor product dictionary of basis functions $\phi(d)\otimes \phi(m)\otimes \phi(x)$. As such, $\gamma_0(d,m,x)=\langle \gamma_0, \phi(d)\otimes \phi(m)\otimes \phi(x)\rangle_{\mathcal{H}} $ and $\|\phi(d)\otimes \phi(m)\otimes \phi(x)\|_{\mathcal{H}}=\|\phi(d)\|_{\mathcal{H}_{\mathcal{D}}} \|\phi(m)\|_{\mathcal{H}_{\mathcal{M}}} \|\phi(x)\|_{\mathcal{H}_{\mathcal{X}}}$. We place regularity conditions on this RKHS construction in order to prove our decoupling result. Anticipating Section 10 of rahul2024supplement, we include conditions for an outcome RKHS in parentheses.

Commonly used kernels are continuous and bounded. Measurability is a similarly weak condition. The characteristic property means that different distributions will have different embeddings in the RKHS. For example, the indicator kernel is characteristic over a discrete domain, while the exponentiated quadratic kernel is characteristic over a continuous domain.

See Section 13 of rahul2024supplement for the full proof. The quantity $\mu_{m}(d,x)=\int \phi(m)\mathrm{d}P(m|d,x)$ embeds the conditional distribution $P(m|d,x)$ as an element of the RKHS $\mathcal{H}_{\mathcal{M}}$, which is a popular technique in the RKHS literature. It satisfies, for $f\in\mathcal{H}_{\mathcal{M}}$, $\langle f, \mu_{m}(d,x) \rangle_{\mathcal{H}_{\mathcal{M}}}=\int f(m) \mathrm{d}P(m|d,x)$.

Theorem (ref) decouples $\gamma_0(d',m,x)$, $P(m|d,x)$, and $P(x)$, providing a blueprint for estimation that avoids density estimation and sequential integration. Our estimator will be $\hat{\theta}^{ME}(d,d')=\langle \hat{\gamma}, \phi(d')\otimes \hat{\mu}_{m,x}(d) \rangle_{\mathcal{H}}$ where $ \hat{\mu}_{m,x}(d) =n^{-1}\sum_{i=1}^n\{\hat{\mu}_m(d,X_i)\otimes \phi(X_i)\}$. The estimator $\hat{\gamma}$ is a standard kernel ridge regression. The estimator $\hat{\mu}_m(d,x)$ is an appropriately generalized kernel ridge regression. We combine them by averaging and taking the product.

We derive this algorithm in Section 13 of rahul2024supplement. We give theoretical values for $(\lambda,\lambda_1)$ that optimally balance bias and variance in Theorem (ref) below. Section 16 gives practical tuning procedures with closed form solutions to empirically balance bias and variance, one of which is asymptotically optimal. We formally define the operator space $\mathcal{L}_2(\mathcal{H}_{\mathcal{M}},\mathcal{H}_{\mathcal{D}}\otimes \mathcal{H}_{\mathcal{X}})$ below.

Uniform consistency with finite sample rates

Towards a guarantee of uniform consistency, we place regularity conditions on the original spaces. Anticipating Section 10 of rahul2024supplement, we include conditions for the outcome in parentheses.

A Polish space is a separable and completely metrizable topological space. Random variables supported in a Polish space may be discrete or continuous and even texts, graphs, or images. Boundedness of the outcome $Y$ can be relaxed. Next, we place assumptions on the smoothness of the regression $\gamma_0$ and the effective dimension of $\mathcal{H}$ the sense of (ref).

See Section 14 of rahul2024supplement for alternative ways of writing and interpreting Assumption (ref) in the tensor product space $\mathcal{H}$. We place similar conditions on the conditional kernel embedding $\mu_m(d,x)$, which is a generalized regression. We articulate this assumption abstractly for the conditional kernel embedding $\mu_{a}(b)=\int \phi(a)\mathrm{d}P(a|b)$ where $a\in\mathcal{A}_{\ell}$ and $b\in\mathcal{B}_{\ell}$. As such, all one has to do is specify $\mathcal{A}_{\ell}$ and $\mathcal{B}_{\ell}$ to specialize the assumption. For $\mu_m(d,x)$, $\mathcal{A}_1=\mathcal{M}$ and $\mathcal{B}_1=\mathcal{D}\times \mathcal{X}$. We parametrize the effective dimension and smoothness of $\mu_{a}(b)$ by $(b_{\ell},c_{\ell})$.

Formally, define the conditional expectation operator $E_{\ell}:\mathcal{H}_{\mathcal{A}_{\ell}}\rightarrow\mathcal{H}_{\mathcal{B}_{\ell}}$, $f(\cdot)\mapsto E\{f(A_{\ell})|B_{\ell}=\cdot\}$. By construction, $E_{\ell}$ encodes the same information as $\mu_{a}(b)$ since $$ \mu_{a}(b)=\int \phi(a)\mathrm{d}P(a|b) =[E_{\ell}\{\phi(\cdot)\}](b) =[E_{\ell}^* \{\phi(b)\}](\cdot),\quad a\in\mathcal{A}_{\ell},\quad b\in\mathcal{B}_{\ell}, $$ where $E_{\ell}^*$ is the adjoint of $E_{\ell}$. We denote the space of Hilbert--Schmidt operators between $\mathcal{H}_{\mathcal{A}_{\ell}}$ and $\mathcal{H}_{\mathcal{B}_{\ell}}$ by $\mathcal{L}_2(\mathcal{H}_{\mathcal{A}_{\ell}},\mathcal{H}_{\mathcal{B}_{\ell}})$. Here, $\mathcal{L}_2(\mathcal{H}_{\mathcal{A}_{\ell}},\mathcal{H}_{\mathcal{B}_{\ell}})$ is an RKHS in its own right, for which we place smoothness and effective dimension assumptions in the sense of (ref). See grunewalder2013smooth and singh2019kernel for details.

Just as we place approximation assumptions for $\gamma_0$ in terms of $\mathcal{H}$, which provides the features onto which we project $Y$, we place approximation assumptions for $E_{\ell}$ in terms of $\mathcal{H}_{\mathcal{B}_{\ell}}$, which provides the features $\phi(B_{\ell})$ onto which we project $\phi(A_{\ell})$. Under these conditions, we arrive at our first main theoretical guarantee.

See Section 14 of rahul2024supplement for the proof. By Proposition (ref), the quantities in Definition (ref) are uniformly consistent with the same rate, which combines optimal rates for standard fischer2017sobolev and generalized li2022optimal kernel ridge regression in RKHS norm. Section 14 gives the exact finite sample rate and the explicit specialization of Assumption (ref). The rate is at best $n^{-1/4}$ when $(c,c_1)=2$ and $(b,b_1)\rightarrow \infty$, i.e. when $(\gamma_0,\mu_m)$ are very smooth with finite effective dimensions. The rates reflect the challenge of a $\sup$ norm guarantee, which is much stronger than an $\mathbb{L}^2$ norm guarantee and is useful for policymakers who may be concerned about each treatment value. See (ref) to specialize these rates for Sobolev spaces.

Time-varying response curves

Robins' g-formula

So far we have considered the effect of a single treatment $D$. Next, we consider the effect of a sequence of time-varying treatments $D_{1:T}=d_{1:T}$ on the counterfactual outcome $Y^{(d_{1:T})}$. If the sequence of treatment values $d_{1:T}$ is observed in the data, this problem may be called on-policy planning; if not, it may be called off-policy planning.

Whereas much of the semiparametric literature restricts $d_t$ to be discrete, we allow $d_t$ to be continuous and consider a nonparametric approach to time-varying response curves gill2001causal. For example, in Section (ref), we estimate the effect of $d_1$ class hours in year one, and $d_2$ class hours in year two, on counterfactual employment, with treatment-confounder feedback. In the spirit of off-policy planning, we consider a distribution shift from $P$ to $\tilde{P}$.

In seminal work, robins1986new states sufficient conditions under which time-varying responses can be measured from the outcome $Y$, treatments $D_{1:T}$, and covariates $X_{1:T}$. We refer to this collection of conditions as sequential selection on observables. We modestly extend the classic identification result by considering incremental responses.

See Section 12 of rahul2024supplement for the identifying assumptions and proof of our extension. We consider a fully nonparametric g-formula with possibly continuous treatments that allows for distribution shift. Lemma (ref) handles auxiliary Markov restrictions as special cases, e.g. if covariates follow a Markov process, then $\theta_0^{GF}$ simplifies by setting $P\{x_t|d_{1:(t-1)},x_{1:(t-1)}\}=P(x_t|d_{t-1},x_{t-1})$.

Sequential kernel embedding

Similar to Lemma (ref), Lemma (ref) identifies each time-varying response as a sequential integral of the form $\int \gamma_0(d_{1:T},x_{1:T})\mathrm{d}Q$ for the distribution $Q=P(x_1) \prod_{t=2}^{T} P\{x_t|d_{1:(t-1)},x_{1:(t-1)}\}$ or $Q=\tilde{P}(x_1) \prod_{t=2}^{T} \tilde{P}\{x_t|d_{1:(t-1)},x_{1:(t-1)}\}$. As before, components of the sequential integral are coupled and therefore challenging to estimate, since $x_1$ appears in $\gamma_0(d_{1:T},x_{1:T})$, $P\{x_t|d_{1:(t-1)},x_{1:(t-1)}\}$, and $P(x_1)$. As in Section (ref), we construct an appropriate RKHS to decouple these components, then to encode $Q$ by a sequential kernel embedding. With these techniques, we again reduce sequential causal inference into the combination of kernel ridge regressions. For clarity, we present the algorithm with $T=2$, and we define $\omega_0(d_1,d_2;x_1)=\int \gamma_0(d_1,d_2,x_1,x_2)\mathrm{d}P(x_2|d_1,x_1)$ so that $\theta_0^{GF}(d_1,d_2)=\int \omega_0(d_1,d_2;x_1)\mathrm{d}P(x_1)$. We consider $T>2$ in Section 11 of rahul2024supplement, which also showcases the role of Markov assumptions.

To construct the RKHS for $\gamma_0$, we define RKHSs for each treatment $D_t$ and each covariate $X_t$. Using identical notation as Section (ref), we assume the regression $\gamma_0$ is an element of the RKHS $\mathcal{H}$ with the kernel $k(d_1, d_2, x_1, x_2; d_1', d_2', x_1', x_2') = k_{\mathcal{D}}(d_1, d_1') k_{\mathcal{D}}(d_2, d_2') k_{\mathcal{X}}(x_1, x_1') k_{\mathcal{X}}(x_2, x_2')$, i.e. $\gamma_0\in\mathcal{H}=\mathcal{H}_{\mathcal{D}}\otimes\mathcal{H}_{\mathcal{D}}\otimes \mathcal{H}_{\mathcal{X}} \otimes \mathcal{H}_{\mathcal{X}}$. As such, $\gamma_0(d_1,d_2,x_1,x_2)=\langle \gamma_0, \phi(d_1)\otimes\phi(d_2) \otimes \phi(x_1)\otimes \phi(x_2)\rangle_{\mathcal{H}} $ and $\|\phi(d_1)\otimes \phi(d_2)\otimes \phi(x_1)\otimes \phi(x_2)\|_{\mathcal{H}}=\|\phi(d_1)\|_{\mathcal{H}_{\mathcal{D}}} \|\phi(d_2)\|_{\mathcal{H}_{\mathcal{D}}} \|\phi(x_1)\|_{\mathcal{H}_{\mathcal{X}}} \|\phi(x_2)\|_{\mathcal{H}_{\mathcal{X}}}$. Under regularity conditions on this RKHS construction, we prove an analogous decoupling result.

See Section 13 of rahul2024supplement for the proof. In $\theta_0^{GF}$, the conditional kernel embedding of $P(x_2|d_1,x_1)$ is $\mu_{x_2}(d_1,x_1)=\int \phi(x_2)\mathrm{d}P(x_2|d_1,x_1)$, and it satisfies $\langle f,\mu_{x_2}(d_1,x_1) \rangle_{\mathcal{H}_{\mathcal{X}}}= \int f(x_2)\mathrm{d}P(x_2|d_1,x_1)$.

As before, this decoupling is a blueprint for estimation. For example, our estimator will be $\hat{\theta}^{GF}(d_1,d_2)=\langle \hat{\gamma}, \phi(d_1)\otimes\phi(d_2)\otimes \hat{\mu}_{x_1,x_2}(d_1) \rangle_{\mathcal{H}}$ where $ \hat{\mu}_{x_1,x_2}(d_1)=n^{-1}\sum_{i=1}^n\{\phi(X_{1i})\otimes \hat{\mu}_{x_2}(d_1,X_{1i}) $. Here, $\hat{\gamma}$ is a kernel ridge regression, $\hat{\mu}_{x_2}(d_1,x_1)$ is a generalized kernel ridge regression, and we combine them by averaging and taking the product.

We derive these algorithms in Section 13 of rahul2024supplement. We give theoretical values for $(\lambda,\lambda_4,\lambda_5)$ that optimally balance bias and variance in Theorem (ref) below. Section 16 gives practical tuning procedures with closed form solutions to empirically balance bias and variance, one of which is asymptotically optimal. Note that $\hat{\theta}^{DS}$ requires observations of the treatments and covariates from the alternative population $\tilde{P}$.

Uniform consistency with finite sample rates

Towards a guarantee of uniform consistency, we place regularity conditions on the RKHSs and original spaces via Assumptions (ref) and (ref). We also assume the regression $\gamma_0$ is smooth and quantify the effective dimension of $\mathcal{H}$ via Assumption (ref). For the conditional kernel embeddings $\mu_{x_2}(d_1,x_1)$ and $\nu_{x_2}(d_1,x_1)$, we place further smoothness and effective dimension conditions via Assumption (ref). With these assumptions, we arrive at our next main result.

See Section 14 of rahul2024supplement for the proof, exact finite sample rates, and explicit specializations of Assumption (ref). As before, these rates are at best $n^{-1/4}$ when $(c,c_4,c_5)=2$ and $(b,b_4,b_5)\rightarrow \infty$. See (ref) for the Sobolev special case.

Simulations and application

Simulations

We evaluate our estimators on various nonparametric designs. For each nonparametric design and sample size, we implement 100 simulations and calculate mean square error (MSE) with respect to the true causal function. Figure (ref) visualizes results, where a lower MSE is desirable. See Section 17 of rahul2024supplement for the data generating processes and implementation details.

The mediated response design huber2018direct involves learning the nonlinear causal function $\theta_0^{ME}(d,d')=0.3d'+0.09d+0.15 dd'+0.25 (d')^3$. A single observation is a tuple $(Y,D,M,X)$ for the outcome, treatment, mediator, and covariates where $Y,D,M,X\in\mathbb{R}$. In addition to our estimator (RKHS, white), we implement the estimator of huber2018direct (IPW, checkered gray), which involves Nadaraya--Watson density estimation en route to generalized inverse propensity weighting. By the Wilcoxon rank sum test, RKHS significantly outperforms IPW at all sample sizes, with p values below $10^{-3}$.

Next, we consider a time-varying dose response design, extending a time-fixed design colangelo2020double. The nonlinear causal function is $\theta_0^{GF}(d_1,d_2)=0.6d_1+0.5d_1^2+1.2d_2+d_2^2$. A single observation is a tuple $(Y,D_{1:2},X_{1:2})$ for the outcome, treatments, and covariates where $Y,D_t\in\mathbb{R}$ and $X_t\in\mathbb{R}^{100}$. See Section 17 of rahul2024supplement for low and moderate dimensional settings. Our machine learning approach for time-varying response curves is uniformly consistent and allows for nonlinearity, dependence over time, and effect modification, which appears to be new.

To illustrate why treatment-confounder feedback and effect modification matter, we compare $\hat{\theta}^{GF}(d_1,d_2)$ \{RKHS(GF)\!, white\} with estimators that ignore these complexities to various degrees. Using the dose response estimator of singh2020kernel \{RKHS(ATE)\!, checkered white\}, we take $D_2$ to be the treatment and misclassify $D_1$ as a covariate. Using the heterogeneous response estimator of singh2020kernel \{RKHS(CATE)\!, lined white\}, we take $D_2$ to be the treatment and misclassify $D_1$ as the subcovariate with meaningful heterogeneity. We also implement the estimator of lewis2020double (SNMM, gray), which is a machine learning approach with linearity, Markov, and no-effect-modification assumptions that do not hold in this setting. By the the Wilcoxon rank sum test, RKHS(GF) significantly outperforms the alternatives at $n=5000$ with p value below $10^{-3}$. The ability of RKHS(GF) to capture treatment-confounder feedback, and effect modification by time-varying confounders, helps when the sample size is large enough.

Application: US Job Corps

We estimate the mediated and time-varying responses of the Job Corps, the largest job training program for disadvantaged youth in the US. The Job Corps serves about 50,000 participants annually, and it is free for individuals who meet low income requirements. Access to the program was randomized from November 1994 to February 1996; see schochet2008does for details. Though access to the program was randomized, individuals could decide whether to participate and for how many hours over multiple years. We assume that, conditional on the observed covariates, those decisions were as good as random in the sense formalized in Section 12 of rahul2024supplement.

First, we consider employment to be a possible mechanism through which class hours affect arrests, under the identifying assumptions of flores2012estimating,huber2018direct. The covariates $X\in\mathbb{R}^{40}$ are measured at baseline; the treatment $D\in\mathbb{R}$ is total hours spent in academic or vocational classes in the first year after randomization; the mediator $M\in\mathbb{R}$ is the proportion of weeks employed in the second year after randomization; and the outcome $Y\in\mathbb{R}$ is the number of times an individual is arrested by police in the fourth year after randomization. We use the same covariates $X\in\mathbb{R}^{40}$ and sample as colangelo2020double, with $n=2,913$ observations. In Figure (ref), $\theta_0^{TE}(d,d')$, $\theta_0^{DE}(d,d')$, and $\theta_0^{IE}(d,d')$ are the total, direct, and indirect responses, respectively, of $d'$ class hours relative to $d$ class hours on arrests. In particular, $\theta_0^{IE}(d,d')$ estimates the extent to which the response is mediated by the mechanism of employment.

At best, the total response of receiving 1,600 class hours (40 weeks) versus 480 class hours (12 weeks) may be a reduction of about 0.1 arrests. The direct response estimate, of class hours on arrests, mirrors the total response estimate. Our indirect response estimate of class hours on arrests, as mediated through employment, is essentially zero. Our results extend the findings of huber2018direct, allowing both $(d,d')$ to vary. It appears that the effect of class hours on arrests is direct; there may be benefits of the training program that are not explained by employment alone. These benefits, however, may require many class hours.

Next, we evaluate the time-varying response of job training on employment. Here, $X_1\in\mathbb{R}^{65}$ are covariates at baseline; $D_1\in\mathbb{R}$ is the total class hours in the first year; $X_2\in\mathbb{R}^{30}$ are covariates observed at the end of the first year; $D_2\in \mathbb{R}$ is the total class hours in the second year; and $Y\in\mathbb{R}$ is the proportion of weeks employed in the fourth year. The covariates and the sample of $n=3,141$ observations we use are similar to colangelo2020double. The time-varying response $\theta^{GF}_0(d_1,d_2)$ is the counterfactual mean employment given $d_1$ class hours in year one and $d_2$ class hours in year two; $\theta^{GF,\nabla}_0(d_1,d_2)$ is the increment of counterfactual mean employment given $d_1$ class hours in year one and incrementally more than $d_2$ class hours in year two. Figure (ref) visualizes the time-varying response estimate and its derivative with respect to the second dose.

The effect of training on employment appears to be positive when the duration of training is relatively brief. At best, the response to receiving job training appears to be 64% employment, compared to receiving no class hours at all which gives 56% employment. The maximum response is achieved by 480--1280 class hours (12--32 weeks) in year one and 0--480 (0--12 weeks) in year two. There is another local maximum of counterfactual employment achieved by 1200 class hours (30 weeks) in both years. Class hours in year one and year two may be complementary at low levels, as visualized by the incremental response. The large plateau in counterfactual employment suggests that a successful yet cost effective policy may be 480 class hours (12 weeks) in the first year and an optional, brief follow up in the second year.

In summary, under standard identifying assumptions, we find that the US Job Corps may provide two distinct benefits: reducing arrests and increasing employment, under different durations of class hours. Many class hours in the first year may directly decrease arrests in the fourth year, while few class hours in the first and second years may significantly increase employment in the fourth year. Section 18 of rahul2024supplement provides implementation details and verifies that our results are robust to the sample choice.

Discussion

Previous methods using kernels for continuous treatments singh2020kernel do not handle treatment-confounder feedback, and therefore cannot analyze the later rounds of Job Corps surveys. We propose the sequential kernel embedding to do so. Whereas survey data for mediation analysis were previously available huber2018direct, we clean additional survey data for time-varying analysis from raw files schochet2008does. By providing clean data, we enable empirical analysis of the later rounds of Job Corps surveys, where class hours in different years may be viewed as a sequence of time-varying continuous treatments. Future work may apply the sequential embedding in dynamic programming for optimal policy estimation.

\spacingset{1}

\DeclareRobustCommand{\VAN}[2]{#2} \DeclareRobustCommand{\VAN}[2]{#1}

\spacingset{1.65}