Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
49,699 characters · 29 sections · 44 citation commands
ScoreMatchingRiesz: Score Matching for Debiased Machine Learning and Policy Path Estimation
Estimating the Riesz representer is a central problem in debiased machine learning for causal inference and structural parameter estimation Chernozhukov2022automaticdebiased. For example, in average treatment effect (ATE) estimation, the Riesz representer corresponds to inverse propensity score weights, and accurate estimation of this representer is essential for accurate ATE estimation. Score matching refers to methods for estimating the score, the derivative of a log density with respect to a variable. Score matching has recently attracted renewed attention through the diffusion model literature Song2020generativemodeling.
Riesz representer estimation and score matching are closely related, but leveraging this link for debiased machine learning requires more than simply importing existing score estimators. For example, in average marginal effect (AME) estimation, the Riesz representer coincides with the (negative) score of the joint density of $(D,Z)$ with respect to the treatment coordinate $D$ given covariates $Z$. Moreover, Hyv\"arinen score matching Hyvarinen2005estimationof yields an objective that is closely related to Riesz regression Chernozhukov2021automaticdebiased. However, Hyv\"arinen-type objectives can be numerically challenging in high dimensions because they involve divergence computations. In applications where the representer is written in terms of density ratios Kato2025directbias, one may instead turn to density ratio estimation (DRE), but flexible DRE is known to suffer from severe overfitting Kato2021nonnegativebregman,Rhodes2020telescopingdensityratio. While denoising score matching (DSM) Vincent2011aconnection and DRE-$\infty$ Choi2022densityratio have been developed to stabilize score and ratio estimation, their use as Riesz representer estimators within Neyman orthogonal scores calls for an explicit construction and an analysis.
Motivated by these relationships, this paper develops score matching-based methods for Riesz representer estimation, which we call ScoreMatchingRiesz. The key idea is to use score models as a Riesz representer estimation module: we construct the Riesz representer either directly from a score field or via density ratios recovered by integrating a score along a bridge of intermediate distributions. Our framework involves two types of scores. The data score is the derivative of the log density with respect to the data variable, and it is estimated by diffusion-style score matching. The time score is the derivative of the log density along a continuum of intermediate distributions, and it is estimated via infinitesimal classification-based score matching (the continuum limit underlying DRE-$\infty$). We refer to ScoreMatchingRiesz designed for data-score estimation as Data-ScoreMatchingRiesz and that designed for time-score estimation as Time-ScoreMatchingRiesz.
In AME estimation, the Riesz representer can be estimated directly from the data score. In other applications such as APE and ATE estimation, however, estimating the Riesz representer requires estimating density ratios, and time-score estimation provides a general way to recover these ratios through a bridge between endpoint distributions. In some restricted settings, the data-score construction is sufficient on its own. For example, in APE estimation under a known pushforward policy, the relevant density ratio can be expressed as an integral of the data score along a known path, which yields the representer without explicitly learning an endpoint-to-endpoint ratio. Joint training of data and time scores can further stabilize the estimation process across these regimes. See Table (ref).
This score-based representer construction also provides a concrete link between AME and APE through a continuum of policy effects. In this study, we define the policy path as a parameter of interest and propose an estimator. In shift-policy settings, the APE representer can be constructed from the AME representer by integrating the score along the policy path. Using this relationship, we estimate intermediate policy effects as the policy shifts continuously.
In summary, our contributions are as follows:
There are additional contributions, mainly in the appendices, due to space constraints. For example, we discuss when AME and APE are seamlessly connected via scores. In Appendix (ref), we provide an application of our method to the local projection in economic analysis Jorda2005estimationand.
\paragraph{Related work.} Constructing (asymptotically) efficient estimators is a central task in causal inference VanderVaart1998asymptoticstatistics, and various unified perspectives on this task have been established. Among them, the debiased machine learning framework formulates the problem through Neyman orthogonal score equations together with cross-fitting Chernozhukov2018doubledebiased. Various methods for Riesz representer estimation have been proposed, including Riesz regression Chernozhukov2021automaticdebiased, nearest neighbor matching Kato2025nearestneighbor,Lin2023estimationbased, and covariate balancing Zhao2019covariatebalancing,BrunsSmith2025augmentedbalancing.
The goal of DRE is to estimate the ratio between two probability density functions. Many methods have been proposed for DRE Sugiyama2012densityratio, but flexible DRE is known to suffer from severe overfitting, called train loss hacking Kato2021nonnegativebregman or the density-chasm Rhodes2020telescopingdensityratio. One approach to mitigating this overfitting is to avoid estimating a single log density ratio in one step and instead introduce a sequence of intermediate density ratios Rhodes2020telescopingdensityratio. Building on this line of work, Choi2022densityratio proposes an infinitesimal classification approach, which takes a continuum limit of telescoping ratios and learns the resulting score representation via score matching. We use this idea through Time-ScoreMatchingRiesz as a representer-estimation primitive when the Riesz representer is ratio-based.
As we and prior work note Kato2025directbias,Kato2026rieszrepresenter,Benmichael2021balancingact, Riesz representer estimation admits complementary views via score matching and DRE. As a consequence, the instability issues known for Hyv\"arinen score matching and the overfitting issues documented for flexible DRE can also arise in Riesz representer estimation. We make these connections operational by constructing Riesz representers from data scores and time scores, and by providing convergence guarantees in the form needed for debiased machine learning.
For detailed related work, see Appendix (ref).
Let $W=(X,Y)$ consist of a regressor $X\in{\mathcal{X}}\subseteq \mathbb{R}^d$ and a scalar outcome $Y\in{\mathcal{Y}}\subseteq \mathbb{R}$. Assume $(X,Y)\sim P_0$, and define the regression function $\gamma_0(x)\coloneqq {\mathbb{E}}[Y\mid X=x]$. We observe $\{(X_i,Y_i)\}_{i=1}^n$, where $(X_i,Y_i)$ are i.i.d. copies of $(X,Y)$. Our target parameter is \[ \theta_0 \coloneqq {\mathbb{E}}\big[m(W,\gamma_0)\big], \] where $m(W,\gamma)$ is a functional of the data and a regression function $\gamma\colon{\mathcal{X}}\to{\mathcal{Y}}$. We assume that the map $\gamma \mapsto {\mathbb{E}}\big[m(W,\gamma)\big]$ is linear and continuous on $L_2(P_X)$ (the $L_2$ space under the marginal law of $X$).
By the Riesz representation theorem, there exists a function $\alpha_0\in L_2(P_X)$ such that, for all $\gamma\in L_2(P_X)$, ${\mathbb{E}}[m(W,\gamma)] = {\mathbb{E}}[\alpha_0(X)\gamma(X)]$ holds. We call $\alpha_0$ the Riesz representer associated with $m$. Let $\eta_0\coloneqq(\gamma_0,\alpha_0)$ denote the nuisance parameters. For a candidate pair $\eta=(\gamma,\alpha)$ and scalar $\theta$, define the Neyman orthogonal score $\psi(W;\eta,\theta) \coloneqq m(W,\gamma) + \alpha(X)\bigl(Y-\gamma(X)\bigr) - \theta$. Then ${\mathbb{E}}[\psi(W;\eta_0,\theta_0)]=0$ and the moment condition is orthogonal to first-order perturbations of $\eta$ at $\eta_0$ under regularity conditions Chernozhukov2018doubledebiased.
An estimator $\widehat{\theta}$ is obtained by plugging in nuisance estimates $\widehat{\eta}=(\widehat{\gamma},\widehat{\alpha})$ and solving \[\frac{1}{n}\sum_{i=1}^n \psi\left(W_i;\widehat{\eta},\widehat{\theta}\right) = 0;\] equivalently, the estimator is given by \[ \widehat{\theta} = \frac{1}{n}\sum_{i=1}^n \Big(m(W_i,\widehat{\gamma})+\widehat{\alpha}(X_i)(Y_i-\widehat{\gamma}(X_i))\Big). \] Assume that $\widehat{\eta}$ satisfies the Donsker condition or is constructed via cross-fitting. Then, under suitable convergence rate conditions on $\widehat{\eta}$, we have $\sqrt{n}\left(\widehat{\theta}-\theta_0\right) \xrightarrow{d} {\mathcal{N}}(0,V)$, where $V > 0$ is the efficient asymptotic variance.
\paragraph{Notation and assumptions.} For simplicity, we assume that the components of $X$ are continuous, except in specific cases such as the treatment $D$ in ATE estimation. We also assume that each continuous random variable $R$ admits a probability density function $p_0(r)$ under $P_0$. We write $R\sim p_0$. Expectations are taken under $P_0$ unless stated otherwise; when needed, we write ${\mathbb{E}}_{R\sim p}[\cdot]$ for $R$ with density $p$. For a pdf $p_0(x)$, we define the data score as $s^{\text{data}}_0(x) \coloneqq \nabla_x\log p_0(x)$. For a one-parameter family of pdfs $\{p_t(x)\}_{t\in{\mathcal{T}}}$, we define the time score as $s^{\text{time}}_t(x) \coloneqq \partial_t\log p_t(x)$. We use $t$, $\delta$, and $\sigma$ for indexing the bridge, policy-shift, and the Gaussian noise level, but they become mathematically equivalent under certain assumptions.
We illustrate how the Riesz representer appears in applications by presenting examples of AME and APE estimation. Note that our method can also be applied to more general cases such as ATE estimation (see Appendix (ref)).
\paragraph{AME estimation.} Let $X=(D,Z)$ with continuous $D\in[-1,1]$. The AME is \[ \theta^{\mathrm{AME}}_0 \coloneqq {\mathbb{E}}\bigl[\partial_d \gamma_0(D,Z)\bigr], \quad m^{\mathrm{AME}}(W,\gamma)\coloneqq \partial_d \gamma(D,Z). \] Under standard integration-by-parts conditions (including suitable boundary behavior), the Riesz representer is the negative score of the joint density of $(D,Z)$ with respect to $d$ and given as \[ \alpha^{\mathrm{AME}}_0(D,Z)\coloneqq -\partial_d \log p_0(d,Z)\big|_{d=D}. \]
\paragraph{APE estimation.} Let $P_1$ and $P_{-1}$ be two counterfactual distributions of $X$ with densities $p_1(x)$ and $p_{-1}(x)$, both absolutely continuous with respect to $P_0$ (density $p_0$). The APE is \[ \theta^{\mathrm{APE}}_0 \coloneqq {\mathbb{E}}_{X\sim p_1}[\gamma_0(X)]-{\mathbb{E}}_{X\sim p_{-1}}[\gamma_0(X)]. \] This can be written as $\theta^{\mathrm{APE}}_0={\mathbb{E}}[m^{\mathrm{APE}}(W,\gamma_0)]$ with $m^{\mathrm{APE}}(W,\gamma)\coloneqq \gamma(X) \frac{p_1(X)-p_{-1}(X)}{p_0(X)}$. Hence, the Riesz representer is given as \[ \alpha^{\mathrm{APE}}_0(X)\coloneqq \frac{p_1(X)-p_{-1}(X)}{p_0(X)}. \]
This section reviews Riesz regression Chernozhukov2021automaticdebiased, score matching Song2021denoisingdiffusion, and DRE-$\infty$ Choi2022densityratio. Score matching can be carried out via Hyv\"arinen score matching or denoising score matching (DSM). In AME estimation, Riesz regression coincides with Hyv\"arinen score matching and can also be interpreted as a special case of LSIF. Accordingly, we first introduce Riesz regression (Hyv\"arinen score matching) and then DSM and DRE-$\infty$.
Throughout this subsection, let $X\in\mathbb{R}^d$ be a continuously distributed random vector with density $p_0(x)$ under $P_0$. Recall that we defined the data score of $p_0$ as \[ \text{(Data score)}\ \ s^{\text{data}}_0(x)\coloneqq \nabla_x\log p_0(x)\in\mathbb{R}^d. \] The Riesz representer in AME estimation directly corresponds to the data score as $\alpha^{\mathrm{AME}}_0(D,Z)=-\partial_d\log p_0(D,Z)$.
We introduce a parametric family $\{p_t(x)\}_{t\in{\mathcal{T}}}$ such that $p_t(x)=p_0(x)$ at $t=0$. This family is specified by the researcher based on applications. We refer to the parameter $t$ as “time.” Using this family, we define the time score as \[ \text{(Time score)}\ \ s^{\text{time}}_t(x)\coloneqq \partial_t \log p_t(x)\in\mathbb{R}. \]
We first consider the estimation of the data score $s^{\text{data}}_0$. Let $s^{\text{data}}\colon{\mathcal{X}}\to\mathbb{R}^d$ be a candidate score function (e.g., a neural network $s^{\text{data}}_{\theta}$). A natural population criterion is the mean squared error (MSE) between $s$ and the true score $s^{\text{data}}_0$: ${\mathcal{R}}^\dagger(s^{\text{data}}) \coloneqq \frac12{\mathbb{E}}_{X\sim p_0}\big[\|s^{\text{data}}(X)-s^{\text{data}}_0(X)\|_2^2\big]$. This objective is not directly computable because $s^{\text{data}}_0$ is unknown. However, under standard integration-by-parts conditions, minimizing ${\mathcal{R}}^\dagger$ is equivalent to minimizing
where $\nabla_x\cdot s^{\text{data}}(x)\coloneqq \sum_{j=1}^d \partial_{x_j}s^{\text{data}}_j(x)$. This estimation approach is called Hyv\"arinen score matching, a special case of LSIF Kanamori2009aleastsquares and Riesz regression Chernozhukov2021automaticdebiased (See Appendix (ref)).
In high dimensions, directly optimizing ((ref)) can be numerically challenging because it involves the divergence $\nabla_x\cdot s_{\theta}(X)$, which requires Jacobian-trace computations. A widely used alternative is DSM Vincent2011aconnection, which replaces the (possibly rough) target score $s^{\text{data}}_0$ by the score of a Gaussian-smoothed density. We explain the details in Section (ref). In this study, we refer to DSM-based Riesz representer estimation as Data-ScoreMatchingRiesz.
Let $q(x)$ and $p(x)$ be two densities on ${\mathcal{X}}$, and assume $q$ is absolutely continuous with respect to $p$ on ${\mathcal{X}}$. Define the density ratio as $r(x)\coloneqq q(x) / p(x)$. Assume we can draw i.i.d. samples from $q$ and $p$.
\paragraph{Bridge distributions and telescoping ratios.} Following Choi2022densityratio, construct a continuum of bridge distributions $\{p_t\}_{t\in[0,1]}$ by sampling \[ X_t \coloneqq \beta^{(1)}(t)X_0+\beta^{(2)}(t)X_1, \] where $X_0\sim q$, $X_1\sim p$ are independent, and $\beta^{(1)},\beta^{(2)}:[0,1]\to[0,1]$ are $C^2$ functions satisfying $\beta^{(1)}(0)=1$, $\beta^{(2)}(0)=0$, $\beta^{(1)}(1)=0$, $\beta^{(2)}(1)=1$. Let $p_t$ denote the law of $X_t$ (and write $p_t(x)$ for its density when it exists).
Using intermediate ratios, we have for any integer $T\ge1$, $r(x)=\frac{q(x)}{p(x)}=\prod_{k=1}^T \frac{p_{(k-1)/T}(x)}{p_{k/T}(x)}$. Taking logs gives $\log r(x)=\sum_{k=1}^T \log \frac{p_{(k-1)/T}(x)}{p_{k/T}(x)}$. In the continuum limit, the following result holds:
The time score is $s^{\text{time}}_t(x)=\partial_t\log p_t(x)$.
\paragraph{Time score matching objective.} Let $s^{\text{time}}_{\theta, t}(x)$ be a time score model. A natural population loss is ${\mathcal{R}}^\dagger(s^{\text{time}}_{\theta}) \coloneqq \int_0^1 {\mathbb{E}}_{X_t\sim p_t}\left[ \lambda(t)\Bigl(\partial_t\log p_t(X_t)-s^{\text{time}}_{\theta, t}(X_t)\Bigr)^2 \right]{\mathrm{d}} t$, where $\lambda:[0,1]\to\mathbb{R}_+$ is a weighting function. The true time score is unknown, but Choi2022densityratio shows that (up to an additive constant that does not depend on $\theta$) this objective is equivalent to an integration-by-parts objective:
In the integral term, $\partial_t s^{\text{time}}_{\theta, t}(x)$ denotes the partial derivative with respect to $t$ while holding $x$ fixed.
\paragraph{Recovering the density ratio.} After fitting $s^{\text{time}}_{\theta,t}$, the ratio can be estimated by numerical integration as $\widehat{\log r}(x) \coloneqq \int_{1}^{0} s^{\text{time}}_{\theta, t}(x) {\mathrm{d}} t$ and $\widehat{r}(x)\coloneqq \exp\left(\widehat{\log r}(x)\right)$.
We first define Data-ScoreMatchingRiesz, which is an application of DSM. Let $\sigma>0$ and define a noising kernel \[ q_\sigma(\tilde x\mid x)={\mathcal{N}}(\tilde x;x,\sigma^2 I_d), \quad \tilde X = X + \sigma \varepsilon,\ \ \varepsilon\sim{\mathcal{N}}(0,I_d). \] Let $p_\sigma$ denote the marginal density of $\tilde X$, i.e. the convolution $p_\sigma = p_0 * {\mathcal{N}}(0,\sigma^2 I_d)$. Data-ScoreMatchingRiesz fits a model $s^{\text{data}}_{\theta, \sigma}(\tilde x)$ by minimizing \[{\mathcal{R}}^{\mathrm{DSM}\dagger}(s^{\text{data}}_{\theta, \sigma}) \coloneqq {\mathbb{E}}_{\sigma\sim \pi}{\mathbb{E}}_{X\sim p_0}{\mathbb{E}}_{\tilde X\sim q_\sigma(\cdot\mid X)}\Big[ \lambda(\sigma) \bigl\| s^{\text{data}}_{\theta, \sigma}(\tilde X) -\nabla_{\tilde x}\log q_\sigma(\tilde X\mid X) \bigr\|_2^2\Big],\] where $\pi$ is a distribution over noise levels and $\lambda(\sigma)\ge 0$ is a weighting function. For Gaussian noise, $\nabla_{\tilde x}\log q_\sigma(\tilde x\mid x) = -\frac{\tilde x-x}{\sigma^2} = -\frac{\varepsilon}{\sigma}$ holds. Hence, an equivalent feasible form of ${\mathcal{R}}^{\mathrm{DSM}\dagger}(s^{\text{data}}_{\theta, \sigma})$ is
It is known that for each fixed $\sigma$, the population minimizer $\theta^*$ satisfies $s^{\text{data}}_{\theta^*, \sigma}(\cdot) \coloneqq \nabla_{\tilde x}\log p_\sigma(\cdot)$.
This is the formulation most commonly used in diffusion models Song2021denoisingdiffusion: one trains a single network on a continuum of noise levels $\sigma$, which can be interpreted as a “time” variable, and then evaluates (or extrapolates) the score near the small-noise regime.
The (time) score matching approach above extends directly to Riesz representer estimation when the representer can be expressed through density ratios.
In AME estimation with continuous $D$, the Riesz representer is $\alpha^{\mathrm{AME}}_0(D,Z)= -\partial_d\log p_0(d,Z)\big|_{d=D}$. This Riesz representer can be interpreted as a time score.
\paragraph{Time score interpretation.} For intuition, consider the affine bridge $D_t=\beta^{(1)}(t)+\beta^{(2)}(t)D$ with $\beta^{(1)}(0)=0$, $\beta^{(2)}(0)=1$, and derivatives $\partial_t\beta^{(1)}(0)=1$, $\partial_t\beta^{(2)}(0)=0$. Let $p_t(d,z)$ be the density of $(D_t,Z)$. A direct change-of-variables calculation yields, for $t$ where a density exists,
Evaluating at $t=0$ gives $\partial_t\log p_t(d,z)\big|_{t=0}=-\partial_d\log p_0(d,z)$, so the AME representer can be viewed as a time score at $t=0$ for this local translation-like bridge. This observation connects AME estimation to the time-score framework used for ratio-based representers.
In APE estimation, the Riesz representer is the difference of two density ratios: $\alpha^{\mathrm{APE}}_0(x)=\frac{p_1(x)}{p_0(x)}-\frac{p_{-1}(x)}{p_0(x)}$. Thus it suffices to estimate $p_1/p_0$ and $p_{-1}/p_0$.
\paragraph{Bridge construction and time score training.} Construct a bridge family $\{p_t\}_{t\in[-1,1]}$ for $X$ using independent draws $X_0\sim p_0$, $X_1\sim p_1$, $X_{-1}\sim p_{-1}$ and the same style of interpolation as in the ATE case: \[ X_t \coloneqq
\] We train $s^{\text{time}}_{\theta, t}(x)$ by minimizing
\paragraph{Constructing $\widehat{\alpha}^{\mathrm{APE}}$.} After training, we obtain $\widehat{\log\frac{p_1}{p_0}}(x)\coloneqq \int_{0}^{1} s^{\text{time}}_{\theta, t}(x) {\mathrm{d}} t$ and $\widehat{\log\frac{p_{-1}}{p_0}}(x)\coloneqq \int_{0}^{-1} s^{\text{time}}_{\theta, t}(x) {\mathrm{d}} t$. Then, we estimate the Riesz representer as \[\widehat{\alpha}^{\mathrm{APE}}(x) \coloneqq \exp\left(\widehat{\log\frac{p_1}{p_0}}(x)\right) - \exp\left(\widehat{\log\frac{p_{-1}}{p_0}}(x)\right).\]
To stabilize training, we jointly train the data score while training the time score, following the suggestion in Choi2022densityratio, in a multitask learning manner. For details, see Appendix (ref).
The procedure for estimating the parameter of interest is as follows: (1) estimate the nuisance parameters $\eta_0 = (\gamma_0, \alpha_0)$; (2) plug in the estimates $\widehat{\eta} = (\widehat{\gamma}, \widehat{\alpha})$ into the Neyman score and solve for $\widehat{\theta}$ from $\frac{1}{n}\sum_{i=1}^n \psi\bigl(W_i;\widehat{\eta},\widehat{\theta}\bigr)=0$.
If the estimators of the nuisance parameters do not satisfy the Donsker condition, we employ cross-fitting to control the empirical process term Klaassen1987consistentestimation,Zheng2011crossvalidatedtargeted,Chernozhukov2018doubledebiased.
ScoreMatchingRiesz admits a generalized formulation based on Bregman divergences, which makes the choice of the loss and the link explicit. Under suitable choices, the implied ratio-based representer satisfies an automatic covariate balancing property that directly reduces the error for the Neyman orthogonal score Kato2026rieszrepresenter,BrunsSmith2025augmentedbalancing. Due to space constraints, we focus on Time-ScoreMatchingRiesz; the same argument applies to Data-ScoreMatchingRiesz.
Let $g\colon\mathbb{R}\to\mathbb{R}$ be differentiable and strictly convex, and let $D_g(u\|v)\coloneqq g(u)-g(v)-g'(v)(u-v)$ denote the Bregman divergence. For the bridge $\{p_t\}_{t\in[0,1]}$ with time score $s^{\text{time}}_{0,t}(x)\coloneqq \partial_t\log p_t(x)$, define the weighted integrated divergence \[\mathrm{BD}_g^\dagger\left(s^{\text{time}}_{0}\middle\| s^{\text{time}}_{\theta}\right) \coloneqq \int_0^1 {\mathbb{E}}_{X_t\sim p_t}\Bigl[\lambda(t)\, D_g\bigl(s^{\text{time}}_{0,t}(X_t)\,\big\|\, s^{\text{time}}_{\theta,t}(X_t)\bigr)\Bigr]{\mathrm{d}} t,\] where $s^{\text{time}}_{\theta,t}$ is a differentiable model and $\lambda(t)\ge 0$ is a weighting function. If the model is well specified, minimizing $\mathrm{BD}_g^\dagger\left(s^{\text{time}}_{0}\middle\| s^{\text{time}}_{\theta}\right)$ recovers the true time score.
For $g(u)=u^2/2$, $\mathrm{BD}_g(s^{\text{time}}_{\theta})$ reduces to the squared-loss objective used in Time-ScoreMatchingRiesz.
Let $r_0(x)\coloneqq q(x)/p(x)$ denote the endpoint ratio. Time-ScoreMatchingRiesz constructs a log-ratio model by integrating the time score and applying an exponential link: $\ell_{\theta}(x)\coloneqq -\int_0^1 s^{\text{time}}_{\theta,t}(x){\mathrm{d}} t$ and $r_\theta(x)\coloneqq \exp\left(\ell_{\theta}(x)\right)$. Since $r_0$ satisfies the covariate-balancing identity ${\mathbb{E}}_{X\sim p}\bigl[r_0(X)\varphi(X)\bigr]={\mathbb{E}}_{X\sim q}\bigl[\varphi(X)\bigr]$ for all $\varphi\in L_2(p)$, ratio-based representers in ScoreMatchingRiesz can be read as balancing weights.
\paragraph{Loss, link, and model choice.} Under matched choices of a Bregman loss and a link, generalized Riesz regression yields (approximate) sample covariate balancing as a first-order condition Kato2026rieszrepresenter. In our setting, choosing the generator $g$ in $\mathrm{BD}_g(s^{\text{time}}_{\theta})$ and the link $\ell_{\theta}$ determines them. This viewpoint suggests selecting the score model class and $g$ so that the induced ratio $r_\theta$ can balance the same features that drive the outcome regression.
\paragraph{Neyman error minimization.} For the orthogonal estimator $\widehat{\theta}$, the leading remainder is the error term ${\mathbb{E}}[(\widehat\alpha-\alpha_0)(\widehat\gamma-\gamma_0)]$. When a Riesz representer estimator achieves balancing for a feature map spanning the regression function, this error term is automatically reduced, and in linear settings it can vanish exactly Kato2026rieszrepresenter,BrunsSmith2025augmentedbalancing. Appendix (ref) records the corresponding balancing identities and a concise sample-level statement.
We introduce a policy path---a continuum of policy effects indexed by a scalar policy intensity $\delta$. The path provides a continuous representation of finite-shift policy effects (APEs). Moreover, only under a known pushforward structure does the $\delta\to0$ limit recover the AME.
Let $\{P_\delta\}_{\delta\in[-\delta_{\max},\delta_{\max}]}$ be counterfactual laws of the regressor $X\in{\mathcal{X}}$ with densities $\{p_\delta\}$. A policy is called a known pushforward if there exists a known measurable map $T_\delta:{\mathcal{X}}\to{\mathcal{X}}$ such that \[ X_\delta = T_\delta(X)\quad\text{for }X\sim P_0, \] equivalently $P_\delta=(T_\delta)_\# P_0$, where $(T_\delta)_\# P_0$ denotes the pushforward distribution (formal definition and density formulas are given in Appendix (ref)).
Define the (symmetric) policy-path parameter
where $\gamma_0(x)\coloneqq {\mathbb{E}}[Y\mid X=x]$. For any target shift $\delta^\star$, the APE at that shift is simply the point on the path: $\theta_0^{\mathrm{APE}}(\delta^\star)=\theta_0(\delta^\star)$. Thus, $\delta\mapsto\theta_0(\delta)$ is a continuous representation of finite policy effects.
When $p_{\pm\delta}\ll p_0$, we can rewrite $\theta_0(\delta)$ under $P_0$ using density ratios $r_{\pm\delta}(x)\coloneqq p_{\pm\delta}(x)/p_0(x)$:
Hence, estimating the path reduces to estimating the path representer $\alpha_{0,\delta}$, i.e., (a difference of) density ratios.
Assume $\delta\mapsto p_\delta(x)$ is differentiable at $\delta=0$ for a.e.\ $x$ and define the policy (time) score at the baseline $s^{\mathrm{pol}}_0(x)\coloneqq \partial_\delta\log p_\delta(x)\big|_{\delta=0}$. Then a standard differentiation argument yields $\theta_0'(0) = 2\,{\mathbb{E}}\bigl[\gamma_0(X)\,s^{\mathrm{pol}}_0(X)\bigr]$. Crucially, $s^{\mathrm{pol}}_0$ depends on how the policy perturbs the law of $X$ and is not determined by the observational data score $s^{\mathrm{data}}_0(x)=\nabla_x\log p_0(x)$ in general. Therefore, without additional structure, the infinitesimal limit $\theta_0(\delta)/(2\delta)$ is not canonically tied to the AME.
\paragraph{AME limit for shift policies under known pushforward.} Suppose the policy is a known pushforward $P_\delta=(T_\delta)_\#P_0$ with $T_0=\mathrm{Id}$. Then, $\theta_0(\delta) = {\mathbb{E}}\bigl[\gamma_0(T_{+\delta}(X))-\gamma_0(T_{-\delta}(X))\bigr]$ holds, so the local behavior of the path is governed by the geometry of $T_\delta$. In the canonical translation (shift) policy $X=(D,Z)$, $T_\delta(d,z)=(d+\delta,z)$, we have $\theta_0(\delta)={\mathbb{E}}[\gamma_0(D+\delta,Z)-\gamma_0(D-\delta,Z)]$ and, under mild smoothness, $\lim_{\delta\downarrow0}\frac{\theta_0(\delta)}{2\delta} = {\mathbb{E}}[\partial_d\gamma_0(D,Z)] =: \theta_0^{\mathrm{AME}}$. Thus, the policy path recovers AME as its $\delta\to0$ limit only in the structured (known-pushforward) regime. Appendix (ref) further shows that the same structure allows construction of $\alpha_{0,\delta}$ (hence APEs and the full path) from the data score.
\paragraph{Coincidence of the Data Score and Time Score} A one-parameter policy family $\{p_\delta\}$ itself provides a natural “time” index. Define the time score along the policy as $s^{\mathrm{time}}_\delta(x)\coloneqq \partial_\delta\log p_\delta(x)$. Then $s^{\mathrm{time}}_\delta$ is exactly the policy score. For translation pushforwards,
so the time score along the policy coincides with the treatment-direction data score evaluated along the known path. This is the basic reason why, under known pushforwards, a single data-score model on $P_0$ can generate ratios for all $\delta$ via one-dimensional integration (Appendix (ref)).
This section provides theoretical analysis related to our proposed ScoreMatchingRiesz. Because our nuisance components are constructed from either the data score $\nabla_x \log p(x)$ or the time score $\partial_t \log p_t(x)$, we focus on $L_2$-type guarantees for these scores.
This subsection records $L_2$-type score estimation rates from diffusion-model-based (denoising) score matching, as established by Oko2023diffusionmodels. Throughout, let $p_0$ be a probability density on a bounded domain ${\mathcal{X}}\subset\mathbb{R}^d$ (e.g., ${\mathcal{X}}=[-1,1]^d$), and let $X\sim p_0$. Recall that for $\sigma>0$, we defined the smoothed density $p_\sigma \coloneqq p_0 * {\mathcal{N}}(0,\sigma^2 I_d)$ (so $\tilde X_\sigma \sim p_\sigma$), and its (data) score $\nabla_x\log p_\sigma(x)$ for $x\in{\mathcal{X}}$. Let $\widehat s^{\text{data}}_\sigma(x)$ denote an estimator of this score.
\paragraph{Additional notations and assumptions.} For a function $g\colon{\mathcal{X}}\to\mathbb{R}^d$ and a density $p$ on ${\mathcal{X}}$, define $\|g\|_{L_2(p)}^2 \coloneqq {\mathbb{E}}_{X\sim p}\bigl[\|g(X)\|_2^2\bigr]$. For a score field $g(\cdot,\sigma)$ depending on $\sigma\in[\sigma_{\min},\sigma_{\max}]$, define the integrated $L_2$ norm $\|g_{\cdot}\|_{\nu}^2 \coloneqq \int_{\sigma_{\min}}^{\sigma_{\max}} \|g_\sigma(\cdot)\|_{L_2(p_\sigma)}^2 {\mathrm{d}} \sigma$, where $0<\sigma_{\min}<\sigma_{\max}<\infty$. We write $a_n\lesssim b_n$ if $a_n\le C b_n$ for a universal constant $C>0$ and all large $n$. We write $\mathrm{polylog}(n)$ for a factor bounded by $(\log n)^c$ for some $c>0$. Let $B^{s}_{p,q}({\mathcal{X}})$ denote a Besov space on ${\mathcal{X}}$ with smoothness $s>0$ and integrability parameters $p,q\in[1,\infty]$. We assume $p_0\in B^{s}_{p,q}({\mathcal{X}})$.
\paragraph{$L_2$ rate.} Oko2023diffusionmodels shows the following result, which is directly applicable to Data-ScoreMatchingRiesz.
In AME estimation with $X=(D,Z)$ and continuous $D$, the Riesz representer is $\alpha^{\mathrm{AME}}_0(D,Z)=-\partial_d\log p_0(D,Z)$ (equivalently, $-\{s^{\text{data}}_0(D,Z)\}_d$).
We next give an $L_2$ rate for Time-ScoreMatchingRiesz. Since the formulation varies across applications, we consider the following representative problem.
\paragraph{Additional notations and assumptions.} Let $q$ and $p$ be two probability densities on ${\mathcal{X}}$ such that $q\ll p$. Define the density ratio as $r(x)\coloneqq \frac{q(x)}{p(x)}$ for $x\in{\mathcal{X}}$. Let $\{p_t\}_{t\in[0,1]}$ be a family of bridge densities connecting $q$ and $p$: $p_0=q$ and $p_1=p$. We assume $t\mapsto p_t(x)$ is differentiable for a.e.\ $x$, and define the time score $s^{\text{time}}_t(x)\coloneqq \partial_t\log p_t(x)$ for $(x,t)\in{\mathcal{X}}\times(0,1)$. Write $X_t\sim p_t$.
Let $\lambda:[0,1]\to\mathbb{R}_+$ be a measurable weighting function and define $\lambda_{\min}\coloneqq \inf_{t\in[0,1]}\lambda(t)$. When $\lambda_{\min}>0$, define the weighted $L_2$ norm for a family $\{f_t\}_{t\in[0,1]}$ with $f_t\colon{\mathcal{X}}\to\mathbb{R}$ by $\|f\|_{\mu_\lambda}^2 \coloneqq \int_0^1 {\mathbb{E}}_{X_t\sim p_t}\bigl[\lambda(t) f_t(X_t)^2\bigr] {\mathrm{d}} t$. We write $\mu_1$ for the special case $\lambda(t)\equiv 1$.
Let $P_n$ denote an effective complexity parameter for ${\mathcal{F}}_n$ (e.g., number of parameters or an upper bound thereof), used to express approximation and generalization rates.
Assumption (ref) is a technical device to transfer $L_2$ control along the bridge to downstream bounds for ratios and representers. Appendix (ref) summarizes practical bridge and weighting choices that mitigate violations of overlap.
\paragraph{$L_2$ rate.} Let ${\mathcal{R}}(\cdot)$ denote the population time-score-matching objective in ((ref)). Under regularity conditions (as in Lemma (ref)), $s^{\text{time}}_t$ is the population minimizer of ${\mathcal{R}}$. Let $\widehat{{\mathcal{R}}}_n(\cdot)$ be an empirical counterpart of ${\mathcal{R}}(\cdot)$ computed from samples from $q$ and $p$ and Monte Carlo bridge draws as in Section (ref). Let ${\mathcal{F}}_n$ be a class of time-score models $s^{\text{time}}_{\theta, t}(x)$ (e.g., neural networks) that are differentiable in $t$. Define the empirical risk minimizer as $\widehat s \in \arg\min_{s\in{\mathcal{F}}_n} \widehat{{\mathcal{R}}}_n(s)$. Then, the following result holds.
Integrating an estimated time score $\widehat s$ yields a plug-in estimator of $\log r$ and hence of $r$. Appendix (ref) records the corresponding $L_2$ rate for the Riesz representer.
We recall the asymptotic efficiency result for $\widehat{\theta} = \frac{1}{n}\sum_{i=1}^n \Big(m(W_i,\widehat{\gamma})+\widehat{\alpha}(X_i)(Y_i-\widehat{\gamma}(X_i))\Big)$. Assume that both $\widehat{\gamma}$ and $\widehat{\alpha}$ satisfy the Donsker condition or are constructed via cross-fitting. Assume that the following conditions hold: (i) $\|\widehat{\alpha} - \alpha_0\|_2 = o_p(1)$; (ii) $\|\widehat{\gamma} - \gamma_0\|_2 = o_p(1)$; (iii) $\|\widehat{\alpha} - \alpha_0\|_2\|\widehat{\gamma} - \gamma_0\|_2 = o_p(1/\sqrt{n})$. Then, it holds that $\sqrt{n}\left(\widehat{\theta}-\theta_0\right) \xrightarrow{d} {\mathcal{N}}(0,{\mathbb{E}}[\psi(W;\eta_0,\theta_0)^2])$ as $n\to\infty$.
We compare Data-ScoreMatchingRiesz (Data-SMR) and Time-ScoreMatchingRiesz (Time-SMR), and we use Riesz regression Chernozhukov2021automaticdebiased,Kanamori2009aleastsquares as a baseline. Note that Riesz regression is identical to Hyv\"arinen score matching in AME estimation and LSIF in APE estimation, while its dual corresponds to stable balancing weights Zubizarreta2015stableweights. Therefore, it is sufficient to compare our method only with Riesz regression because it is equivalent to various other methods. For details on these equivalences, see BrunsSmith2025augmentedbalancing and Kato2026rieszrepresenter. All methods use the same MLP architecture for $\widehat\gamma$ and two-fold cross-fitting. Additional results are reported in Appendices (ref) and (ref).
\paragraph{Data generating process.} Throughout, $X=(X^{(1)},X^{(2)},X^{(3)})\in\mathbb{R}^3$ follows $X\sim {\mathcal{N}}(0,\Sigma)$ with \[ \Sigma=
. \] We set $D\coloneqq X^{(1)}$ and $Z\coloneqq (X^{(2)},X^{(3)})$. The outcome is $Y=\mu(X)+\varepsilon$ with $\varepsilon\sim{\mathcal{N}}(0,1)$ and $\mu(x) = 1+x^{(1)}+0.1(x^{(1)})^2+2\sin(x^{(1)})+x^{(2)}+x^{(1)}x^{(2)}+(x^{(3)})^2+(x^{(3)})^3$.
\paragraph{Targets.} We estimate the AME $\theta_0^{\mathrm{AME}}={\mathbb{E}}[\partial_d\gamma_0(D,Z)]$ and the APE under the shift policy at $\delta=1$, $\theta_0^{\mathrm{APE}}={\mathbb{E}}[\mu(D+1,Z)]-{\mathbb{E}}[\mu(D-1,Z)]$. The ground truths are computed by population Monte Carlo.
\paragraph{Evaluation metrics.} Across $200$ trials, we report (i) bias and MSE of $\widehat\theta$ and (ii) empirical coverage of 95% confidence intervals. For the shift policy, we also estimate the full policy path $\delta\mapsto\theta_0(\delta)$ on a grid (Section (ref)) and plot pointwise 95% bands, the true path, and the local linear approximation $2\delta\,\theta_0^{\mathrm{AME}}$.
\paragraph{Results.} Tables (ref) and (ref) show that Data-SMR is the most accurate and stable method, achieving the smallest MSE for both AME (0.161) and APE at $\delta=1$ (0.486) with near-nominal coverage, whereas Time-SMR is less accurate and exhibits marked undercoverage for APE (75%). Riesz regression is substantially less stable, especially for APE where it suffers severe variance amplification (MSE 74.918). Figure (ref) further suggests that the policy path deviates from the local linear approximation $2\delta,\theta_0^{\mathrm{AME}}$ as $\delta$ increases, highlighting the need for stable representer estimation to recover the nonlinear curve.
We proposed ScoreMatchingRiesz to stably estimate the Riesz representer, which also provides a unified interpretation across causal targets. In practice, although Data-ScoreMatchingRiesz shows preferable empirical performance, there are applications where we cannot use Data-ScoreMatchingRiesz and only Time-ScoreMatchingRiesz is available; thus, we need to choose the method based on the application. We define the policy path that bridges the AME and APE, and we demonstrated that this path can be estimated using the score. These connections suggest that score matching is a general tool for causal inference.
\onecolumn