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.
23,920 characters · 11 sections · 56 citation commands
Nearest Neighbor Matching as Least Squares Density Ratio Estimation and Riesz Regression
Nearest-neighbor matching is a popular method for treatment effect estimation from observational data. Lin2023estimationbased shows that nearest-neighbor (NN) matching can be derived from density ratio estimation (DRE), providing a novel density ratio estimator.
In this study, we show that the density ratio estimator proposed in Lin2023estimationbased is equivalent to Least-Squares Importance Fitting (LSIF), proposed by Kanamori2009aleastsquares, with a specific kernel function. Furthermore, we show that nearest-neighbor matching can be viewed as a specific implementation of Riesz regression, proposed by Chernozhukov2024automaticdebiased, by demonstrating the equivalence between LSIF and Riesz regression. Our results imply the following: (i) we can directly apply various results established for LSIF to NN matching, such as the choice of alternative kernels, convergence rates, and possible extensions; (ii) we confirm the generality of LSIF and Riesz regression for a variety of problems by identifying their equivalence with NN matching.
In Section (ref), we formulate our problem. In Section (ref), we review the arguments in Lin2023estimationbased. In Section (ref), we show that the DRE method proposed in Lin2023estimationbased is equivalent to LSIF, proposed in Kanamori2009aleastsquares. In Section (ref), we show that NN matching is equivalent to Riesz regression, proposed in Chernozhukov2024automaticdebiased, by establishing the equivalence between Riesz regression and LSIF. The results in this study are partially based on our previous work Kato2025directbias and Kato2025directdebiased.
Let $X \in {\mathcal{X}} \subseteq {\mathbb{R}}^d$ be $d$-dimensional covariates with a covariate space ${\mathcal{X}}$, $D \in \{1, 0\}$ be a treatment indicator, and $Y_i \in {\mathcal{Y}} \subseteq {\mathbb{R}}$ be an outcome with an outcome space ${\mathcal{Y}}$. Let $Y(1)$ and $Y(0)$ be the potential outcomes of treated and control units, where $Y = D \cdot Y(1) + (1 - D) \cdot Y(0)$ holds. We define the Average Treatment Effect (ATE) as \[\tau_0 \coloneqq \mathbb{E}\left[Y(1) - Y(0)\right].\] We observe samples $\{(X_i, D_i, Y_i)\}^n_{i=1}$, where $(X_i, D_i, Y_i)$ is an i.i.d. copy of $(X, D, Y)$. Our goal is to estimate $\tau_0$ using these observations.
\paragraph{Assumptions and notation.} This section describes the assumptions and notation. We assume that the covariates in the treated and control groups have densities. Let $f(x \mid D = 1)$ denote the probability density of covariates in the treated group, $f(x \mid D = 0)$ the probability density in the control group, and $f(x)$ the marginal covariate density. We also denote the joint density of $(X, D)$ by $f(x, d)$. Let us define the propensity score as $e_0(X) \coloneqq P(D = 1\mid X) = \frac{f(x, 1)}{f(x)}$ and assume that there exists a constant $\epsilon \in (0, 1/2)$ independent of $n$ such that $\epsilon < e_0(X) < 1 - \epsilon$. We assume that $Y(1)$ and $Y(0)$ are independent of $D$ conditional on $X$ (unconfoundedness). We also assume that $X$, $Y(1)$, and $Y(0)$ are bounded. These assumptions may be relaxed, but we adopt them for simplicity without loss of generality in our claims. We write $\mathbbm{1}(\cdot)$ for the indicator function. Let $n_1 = \sum^n_{i=1}\mathbbm{1}(D_i = 1)$ and $n_0 = \sum^n_{i=1}\mathbbm{1}(D_i = 0)$.
Lin2023estimationbased first proposes their own density ratio estimation method. Let $X \sim \nu_0$ and $Z \sim \nu_1$ be independent with absolutely continuous laws on $\mathbb{R}^d$. We observe i.i.d. samples $\{X_i\}_{i=1}^{N_0}$ and $\{Z_j\}_{j=1}^{N_1}$ and aim to estimate the density ratio \[ r_0(x) \coloneqq \frac{f_{(1)}(x)}{f_{(0)}(x)}, \] where $f_{(1)}$ and $f_{(0)}$ are the probability densities of $\nu_1$ and $\nu_0$, respectively. For $M \in \{1,\dots,N_0\}$ and $z \in \mathbb{R}^d$, let ${\mathcal{X}}_{(M)}(z)$ be the $M$-th nearest neighbor of $z$ in $\{X_i\}_{i=1}^{N_0}$ under a given metric $\|\cdot\|$. Define the catchment area of $x$ as \[ A_M(x) \coloneqq \bigl\{z : \|x - z\| \le \|{\mathcal{X}}_{(M)}(z) - z\|\bigr\}, \] and the matched-times count as \[ K_M(x) \coloneqq \sum_{j=1}^{N_1} \mathbbm{1}\big(Z_j \in A_M(x)\big). \] Lin2023estimationbased propose the one-step estimator \[ \widehat r_M(x) = \frac{N_0}{N_1} \frac{K_M(x)}{M}, \] and show that it is computationally efficient and rate-optimal for Lipschitz densities.
The $M$–NN matching estimator of Abadie2006largesample imputes the missing potential outcome of unit $i$ by averaging the observed outcomes of its $M$ nearest neighbors in the opposite treatment arm. Let \[ J_M(i)\ \subset\ \{1,\ldots,n\} \] be the index set of the $M$ nearest neighbors of $X_i$ among the units with $D_j = 1 - D_i$. The imputed outcomes are \[ \widehat Y_i(0) =
\qquad \widehat Y_i(1) =
\] The NN matching estimator for the average treatment effect (ATE) is \[ \widehat\tau_M \coloneqq \frac{1}{n} \sum_{i=1}^n \left(\widehat Y_i(1) - \widehat Y_i(0)\right). \] Introduce the matched-times count (how many times unit $i$ is used as a match by units in the opposite group) as \[ K_M(i) \coloneqq \sum_{j=1,\ D_j = 1 - D_i}^n \mathbbm{1}\Big(i \in J_M(j)\Big). \] Then, $\widehat\tau_M$ can be written as follows:
\paragraph{Connection to DRE.} In ATE estimation, the density-ratio estimator proposed in Lin2023estimationbased is given as \[ \frac{n_0}{n_1} \frac{K_M(i)}{M} \xrightarrow{p} \frac{f(X_i\mid D_i = 1)}{f(X_i \mid D_i = 0)} \quad (D_i = 0), \qquad \frac{n_1}{n_0} \frac{K_M(i)}{M} \xrightarrow{p} \frac{f(X_i\mid D_i = 0)}{f(X_i\mid D_i = 1)} \quad (D_i = 1), \] so that, using $n_1 / n_0 \to P(D = 1) / P(D = 0)$, the factor $1 + \frac{K_M(i)}{M}$ is a consistent estimator of the inverse propensity weights: \[ 1 + \frac{K_M(i)}{M} \xrightarrow{p}
\] Substituting these weights into ((ref)) gives the DRE–motivated representation of NN matching.
\paragraph{Bias-corrected NN matching.} Let $\widehat\mu_\omega(x)$ estimate $\mu_\omega(x) \coloneqq \mathbb{E}[Y \mid X = x, D = \omega]$, and set residuals $\widehat R_i \coloneqq Y_i - \widehat\mu_{D_i}(X_i)$ as well as the regression plug-in $\widehat\tau_{\mathrm{reg}} \coloneqq \frac{1}{n} \sum_{i=1}^n \bigl\{\widehat\mu_1(X_i) - \widehat\mu_0(X_i)\bigr\}$. The bias-corrected $M$–NN estimator of Abadie2011biascorrected can be written as
With $M \to \infty$ and $M \log n / n \to 0$, the weights $1 + \frac{K_M(i)}{M}$ consistently estimate the inverse propensity scores, so $\widehat\tau_{bc,M}$ is doubly robust and achieves the semiparametric efficiency bound under suitable smoothness and overlap conditions Lin2023estimationbased.
LSIF aims to solve the least-squares problem \[ r^* \coloneqq \operatorname*{arg\,min}_{r \in {\mathcal{R}}} {\mathbb{E}}_{f_{(0)}} \big[\big(r_0(X) - r(X)\big)^2\big], \] where ${\mathcal{R}}$ denotes a hypothesis class for $r_0$, and ${\mathbb{E}}_{f_{(0)}}$ denotes expectation with respect to the density $f_{(0)}(x)$. Using the identity \[ {\mathbb{E}}_{f_{(0)}} \big[(r_0 - r)^2\big] = {\mathbb{E}}_{f_{(0)}} \big[r(X)^2\big] - 2{\mathbb{E}}_{f_{(1)}} \big[r(X)\big] + \text{const}, \] we obtain the equivalent optimization that does not involve $r_0$: \[ r^* = \operatorname*{arg\,min}_{r \in {\mathcal{R}}} \Big\{{\mathbb{E}}_{f_{(0)}} \big[r(X)^2\big] - 2{\mathbb{E}}_{f_{(1)}} \big[r(X)\big]\Big\}, \] where the constant term is irrelevant for optimization.
We consider a linear-in-parameters density-ratio model; that is, with a basis $\Phi(x) \in \mathbb{R}^d$ we approximate $r_0$ as \[ r_\beta(x) = \beta^\top \Phi(x), \] where $\beta \in \mathbb{R}^d$. We estimate $\beta$ by ridge regression: \[ \widehat{\beta} \coloneqq \operatorname*{arg\,min}_{\beta \in {\mathbb{R}}^d} \left\{ \frac{1}{N_0} \sum_{i=1}^{N_0} r_\beta(X_i)^2 - 2\frac{1}{N_1} \sum_{j=1}^{N_1} r_\beta(Z_j) + \frac{\lambda}{2} \|\beta\|_2^2 \right\}, \] where $\lambda \ge 0$ is a regularization parameter, $\{X_i\}_{i=1}^{N_0} \sim f_{(0)}$, and $\{Z_j\}_{j=1}^{N_1} \sim f_{(1)}$.
This program has the closed-form solution \[ \widehat\beta = \operatorname*{arg\,min}_{\beta \in \mathbb{R}^d} \left\{\frac{1}{2} \beta^\top \widehat H \beta - \beta^\top \widehat h + \frac{\lambda}{2} \|\beta\|_2^2\right\} = \left(\widehat H + \lambda I\right)^{-1} \widehat h, \] where
We continue considering the DRE problem in Section (ref). We show the equivalence between Lin2023estimationbased's DRE method and LSIF, which also implies that NN matching can be justified via LSIF.
We aim to estimate $r_0(c)$ for a fixed point of interest $c \in \{Z_j\}^{N_1}_{j=1}$. Let us define the following $1$-dimensional basis function: \[\Phi(x) = \Phi_c(x) = \mathbbm{1}\big(x \in A_M(c)\big),\] where recall that $A_M(x) \coloneqq \bigl\{z : \|x - z\| \le \|{\mathcal{X}}_{(M)}(z) - z\|\bigr\}$, and ${\mathcal{X}}_{(M)}(z)$ is the $M$-th nearest neighbor of $z$ in $\{X_i\}_{i=1}^{N_0}$ under a given metric $\|\cdot\|$.
Then, we define a linear-in-parameters density-ratio model as \[\widehat{r}(c) = \beta \Phi_c(c),\] where $\beta$ is a scalar. Here, recall that LSIF gives \[\widehat\beta = \left(\widehat H + \lambda I\right)^{-1} \widehat h,\] where
Recall that $K_M(c) \coloneqq \sum_{j=1}^{N_1} \mathbbm{1}\big(Z_j \in A_M(c)\big)$. Therefore, when $\lambda = 0$, we have \[\widehat\beta = \left(\widehat H + \lambda I\right)^{-1} \widehat h = \frac{N_0}{N_1} \cdot \frac{K_M(c)}{M},\] and $\widehat{r}(c) = \widehat{\beta} \Phi_c(c) = \frac{N_0}{N_1} \cdot \frac{K_M(c)}{M}$. This estimator coincides with $\widehat r_M(c)$ derived in Lin2023estimationbased. Thus, we show that the DRE method in Lin2023estimationbased is equivalent to LSIF with the basis function $\widehat{r}(c) = \beta \Phi_c(c)$.
Theorem (ref) shows that the density-ratio estimator of Lin2023estimationbased is an LSIF solution for a particular basis function. Hence, general properties of LSIF provide an alternative route to their results. In particular, we can interpret the density-ratio estimator of Lin2023estimationbased as LSIF with a specific kernel function.
\paragraph{Connection to existing DRE methods.} This equivalence result implies that we can leverage various results established for DRE. In particular, Sugiyama2011densityratio shows that a class of DRE methods is equivalent from the viewpoint of Bregman divergence. This class includes Kullback–Leibler divergence minimization Nguyen2007estimatingdivergence,Sugiyama2008directimportance and learning from positive and unlabeled data Lancaster1996casecontrolstudies,duPlessis2015convexformulation,Kato2019learningfrom. For example, we can apply convergence rate results shown in Kanamori2012statisticalanalysis for more general reproducing kernel Hilbert spaces (RKHS). We can also incorporate the covariate shift setting, as investigated in Uehara2020offpolicy, Kato2024doublyrobust, and Chernozhukov2025automaticdebiased. When applying neural networks, overfitting can be prevented by using regularization techniques such as those proposed in Rhodes2020 and Kato2021nonnegativebregman.
Note that, since ${\mathcal{R}}$ is restricted to $\{\beta\Phi_c(x)\}$, for $r^* \coloneqq r^*_c \coloneqq \operatorname*{arg\,min}_{r_c \in \{\beta\Phi_c(x)\}} {\mathbb{E}}_{f_{(0)}} \big[\big(r_0(X) - r_c(X)\big)^2\big]$, we usually have $r^*_c \neq r_0$. In contrast, existing studies analyze estimators under the assumption that $r_0 \in {\mathcal{R}}$, as shown in Kanamori2012statisticalanalysis. Even when $r_0 \notin {\mathcal{R}}$, in many cases, we can still apply existing results to obtain a bound on the mean squared error (MSE). For example, by applying Theorem 2 in Kanamori2012statisticalanalysis, we can bound $\mathbb{E}\left[\big(r^*_c(X) - \widehat{r}_c(X)\big)^2\right]$ for $\widehat{r}_c = \widehat{\beta}\Phi_c(x)$. In addition, we can assess how large a sample size is needed to achieve $r^*_c(c) \approx r_0(c)$.
We can also derive an MSE bound for $r_0$ over the entire $x \in {\mathcal{X}}$ by considering a hypothetical estimator $\widetilde{r}(x)$ such that, for each $x$, we use $\widehat{r}_x(x) = \widehat{\beta}\Phi_x(x)$ obtained via NN matching. Then, the MSE can be bounded as $\mathbb{E}\left[\big(r_0(X) - \widetilde{r}(X)\big)^2\right] = 2\left(\mathbb{E}\left[\big(r_0(X) - r^*_X(X)\big)^2\right] + \mathbb{E}\left[\big(r^*_X(X) - \widehat{r}_X(X)\big)^2\right]\right)$, which can be evaluated via the existing results.
This section further shows the equivalence among nearest-neighbor matching, LSIF, and Riesz regression. To demonstrate this, we show that the bias-correction term can be estimated using LSIF in a more direct way than the formulation of Lin2023estimationbased. That is, we can directly estimate the inverse propensity scores by interpreting them as density ratios: \[ w_0(1, X) = \frac{1}{e_0(X)} = \frac{f(X)}{f(X, 1)},\qquad w_0(0, X) \coloneqq \frac{1}{1 - e_0(X)} = \frac{f(X)}{f(X, 0)}. \] Note that the LSIF formulation in this section is nearly the same as the one in Section (ref), and corresponds to the expressions involving $1 + \frac{K_M(i)}{M}$ in Lin2023estimationbased. We consider this formulation more straightforward than that in Lin2023estimationbased.
We first consider LSIF to estimate $w_0(1, \cdot)$. Note that we separately estimate $w_0(1, \cdot)$ and $w_0(0, \cdot)$, and $w_0(0, \cdot)$ can be estimated in a similar way.
We aim to estimate $w_0(1, c)$ for a fixed point of interest $c = X_i$ such that $D_i = 1$. Following LSIF, we can estimate an empirical version of the following least-squares problem:
where the expectation is taken under the denominator $f(X, 1)$ for the squared term and the numerator $f(X)$ for the linear term (dropping constants).
Since our interest lies only in the point $c$, we use the following one-dimensional basis function again: \[ \Phi(x) = \Phi_c(x) = \mathbbm{1}\big(x \in A_M(c)\big). \]
We consider LSIF with a linear-in-parameter model \[ w(1, c) = \theta \Phi_c(c), \] where $\theta$ is a scalar. Then, the estimator is given as \[ \widehat w(1, c) = \widehat\theta^\top \Phi(c) \] with the estimated parameter \[ \widehat\theta \coloneqq \operatorname*{arg\,min}_{\theta \in \mathbb{R}} \left\{\frac{1}{2} \theta^\top \widehat H \theta - \theta^\top \widehat h + \frac{\lambda}{2} \|\theta\|_2^2\right\} = \left(\widehat H + \lambda I\right)^{-1} \widehat h, \] where
Here, we have
where recall that \[ K_M(i) \coloneqq \sum_{j=1,\ D_j = 1 - D_i}^n \mathbbm{1}\Big(i \in J_M(j)\Big). \]
Therefore, when $\lambda = 0$, we estimate $\widehat w(1, x)$ as \[ \widehat w(1, x) = 1 + \frac{K_M(i)}{M}. \] Similarly, we can estimate $\widehat w(0, c)$ for a fixed point of interest $c = X_i$ such that $D_i = 0$ by solving an empirical version of the following problem:
Then, the estimator is given as \[ \widehat w(0, c) = 1 + \frac{K_M(i)}{M}. \] Using these estimators, we construct the following inverse propensity score estimator for the ATE: \[ \widehat\tau_M = \frac{1}{n} \left[ \sum_{i : D_i = 1} \left(1 + \frac{K_M(i)}{M}\right) Y_i - \sum_{i : D_i = 0} \left(1 + \frac{K_M(i)}{M}\right) Y_i \right]. \] This estimator is equivalent to ((ref)) proposed in Lin2023estimationbased, which is shown to be equal to the NN matching estimator of Abadie2006largesample. We can also construct the bias-corrected estimator in a similar way. Thus, we have shown the equivalence between LSIF, the method in Lin2023estimationbased, and NN matching in a different way from Section (ref).
In debiased machine learning, the estimation of the Riesz representer plays a core role Chernozhukov2022automaticdebiased. To estimate the Riesz representer effectively, Riesz regression has been developed by Chernozhukov2024automaticdebiased. In this context, Riesz regression is essentially the same as LSIF in Kanamori2009aleastsquares.
Let $W = (Y, Z)$, where $Z$ is a regressor and $\gamma_0 \coloneqq \mathbb{E}\left[Y \mid Z\right]$ is a regression function. We consider the target parameter to be a linear functional of the regression function, $\tau_0 = {\mathbb{E}}[m(W, \gamma_0)]$, with Riesz representer $\alpha_0$ satisfying \[ {\mathbb{E}}\bigl[m(W, \gamma)\bigr] = {\mathbb{E}}\bigl[\alpha_0(X) \gamma(X)\bigr] \quad \forall \gamma. \] Chernozhukov2024automaticdebiased show that $\alpha_0$ is the minimizer of the Riesz regression risk: \[ \alpha_0 = \operatorname*{arg\,min}_{\alpha} {\mathbb{E}}\left[\alpha(X)^2 - 2 m(W, \alpha)\right], \] leading to the orthogonal, doubly robust score \[ \psi(W, \gamma, \alpha, \tau) = m(W, \gamma) - \tau + \alpha(Z)\bigl(Y - \gamma(Z)\bigr), \] whose expectation is zero if either $\gamma = \gamma_0$ or $\alpha = \alpha_0$.
In ATE estimation, for $Z = (D, X)$, the Riesz representer is given as \[ \alpha_0(Z) = \mathbbm{1}(D = 1) w_0(1, X) - \mathbbm{1}(D = 0) w_0(0, X). \] For this Riesz representer, the Riesz regression objective is
This objective function ((ref)) for Riesz regression can also be obtained by simultaneously solving the problems ((ref)) and ((ref)). Therefore, Riesz regression is equivalent to LSIF, which implies that NN matching is a specific case of Riesz regression.
Our results provide a constructive dictionary in which Riesz regression collapses to LSIF and to the NN density-ratio estimator, yielding an explicit bridge from automatic debiased machine learning to NN matching: \[ \boxed{\text{NN matching}}\ \equiv\ \boxed{\text{Riesz regression}}\ \equiv\ \boxed{\text{LSIF}}. \] This equivalence clarifies why increasing $M$ improves statistical efficiency and why the same estimator inherits both graph-based computational advantages and LSIF’s quadratic-loss simplicity.
The results of this study are partially shown in Kato2025directbias and Kato2025directdebiased. In those studies, further extensions have been proposed. For example, they show that Riesz regression can be generalized by incorporating Bregman divergence, which also bridges covariate balancing and Riesz regression. In fact, the Riesz representer estimated via LSIF or Riesz regression is in a dual relationship with the stable weights of Zubizarreta2015stableweights. In addition, they show the connection between Riesz regression and targeted maximum likelihood estimation (TMLE) in vanderLaan2006targetedmaximum. Combining these findings with the results of this study, we can relate NN matching with LSIF, Riesz regression, covariate balancing, and TMLE. Furthermore, based on the results of BrunsSmith2025augmentedbalancing, the bias-corrected NN estimator can be expressed as a single linear model.