Simplifying debiased inference via automatic differentiation and probabilistic programming

Abstract

We introduce an algorithm that simplifies the construction of efficient estimators, making them accessible to a broader audience. ‘Dimple’ takes as input computer code representing a parameter of interest and outputs an efficient estimator. Unlike standard approaches, it does not require users to derive a functional derivative known as the efficient influence function. Dimple avoids this task by applying automatic differentiation to the statistical functional of interest. Doing so requires expressing this functional as a composition of primitives satisfying a novel differentiability condition. Dimple also uses this composition to determine the nuisances it must estimate. In software, primitives can be implemented independently of one another and reused across different estimation problems. We provide a proof-of-concept Python implementation and showcase through examples how it allows users to go from parameter specification to efficient estimation with just a few lines of code.

1. Introduction

1.1. Motivation

There is a long history of constructing estimators that are asymptotically efficient (‘efficient’ hereafter) in infinite-dimensional models (Levit, 1975; Hasminskii and Ibragimov, 1979; Pfanzagl, 1990; Bickel et al., 1993). Several approaches for doing this permit the use of arbitrary methods for estimating needed nuisance functions. These include one-step estimation (Pfanzagl, 1982), estimating equations (van der Laan and Robins, 2003), targeted learning (van der Laan and Rubin, 2006), and double machine learning (Chernozhukov et al., 2017). These frameworks construct efficient estimators using an object known as the efficient influence function (EIF). Though EIFs have been derived in many examples (Newey, 1990; Bickel et al., 1993; van der Laan and Robins, 2003), computing them in new problems can be challenging, requiring tools from functional analysis and differential geometry. These calculations can represent a bottleneck in scientific discovery, as some novel analyses cannot be performed until an expert derives the EIF.

Investigators have sought to use numerical approaches to overcome this bottleneck. Frangakis et al. (2015) built on the concept of an influence curve from robust statistics (Hampel, 1968) to propose a method using finite differences to approximate nonparametric EIFs. This method applies whenever the data are discrete or the parameter is a smooth functional of the distribution function. Luedtke et al. (2015) provided a modified finite difference approach that correctly approximates the EIF more generally, even when the data are continuous. Van der Laan et al. (2015) and Carone et al. (2018) improved upon this proposal by providing numerical strategies for approximating EIFs in the harder case where the model is semiparametric, rather than nonparametric. In a concurrent line of work, Ichimura and Newey (2022) developed a numerical approach to evaluate the influence function of a given estimator, which coincides with the EIF when the estimator is efficient. Though one pursues numerical approximation of the influence function of an estimator and the other pursues numerical approximation of the EIF of a statistical functional, the proposals in Luedtke et al. (2015) and Ichimura and Newey (2022) address challenges related to the possible non-discreteness of the data in similar ways.

The aforementioned numerical approaches rely on two hyperparameters: a step size $ϵ$ for a finite difference approximation and a bandwidth $\lambda$ for a kernel approximation of a delta function. While both must be small, $ϵ$ should be an order of magnitude smaller than $\lambda$. However, setting $ϵ$ too small can lead to numerical instability. Carone et al. (2018) provide guidance on the choice of hyperparameters through a diagnostic they call $ϵ-\lambda$ plots. Jordan et al. (2022) studied the conditions that a numerical approximation must meet to guarantee that the resulting one-step estimator possesses desirable statistical properties. The complexity of these requirements underscores the need for careful hyperparameter selection in existing numerical approaches, which could limit their adoption. Researchers in other fields have encountered similar challenges when using numerical differentiation (Griewank and Walther, 2008).

Automatic differentiation avoids these challenges (Wengert, 1964; Speelpenning, 1980). In this approach, a function $f$ is written as a composition of differentiable functions with known gradients, which we refer to as primitives. Then, an algorithm can iteratively apply the chain rule to derive the gradient of $f$ at a specified point. Unlike numerical differentiation, automatic differentiation is hyperparameter-free and is always accurate up to working precision. Differentiable programming extends automatic differentiation by identifying a compositional structure for $f$ directly via computer code that can be used to evaluate it, simplifying implementation (Baydin et al., 2018; Blondel and Roulet, 2024). Automatic differentiation and differentiable programming have been used with success in applications (Bischof, 1995; Homescu, 2011), and have played a crucial role in the rapid progress made in deep learning (Goodfellow et al., 2016). However, they have yet to be used to compute a semi- or nonparametric EIF. This may be due to the fact that EIFs correspond to a gradient of a functional defined on an infinite-dimensional statistical manifold, and so it has been unclear what primitives can be chained together to compute them algorithmically.

1.2. Contributions

Our first two contributions are as follows:

• 1In Section 2.1, we provide an automatic differentiation algorithm for computing the EIF of any parameter that can be expressed as a composition of differentiable Hilbert-valued primitives.
• 2In Section 2.2, we establish the differentiability of a variety of primitives. These primitives can be composed to express many parameters of interest.

We then turn to estimation, building on developments in probabilistic programming (van de Meent et al., 2018). Probabilistic programming and finite-dimensional automatic differentiation have been employed for Bayesian inference (Ge et al., 2018; Tran et al., 2018; Stan Development Team, 2023), parametric maximum likelihood estimation (Fournier et al., 2012; Kristensen et al., 2015), M-estimation (Zivich et al., 2022; Kosmidis and Lunardon, 2024), and efficient estimation in high-dimensional parametric models (Agrawal et al., 2024). However, to date there is not a probabilistic programming framework for efficiently estimating a generic smooth parameter that may rely on nonparametric nuisances like density and regression functions. Our remaining contributions address this gap.

• 3In Section 3, we provide an algorithm to estimate any parameter that is expressed as a composition of known primitives and conditions under which this estimator will be efficient.
• 4In Section 1.3, we introduce a proof-of-concept Python package implementing our approach, available at http://github.com/alexluedtke12/pydimple.

We refer to our framework for constructing efficient estimators using automatic differentiation and probabilistic programming as ‘dimple’, short for ‘debiased inference made simple’.

The goal of dimple is similar in spirit to that of automatic debiased machine learning methods: to avoid analytical derivations when constructing efficient estimators (Chernozhukov et al., 2022a,b,c). Those existing methods apply to regression functionals, whose EIFs depend on the functional’s form through nuisance functions known as Riesz representers. Automatic debiased machine learning provides loss functions that can be used to estimate these nuisances without explicitly deriving their form. Compared to these methods, dimple innovates in several directions. It allows the functional to depend on the distribution through non-regression summaries, uses automatic differentiation to compute the EIF, and is integrated into a user-friendly probabilistic programming framework. It also applies to general Hilbert-valued parameters (Luedtke and Chung, 2024).

Beyond simplifying estimation, our automatic differentiation approach and the chain rule that underlies it may be of independent interest for their use in deriving semiparametric efficiency bounds, which are given by the variance of the EIF when it exists (van der Vaart and Wellner, 1989; van der Vaart, 1991a). Our theoretical results show this EIF exists for any parameter expressible as a composition of differentiable primitives, and applying automatic differentiation by hand gives its analytic form. An example of this is given in Section 2.3.

1.3. Illustrations of proposed approach

Our proof-of-concept Python package, pydimple, uses dimple to construct nonparametric efficient estimators. We illustrate it with three examples. The following Python code imports the classes and functions from the package needed to run them:

In our first example, the objective is to estimate the expected density $\psi \left(P\right)≔E_P[p(z)]$, with $p$ the probability density function of $Z~P$.

Example 1:

Expected density (Bickel and Ritov, 1988).

The first line of code defines a distribution $P$ and identifies a dataset dat containing iid draws from it to be used for estimation. The next two lines serve as a blueprint of the statistical estimation problem, formulating it using primitives from pydimple. The fourth line computes a 5-fold cross-fitted estimator, Wald-type standard error, and 95% confidence interval of $\psi \left(P\right)$. Internally, pydimple estimates density functions using kernel density estimation and all other nuisances using lightgbm (Ke et al., 2017), with hyperparameters tuned via Optuna (Akiba et al., 2019).

In our next example, the objective is to estimate the nonparametric $R^2$ criterion, defined as $\psi \left(P\right)≔1-\int {\left\{y-E_P(Y\mid X=x)\right\}}^{2}dP(x,y)/{\text{Var}}_P\left(Y\right)$, where $P$ is a distribution of $(X,Y)$.

Example 2:

Nonparametric $R^2$ criterion (Williamson et al., 2021).

The final example concerns the estimation of the longitudinal G-formula (Robins, 1986; Bang and Robins, 2005; Luedtke et al., 2017; Rotnitzky et al., 2017; Chernozhukov et al., 2022b). This parameter is defined in longitudinal settings where covariate-treatment pairs $(X_t,A_t)$ are collected at times $t=0,1,\dots ,T-1$, and then an outcome $Y$ is collected at the end of the study. For $t\in \{\mathrm{0,1},\dots ,T-1\}$, let $H_t$ denote the history $(X_0,A_0,X_1,A_1,\dots ,X_{t-1},A_{t-1},X_t)$ available just before treatment $A_t$ is administered. Also let $H_T=\left(H_{T-1},A_{T-1},Y\right)$ denote all of the covariate, treatment, and outcome data. The parameter $\psi \left(P\right)$ is defined recursively by letting ${\mu }_{P,T}\left(h_T\right)=y,{\mu }_{P,t}\left(h_t\right)=E_P\left[{\mu }_{P,t+1}\left(H_{t+1}\right)\mid A_t=1,H_t=h_t\right]$ for $t=T-1$, $T-2,\dots ,0$, and $\psi \left(P\right)=E_P\left[{\mu }_{P,0}\left(H_0\right)\right]$. Under causal conditions, $\psi \left(P\right)$ gives the counterfactual mean outcome when everyone receives treatment 1 at all time points. The recursive definition of this parameter can be easily expressed in pydimple via a for loop, as illustrated below.

Example 3:

Longitudinal G-Formula (Robins, 1986).

In Section 4, we evaluate the performance of pydimple in a Monte Carlo study. The results support that dimple can produce asymptotically efficient estimators and valid confidence intervals. At sample sizes of $n\ge 1000$, the confidence intervals have approximately nominal coverage and nearly optimal widths, scaling with the standard deviation of the EIF. These results are not surprising given that dimple uses a variant of one-step estimation, which has been shown to perform well in a variety of settings (Pfanzagl, 1982; Bickel et al., 1993). Details on this estimation framework and its theoretical properties are given in Section 3.1.

2. Automatic differentiation of statistical functionals

2.1. Algorithm and theoretical guarantee

We present an algorithm for automatically differentiating a parameter $\psi$ that maps from a set of distributions $ℳ$ into a Hilbert space ${𝒲}_{\psi }$. This and all Hilbert spaces in this work are real. The set $ℳ$, called the model, may be nonparametric or semiparametric.

Using our approach requires that $\psi$ can be expressed as in Algorithm 1, meaning that, for any $P\in ℳ$, calling that algorithm with input $P$ returns $\psi \left(P\right)$. The $j$-th line of that algorithm computes the evaluation $h_j$ of a primitive ${\theta }_j:ℳ\times {𝒰}_j\to {𝒱}_j$ applied to $P$ and $h_{\text{pa}\left(j\right)}≔\left(h_i:i\in \text{pa}\left(j\right)\right)$. Here, ${𝒰}_j\supseteq {\prod }_{i\in \text{pa}\left(j\right)}{𝒱}_i$ and $\text{pa}\left(j\right)\subseteq [j-1]≔\{\mathrm{1,2},\dots ,j-1\}$, where $\left[0\right]≔\varnothing$. When $\text{pa}\left(j\right)=\varnothing$, we let ${𝒰}_j≔\left\{0\right\}$ and $h_{\text{pa}\left(j\right)}≔0$. This adheres to the following general convention: for any tuple $\left(a_i:i\in \left[j\right]\right)$ and set $𝒮\subseteq \left[j\right]$, $a_{𝒮}≔\left(a_i:i\in 𝒮\right)$ when $𝒮\ne \varnothing$ and $a_{𝒮}≔0$ otherwise.

Algorithm 1.

To use automatic differentiation, there must exist an algorithm of the following form that takes as input a generic $P\in ℳ$ and returns $\psi (P)$

Though Algorithm 1 does not rely on ${𝒱}_j$ having an inner product structure, our automatic differentiation scheme does. Hence, we suppose we can specify an ambient Hilbert space ${𝒲}_j\supseteq {𝒱}_j$ for each $j\in [k-1]$. We also let ${𝒲}_k≔{𝒲}_{\psi }$ and view ${𝒰}_j$ as a subset of the direct sum Hilbert space ${⨁}_{i\in \text{pa}\left(j\right)}{𝒲}_i$, which we take to be a trivial Hilbert space $\left\{0\right\}$ when $\text{pa}\left(j\right)=0$.

To ensure that a chain rule can be applied to differentiate $\psi$, each ${\theta }_j$ must be differentiable in a sense we now introduce. We do this for a generic primitive $\theta :ℳ\times 𝒰\to 𝒱$ with $𝒰\subseteq 𝒯$ and $𝒱\subseteq 𝒲$ for ambient Hilbert spaces $𝒯$ and $𝒲$. We denote the norms on these spaces by $\Vert \cdot {\Vert }_{𝒯}$ and $\Vert \cdot {\Vert }_{𝒲}$. Given $u\in 𝒰$ and $t\in 𝒯$, we let $𝒫(u,𝒰,t)$ be the set of paths $\left\{u_{ϵ}:ϵ\in \left[\mathrm{0,1}\right]\right\}\subseteq 𝒰$ that satisfy ${‖u_{ϵ}-u-ϵt‖}_{𝒯}=o\left(ϵ\right)$. For $u\in 𝒰$, we call ${\check{𝒰}}_u≔\{t\in 𝒯:𝒫(u,𝒰,t)\ne \varnothing \}$ the tangent set of $𝒰$ at $u$ and the $𝒯$-closure of its linear span, ${\dot{𝒰}}_u$, the tangent space of $𝒰$ at $u$. The tangent space ${\dot{ℳ}}_P\subseteq L^2\left(P\right)$ of $ℳ$ at $P$ is similarly defined as the $L^2\left(P\right)$-closure of the linear span of the tangent set ${\check{ℳ}}_P≔\left\{s\in L^2\left(P\right):𝒫(P,ℳ,s)\ne \varnothing \right\}$, where $𝒫(P,ℳ,s)$ is the collection of quadratic mean differentiable univariate submodels $\left\{P_{ϵ}:ϵ\in \left[\mathrm{0,1}\right]\right\}\subseteq ℳ$ with score $s$ at $ϵ=0$ and $P_{ϵ=0}=P$. We call $\theta$ totally pathwise differentiable at $(P,u)$ if there exists a continuous linear operator ${\dot{\theta }}_{P,u}:{\dot{ℳ}}_P\oplus {\dot{𝒰}}_u\to 𝒲$ such that, for all $s\in {\check{ℳ}}_P,t\in {\check{𝒰}}_u,\left\{P_{ϵ}:ϵ\right\}\in 𝒫(P,ℳ,s)$, and $\left\{u_{ϵ}:ϵ\right\}\in 𝒫(u,𝒰,t)$,

$${‖\theta \left(P_{ϵ},u_{ϵ}\right)-\theta (P,u)-ϵ{\dot{\theta }}_{P,u}(s,t)‖}_{𝒲}=o(ϵ).$$ (1)

The differential operator ${\dot{\theta }}_{P,u}$ is unique. We denote its Hermitian adjoint by ${\dot{\theta }}_{P,u}^{*}:𝒲\to {\dot{ℳ}}_P\oplus {\dot{𝒰}}_u$.

Though we have not seen this precise notion of differentiability introduced previously, total pathwise differentiability is closely related to classical notions that apply to functions of a single argument, $u\in 𝒰$ or $P\in ℳ$. In particular, the total pathwise differentiability of $\theta$ at $(P,u)$ implies $\theta (P,\cdot )$ is Hadamard differentiable (Averbukh and Smolyanov, 1967) and $\theta (\cdot ,u)$ is pathwise differentiable (van der Vaart, 1991a), and the converse implication holds under additional regularity conditions — see Lemmas S5 and S6 in the appendix. In the special case where $𝒰$ is the trivial vector space $\left\{0\right\}$, $\theta$ is totally pathwise differentiable if and only if $\nu (\cdot )≔\theta (\cdot ,0)$ is pathwise differentiable. In these cases, ${\dot{\theta }}_P^{*}\left(w\right)=\left({\dot{\nu }}_P^{*}\left(w\right),0\right)$. The map ${\dot{\nu }}_P^{*}:𝒲\to {\dot{ℳ}}_P$ is called the efficient influence operator of $\nu$ at $P$; it can be used to draw inferences about $\nu \left(P\right)$ (Luedtke and Chung, 2024). When $𝒲=\mathbb{R}$ and $w=1,{\dot{\nu }}_P^{*}\left(w\right)$ is known as the EIF of $\nu$ at $P$ (Bickel et al., 1993, Definition 5.2.3).

Algorithm 2.

Automatic differentiation to evaluate the efficient influence operator at $P$

Like Hadamard and pathwise differentiability (Averbukh and Smolyanov, 1967; van der Vaart, 1991b; Kennedy, 2022), total pathwise differentiability satisfies a chain rule — see Lemma S3 in the appendix. The availability of a chain rule suggests a form of automatic differentiation (Linnainmaa, 1970), which we present in Algorithm 2. Initially, this algorithm requires the evaluation of the composition that defines $\psi$. Subsequently, information backpropagates through the indices $j$ of the composition to evaluate the efficient influence operator. This information is derived from the adjoint ${\dot{\theta }}_{j,P,h_{\text{pa}\left(j\right)}}^{*}$ of the differential operator of ${\theta }_j$ at $\left(P,h_{\text{pa}\left(j\right)}\right)$, where $h_{\text{pa}\left(j\right)}$ is as defined as in Algorithm 1 when that algorithm is called with input $P$.

Theorem 1 (Automatic differentiation works). Suppose $\psi$ can be expressed as in Algorithm 1 and ${\theta }_j$ is totally pathwise differentiable at $\left(P,h_{\text{pa}\left(j\right)}\right)$ for all $j\in \left[k\right]$. Then, $\psi$ is pathwise differentiable at $P$ and Algorithm 2 returns the efficient influence operator evaluation ${\dot{\psi }}_P^{*}\left(f_k\right)$.

To obtain the EIF of a real-valued parameter $\psi$, Algorithm 2 can be run with $f_k=1$. When proving Theorem 1, we define maps ${\eta }_j:ℳ\times \prod _{\ell =1}^j{𝒱}_{\ell }\to {𝒲}_{\psi }$ such that ${\eta }_j\left(P,v_{\left[j\right]}\right)$ is the output of a modification of Algorithm 1 that replaces each line $i\le j$ by the assignment $h_i=v_i$ and leaves each line $i>j$ unchanged, so that $h_i={\theta }_i\left(P,h_{\text{pa}\left(i\right)}\right)$. The crux of our proof involves showing that, just before step $j$ of the loop in Algorithm 2, $f_0={\dot{\eta }}_{j,P,h_{\left[j\right]}}^{*}\left(f_k\right)$.

We now discuss two strategies to increase the applicability of Algorithm 2. The first uses that the choice of ambient Hilbert spaces ${𝒲}_1,{𝒲}_2,\dots ,{𝒲}_{k-1}$ can impact whether the primitives used to express $\psi$ are totally pathwise differentiable. Hence, a good strategy is to choose these Hilbert spaces in a way that makes the primitives differentiable under mild conditions — see the next subsection for examples. The second uses that differentiability is a local property, and so Theorem 1 remains true if the representation of $\psi$ in Algorithm 1 is only valid in some Hellinger neighborhood $𝒩\subset ℳ$ of $P$. Concretely, it suffices that calling Algorithm 1 with input $P^{\prime }$ returns $\psi \left(P^{\prime }\right)$ for all $P^{\prime }\in 𝒩$; no such requirement is needed for $P^{\prime }\notin 𝒩$. This relaxation may help identify compositions of primitives to automatically differentiate.

There are many forms of automatic differentiation. They differ in how derivative information is propagated. The two extremes are reverse mode (Linnainmaa, 1970) and forward mode (Wengert, 1964). As exhibited in Algorithm 2, reverse mode backpropagates information through adjoints. Forward mode takes the opposite approach, propagating it forward through differential operators as the function is evaluated. There are also intermediate approaches that propagate information both forward and backward, each of which may have a different time complexity. Because identifying the best such traversal is an NP-complete problem (Naumann, 2008), many implementations of automatic differentiation employ either a simple forward or reverse mode. A general rule of thumb is that reverse mode will be faster than forward mode when a function’s input dimension exceeds its output dimension (Chapter 3 of Griewank and Walther, 2008). In our context, the input dimension is infinite unless $ℳ$ is parametric. Consequently, we pursue reverse mode automatic differentiation in this work. We leave consideration of alternative traversal strategies to future study.

2.2. Useful primitives

Table 1 provides a selection of useful primitives. In Appendix C, we establish that, under conditions, each is totally pathwise differentiable when $ℳ$ is a locally nonparametric model of distributions on $𝒵$. Some of these primitives involve a coarsened random variable $X=𝒞(z)$, where $𝒞:𝒵\to 𝒳$ is a many-to-one map; when $Z~Q\in ℳ$, we let $Q_X$ denote the marginal distribution of $X$. The primitives fall into three distinct categories, which we elaborate on below.

Table 1:

Primitives $\theta$ with domain $ℳ\times 𝒰$. Given the specified ambient Hilbert spaces $𝒯\supseteq 𝒰$ and $𝒲\supseteq \text{Image}\left(\theta \right)$, each is totally pathwise differentiable under conditions given in Appendix C.

Notation: $Q\in ℳ;X\in 𝒳$ and $A\in \left\{\mathrm{0,1}\right\}$ are coarsenings or subvectors of $Z;Q_X,Q_{A,X}$ marginal distribution of $X,(A,X)$ when $Z~Q;L^{\star }(\rho )=$ a set of $\rho$-a.s. equivalence classes, each satisfying a moment condition, where $\rho$ and $Q$ are equivalent measures; $L^{\star }\left({\rho }_X\right)$ is defined similarly but with $Q$ replaced by $Q_X;{\Pi }_P(u)(z)=u(z)-E_P[u(Z)\mid X=x];q_X,p_X\left(q_{Z\mid X},p_{Z\mid X}\right)=\text{marginal}$ (conditional) densities under sampling from $Q,P;\lambda ,{\lambda }_1,{\lambda }_2=\sigma$-finite measures.

	$𝒰$	$𝒯$	$𝒲$	$\theta (P,u)$	${\dot{\theta }}_{P,u}^{*}(w)$
Conditional mean	$L^{\star }(\rho )$	$L^2(Q)$	$L^2\left(Q_X\right)$	$E_P[u(Z)\mid X=\cdot ]$	$\left(\frac{q_X}{p_X}[u-\theta (P,u)]w,\frac{p_{Z\mid X}}{q_{Z\mid X}}w\right)$
r-fold conditional mean	$L^{\star }\left({\rho }^r\right)$	$L^2\left(Q^r\right)$	$L^2\left(Q_X^r\right)$	$E_{P^r}\left[u\left(Z_{\left[r\right]}\right)\mid X_{\left[r\right]}=\cdot \right]$	see appendix
Conditional variance	$L^{\star }(\rho )$	$L^2(Q)$	$L^2\left(Q_X\right)$	${\text{Var}}_P[u(Z)\mid X=\cdot ]$	$\left(\frac{q_X}{p_X}\left[{\Pi }_P(u{)}^2-\theta (P,u)\right]w,2\frac{p_{Z\mid X}}{q_{Z\mid X}}{\Pi }_P(u)w\right)$
Conditional covariance	$L^{\star }(\rho {)}^2$	$L^2(Q{)}^{\oplus 2}$	$L^2\left(Q_X\right)$	${\text{cov}}_P\left[u_1(Z),u_2(Z)\mid X=\cdot \right]$	see appendix
Kernel embedding	$L^{\star }\left({\rho }_X\right)$	$L^2\left(Q_X\right)$	RKHS	$\int K(x,\cdot )u(x)P_X(dx)$	$\left(wu-\int wud{P}_X,\frac{p_X}{q_X}w\right)$
Optimal value	$𝒰$	$𝒯$	$\mathbb{R}$	${\text{min}}_{y\in 𝒴}{F}_y(P,u)$	${\dot{F}}_{{\text{argmin}}_y{F}_y}^{*}(P,u),P,u$
Optimal solution	$𝒰$	$𝒯$	${\mathbb{R}}^d$	${\text{argmin}}_{w\in 𝒲}{F}_w(P,u)$	see appendix
Pathwise diff. parameter	{0}		$𝒲$	$\nu (P)$	$\left({\dot{\nu }}_P^{*}(w),0\right)$
↳ Root-density	{0}		$L^2(\lambda )$	$(dP/d\lambda {)}^{1/2}$	$\left(\frac{w}{2\nu (P)}-\int \frac{w}{2\nu (P)}dP,0\right)$
↳ Conditional density	{0}		$L^2(Q)$	$p_{Y\mid X}$, with $Z=(X,Y)$	$\left({\Pi }_P\left(\frac{q}{p}w{p}_{Y\mid X}\right),0\right)$
↳ Dose-response function	{0}		$L^2(\lambda )$	$\int E_P[Y\mid A=\cdot ,X=x]P_X(dx)$	see appendix
↳ Counterfactual density	{0}		$L^2(\lambda )$	$y\mapsto E_P[p(y\mid A=1,X)]$	see appendix
Hadamard diff. map	$𝒰$	$𝒯$	$𝒲$	$\zeta (u)$	$\left(0,{\dot{\zeta }}_u^{*}(w)\right)$
↳ Inner product	$ℛ\oplus ℛ$		$\mathbb{R}$	${⟨u_1,u_2⟩}_{ℛ}$, with $u=\left(u_1,u_2\right)$	$\left(0,\left(w{u}_2,w{u}_1\right)\right)$
↳ Squared norm	$𝒯$		$\mathbb{R}$	$\Vert u{\Vert }_{𝒯}^2$	$(\mathrm{0,2}wu)$
↳ Differentiable function	${\mathbb{R}}^d$		$\mathbb{R}$	$\zeta (u)$	$(0,w\nabla \zeta (u))$
↳ Pointwise operations	$L^2(\lambda {)}^{\oplus d}$		$L^2(\lambda )$	$x\mapsto f_x\circ \bar{u}(x)$, with $\bar{u}(x)={\left(u_j(x)\right)}_{j=1}^d$	$\left(0,{\left(x\mapsto w(x){\nabla }_j{f}_x\circ \bar{u}(x)\right)}_{j=1}^d\right)$
↳ Bounded affine map	$𝒯$		$𝒲$	$\kappa (u)+c$, with $c\in 𝒲$, bounded linear $\kappa$	$\left(0,{\kappa }^{*}(w)\right)$
↳ Constant map	{0}		$𝒲$	$c$, with $c\in 𝒲$	(0,0)
↳ Coordinate projection	${ℛ}_1\oplus {ℛ}_2$		${ℛ}_1$	$u_1$, with $u=\left(u_1,u_2\right)$	(0, ($w,0$))
↳ Change of measure	$L^2\left({\lambda }_1\right)$		$L^2\left({\lambda }_2\right)$	$u$	($0,\frac{d{\lambda }_2}{d{\lambda }_1}w$)
↳ Lift to new domain	$L^2\left(Q_X\right)$		$L^2(Q)$	$z\mapsto u(x)$	$\left(0,E_Q[w(Z)\mid X=\cdot ]\right)$
↳ Fix binary argument	$L^2\left(Q_{A,X}\right)$		$L^2\left(Q_X\right)$	$x\mapsto u(1,x)$	$\left(0,(a,x)\mapsto \frac{a}{Q(A=1\mid x)}w(x)\right)$

The first consists of maps $\theta$ that depend on both their distributional input $P$ and Hilbert-valued input $u$. A simple example is a conditional mean operator $\theta (P,u)(\cdot )=E_P[u(z)\mid X=\cdot ]$, with $u$ a real-valued function. An $r$-fold conditional mean is a slightly more complex case, corresponding to a conditional variant of the estimand pursued by U-statistics (Hoeffding, 1948). Here, $u$ is a function of $r\ge 2$ independent observations and $\theta (P,u)(\cdot )=E_{P^r}\left[u\left(Z_1,Z_2,\dots ,Z_r\right)\mid \left(X_1,X_2,\dots ,X_r\right)=\cdot \right]$. Conditional variances and covariances are other examples, as are the optimal value and solution of optimization problems whose objective functions depend on the data-generating distribution $P$ and a Hilbert-valued input $u$. Kernel embedding operators $\theta (P,u)=\int K(x,\cdot )u\left(x\right)P_X\left(dx\right)$ represent another example, where $u$ is a bounded function and $K$ is a bounded kernel of a reproducing kernel Hilbert space (RKHS) (Chapter 4.9.1.1 of Berlinet and Thomas-Agnan, 2011). This primitive has been used to construct tests of equality of distributions (Gretton et al., 2012; Muandet et al., 2021) and independence (Gretton et al., 2007).

The second category consists of maps $\theta$ that solely depend on their distribution-valued input, so that there is some pathwise differentiable parameter $\nu :ℳ\to 𝒲$ such that $\theta (P,u)=\nu \left(P\right)$ for all $(P,u)\in ℳ\times 𝒰$. Since $u$ plays no role in determining the value of $\theta (P,u)$, it suffices to assume that $𝒰$ is a trivial vector space $\left\{0\right\}$. Examples of $\nu$ include root density functions, conditional density functions, dose-response functions (Díaz and van der Laan, 2013), and counterfactual density functions (Kennedy et al., 2021), as detailed in Luedtke and Chung (2024). That work also shows that conditional average treatment effect functions (Hill, 2011) are pathwise differentiable.

The third category consists of maps $\theta$ that solely depend on their Hilbert-valued input, so that there is some Hadamard differentiable map $\zeta :𝒰\to 𝒲$ such that $\theta (P,u)=\zeta \left(u\right)$ for all $(P,u)$.

Simple examples are inner products, squared norms, and finite-dimensional differentiable functions. Pointwise operations provide a richer example — for instance, $\zeta$ may take as input functions $u_1$ and $u_2$ and return $x\mapsto f\left(u_1\left(x\right),u_2\left(x\right)\right)$ with $f(a,b)≔[a+b]\text{cos}\left(a\right)$. The pointwise operation may also depend on $x$, as would be the case if $\zeta \left(u_1,u_2\right)\left(x\right)=f_x\left(u_1\left(x\right),u_2\left(x\right)\right)$ with $f_x(a,b)=(a+x)\text{cos}\left(b\right)$.

Bounded affine maps provide further examples of primitives from the third category. Simple special cases are given by constant maps and coordinate projections. Changes of measure are another special case, so that for measures ${\lambda }_2\ll {\lambda }_1$ with $d{\lambda }_2/d{\lambda }_1$ essentially bounded, $\zeta \left(u\right)$ is the isometric embedding of $u\in L^2\left({\lambda }_1\right)$ into $L^2\left({\lambda }_2\right)$; this map facilitates the chaining of primitives ${\zeta }_1:{𝒰}_0\to L^2\left({\lambda }_1\right)$ and ${\zeta }_2:L^2\left({\lambda }_2\right)\to {𝒲}_0$ via ${\zeta }_2\circ \zeta \circ {\zeta }_1$. Primitives ${\zeta }_1:{𝒰}_0\to L^2\left(Q_X\right)$ and ${\zeta }_2:L^2\left(Q\right)\to {𝒲}_0$ can similarly be chained together using a lifting map $\zeta :L^2\left(Q_X\right)\to L^2\left(Q\right)$. Further examples of affine maps are those that fix a binary argument to unity, so that $\zeta \left(u\right)\left(x\right)=u(1,x)$. When paired with a conditional mean operator $(P,u)\mapsto E_P[u(z)\mid (A,X)=\cdot ]$, this makes it possible to express $(P,u)\mapsto E_P[u(z)\mid A=1,X=\cdot ]$ as a composition of totally pathwise differentiable primitives; such conditional expectations appear frequently in parameters arising in causal inference (Robins, 1986). Bounded affine maps $\zeta$ can also be composed with conditional mean primitives to express affine functionals of regressions commonly studied in the literature (Chernozhukov et al., 2018, 2022c), namely $\psi \left(P\right)=E_P\left[\zeta \left(E_P[Y\mid X=\cdot ]\right)(z)\right]$ with $\zeta :L^2\left(Q_X\right)\to L^2\left(Q\right)$.

The choice of ambient Hilbert spaces $𝒯$ and $𝒲$ can impact a primitive’s total pathwise differentiability. For the first five primitives in Table 1 and the conditional density primitive, this choice involves selecting a $Q\in ℳ$ indexing one or more $L^2$ spaces. As detailed in Appendix C, sufficient conditions for each primitive’s total pathwise differentiability at a given $(P,u)$ require one or more of the following to be bounded: $dP/dQ,dQ/dP,d{P}_X/d{Q}_X$, or $d{Q}_X/d{P}_X$. If $Q=P$ then these functions are all equal to unity, and so these boundedness conditions certainly hold. Given this, when composing these primitives to differentiate $\psi$ at $P$ we generally recommend choosing $Q=P$ to index the ambient Hilbert spaces.

In Appendix C.5, we describe how to transform primitives for conditional quantities to those for marginal quantities. This provides a way to derive a marginal mean primitive, for example. We also propose a way to adapt the primitives from nonparametric to semiparametric models, though we caution that this proposal requires projecting onto the tangent space of the model and so implementing it can require specialized expertise.

2.3. Example: nonparametric $R^2$

We apply Algorithm 2 to compute the EIF of the nonparametric $R^2$ parameter at a distribution $P$ in a nonparametric model $ℳ$. To avoid notational overload on $P$, we denote a generic element of $ℳ$ by $P^{\prime }$ when presenting this example.

We express $\psi \left(P^{\prime }\right)≔1-E_{P^{\prime }}\left[{\left\{Y-{\mu }_{P^{\prime }}\left(X\right)\right\}}^2\right]/{\text{Var}}_{P^{\prime }}\left(Y\right)$ as a composition of eight primitives: a constant ${\theta }_1\left(P^{\prime },0\right)=(z\mapsto y)\in L^2\left(P\right)$, variance ${\theta }_2\left(P^{\prime },h_1\right)={\text{Var}}_{P^{\prime }}\left[h_1(z)\right]\in \mathbb{R}$, constant ${\theta }_3\left(P^{\prime },0\right)=(z\mapsto y)\in L^2\left(P\right)$, conditional mean ${\theta }_4\left(P^{\prime },h_3\right)(\cdot )=E_{P^{\prime }}\left[h_3(z)\mid X=\cdot \right]\in L^2\left(P_X\right)$, lifting ${\theta }_5\left(P^{\prime },h_4\right)=\left(z\mapsto h_4\left(x\right)\right)\in L^2\left(P\right)$, pointwise operation ${\theta }_6\left(P^{\prime },h_5\right)=\left(z\mapsto {\left[y-h_5(z)\right]}^2\right)\in L^2\left(P\right)$, mean ${\theta }_7\left(P^{\prime },h_6\right)=E_{P^{\prime }}\left[h_6(z)\right]\in \mathbb{R}$, and differentiable function ${\theta }_8\left(P^{\prime },\left(h_2,h_7\right)\right)=1-h_7/h_2\in \mathbb{R}$. Here, pa(1) = pa(3) = ∅, pa(2) = {1}, pa(4) = {3}, pa(5) = {4}, pa(6) = {5}, pa(7) = {6}, and pa(8) = {2, 7}. The redundancy of the constant maps ${\theta }_1$ and ${\theta }_3$ has no impact on the applicability of Algorithm 2. For the composition of primitives to express $\psi$, it must return $\psi \left(P^{\prime }\right)$ for every $P^{\prime }\in ℳ$. This will hold if each claimed element of $L^2\left(P\right)$ or $L^2\left(P_X\right)$ is a $P$-a.s. equivalence class of functions that satisfies an appropriate second-moment condition. To this end, we suppose all pairs of distributions in $ℳ$ are mutually absolutely continuous, which makes it so, for each $P^{\prime }\in ℳ,P^{\prime }$-a.s. and $P$-a.s. equivalence classes of functions are one and the same. For the moment condition, we suppose there exists $c<\infty$ such that $P^{\prime }\left(\right|Y|\le c)=1$ for all $P^{\prime }\in ℳ$. To avoid dividing by zero, we also suppose ${\text{Var}}_{P^{\prime }}\left(Y\right)>0$ for all $P^{\prime }$.

Figure 1 illustrates the computation graph for our representation of $\psi$ (Linnainmaa, 1970; Bauer, 1974). In it, each primitive is a node. A directed arrow points from node $i$ to node $j$ if $i\in \text{pa}\left(j\right)$, and from $P$ to node $j$ if ${\theta }_j$ depends nontrivially on its distribution-valued argument. Table 2 provides a step-by-step illustration of how Algorithm 2 computes the EIF at $P$. As anticipated in Section 1.2, this manual application of automatic differentiation has yielded an analytical expression for the EIF.

Figure 1:

Computation graph for the nonparametric $R^2$ parameter.

Table 2:

Step-by-step evaluation of Algorithm 2 for the nonparametric $R^2$ parameter when it is expressed as in Figure 1. Each row $j$ lists the values $f_0,f_1,\dots ,f_j$ take just before executing line 3 of the for loop in Algorithm 2. The value of $f_0$ in the bottom row is the efficient influence function.

Notation: ${\mu }_P(x)=E_P[Y\mid X=x]$, ${\sigma }_P^2={\text{Var}}_P(Y)$, and ${\gamma }_P=E_P\left[{\text{Var}}_P(Y\mid X)\right]$.

Step ($j$)	Variables in Algorithm 2
	${\mathit{f}}_8$	${\mathit{f}}_7$	${\mathit{f}}_6$	${\mathit{f}}_5$	${\mathit{f}}_4$	${\mathit{f}}_3$ §	${\mathit{f}}_2$	${\mathit{f}}_1$ §	${\mathit{f}}_0$
8	1	0	0	0	0	−	0	−	0
7		$-\frac{1}{{\sigma }_P^2}$	0	0	0	−	$\frac{{\gamma }_P}{{\sigma }_P^4}$	−	0
6			$-\frac{1}{{\sigma }_P^2}$	0	0	−	$\frac{{\gamma }_P}{{\sigma }_P^4}$	−	$z\mapsto -\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
5				$z\mapsto \frac{2\left[y-{\mu }_P(x)\right]}{{\sigma }_P^2}$	0	−	$\frac{{\gamma }_P}{{\sigma }_P^4}$	−	$z\mapsto -\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
4					0†	−	$\frac{{\gamma }_P}{{\sigma }_P^4}$	−	$z\mapsto -\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
3						−	$\frac{{\gamma }_P}{{\sigma }_P^{+}}$	−	$z\mapsto -\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
2							$\frac{{\gamma }_P}{{\sigma }_P^4}$	−	$z\mapsto -\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
1								−	$z\mapsto \frac{\left[{\left\{y-E_P\left[Y\right]\right]}^2-{\sigma }_P^2\right]{\gamma }_P}{{\sigma }_P^2}-\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$
									$z\mapsto \frac{\left[{\left\{y-E_P\left[Y\right]\right]}^2-{\sigma }_P^2\right]{\gamma }_P}{{\sigma }_P^4}-\frac{{\left[y-{\mu }_P(x)\right]}^2-{\gamma }_P}{{\sigma }_P^2}$

^§${\theta }_3$ and ${\theta }_1$ are constant maps and so have derivative zero. Hence, backpropagating $f_3$ and $f_1$ on steps $j=3$ and $j=1$ has no impact on earlier nodes $(f_i:i<j)$, and so these values need not be stored.

^†${\theta }_5:L^2\left(P_X\right)\to L^2(P)$ is a lifting, so $f_4$ updates as $f_4+=E_P\left[f_5(Z)\mid X=\cdot \right]$. This update equals zero.

In Appendix D, we apply Algorithm 2 to two other parameters: the expected density $\psi \left(P\right)=E_P[p(z)]$ and expected conditional covariance $\psi \left(P\right)=E_P\left[{\text{cov}}_P(A,Y\mid X)\right]$.

Algorithm 3.

Obtain initial estimates of $\psi (P)$ and ${\dot{\psi }}_P^{*}\left(f_k\right)$ using data drawn from $P$

3. Probabilistic programming of efficient estimators

3.1. Estimator construction

Suppose now that we observe independent draws $Z_1,Z_2,\dots ,Z_n$ from an unknown $P\in ℳ$. Our goal is to infer $\psi \left(P\right)$. We do this using a variant of a one-step estimator (Pfanzagl, 1982), which debiases an initial estimator using the efficient influence operator.

Algorithm 3 presents a method to obtain this initial estimate and the evaluation of the efficient influence operator at a user-specified $f_k$. The algorithm employs a forward and backward pass. In the forward pass, it recursively estimates the $P$-dependent evaluations of the operators ${\theta }_j$ that appear in Algorithm 1, ultimately yielding the initial estimate ${\widehat{h}}_k$ of $\psi \left(P\right)$. In the backward pass, it recursively estimates $P$-dependent adjoint evaluations and backpropagates them as in Algorithm 2, yielding an estimate ${\widehat{f}}_0$ of the efficient influence operator evaluation ${\dot{\psi }}_P^{*}\left(f_k\right)$. Black-box machine learning tools may be used to estimate all of the needed nuisances. When $\psi$ is real-valued and $f_k=1,{\widehat{f}}_0$ is an estimate of the EIF. To fix ideas, we focus on this real-valued setting in the remainder of the main text and refer the reader to Appendix B.3 for a discussion of more general settings where $\psi$ takes values in a generic separable Hilbert space, such as an $L^2$ space—these more general settings cover, for example, the case where $\psi$ returns a square integrable dose-response function..

Inspired by one-step estimation, the outputs of Algorithm 3 can be used to construct

$$\widehat{\psi }≔{\widehat{h}}_k+\frac{1}{n}\sum _{i=1}^n{\widehat{f}}_0\left(Z_i\right)$$

and 95% confidence interval $\widehat{\psi }\pm 1.96\widehat{\sigma }/n^{1/2}$, where ${\widehat{\sigma }}^2$ is the empirical variance of $\widehat{f_0}$. In fact, $\widehat{\psi }$ is exactly the traditional one-step estimator $\psi \left(\widehat{P}\right)+\frac{1}{n}\sum _{i=1}^n{\dot{\psi }}_{\widehat{P}}^{*}\left(1\right)\left(Z_i\right)$ under the strong compatibility condition that there exists a $\widehat{P}\in ℳ$ such that ${\widehat{h}}_j={\theta }_j\left(\widehat{P},{\widehat{h}}_{\text{pa}\left(j\right)}\right)$ and ${\widehat{\vartheta }}_j={\dot{\theta }}_{j,\widehat{P},{\widehat{h}}_{\text{pa}\left(j\right)}}^{*}\left({\widehat{f}}_j\right)$ for all $j\in \left[k\right]$. To ease implementation, Algorithm 3 makes a weaker compatibility requirement that allows each ${\widehat{\vartheta }}_{j,0},{\widehat{h}}_j$, and ${\widehat{h}}_{\text{pa}\left(j\right)}$ to be compatible with a ${\widehat{P}}_j$ that may depend on $j$. This dependence makes it so $\widehat{\psi }$ may not be a traditional one-step estimator; the benefit of permitting this is that it allows each primitive to be paired with a nuisance estimation routine that can be implemented once and reused across different parameters $\psi$—see the next subsection for details.

Because $\widehat{\psi }$ is not generally a traditional one-step estimator, we state sufficient conditions for its efficiency here. These conditions concern the same objects as for a traditional one-step estimator (Bickel et al., 1993), namely a remainder ${ℛ}_n≔{\widehat{h}}_k+\int {\widehat{f}}_0(z)P\left(dz\right)-\psi \left(P\right)$ and drift term ${𝒟}_n≔\int \left[{\widehat{f}}_0(z)-f_0(z)\right]\left(P_n-P\right)\left(dz\right)$, with $P_n$ the empirical distribution and $f_0≔{\dot{\psi }}_P^{*}\left(1\right)$ the EIF.

Theorem 2 (Efficiency of $\widehat{\psi }$). Suppose $\psi$ is real-valued, $f_k=1$, and the conditions of Theorem 1 hold. If ${𝒟}_n$ and ${ℛ}_n$ are $o_p\left(n^{-1/2}\right)$, then $\widehat{\psi }$ is an efficient estimator of $\psi \left(P\right)$, in that

$$\widehat{\psi }-\psi \left(P\right)=\frac{1}{n}\sum _{i=1}^n{f}_0\left(Z_i\right)+o_p\left(n^{-1/2}\right).$$

Our Python package implements a cross-fitted counterpart of $\widehat{\psi }$, which we present and study in Appendix B.1. Like $\widehat{\psi }$, this estimator is efficient when a remainder and drift term are negligible (see Theorem S1). The cross-fitted estimator may be preferred over $\widehat{\psi }$ because its drift term is $o_p\left(n^{-1/2}\right)$ even without an empirical process condition (Schick, 1986). This is in contrast to $\widehat{\psi }$, for which such a condition is often used to ensure negligibility of that term (Lemma 19.24 of van der Vaart, 2000).

The compatibility condition in Algorithm 3 helps make ${ℛ}_n=o_p\left(n^{-1/2}\right)$ plausible. To see this, we add and subtract terms and use that $\int {\widehat{\vartheta }}_{j,0}d{\widehat{P}}_j=0$ and ${\widehat{f}}_0=\sum _{j=1}^k{\widehat{\vartheta }}_{j,0}$ to reexpress ${ℛ}_n$ as

$${ℛ}_n=\left[{\widehat{h}}_k-\psi (P)-\sum _{j=1}^k\int {\vartheta }_{j,0}(z)\left({\widehat{P}}_j-P\right)(dz)\right]+\left[\sum _{j=1}^k\int \left[{\vartheta }_{j,0}(z)-{\widehat{\vartheta }}_{j,0}(z)\right]\left({\widehat{P}}_j-P\right)(dz)\right],$$

where ${\vartheta }_{j,0}$ is the first entry of the augmentation term ${\dot{\theta }}_{j,P,h_{\text{pa}\left(j\right)}}^{*}\left(f_j\right)$ from Algorithm 2. If both terms on the right are $o_p\left(n^{-1/2}\right)$, then ${ℛ}_n$ is as well. Moreover, since the above holds for any ${\widehat{P}}_j$ compatible with the nuisances ${\widehat{h}}_j,{\widehat{h}}_{\text{pa}\left(j\right)}$, and ${\widehat{\vartheta }}_{j,0}$, it suffices that there exist such ${\widehat{P}}_j,j\in \left[k\right]$, that make each of the two terms on the right $o_p\left(n^{-1/2}\right)$. By Cauchy-Schwarz, the magnitude of the latter term upper bounds by $\sum _{j=1}^k{‖d{\widehat{P}}_j/dP-1‖}_{L^2\left(P\right)}{‖{\widehat{\vartheta }}_{j,0}-{\vartheta }_{j,0}‖}_{L^2\left(P\right)}$, which is $o_p\left(n^{-1/2}\right)$ under an $n^{-1/4}$-rate condition akin to those arising in the study of traditional one-step estimators. In Appendix B.2, we show that the first term on the right is the remainder in a von Mises expansion, which makes it so it too will be $o_p\left(n^{-1/2}\right)$ under appropriate conditions. Similar arguments have been made previously to motivate one-step estimators when $\psi$ is a generic pathwise differentiable parameter (Robins et al., 2009; Luedtke and Chung, 2024).

For a particular parameter $\psi$, the remainder ${ℛ}_n$ can also be analyzed explicitly. In Appendix E, we provide this analysis for the nonparametric $R^2$ parameter, and we also provide an analogue of Table 2 that helps to visualize Algorithm 3.

3.2. Software implementation

To implement Algorithm 3 in software, it is necessary to convert a computer code representation of the parameter $\psi$ — such as the one given on lines 2–4 of Example 2 — into the compositional form in Algorithm 1. Any computer program, including those with function calls and control flow structures like loops and conditionals, can be represented via its computation graph (Linnainmaa, 1970; Bauer, 1974). When the operations in this graph are totally pathwise differentiable primitives, as they are in Figure 1, this graph yields an expression for $\psi$ as in Algorithm 1. Representing programs via their computation graphs enables user-friendly automatic differentiation, as seen in software like TensorFlow, PyTorch, and JAX (Abadi et al., 2015; Bradbury et al., 2018; Paszke et al., 2019).

Our proof-of-concept package also adopts computation graphs, using a static graph structure (Baydin et al., 2018). When defining the parameter, each primitive call takes dimple objects as inputs and returns another dimple object. At this stage, a dimple object comprises a primitive and a set of pointers to its inputs, serving the roles of ${\theta }_j$ and $\text{pa}\left(j\right)$ in Algorithm 1. The dimple objects can be used to express $\psi$ as in Algorithm 1, and also provide the form of the computation graph. Estimation — as executed, for instance, via ‘estimate(R2)’ in Example 2 — calls Algorithm 3 with this computation graph to populate each dimple object with initial estimates ${\widehat{h}}_j$ in the forward pass, with adjoint estimates ${\widehat{f}}_j$ added in the backward pass.

Primitives can be implemented independently of one another and reused across different parameters $\psi$. In this modular approach, each primitive ${\theta }_j$ in Algorithm 3 requires two associated nuisance estimation routines: one for the forward pass and another for the backward pass. Algorithm 4 details the inputs and outputs of these routines for a generic primitive $\theta :ℳ\times 𝒰\to 𝒲$. Like Algorithm 3, the backward routine makes a compatibility requirement: the estimates of $\theta (P,u)$ and the first entry of ${\dot{\theta }}_{P,u}^{*}\left(w\right)$ must be compatible with a common $\widehat{P}\in ℳ$. There is no need to explicitly exhibit $\widehat{P}$, which in many cases will make it possible to only estimate certain summaries of $P$, such as regression functions or conditional densities. The forward and backward routines may also take hyperparameters as inputs, such as specifications of statistical learning algorithms to use to fit nuisances or folds to be used during cross-fitting.

Algorithm 3 can be executed using these nuisance estimation routines. At step $j$ of the forward pass, the forward routine for ${\theta }_j$ is called with inputs $Z_{𝒥}$ and ${\widehat{h}}_{\text{pa}\left(j\right)}$, yielding ${\widehat{h}}_j$. At step $j$ of the backward pass, the backward routine is called with inputs $Z_{𝒥},{\widehat{h}}_{\text{pa}\left(j\right)},{\widehat{h}}_j$, and ${\widehat{f}}_j$, yielding ${\widehat{\vartheta }}_j$.

Forward and backward routines for the primitives in Table 1 are given in Appendix F. Those routines allow needed regression, density functions, and density ratios to be estimated using flexible statistical learning tools, such as gradient-boosted trees (Ke et al., 2017). Riesz losses (Chernozhukov et al., 2022c) are used when estimating the adjoint of bounded affine maps whose ambient Hilbert spaces depend on $P$.

Algorithm 4.

Requirements of nuisance estimation routines for a generic $\theta :ℳ\times 𝒰\to 𝒲$

4. Simulation study

A simulation study was conducted to verify whether a handful of lines of pydimple code truly suffice to construct performant one-step estimators for the examples in Section 1.3. Such a finding would be important given that existing implementations of one-step estimators in these problems can consist of hundreds of lines of code (e.g., Luedtke et al., 2017; Williamson et al., 2021). For a given Monte Carlo replicate of our study, a dataset dat was generated as an iid sample of $n\in \left\{\mathrm{250,1000,4000,16000}\right\}$ draws from a distribution $P$, and then the pydimple code from Section 1.3 was executed to perform estimation. A total of 1000 Monte Carlo replicates were performed for each scenario. Code for this study is available at http://github.com/alexluedtke12/pydimple.

For the expected density parameter, we followed Carone et al. (2018) in letting $P=\text{Beta}\left(\mathrm{3,5}\right)$. For the nonparametric $R^2$ parameter, we followed Williamson et al. (2021) in letting $P$ be the distribution of $(X,Y)$ with $X=\left(X_1,X_2\right)$ a pair of independent Unif[−1, 1] random variables and $Y\mid X~N\left(25X_1^2/\mathrm{9,1}\right)$. For the longitudinal G-formula example, we used the data-generating process from Simulation 1 of Luedtke et al. (2017).

Table 3 reports the results. The coverage of 95% confidence intervals returned by pydimple approximately attains the nominal value for all scenarios with $n\ge 1000$. Coverage tends to be below nominal when $n=250$, taking a minimal value of 79%. The widths of the confidence intervals are close to the asymptotically optimal width of $2\times 1.96\times \sigma /n^{1/2}$, with $\sigma$ the standard deviation of the EIF. For each expected density and nonparametric $R^2$ scenario, the variance of the estimators always approximately equals the efficient variance, ${\sigma }^2/n$. In the longitudinal G-formula scenario, this variance is much larger than the efficient variance at smaller sample sizes but the relative difference diminishes as $n$ grows. The inflated variance appears to be partially driven by outliers arising from inverse propensity weighting in the EIF. Bias is minimal in all scenarios. Together these results support our theoretical findings that dimple can produce asymptotically efficient estimators under reasonable conditions (Theorem 2).

Table 3:

Performance of pydimple package in the examples from Section 1.3. Assessed via coverage of 95% confidence intervals, relative $\text{width}=n^{1/2}\times (\text{mean of CI})/(2\times 1.96\times \text{standard deviation of EIF})$, $\text{relative variance}=n\times (\text{variance of estimator})/(\text{variance of EIF})$, and squared bias divided by the mean squared error. Asymptotically, an efficient estimator would have these four quantities equal to the values in the Target Value row.

	$\mathit{n}$	Coverage	Rel. Width	Rel. Variance	Bias2/MSE
Target Value		≥ 95%	≤ 1.00	≤ 1.00	= 0.00
Expected Density	250	91%	0.94	1.13	0.04
	1000	92%	0.96	0.95	0.07
	4000	93%	0.97	1.06	0.02
	16000	95%	0.98	0.98	0.03
Nonparametric $R^2$	250	87%	0.89	1.12	0.10
	1000	92%	0.94	1.05	0.05
	4000	94%	0.97	1.03	0.00
	16000	93%	0.98	1.10	0.01
Longitudinal G-Formula	250	79%	1.09	4.44	0.12
	1000	93%	1.28	2.01	0.02
	4000	94%	1.25	2.10	0.01
	16000	94%	1.06	1.26	0.00

On the author’s 2021 MacBook Pro, the average time required to evaluate a single Monte Carlo replicate for Examples 1, 2, and 3 was about 3, 20, and 110 seconds respectively when $n=1000$.

5. Discussion

Our results suggest researchers should refocus their efforts when constructing an efficient estimator of a new parameter. Traditionally, this process involves labor-intensive calculations to derive the EIF. Our approach avoids this task for parameters expressible as compositions of known primitives. Yet, we recognize that this method will not apply to all new parameters, particularly those that cannot be represented in this form. In such cases, we propose a shift in focus: rather than immediately studying the differentiability of the parameter itself, researchers should aim to express it as a composition involving both known and novel primitives. By establishing the differentiability of each new primitive, they can deduce the differentiability of the parameter itself. This modular approach can both simplify the immediate analysis and the study of other parameters in future research. An example of a primitive that would be useful for ongoing research—but is missing from Table 1—is one that returns a solution to an infinite-dimensional inverse problem (Bennett et al., 2023; Cui et al., 2024; Smucler et al., 2025).

As noted in Carone et al. (2018), any algorithmic approach to statistical inference inevitably has potential for misuse. Per Theorems 1 and 2, the validity of our method depends on differentiability and nuisance estimation conditions. While an expert can interpret and verify these conditions, non-specialists may employ a software implementation without doing so. This may lead to an overconfidence in the inferential conclusions reached, especially in high-dimensional problems or other settings where these conditions are unlikely to be met. To mitigate this risk, future software could provide a list of conditions sufficient for ensuring the differentiability of each primitive and the rates of convergence required on nuisance estimators for efficiency. Though Appendix C already gives sufficient conditions for the differentiability of the primitives in Table 1, work is still needed to figure out how to have software return this information in an accessible manner.

The composition of primitives used to express the parameter $\psi$ will generally be nonunique. The same holds in traditional automatic differentiation settings, where, for example, the function $f\left(x\right)=2x^2+1$ can be expressed as $g\circ h$ with $g\left(x\right)=2x+1$ and $h\left(x\right)=x^2,g\left(x\right)=\text{cosh}\left(x\right)$ and $h\left(x\right)=2{\text{sinh}}^{-1}\left(x\right)$, or $g\left(x\right)=2x^{1/2}+1$ and $h\left(x\right)=x^4$. In traditional settings, choosing different compositions can impact the conditions under which the chain rule can be applied—for example, $g\left(x\right)=2x^{1/2}+1$ is not differentiable at 0, even though $f\left(x\right)=2x^2+1$ is. As indicated by the differentiability conditions in Theorem 1, the applicability of the chain rule also depends on the choice of representation for $\psi$ in our statistical setting. Beyond that, this choice can also impact the nuisances that need to be estimated. For example, suppose $\psi \left(P\right)=E_P[Z{]}^2$ is expressed as $⟨z\mapsto z,dP/d\lambda {⟩}_{L^2\left(\lambda \right)}^2$ using inner product and density primitives. Algorithm 3 would estimate the nuisance $dP/d\lambda$ for this representation and the efficiency of the one-step estimator would rely on an $n^{-1/4}$ rate condition. Estimation of this nuisance, and the associated rate condition, would be avoided if $\psi$ were instead expressed using a marginal mean primitive and square function. Algorithmically finding alternative representations of a parameter that weaken nuisance estimation requirements is an important area for future work.

There are also other interesting avenues for future work. First, though our algorithms apply to arbitrary models, the particular nuisance estimation routines given in Appendix F are designed for nonparametric settings. Future work should seek to identify semiparametric model restrictions that allow for modular nuisance estimation. Second, while we focused on algorithmically constructing one-step estimators, it would be worth considering whether this can be done for other estimation strategies, such as targeted minimum loss-based estimation (van der Laan and Rubin, 2006; van der Laan and Gruber, 2016). Third, it would be interesting to investigate whether automatic differentiation can be used to obtain higher-order influence functions (Robins et al., 2008).

Dimple provides a means to rapidly translate new developments in debiased estimation from theory to practice. While our proof-of-concept software uses a cross-fitted one-step estimator, future implementations should incorporate recent proposals like iteratively debiasing nuisance function estimators for added robustness (Rotnitzky et al., 2017; Luedtke et al., 2017), employing separate data subsamples to estimate different nuisance functions (Newey and Robins, 2018), and enforcing known shape constraints on the parameter of interest (Westling and Carone, 2020; Ham et al., 2024).

This mirrors the approach taken by other differentiable and probabilistic programming frameworks, where new algorithms, architectures, and regularization techniques are routinely integrated into existing software (Bradbury et al., 2018; Paszke et al., 2019; Stan Development Team, 2023). This practice ensures that state-of-the-art methods are widely accessible, a model we strive to emulate.