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.
198,265 characters · 55 sections · 161 citation commands
A Unified Framework for Debiased Machine Learning: Riesz Representer Fitting under Bregman Divergence
The Riesz representer plays a crucial role in debiased machine learning for a variety of causal and structural parameter estimation problems Chen2015sievesemiparametric,Chernozhukov2022automaticdebiased, including Average Treatment Effect (ATE) estimation Imbens2015causalinference, Average Marginal Effect (AME) estimation, Average Policy Effect (APE) estimation, and covariate shift adaptation Shimodaira2000improvingpredictive. The Riesz representer arises from the Riesz representation theorem for a (typically linear) parameter functional and has a close connection to semiparametric efficiency bounds Newey1994theasymptotic. In particular, by using the Riesz representer appropriately, we can construct semiparametric efficient estimators that are asymptotically linear with the efficient influence function, which is also referred to as a Neyman orthogonal score Chernozhukov2018doubledebiased.
In many applications, the Riesz representer admits a closed-form expression in terms of other nuisance objects, but estimating it well is nontrivial. For example, in ATE estimation, the Riesz representer can be written using the inverse propensity score. A straightforward approach is to estimate the propensity score and then construct the Riesz representer by taking its inverse. In covariate shift adaptation, the Riesz representer is given by a density ratio, the ratio of two probability density functions (pdfs). A straightforward approach is to estimate the two pdfs and then take their ratio. However, these approaches are not necessarily designed to minimize the estimation error of the Riesz representer itself, and their performance for the task of Riesz representer estimation is not guaranteed a priori.
To address this issue, end-to-end approaches for Riesz representer estimation have been explored, such as Riesz regression Chen2014sieveinference,Chernozhukov2021automaticdebiased. In addition, many application-specific methods have been developed. For example, in ATE estimation, entropy balancing weights Hainmueller2012entropybalancing, stable balancing weights Zubizarreta2015stableweights, tailored loss minimization Zhao2019covariatebalancing, and calibrated estimation Tan2019regularizedcalbrated have been proposed. In covariate shift adaptation, direct density ratio estimation methods have also been proposed Sugiyama2012densityratio. These developments suggest that “direct” estimation of the Riesz representer can be both statistically and practically advantageous.
This study provides a unified framework for estimating the Riesz representer by fitting a Riesz representer model to the true Riesz representer under a Bregman divergence. Our framework not only accommodates various existing methods as special cases, but also yields new algorithms, clarifies dual balancing interpretations, and provides convergence rate analysis and other theoretical and practical implications.
We first describe a baseline setup that covers many causal and structural parameters with a single i.i.d.\ sample. Extensions to multi-sample settings, such as covariate shift adaptation, are presented in Section (ref).
We denote the observation by \(W \coloneqq (X,Y)\), where \(Y\in{\mathcal{Y}}\) is an outcome and \(X\in{\mathcal{X}}\) is a regressor vector. Here, ${\mathcal{Y}} \subseteq {\mathbb{R}}$ and ${\mathcal{X}} \subseteq {\mathbb{R}}^k$ are outcome and ($k$-dimensional) regressor spaces, respectively. Let $P_0$ be the distribution that generates $W$. We observe $n$ i.i.d.\ copies of $W$, denoted as \[ {\mathcal{D}} \coloneqq \big\{W_i\big\}^n_{i=1} = \big\{(X_i, Y_i)\big\}^n_{i=1}. \] We denote the regression function by $\gamma_0(x) \coloneqq {\mathbb{E}}_{P_0}\left[Y\mid X=x\right]$. We drop the subscript $P_0$ when the dependence is clear.
\paragraph{Parameter functional.} Our goal is to estimate a parameter of interest of the form \[ \theta_0 \coloneqq {\mathbb{E}}\left[m(W,\gamma_0)\right], \] where $m(W,\gamma)$ is a functional that depends on $W$ and a regression function $\gamma\colon {\mathcal{X}} \to {\mathcal{Y}}$. The map $\gamma \mapsto m(W,\gamma)$ is defined for generic $\gamma$, not limited to the true regression function $\gamma_0$. By changing $m$, we can recover many parameters as special cases, including ATE, AME, and APE.
\paragraph{Riesz representer.} Let ${\mathcal{H}} \coloneqq \left\{\gamma\colon {\mathcal{X}}\to{\mathbb{R}} \mid {\mathbb{E}}\left[\gamma(X)^2\right]<\infty\right\}$ be the Hilbert space $L_2(P_X)$ with inner product $\langle f,g\rangle \coloneqq {\mathbb{E}}\left[f(X)g(X)\right]$. For simplicity, we assume that the map \[ \gamma \longmapsto {\mathbb{E}}\left[m(W,\gamma)\right] \] is linear and continuous on ${\mathcal{H}}$ . Equivalently, there exists a constant $C>0$ such that \[ {\mathbb{E}}\left[m(W,\gamma)\right]^2 \leq C\,{\mathbb{E}}\left[\gamma(X)^2\right] \] holds for all $\gamma \in {\mathcal{H}}$. By the Riesz representation theorem, there exists a function $\alpha_0\in {\mathcal{H}}$ such that \[ {\mathbb{E}}\left[m(W,\gamma)\right]={\mathbb{E}}\left[\alpha_0(X)\gamma(X)\right] \] for all $\gamma\in{\mathcal{H}}$.
The function $\alpha_0$ is referred to as the Riesz representer Chen2015sievesemiparametric,Chernozhukov2022automaticdebiased. As discussed below, it plays an important role in constructing efficient estimators. In ATE estimation, the Riesz representer has also been referred to as the bias-correction term or clever covariates vanderLaan2006targetedmaximum,Schuler2024introductionmodern.
\paragraph{Neyman orthogonal score.} Let $\eta_0 \coloneqq (\gamma_0,\alpha_0)$ be the nuisance parameter, where $\alpha_0$ is the Riesz representer associated with $m$. Define the Neyman orthogonal score by
It holds that \[ {\mathbb{E}}\left[\psi\left(W;\eta_0,\theta_0\right)\right]=0, \] which yields an estimating equation for $\theta_0$.
Moreover, $\psi$ is Neyman orthogonal at $\eta_0$ in the sense that the Gateaux derivative with respect to $\eta$ vanishes: \[ \partial_{\eta}{\mathbb{E}}\left[\psi\left(W;\eta,\theta_0\right)\right]\big|_{\eta=\eta_0}=0. \] To see this, note that the derivative with respect to $\alpha$ is ${\mathbb{E}}\left[Y-\gamma_0(X)\right]=0$, and the derivative with respect to $\gamma$ cancels by the Riesz identity ${\mathbb{E}}\left[m(W,h)\right]={\mathbb{E}}\left[\alpha_0(X)h(X)\right]$. Orthogonality implies that first-order errors from estimating $\eta_0$ do not affect the asymptotic distribution of the final estimator of $\theta_0$, provided cross fitting (or a Donsker condition) and mild convergence rate conditions on the nuisance estimators hold Chernozhukov2018doubledebiased,Chernozhukov2022automaticdebiased.
A useful identity that follows from the definitions is \[ {\mathbb{E}}\left[\psi\left(W;\eta,\theta_0\right)\right]={\mathbb{E}}\left[\left(\alpha_0(X)-\alpha(X)\right)\left(\gamma(X)-\gamma_0(X)\right)\right], \] which makes explicit that the score drift is second order in the product of nuisance estimation errors.
\paragraph{Estimator of the parameter of interest.} We consider a base estimator for $\theta_0$, which we call the Augmented Riesz Weighted (ARW) estimator. Let $\widehat{\eta}\coloneqq (\widehat{\gamma},\widehat{\alpha})$ be an estimator of $\eta_0$. Replacing the moment condition ${\mathbb{E}}\left[\psi\left(W;\eta_0,\theta_0\right)\right]=0$ with its empirical analogue yields an estimator $\widehat{\theta}^{\text{ARW}}$ that satisfies \[ \frac{1}{n}\sum^n_{i=1}\psi\left(W_i;\widehat{\eta},\widehat{\theta}^{\text{ARW}}\right)=0. \] Equivalently, \[ \widehat{\theta}^{\text{ARW}} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(m\left(W_i,\widehat{\gamma}\right) + \widehat{\alpha}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right)\Big). \] This estimator is the canonical one-step or estimating-equation estimator built from the orthogonal score. In many semiparametric models, it is asymptotically equivalent to other efficient constructions (e.g., estimating equation methods and TMLE) under standard conditions VanderVaart2002semiparametricstatistics,Schuler2024introductionmodern.
In ATE estimation, $\widehat{\theta}^{\text{ARW}}$ coincides with the augmented inverse probability weighting (AIPW) estimator.
\paragraph{Alternative estimators.} In addition to the ARW estimator, we also consider the following three estimators: the regression adjustment (RA) estimator, the Riesz weighted (RW) estimator, and the targeted maximum likelihood estimator (TMLE):
The RA estimator is motivated by the definition $\theta_0={\mathbb{E}}\left[m(W,\gamma_0)\right]$ and is also referred to as a plug-in or direct method estimator. The RW estimator is motivated by the identity $\theta_0={\mathbb{E}}\left[\alpha_0(X)\gamma_0(X)\right]={\mathbb{E}}\left[\alpha_0(X)Y\right]$, where the second equality follows from ${\mathbb{E}}\left[Y-\gamma_0(X)\mid X\right]=0$. In ATE estimation, the RW estimator coincides with the inverse probability weighting (IPW) estimator.
This study proposes generalized Riesz regression, a general approach for estimating the Riesz representer by Bregman divergence minimization. By varying the Bregman divergence (equivalently, the convex generator) and the link function of the Riesz representer model, generalized Riesz regression recovers a range of existing objectives and yields new ones. We further show that appropriate loss--link pairs induce balancing weights automatically, a phenomenon we call automatic regressor balancing as a generalization of covariate balancing. We then clarify how regressor balancing connects to the orthogonal score, yielding automatic Neyman orthogonalization and automatic Neyman error minimization. Finally, we provide convergence rate results for generalized Riesz regression with RKHS models and neural networks.
In summary, our main contributions are:
These contributions not only provide new methodological and theoretical findings, but also connect existing studies that have developed related ideas in parallel.
\paragraph{Generalized Riesz Regression.} We formulate Riesz representer estimation as a problem of fitting a Riesz representer model to the true Riesz representer under a Bregman divergence Bregman1967relaxationmethod in Section (ref). We refer to Riesz representer fitting under a Bregman divergence as generalized Riesz regression.\footnote{We may also refer to it as Bregman--Riesz regression, direct bias-correction term estimation, generalized tailored loss minimization, or generalized covariate balancing (Remark (ref)). However, we adopt the term generalized Riesz regression because the choice of loss function is closely related to the choice of link function.} The Bregman divergence includes various discrepancy measures, such as squared distance and Kullback--Leibler (KL) divergence, as special cases. Although the true Riesz representer is unknown, we derive an objective function that does not involve it explicitly and can be approximated using only observations. Therefore, we can train the Riesz representer model within an empirical risk minimization framework.
With the squared loss, the Bregman divergence minimization problem aligns with Riesz regression Chernozhukov2021automaticdebiased. With the KL divergence loss, it aligns with tailored loss minimization Zhao2019covariatebalancing. In addition, BrunsSmith2025augmentedbalancing shows that stable balancing weights arise as dual solutions of Riesz regression, while Zhao2019covariatebalancing shows that entropy balancing weights arise as dual solutions of tailored loss minimization. Our formulation places these correspondences within a single loss-based framework. Moreover, subsequent works show that nearest-neighbor matching and score matching can also be viewed as special cases within this perspective Kato2025nearestneighbor,Kato2025scorematchingriesz.
\paragraph{Automatic Regressor Balancing.} We show that generalized Riesz regression admits a dual interpretation as a regressor balancing problem under specific loss--link choices. For example, generalized Riesz regression with squared distance becomes identical to stable balancing weights when we use a linear link for the Riesz representer model. In contrast, generalized Riesz regression with a KL-type loss becomes identical to entropy balancing weights when we use a logistic link. We make these correspondences explicit in Section (ref). We refer to this property as automatic regressor balancing. Covariate balancing, which targets balance of the covariates $Z$ in settings with $X=(D,Z)$, appears as a special case.
\paragraph{Automatic Neyman Orthogonalization.} Regressor balancing has direct implications for the orthogonal score. To illustrate, suppose $\widehat{\alpha}$ satisfies an exact balancing identity for the function $\gamma_0$, \[ \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)\gamma_0(X_i)=\frac{1}{n}\sum^n_{i=1}m\left(W_i,\gamma_0\right). \] Then the RW estimator can be rewritten as an infeasible ARW estimator that uses $\gamma_0$: \[ \widehat{\theta}^{\text{RW}} =\frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)Y_i =\frac{1}{n}\sum^n_{i=1}\Big(m\left(W_i,\gamma_0\right)+\widehat{\alpha}(X_i)\left(Y_i-\gamma_0(X_i)\right)\Big) \eqqcolon \widetilde{\theta}^{\text{ARW}}. \] This identity clarifies that (exact) balancing enforces orthogonalization at the level of the estimating equation. Section (ref) formalizes and extends this idea to approximate balancing and more general loss--link pairs.
\paragraph{Automatic Neyman Error Minimization.} A central object in debiased machine learning is the empirical mean of the orthogonal score with estimated nuisances. For a candidate nuisance pair $\widehat{\eta}=(\widehat{\gamma},\widehat{\alpha})$, consider the empirical drift at $\theta_0$, \[ \frac{1}{n}\sum^n_{i=1}\psi\left(W_i;\widehat{\eta},\theta_0\right). \] Using ${\mathbb{E}}\left[\alpha_0(X)\left(Y-\gamma_0(X)\right)\right]=0$ and the Riesz identity ${\mathbb{E}}\left[m(W,\gamma)\right]={\mathbb{E}}\left[\alpha_0(X)\gamma(X)\right]$, the population counterpart satisfies \[ {\mathbb{E}}\left[\psi\left(W;\widehat{\eta},\theta_0\right)\right] ={\mathbb{E}}\left[\left(\alpha_0(X)-\widehat{\alpha}(X)\right)\left(\widehat{\gamma}(X)-\gamma_0(X)\right)\right]. \] Thus, controlling this drift amounts to controlling a second-order product of nuisance errors, which is the key robustness mechanism behind orthogonal scores.
We show that generalized Riesz regression can be interpreted as directly targeting such score drift through its dual balancing structure. In this sense, generalized Riesz regression not only estimates the Riesz representer $\alpha_0$, but also implicitly targets the orthogonal score itself. We refer to this property as automatic Neyman error minimization.
\paragraph{Convergence Rate Analysis.} In Section (ref), we establish convergence rates of the estimated Riesz representer model toward the true Riesz representer for RKHS regression and neural networks. In particular, we show minimax-optimal rates under standard smoothness and complexity assumptions.
\paragraph{Insights for Practitioners.} Our framework also provides practical guidance. First, the loss--link pair determines both the geometry of fitting (via the Bregman divergence) and the implicit balancing behavior (via the dual). This clarifies when KL-type objectives naturally yield nonnegative, approximately normalized weights, and when squared-loss objectives tend to yield more stable (variance-controlled) weights. Second, our analysis highlights that the choice of Riesz representer loss interacts with the approximation properties of the first-step regression learner for $\gamma_0$. Third, the dual balancing view provides diagnostics: imbalance of regressor moments associated with the learned representation can be used to assess whether the learned Riesz representer is likely to deliver small score drift.
\paragraph{Contents of this Study.} In the following sections, we introduce our general setup (Section (ref)) and then propose our generalized Riesz regression method (Section (ref)). In Section (ref), we define and discuss the automatic covariate balancing property, and in Section (ref), we present the automatic Neyman orthogonalization property. In Section (ref), we summarize insights for practitioners. In Section (ref), we present a convergence rate analysis of generalized Riesz regression.
Appendix (ref) provides a detailed review of related work. In Appendix (ref), we explain the relationship between density ratio estimation and Riesz regression Kato2025rieszregression. In Appendix (ref), we introduce extensions of our framework that are developed in subsequent works Kato2025nearestneighbor,Kato2025scorematchingriesz. In Appendix (ref), we provide Bayesian interpretations of our proposed methods. The remaining appendices contain remarks and proofs for the main text.
This section briefly reviews related work and applications of debiased machine learning. For a more detailed review, see Appendix (ref). Our framework is primarily built on results from Riesz regression, balancing weights, and density ratio estimation. In particular, Bregman divergences have already been applied to density ratio estimation in Sugiyama2011densityratio, which also discusses the duality between empirical risk minimization and moment matching. Our framework generalizes this viewpoint to a broader class of debiased machine learning problems and clarifies how Riesz representer estimation, balancing, and orthogonal scores fit into a single conceptual and algorithmic pipeline.
\paragraph{General Theory.} Many targets in causal and structural inference can be written as linear functionals of a regression function $\gamma_0$ Newey1994theasymptotic. The Riesz representer $\alpha_0$ is the element that represents this functional as an $L_2(P_X)$ inner product, and it enters the efficient influence function through the orthogonal score ((ref)). Therefore, estimating $\alpha_0$ accurately is central to constructing efficient estimators Chernozhukov2022automaticdebiased. Our generalized Riesz regression views $\alpha_0$ estimation as a divergence-minimizing fitting problem, and its dual formulation yields balancing weights that control the drift of the orthogonal score.
We provide representative examples below, along with the corresponding Riesz representers and Neyman orthogonal scores.
\paragraph{ATE Estimation.} Let the regressor be $X\coloneqq (D,Z)$, where $D\in\left\{0,1\right\}$ is a treatment indicator and $Z\in{\mathcal{Z}}$ is a covariate vector. Following the Neyman--Rubin framework Neyman1923surapplications,Rubin1974estimatingcausal, let $Y(1),Y(0)\in{\mathcal{Y}}$ denote potential outcomes. The ATE is \[ \theta^{\text{ATE}}_0 \coloneqq {\mathbb{E}}\left[m^{\text{ATE}}\left(W,\gamma_0\right)\right], \qquad m^{\text{ATE}}\left(W,\gamma\right)\coloneqq \gamma(1,Z)-\gamma(0,Z). \] To identify the ATE, we assume unconfoundedness and overlap: $(Y(1),Y(0))$ is independent of $D$ given $Z$, and there exists $\epsilon\in(0,1/2)$ such that $\epsilon<e_0(Z)<1-\epsilon$ almost surely, where $e_0(Z)\coloneqq {\mathbb{P}}\left(D=1\mid Z\right)$. We also assume suitable moment conditions, such as ${\mathbb{E}}\left[Y(d)^2\right]<\infty$ for $d\in\left\{0,1\right\}$.
In ATE estimation, the Riesz representer is \[ \alpha^{\text{ATE}}_0(X)\coloneqq \frac{D}{e_0(Z)}-\frac{1-D}{1-e_0(Z)}. \] This term is referred to by various names across different methods. In the semiparametric inference literature, it is called a bias-correction term Schuler2024introductionmodern. In TMLE, it is called the clever covariate vanderLaan2006targetedmaximum. In the DML literature, it is called the Riesz representer Chernozhukov2022automaticdebiased. The components $1/e_0(Z)$ and $1/\left(1-e_0(Z)\right)$ are also referred to as inverse propensity scores or balancing weights Horvitz1952generalization,Hainmueller2012entropybalancing. For its estimation, various methods have been proposed, including covariate balancing Imai2013covariatebalancing,Kallus2020generalizedoptimal,Kong2023covariatebalancing and other approaches Lee2010improvingpropensity.
The orthogonal score is \[ \psi^{\text{ATE}}\left(W;\eta,\theta\right)\coloneqq m^{\text{ATE}}\left(W,\gamma\right)+\alpha^{\text{ATE}}(X)\left(Y-\gamma(X)\right)-\theta, \] where $\eta^{\text{ATE}}_0\coloneqq \left(\gamma_0,\alpha^{\text{ATE}}_0\right)$. Solving $\frac{1}{n}\sum^n_{i=1}\psi^{\text{ATE}}\left(W_i;\widehat{\eta},\theta\right)=0$ yields \[ \widehat{\theta}^{\text{ATE}} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}^{\text{ATE}}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right)+m^{\text{ATE}}\left(W_i,\widehat{\gamma}\right)\Big). \] This is the AIPW (doubly robust) estimator.
\paragraph{AME Estimation.} Let $X=(D,Z)$, where $D$ is a scalar continuous treatment. The AME is defined as \[ \theta^{\text{AME}}_0 \coloneqq {\mathbb{E}}\left[\partial_d \gamma_0(D,Z)\right], \qquad m^{\text{AME}}\left(W,\gamma\right)\coloneqq \partial_d\gamma(D,Z). \] Assume that $X$ admits a continuously differentiable pdf $f_0(D,Z)$ in $d$ and that boundary terms vanish so that integration by parts is valid. Then, \[ {\mathbb{E}}\left[\partial_d\gamma(D,Z)\right] = -{\mathbb{E}}\left[\gamma(D,Z)\,\partial_d\log f_0(D,Z)\right], \] and the Riesz representer is the negative score \[ \alpha^{\text{AME}}_0(X)\coloneqq -\partial_d\log f_0(D,Z). \] The orthogonal score is \[ \psi^{\text{AME}}\left(W;\eta,\theta\right)\coloneqq m^{\text{AME}}\left(W,\gamma\right)+\alpha^{\text{AME}}(X)\left(Y-\gamma(X)\right)-\theta, \] with $\eta^{\text{AME}}_0\coloneqq \left(\gamma_0,\alpha^{\text{AME}}_0\right)$. Solving $\frac{1}{n}\sum^n_{i=1}\psi^{\text{AME}}\left(W_i;\widehat{\eta},\theta\right)=0$ yields \[ \widehat{\theta}^{\text{AME}} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}^{\text{AME}}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right)+m^{\text{AME}}\left(W_i,\widehat{\gamma}\right)\Big). \] In Riesz regression for AME settings, the objective coincides with Hyv\"arinen score matching and (in certain linear models) with least-squares importance fitting (LSIF)-type density ratio objectives Hyvarinen2005estimationof,Kanamori2009aleastsquares.
\paragraph{APE Estimation.} We consider the average effect of a counterfactual shift in the distribution of the regressor $X$ from a known $P_{-1}$ to another $P_1$, under the assumption that $\gamma_0$ is invariant to the distribution of $X$. We define the APE as \[ \theta^{\text{APE}}_0 \coloneqq \int \gamma_0(x)\,{\mathrm{d}} \mu(x), \qquad \mu(x)\coloneqq P_1(x)-P_{-1}(x). \] The linear functional is \[ m^{\text{APE}}\left(W,\gamma\right)\coloneqq \int \gamma(x)\,{\mathrm{d}} \mu(x), \] which does not depend on $W$ once $\gamma$ is fixed. Suppose $P_0$, $P_1$, and $P_{-1}$ admit pdfs $p_0$, $p_1$, and $p_{-1}$, and $P_1$ and $P_{-1}$ are absolutely continuous with respect to $P_0$. Then the Riesz representer is \[ \alpha^{\text{APE}}_0(X)\coloneqq \frac{p_1(X)-p_{-1}(X)}{p_0(X)}. \] The orthogonal score is \[ \psi^{\text{APE}}\left(W;\eta,\theta\right)\coloneqq m^{\text{APE}}\left(W,\gamma\right)+\alpha^{\text{APE}}(X)\left(Y-\gamma(X)\right)-\theta, \] with $\eta^{\text{APE}}_0\coloneqq \left(\gamma_0,\alpha^{\text{APE}}_0\right)$. Solving $\frac{1}{n}\sum^n_{i=1}\psi^{\text{APE}}\left(W_i;\widehat{\eta},\theta\right)=0$ yields \[ \widehat{\theta}^{\text{APE}} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}^{\text{APE}}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right)+m^{\text{APE}}\left(W_i,\widehat{\gamma}\right)\Big). \]
\paragraph{Covariate Shift Adaptation.} Covariate shift refers to settings where the marginal distribution of the regressor $X$ changes across populations, while the conditional distribution of $Y$ given $X$ remains invariant Shimodaira2000improvingpredictive,Reddi2015doublyrobust. Let $(X,Y)\sim P_0$ denote the source (labeled) distribution with regressor pdf $p_0(x)$, and let $\widetilde{X}\sim P_{X,1}$ denote the target (unlabeled) regressor distribution with pdf $p_1(x)$. Assume ${\mathbb{E}}\left[Y\mid X=x\right]=\gamma_0(x)$ holds under both populations and $P_{X,1}$ is absolutely continuous with respect to $P_{X,0}$. The target parameter is \[ \theta^{\text{CS}}_0 \coloneqq {\mathbb{E}}_{P_{X,1}}\left[\gamma_0(\widetilde{X})\right] ={\mathbb{E}}_{P_0}\left[r_0(X)\gamma_0(X)\right], \qquad r_0(X)\coloneqq \frac{p_1(X)}{p_0(X)}. \] In this setting, the Riesz representer associated with the linear functional $\gamma\mapsto {\mathbb{E}}_{P_{X,1}}\left[\gamma(\widetilde{X})\right]$, when represented as an $L_2(P_{X,0})$ inner product, is the density ratio $r_0$.
With independent samples $\big\{(X_i,Y_i)\big\}^n_{i=1}\sim P_0$ and $\big\{\widetilde{X}_j\big\}^m_{j=1}\sim P_{X,1}$, an orthogonal moment condition is \[ {\mathbb{E}}_{P_{X,1}}\left[\gamma(\widetilde{X})\right]+{\mathbb{E}}_{P_0}\left[r(X)\left(Y-\gamma(X)\right)\right]-\theta=0, \] and the corresponding estimator takes the doubly robust form \[ \widehat{\theta}^{\text{CS}} \coloneqq \frac{1}{m}\sum^m_{j=1}\widehat{\gamma}(\widetilde{X}_j) +\frac{1}{n}\sum^n_{i=1}\widehat{r}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right). \]
\paragraph{Density Ratio Estimation.} We also note that density ratios are important in many tasks, such as learning with noisy labels Liu2014classificationwith, anomaly detection Smola2009relativenovelty,Hido2008inlierbased,Abe2019anomalydetection,Nam2015directdensityratio,Kato2021nonnegativebregman, two-sample testing Keziou2005testof,Sugiyama2011leastsquarestwosample, and change point detection Kawahara2009changepointdetection. Learning from positive and unlabeled data can also be interpreted as an application of density ratio estimation Kato2019learningfrom. Therefore, density ratio estimation has been studied as an independent task in machine learning Sugiyama2012densityratio.
Sugiyama2008directimportance considers covariate shift adaptation using importance weights estimated by LSIF, which can be interpreted as Riesz regression for density ratio estimation Kato2025rieszregression. Chernozhukov2025automaticdebiased and Kato2024doubledebiasedcovariateshift investigate efficient estimation of parameters under covariate shift from different perspectives.
Kato2025nearestneighbor points out that the nearest neighbor matching-based density ratio estimation method proposed in Lin2023estimationbased is a special case of LSIF for density ratio estimation Kanamori2009aleastsquares. Since LSIF can be interpreted as Riesz regression, nearest neighbor matching-based ATE estimation can also be interpreted as ATE estimation via Riesz regression. Kato2025scorematchingriesz proposes a Riesz representer estimation method based on score matching in diffusion models Hyvarinen2005estimationof,Song2020generativemodeling,Song2020improvedtechniques.
\paragraph{Notations and Assumptions.} If nested parentheses appear as $f\left(\left(\cdot\right)\right)t\right)}$, we often drop one layer when no confusion arises. For example, when $X=(D,Z)$, we write $f(X)=f(D,Z)$ instead of $f\left(\left(D,Z\right)\right)Z\right)}$. Let ${\mathbb{E}}$ denote expectation under $P_0$ unless specified otherwise. We use the subscript ${}_0$ to denote parameters under $P_0$.
In this study, we propose Riesz representer estimation methods by directly fitting a Riesz representer model to the true value under the Bregman divergence, which is a general discrepancy measure that includes the squared loss and the KL divergence as special cases. We refer to our method as generalized Riesz regression. The term generalized Riesz regression reflects the fact that the choice of loss function is closely connected to the choice of link function from the viewpoint of covariate balancing. We explain this viewpoint in Section (ref) and refer to it as automatic covariate balancing. This section provides a general formulation, and we introduce applications of generalized Riesz regression in Section (ref).
This study fits a Riesz representer model $\alpha \colon {\mathcal{X}} \to {\mathcal{A}}$ to the true Riesz representer $\alpha_0(X)$ under a Bregman divergence, where ${\mathcal{A}}\subset {\mathbb{R}}$ is the Riesz representer space. Let $g\colon {\mathcal{A}} \to {\mathbb{R}}$ be a differentiable and strictly convex function on ${\mathcal{A}}$. As discussed in Section (ref), this function $g$ corresponds to the objective (loss) function in covariate balancing. We refer to this function $g$ as the Bregman--loss function or the loss function.
Given $x \in {\mathcal{X}}$, the Bregman divergence between the scalar values $\alpha_0(x)$ and $\alpha(x)$ is defined as \[\text{BD}^\dagger_g\big(\alpha_0(x)\mid \alpha(x)\big) \coloneqq g(\alpha_0(x)) - g(\alpha(x)) - \partial g(\alpha(x)) \big(\alpha_0(x) - \alpha(x)\big),\] where $\partial g$ denotes the derivative of $g$. We then define the average Bregman divergence as \[\text{BD}^\dagger_g\big(\alpha_0\mid \alpha\big) \coloneqq \mathbb{E}\Big[g(\alpha_0(X)) - g(\alpha(X)) - \partial g(\alpha(X)) \big(\alpha_0(X) - \alpha(X)\big)\Big].\] We define the population target as \[\alpha^* = \operatorname*{arg\,min}_{\alpha\in {\mathcal{H}}} \text{BD}^\dagger_g\big(\alpha_0\mid \alpha\big),\] where ${\mathcal{H}}$ denotes models for $\alpha_0$. If $\alpha_0 \in {\mathcal{H}}$, then $\alpha^* = \alpha_0$ holds.
Although $\alpha_0$ is unknown, we can define an equivalent optimization problem that does not involve $\alpha_0$: \[\alpha^* = \operatorname*{arg\,min}_{\alpha\in {\mathcal{H}}} \text{BD}_g\big(\alpha\big),\] where \[\text{BD}_g\big(\alpha\big) \coloneqq \mathbb{E}\Big[ - g(\alpha(X)) + \partial g(\alpha(X)) \alpha(X) - m\big(W, (\partial g) \circ \alpha\big)\Big].\] Here, we use the linearity of $m$ and the Riesz representation theorem, which imply that \[ {\mathbb{E}}\Big[\partial g(\alpha(X))\alpha_0(X)\Big]={\mathbb{E}}\Big[m\big(W, (\partial g) \circ \alpha\big)\Big]. \]
We estimate the Riesz representer $\alpha_0$ by minimizing an empirical Bregman divergence:
where $J(\alpha)$ is a regularization function, and \[ \widehat{\text{BD}}_g(\alpha) \coloneqq \frac{1}{n}\sum^n_{i=1}\Big( - g(\alpha(X_i)) + \partial g(\alpha(X_i)) \alpha(X_i) - m\big(W_i, (\partial g)\circ \alpha\big)\Big). \] The choice of the regularization function is important because Riesz representer estimation is known to exhibit a characteristic overfitting phenomenon, often described as train-loss hacking or the density chasm. For details, see Section (ref).
By choosing different $g$, we obtain various objectives for Riesz representer estimation, including Riesz regression. Specifically, we obtain the following divergences (loss functions) as special cases of the Bregman divergence:
See also Table (ref) for a summary. Appendix (ref) provides a more detailed relationship between Riesz representer estimation and density ratio estimation, and Appendix (ref) provides more detailed relationships among the existing methods for each loss.
We refer to our method as SQ-Riesz when using the squared loss, UKL-Riesz when using the UKL divergence, BKL-Riesz when using the BKL divergence, BP-Riesz when using the BP divergence, and PU-Riesz when using the PU learning loss. We explain these special cases in detail in the following subsection
Let $C \in {\mathbb{R}}$ be a constant. We consider the following convex function: \[g^{\text{SQ}}(\alpha) = (\alpha - C)^2.\] This choice of convex function is motivated by the squared loss. The choice of $C$ depends on the researcher. We propose choosing $C$ so that the automatic covariate balancing property holds, see Section (ref). The derivative of $g^{\text{SQ}}(\alpha)$ with respect to $\alpha$ is given as
Under this choice of $g$, the Bregman divergence objective is given as \[\text{BD}_g\big(\alpha\big) \coloneqq \mathbb{E}\Big[\alpha(X)^2 - 2m\big(W, \big(\alpha(\cdot) - C\big)\big)g)\Big]}.\] Then, the estimation problem can be written as
where \[\widehat{\text{BD}}_{g^{\text{SQ}}}\big(\alpha\big) \coloneqq \frac{1}{n}\sum^n_{i=1}\left( \alpha(X_i)^2 - 2m\big(W_i, \big(\alpha(\cdot)\big)\big)g)}\right).\] Here, for simplicity, we drop constant terms that are irrelevant for the optimization and use the linearity of $m$ for $2(\alpha(\cdot) - C)$\footnote{The original Bregman divergence objective using $g(\alpha) = (\alpha - C)^2$ in ((ref)) is given as
}. This estimation method corresponds to Riesz regression in debiased machine learning Chernozhukov2021automaticdebiased and least-squares importance fitting (LSIF) in density ratio estimation Kanamori2009aleastsquares. Moreover, if we define ${\mathcal{H}}$ appropriately, we can recover nearest neighbor matching, as pointed out in Kato2025nearestneighbor, which extends the argument in Lin2023estimationbased.
Next, we consider a KL-divergence-motivated convex function. Let $C < \inf_x |\alpha(x)|$ be a constant. We define \[g^{\text{UKL}}(\alpha) = (|\alpha| - C)\log\left(|\alpha| - C\right) - |\alpha|.\] The choice of $C$ depends on the researcher. We propose choosing $C$ so that the automatic covariate balancing property holds, see Section (ref). The derivative of $g^{\text{UKL}}(\alpha)$ with respect to $\alpha$ is given as
Under this choice of $g$, the Bregman divergence objective is given as follows\footnote{ This Bregman divergence objective is derived as follows:
}: \[\text{BD}_{g^{\text{UKL}}}\big(\alpha\big) \coloneqq \mathbb{E}\Big[C\log\left(|\alpha(X)| - C\right) + |\alpha(X)| - m\Big(W, \operatorname{sign}\big(\alpha(\cdot)\big) \log\left(|\alpha(\cdot)| - C\right)\Big)\Big].\]
We estimate $\alpha_0$ by minimizing the empirical objective: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha \in {\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{UKL}}}\big(\alpha\big) + \lambda J(\alpha), \] where
In the next subsection, we also introduce the BKL divergence as a KL-divergence-motivated divergence, but the UKL divergence more closely corresponds to the standard KL divergence. The equivalent formulation is known as KLIEP in density ratio estimation. Note that KLIEP is a constrained formulation that is equivalent to UKL divergence minimization, and this equivalence is also known as Silverman's trick Silverman1978densityratios,KatoMinami2023unifiedperspective. This constrained formulation can also be interpreted as a dual formulation. In ATE estimation, tailored loss minimization corresponds to UKL minimization, whose dual yields entropy balancing weights Hainmueller2012entropybalancing.
We introduce the BKL divergence and BKL-Riesz, which are motivated by the KL divergence and logistic regression. Let $C < \inf_x |\alpha(x)|$ be a constant. We define \[g^{\text{BKL}}(\alpha) \coloneqq (|\alpha| - C)\log \big(|\alpha| - C\big) - (|\alpha| + C)\log(|\alpha| + C).\] The choice of $C$ depends on the researcher. We propose choosing $C$ so that the automatic covariate balancing property holds, see Section (ref).
Under this choice of $g$, the Bregman divergence objective is given as follows\footnote{ This Bregman divergence objective is derived as follows:
}: \[\text{BD}_{g^{\text{BKL}}}\big(\alpha\big) \coloneqq \mathbb{E}\left[ C\log \left(\frac{|\alpha(X)| - C}{|\alpha(X)| + C}\right) - m\left(W, \operatorname{sign}\big(\alpha(\cdot)\big) \log\left(\frac{|\alpha(\cdot)| - C}{|\alpha(\cdot)| + C}\right)\right)\right]\right)}}.\]
We estimate $\alpha_0$ by minimizing the empirical objective: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha \in {\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{BKL}}}\big(\alpha\big) + \lambda J(\alpha), \] where
In ATE estimation, this formulation corresponds to MLE for a logistic model of the propensity score. In density ratio estimation, this formulation corresponds to a logistic regression approach, where we classify two datasets using a logistic model and then take the ratio to obtain a density ratio estimator. For details, see Section (ref).
Basu's power (BP) divergence bridges the squared loss and KL divergence Basu1998robustandefficient. Let $C < \inf_x |\alpha(x)|$ be a constant. Based on the BP divergence, we introduce the following function: \[g^{\text{BP}}(\alpha) \coloneqq \frac{\big(|\alpha| - C\big)^{1 + \omega} - \big(|\alpha| - C\big)}{\omega} - |\alpha|.\] The derivative is given as \[\partial g^{\text{BP}}(\alpha) = \left(1 + \frac{1}{\omega}\right)\operatorname{sign}(\alpha)\Big(\big(|\alpha| - C\big)^{\omega} - 1\Big).\] Using this function in the Bregman divergence yields a BP-motivated loss and the corresponding objective for BP-Riesz regression. The choice of $C$ depends on the researcher. We propose choosing $C$ so that the automatic covariate balancing property holds, see Section (ref).
Under this choice of $g$, the Bregman divergence objective is given as follows\footnote{ This Bregman divergence objective is derived from
}:
We estimate $\alpha_0$ by minimizing the empirical objective: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha \in {\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{BP}}}\big(\alpha\big) + \lambda J(\alpha), \] where
Basu's power divergence bridges the squared loss and the (U)KL divergence. When $\omega = 1$, BP-Riesz regression reduces to SQ-Riesz regression, while when $\omega \to 0$, BP-Riesz regression reduces to UKL-Riesz regression. This follows because \[\lim_{\omega \to 0}\frac{\big(|\alpha| - C\big)^\omega - 1}{\omega} = \log\big(|\alpha| - C\big).\] BP-Riesz regression plays an important role in robust estimation of the Riesz representer. UKL-Riesz regression implicitly assumes exponential or sigmoid models for the Riesz representer. If the model is misspecified, the estimation accuracy can deteriorate. As Sugiyama2012densityratio notes, SQ-Riesz regression is more robust to outliers, while UKL-Riesz regression can perform well under correct specification. BP-Riesz regression provides an intermediate objective between these two extremes. In addition, BP-Riesz regression is useful for understanding the automatic covariate balancing property.
We introduce PU learning loss and PU-Riesz, which are motivated by PU learning. Let $C < \inf_x |\alpha(x)|$ be some constant. We define $g^{\text{PU}}$ as \[g^{\text{PU}}(\alpha) \coloneqq \widetilde{C}\log\left(1-|\alpha|\right) + \widetilde{C}|\alpha|\Big(\log\left(|\alpha|\right)-\log\left(1-|\alpha|\right)\Big)\] for some $\widetilde{C} \in {\mathbb{R}}$, and we restrict $\alpha$ to take values in $(0, 1)$. The choice of $\widetilde{C}$ depends on the researcher. It corresponds to the class prior in PU learning and plays a role that differs from the parameter $C$ in the other loss functions. The derivative of $g^{\text{PU}}(\alpha)$ with respect to $\alpha$ is given as
Under this choice of $g$, the Bregman divergence objective is given as follows: \[\text{BD}_{g^{\text{PU}}}\big(\alpha\big) \coloneqq \mathbb{E}\left[ - \widetilde{C}\log\left(1-|\alpha(X)|\right) - m\left(W, \widetilde{C}\operatorname{sign}(\alpha)\Big(\log\left(|\alpha(\cdot)|\right)-\log\left(1-|\alpha(\cdot)|\right)\Big)\right)\right].\]
Then, we estimate $\alpha_0$ by minimizing the empirical objective: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha \in {\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{PU}}}\big(\alpha\big) + \lambda J(\alpha), \] where
PU learning is a classical problem. For example, Lancaster1996casecontrolstudies studies this problem under a stratified sampling scheme Wooldridge2001asymptoticproperties. duPlessis2015convexformulation rediscovers this formulation and calls it unbiased PU learning. Kato2019learningfrom points out the relationship between PU learning and density ratio estimation, and Kato2021nonnegativebregman shows that PU learning is a special case of density ratio model fitting under a Bregman divergence. Our results further generalize these results. Note that PU learning in these settings and our setting is called case-control PU learning. There is also another formulation called censoring PU learning Elkan2008learningclassifiers. Kato2025puate considers ATE estimation in a PU learning setup and applies our method in their study.
This section formalizes the automatic regressor balancing phenomenon: under suitable loss--link choices and linear-in-parameters modeling, the generalized Riesz regression estimator satisfies (approximate) moment-balance conditions as a direct consequence of first-order optimality (KKT) for the regularized ERM problem.
A key point throughout is that balancing is not an additional constraint we impose. Rather, it is an implicit constraint that appears whenever we parametrize the score $u=\partial g\circ \alpha$ linearly and solve a convex (or approximately solved) empirical Bregman problem.
Let ${\bm{\phi}}=(\phi_1,\dots,\phi_p)^\top$ be basis functions ${\bm{\phi}}\colon{\mathcal{X}}\to{\mathbb{R}}^p$. We consider a linear-in-parameter model for the Riesz representer of the form
where $\zeta^{-1}$ is a (possibly $x$-dependent) inverse link function. This introduction of a link function is motivated by generalized linear models.
\paragraph{Examples.} For example, we can approximate the Riesz representer by \[\alpha_{{\bm{\beta}}}(X) \coloneqq {\bm{\phi}}(X)^\top {\bm{\beta}},\] which corresponds to using a linear link function for $\zeta^{-1}$. This linear specification can be applied in many settings, including ATE estimation and density ratio estimation.
We can improve estimation accuracy by incorporating additional modeling assumptions. For example, in ATE estimation, we can approximate the propensity score $e_0(Z) = P(D = 1\mid Z)$ by a logistic model, \[e_{{\bm{\beta}}}(Z) \coloneqq \frac{1}{1 + \exp\Big(- {\bm{\phi}}(Z)^\top {\bm{\beta}}\Big)},\] where ${\bm{\phi}}\colon {\mathcal{Z}} \to {\mathbb{R}}^p$ is a basis function, and ${\bm{\beta}}$ is the corresponding parameter. Note that in this case, we consider a basis function that receives $Z$ not $X$, or we can interpret that ${\bm{\phi}}(D, Z)$ only depends on $Z$ and is independent of $D$. Plugging this propensity score model into the Riesz representer for ATE, we can approximate the Riesz representer by \[\alpha^{\text{ATE}}_{{\bm{\beta}}}(X) \coloneqq \frac{D}{e_{{\bm{\beta}}}(Z)} - \frac{1 - D}{1 - e_{{\bm{\beta}}}(Z)}.\] In such cases, we define $\zeta^{-1}$ so that
Note that we can also model the Riesz representer as
by including $D$ in the basis function. The choice of basis functions depends on the heterogeneity of $\gamma_0(X)$. As we discuss in Section (ref), if $\gamma_0(x)$ is constant for all $x$, ((ref)) may be more appropriate. In contrast, if $\gamma_0(x)$ varies across $x$, ((ref)) may be more appropriate.
Similarly, in covariate shift adaptation, we can model the density ratio as \[\alpha^{\text{CS}}_{{\bm{\beta}}}(X) = \exp\Big(- {\bm{\phi}}(X)^\top {\bm{\beta}}\Big).\]
Automatic balancing arises when the score \[ u_{{\bm{\beta}}}(x)\coloneqq (\partial g)\big(\alpha_{{\bm{\beta}}}(x)\big) \] is linear in the coefficients ${\bm{\beta}}$. To express this cleanly, define feature functions \[ \widetilde\phi_j(x)\coloneqq \widetilde g\left(x,\phi_j(x)\right), \qquad j=1,\dots,p, \] and the linear index
The canonical way to ensure ((ref)) is to choose the link so that
When $g$ is strictly convex and differentiable on its domain, ((ref)) is equivalent to $\alpha_{{\bm{\beta}}}(x)=(\partial g)^{-1}\big(u_{{\bm{\beta}}}(x)\big)$ (possibly branchwise; see Section (ref)).
\paragraph{Interpretation.} Equation ((ref)) says that ((ref)) is best viewed as a generalized linear model in dual coordinates $u=\partial g(\alpha)$. The balancing statements below will be written in terms of the features $\widetilde\phi_j$ that index $u_{{\bm{\beta}}}$.
Recall the empirical Bregman objective (Section (ref))
We estimate ${\bm{\beta}}$ by penalized ERM
and set $\widehat{\alpha}\coloneqq \alpha_{\widehat{{\bm{\beta}}}}$.
\paragraph{Imbalance Gap Functional.} For each $j=1,\dots,p$, let us define the sample imbalance gap functional as
We also define
The quantity $\widehat\Delta_j(\alpha)$ measures mismatch between the weighted empirical moment of $\widetilde\phi_j$ and the empirical target moment induced by $m$.
\paragraph{Interpretation and limitations.} Theorem (ref) is a first-order optimality statement on the training sample used in ((ref)). If cross fitting is used, the same equalities and inequalities generally do not hold on the held-out fold, and imbalance becomes a generalization/diagnostic object rather than an exact constraint. Likewise, if ((ref)) is solved only approximately, or the model is nonconvex, then ((ref)) holds only up to an optimization residual.
\paragraph{When does a constrained “balancing program” coincide with a true dual?} In the linear-score settings emphasized here, for example, SQ-Riesz with a linear link or UKL-Riesz with a compatible log-type link, the map ${\bm{\beta}}\mapsto \widehat{\mathrm{BD}}_g(\alpha_{{\bm{\beta}}})$ is convex, and the penalty $\|{\bm{\beta}}\|_a^a$ is convex for $a\ge 1$. Under standard constraint qualifications, one can derive an explicit Fenchel/Lagrange dual whose constraints are exactly moment-balance inequalities, cf.\ Proposition (ref).
\paragraph{When should the “dual” be interpreted only as a KKT characterization?} Outside the convex or linear regime, for example, when $\alpha$ is represented by a neural network and trained by nonconvex ERM, a constrained balancing program is best interpreted as a KKT-style characterization rather than a literal dual. In such cases, the fitted $\widehat\alpha$ need not satisfy ((ref))--((ref)) exactly due to (i) nonconvexity and (ii) optimization error. Practically, $\max_j|\widehat\Delta_j(\widehat\alpha)|$ is best viewed as a diagnostic of approximate balance.
When ((ref)) is convex, generalized Riesz regression can be viewed as selecting the minimum-$g$ weights among approximately balancing solutions. We state the cleanest form for $\ell_1$ regularization, which yields explicit moment constraints.
Problem ((ref)) is a generalized balancing program, with the generator $g$ determining the geometry of the weights. Exact balance at $\lambda=0$ corresponds to feasibility of the constraints, which need not hold in finite samples or under domain restrictions on $\alpha$; see Section (ref) for a detailed feasibility discussion.
\paragraph{Loss invariance under exact balance.} When exact balance holds on the same dictionary, the final identity ((ref)) depends on balancing, not on which strictly convex generator $g$ selected $\widehat{\alpha}$ among feasible solutions. Thus, in the exact-balance regime, different loss--link pairs that lead to the same balancing equations yield the same RW estimator on that training sample. This “loss invariance” disappears in inexact-balance regimes and is a primary reason loss choice matters in practice (discussed below).
\paragraph{Exact balancing is a sample property for the chosen features.} Even when $\lambda=0$ and $\mathcal{F}_0\neq\emptyset$, exact balancing does not imply $\widehat\alpha=\alpha_0$ pointwise; it only implies moment matching on the chosen dictionary $\{\widetilde\phi_j\}_{j=1}^p$.
\paragraph{Cross fitting breaks exact sample balance.} Exact balancing is derived from KKT conditions on the same sample used to fit $\widehat\alpha$. If $\widehat\alpha$ is estimated on a training fold and evaluated on a separate fold, the evaluation-fold imbalance is generally nonzero. In cross-fitted inference, imbalance is best treated as a diagnostic that should generalize.
\paragraph{Nonconvex ERM does not guarantee KKT satisfaction.} When $\widehat\alpha$ is obtained by nonconvex optimization, for example, neural networks, local minima and early stopping can prevent exact KKT satisfaction even on the training sample. This does not invalidate the objective, but it changes balancing from a theorem-level implication to an empirical property that can be checked.
To make feasibility explicit, define the feasible set of balancing weights at tolerance level $\lambda\ge 0$:
where $\mathrm{dom}(g)$ denotes the domain of the generator, including any sign restrictions induced by the loss--link choice. By construction, $\mathcal{F}_{\lambda_2}\supseteq \mathcal{F}_{\lambda_1}$ whenever $\lambda_2\ge \lambda_1$, so relaxing $\lambda$ weakly enlarges the feasible set.
\paragraph{Exact balancing at $\lambda=0$ may be infeasible.} Exact sample balancing corresponds to $\mathcal{F}_0\neq\emptyset$. Whether $\mathcal{F}_0$ is nonempty depends on the feature map, the sample size, and the domain restrictions on $\alpha$. In the unconstrained case $\mathrm{dom}(g)={\mathbb{R}}$, exact balancing is a system of $p$ linear equations in $n$ unknowns and is generically solvable when $p\le n$ and the corresponding design matrix has full rank. In contrast, if $\mathrm{dom}(g)$ enforces constraints such as $\alpha_i>0$, for example, density ratios, or branchwise sign restrictions, for example, ATE, feasibility additionally requires that the target moments lie in an appropriate convex cone, and infeasibility can occur even when $p\le n$.
\paragraph{Implications for the primal problem.} When the balancing formulation is a true convex dual of the primal, Proposition (ref), infeasibility at $\lambda=0$ signals that the corresponding unregularized primal problem is ill-behaved: it may fail to attain a minimizer or may drive parameters toward the boundary of $\mathrm{dom}(g)$, producing extreme weights (overfitting). Regularization remedies this by enlarging $\mathcal{F}_\lambda$ or enforcing coercivity of the objective.
Condition ((ref)) is most transparently satisfied by choosing the link as a (branchwise) inverse derivative of $g$. Below we list the main choices used in this paper, writing $C\ge 0$ for a shift parameter and $\xi(x)\in\{0,1\}$ for a known branch selector, for example, $\xi(x)=D$ in ATE:
The regularization parameter $\lambda$ plays two conceptually distinct roles:
Accordingly, in applications where $\lambda=0$ leads to unstable or infeasible balancing, it is natural to treat $\lambda$ as a tuning parameter that trades off balance and stability, while monitoring $\max_j|\widehat\Delta_j(\widehat\alpha)|$ as an interpretable diagnostic.
This appendix connects three perspectives on representer fitting: (i) the sieve Riesz equations in a Hilbert space, (ii) Bregman projection geometry, and (iii) the KKT balancing conditions in generalized Riesz regression.
One can estimate the Riesz representer by directly solving the equations implied by the Riesz representation theorem. This approach is a variant of moment matching, often called the sieve Riesz representer, and it has been used in the semiparametric and sieve inference literature, see, e.g., Chen2015sievesemiparametric and Chen2015sievewald. This viewpoint is closely connected to exact balancing. Understanding this approach clarifies how the choice of the loss function $g$ in the Bregman divergence determines the geometry of the representer fit.
For the linear functional $\gamma\mapsto {\mathbb{E}}\left[m(W,\gamma)\right]$, the Riesz representer $\alpha_0\in{\mathcal{H}}$ satisfies ${\mathbb{E}}\left[m(W,\gamma)\right]=\langle \alpha_0,\gamma\rangle$ for all $\gamma\in{\mathcal{H}}$. On a finite-dimensional sieve space $H_p\coloneqq\mathrm{span}\{\phi_1,\ldots,\phi_p\}$, the sieve Riesz representer $\alpha_p\in H_p$ is characterized by linear equations $\langle \alpha_p,\phi_j\rangle={\mathbb{E}}\left[m(W,\phi_j)\right]$. Our KKT balancing equations are the empirical counterpart, and the Bregman geometry viewpoint is developed in Appendix (ref).
\paragraph{Riesz representer as a linear Equation in a Hilbert space} Let ${\mathcal{H}}\coloneqq L_2(P_X)$ with inner product $\langle f,g\rangle \coloneqq {\mathbb{E}}\left[f(X)g(X)\right]$. For the linear map $\gamma\mapsto {\mathbb{E}}\left[m(W,\gamma)\right]$ (Section (ref)), the Riesz representation theorem yields $\alpha_0\in{\mathcal{H}}$ such that
Restricting to a finite-dimensional sieve space ${\mathcal{H}}_p\coloneqq \mathrm{span}\{\phi_1,\ldots,\phi_p\}$, the sieve Riesz representer $\alpha_p\in{\mathcal{H}}_p$ is characterized by
Writing $\alpha_p(x)={\bm{\phi}}(x)^\top{\bm{\beta}}$ with ${\bm{\phi}}\coloneqq (\phi_1,\ldots,\phi_p)^\top$, ((ref)) becomes the linear system
\paragraph{Bregman objectives, dual variables, and a common projection geometry} Recall the pointwise Bregman divergence \[ \mathrm{BD}^\dagger_g\left(\alpha_0(x)\mid \alpha(x)\right) \coloneqq g(\alpha_0(x)) - g(\alpha(x)) - \partial g(\alpha(x))\big(\alpha_0(x)-\alpha(x)\big), \] and the population target $\alpha^*\coloneqq \arg\min_{\alpha\in{\mathcal{H}}}{\mathbb{E}}\left[\mathrm{BD}^\dagger_g\left(\alpha_0(X)\mid \alpha(X)\right)\right]$. A first-order characterization of Bregman projections is: if ${\mathcal{H}}$ is convex and $\alpha^*$ is an interior minimizer, then
with equality along feasible smooth directions. Equation ((ref)) makes clear that all losses share the same underlying $L_2(P_X)$ geometry; what changes across losses is the dual coordinate $\partial g(\alpha)$.
A convenient reparameterization uses the convex conjugate $g^*$ and the dual variable
When $g$ is strictly convex and differentiable, the Fenchel--Young identity implies $g^*(u)=\alpha u-g(\alpha)$ at $u=\partial g(\alpha)$, and the population objective can be written as
up to an additive constant independent of $\alpha$.
\paragraph{Finite-dimensional dual models and KKT.} Consider a model class specified in dual coordinates as
(possibly branchwise to enforce sign restrictions). For the penalized empirical objective \[ \widehat{\bm{\beta}} \in \arg\min_{{\bm{\beta}}\in{\mathbb{R}}^p} \Big\{{\mathbb{E}}_n\left[g^*(u_{\bm{\beta}}(X))\right] - {\mathbb{E}}_n\left[m(W,u_{\bm{\beta}})\right] + \tfrac{\lambda}{a}\|{\bm{\beta}}\|_a^a\Big\}, \] the KKT conditions yield
where $\widehat{\alpha}=\alpha_{\widehat{\bm{\beta}}}$. When $\lambda=0$, ((ref)) reduces exactly to the empirical sieve Riesz equations, the finite-sample analogue of ((ref)).
We now specialize the generic “automatic balancing” equations to the familiar covariate balancing language in treatment-effect problems, and we clarify an important loss--link implication:
\paragraph{ATE balancing with $Z$-only features recovers covariate balance.} In the ATE setting $X=(D,Z)$ with representer $\alpha^{\mathrm{ATE}}(D,Z)=D/e_0(Z)-(1-D)/(1-e_0(Z))$, a common choice in practice is to model the propensity using only covariates $Z$ and to take $\widetilde\phi_j(X)=\phi_j(Z)$. In this case, $m(W,\phi_j(Z))=0$ for ATE-type $m$, because $\phi_j(Z)$ does not vary with $D$, so the exact-balance equations $\widehat\Delta_j(\widehat\alpha)=0$ become \[ \frac{1}{n}\sum^n_{i=1} \widehat\alpha(D_i,Z_i)\,\phi_j(Z_i)=0, \qquad j=1,\dots,p, \] i.e. \[ \frac{1}{n}\sum^n_{i=1} \frac{D_i}{\widehat e(Z_i)}\,\phi_j(Z_i) = \frac{1}{n}\sum^n_{i=1} \frac{1-D_i}{1-\widehat e(Z_i)}\,\phi_j(Z_i), \qquad j=1,\dots,p. \] These are exactly the usual covariate balancing conditions: the weighted covariate moments in the treated and control groups match for the chosen dictionary.
\paragraph{Sigmoid propensity modeling implies a log-link representer.} If we commit to the logistic propensity model $e_{{\bm{\beta}}}(Z)=\Lambda({\bm{\phi}}(Z)^\top{\bm{\beta}})$, then the induced ATE representer is exactly the branchwise log-link form ((ref)). Therefore, the dual-score linearity requirement ((ref)) reduces to a compatibility condition between (i) the chosen generator $g$ and (ii) the log-link induced by the sigmoid propensity model.
\paragraph{UKL-Riesz is loss--link compatible for ATE under a sigmoid propensity model.} Under the shifted UKL generator with $C=1$, the dual score $\partial g^{\mathrm{UKL}}(\alpha)$ becomes linear in the logistic index, indeed, equal to $-{\bm{\phi}}(Z)^\top{\bm{\beta}}$ on both branches. Consequently, UKL-Riesz with the sigmoid-induced log-link enjoys automatic balancing through Theorem (ref). A concise derivation is given in Appendix (ref).
\paragraph{Why logistic MLE (BKL) is not the ATE-balancing choice under the same sigmoid link.} Standard logistic regression MLE corresponds to the Bernoulli likelihood and therefore to the BKL-type generator. Under the same sigmoid-induced ATE representer ((ref)), the BKL dual score is not linear in the index ${\bm{\phi}}(Z)^\top{\bm{\beta}}$, so the KKT balancing mechanism does not align with ATE balancing. This is consistent with the “estimand-driven loss selection” message of Zhao2019covariatebalancing: under logistic modeling, the Bernoulli likelihood corresponds to overlap-style weighting objectives rather than ATE. Within our framework, the resolution is simple: either (i) change the loss to UKL to match the sigmoid-induced log-link for ATE balancing, or (ii) keep the Bernoulli likelihood but reinterpret the target estimand/weighting scheme accordingly.
\paragraph{Why restricting propensity modeling to $Z$ can be undesirable for orthogonalization.} Using $Z$-only features is natural for propensity modeling and yields classical covariate balance, but it can be misaligned with the goal of automatic Neyman orthogonalization: the regression function $\gamma_0(X)={\mathbb{E}}\left[Y\mid D,Z\right]$ is a function of $X=(D,Z)$, and approximating it well typically requires treatment-specific components, for example, separate bases for $D=1$ and $D=0$, or interactions $D\cdot \phi(Z)$. If the balancing dictionary only depends on $Z$, then exact balancing only enforces orthogonality against a narrow outcome model class and may fail under heterogeneous treatment effects. The detailed orthogonality implications and remedies, for example, richer dictionaries, augmentation/TMLE, and inexact-balance regimes, are developed in Section (ref).
This section explains why balancing is not merely a diagnostic for weighting, but a structural device that yields (approximate) Neyman orthogonality and controls the leading bias of debiased estimators. We first treat the benchmark case of exact balancing (typically $\lambda=0$ on the training sample), and then generalize to inexact balancing (regularization $\lambda>0$, cross fitting, or optimization residuals). We also highlight the special linear--linear case in which balancing collapses to regression (OLS/ridge), and we compare the targeting directions of generalized Riesz regression and TMLE.
In Section (ref), we defined the following four estimators for the estimation of the parameter of interest:
In Section (ref), we defined the Neyman error as \[ \text{NeymanError} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}(X_i)\big(Y_i-\widehat{\gamma}(X_i)\big) + m\left(W_i,\widehat{\gamma}\right) - m\left(W_i,\gamma_0\right)\Big), \] which is the discrepancy between an estimated Neyman orthogonal score and the true Neyman orthogonal score.
In Section (ref), for functions $f\colon {\mathcal{X}} \to {\mathbb{R}}$ we defined the imbalance gap as \[ \widehat\Delta\left(\alpha,f\right) \coloneqq \frac{1}{n}\sum^n_{i=1} \Big( \alpha(X_i)f(X_i) - m\left(W_i,f\right) \Big). \]
Let $\varepsilon_i\coloneqq Y_i-\gamma_0(X_i)$. Using $\widehat\Delta$, the sample Neyman error admits the decomposition
The last term is a weighted noise term. Under cross fitting (or other sample-splitting schemes), conditional on the training data used to estimate $\widehat{\alpha}$, the weights are independent of $\varepsilon_i$, so this term has conditional mean zero. Therefore, the leading deterministic component of the sample drift is \[ \text{NeymanError}^\dagger \coloneqq \widehat\Delta\left(\widehat{\alpha},\gamma_0\right) - \widehat\Delta\left(\widehat{\alpha},\widehat{\gamma}\right) = -\widehat\Delta\left(\widehat{\alpha},\widehat{\gamma}-\gamma_0\right), \] which we call the pseudo Neyman error.
\paragraph{Imbalance control over a working regression class.} Suppose that $\widehat\alpha$ is obtained by generalized Riesz regression with the KKT bounds in Theorem (ref), and suppose that $\widehat{\gamma}$ is estimated in a $p$-dimensional working space, \[ \widehat{\gamma}(x)={\bm{\phi}}(x)^\top\widehat{{\bm{\rho}}}, \qquad {\bm{\phi}}(x)=\left(\phi_1(x),\dots,\phi_p(x)\right)^\top. \] If the balancing dictionary coincides with the basis (so $\widetilde\phi_j=\phi_j$), then linearity yields \[ \widehat\Delta\left(\widehat{\alpha},\widehat{\gamma}\right) = \sum_{j=1}^p \widehat{\rho}_j\,\widehat\Delta_j\left(\widehat{\alpha}\right). \] Consequently, Theorem (ref) implies the bound \[ \left|\widehat\Delta\left(\widehat{\alpha},\widehat{\gamma}\right)\right| \le \lambda\sum^p_{j=1}\left|\widehat{\rho}_j\right|\left|\widehat\beta_j\right|^{a-1}, \] and, in particular, for $a=1$, \[ \left|\widehat\Delta\left(\widehat{\alpha},\widehat{\gamma}\right)\right| \le \lambda\|\widehat{{\bm{\rho}}}\|_1. \]
\paragraph{Without approximation error for $\gamma_0$.} If $\gamma_0$ belongs to the linear space spanned by ${\bm{\phi}}$, there exists ${\bm{\rho}}_0\in{\mathbb{R}}^p$ such that $\gamma_0(x)={\bm{\phi}}(x)^\top{\bm{\rho}}_0$. In this case, the same argument yields \[ \left|\widehat\Delta\left(\widehat{\alpha},\gamma_0\right)\right| \le \lambda\sum^p_{j=1}\left|\rho_{0,j}\right|\left|\widehat\beta_j\right|^{a-1}. \] Thus, automatic regressor balancing directly controls the pseudo Neyman error by controlling the imbalance gap on the working space.
\paragraph{With approximation error for $\gamma_0$.} When $\gamma_0\notin\mathrm{span}\left\{\phi_1,\dots,\phi_p\right\}$, balancing should be interpreted as shrinking the orthogonality defect on a working regression class. If $\gamma_0$ is well approximated by some $\gamma_\phi\in\mathrm{span}\left\{\widetilde\phi_j\right\}$ and $\widehat\alpha$ approximately balances the same span, then \[ \frac{1}{n}\sum^n_{i=1}\widehat\alpha(X_i)\gamma_\phi(X_i)\approx \frac{1}{n}\sum^n_{i=1}m\left(W_i,\gamma_\phi\right) \] reduces the component of the drift driven by the working approximation, leaving only the approximation residual $\gamma_0-\gamma_\phi$ and the representer error $\widehat\alpha-\alpha_0$. This is the sense in which balancing is an automatic Neyman error control mechanism: it targets the part of the score error that is linear in the regression approximation.
This second-order structure is discussed in Zhao2019covariatebalancing for the case where UKL-Riesz regression with a basis function ${\bm{\phi}}\colon {\mathcal{Z}} \to {\mathbb{R}}^p$ is used in ATE estimation. Our results generalize that finding to more general models, losses, and parameters of interest.
\paragraph{Neyman error decomposition and the role of imbalance.} At the sample level, the leading error of RW, relative to a target orthogonal score, is governed by imbalance. For $\gamma\in\Gamma_p$ written as $\gamma=\sum_{j=1}^p\rho_j\widetilde\phi_j$,
Therefore, \[ \left| \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)\gamma(X_i) - \frac{1}{n}\sum^n_{i=1}m\left(W_i,\gamma\right) \right| \le \|{\bm{\rho}}\|_1\max_{j\le p}\left|\widehat\Delta_j\left(\widehat{\alpha}\right)\right|. \] Combining this with Theorem (ref) shows that $\lambda$ directly controls the size of the sample orthogonality defect over coefficient-bounded outcome models. This is one concrete sense in which generalized Riesz regression performs automatic Neyman error control and minimization: it chooses $\widehat{\alpha}$ so that the empirical Riesz equations hold approximately on the dictionary, thereby controlling the leading drift term over the corresponding working class.
\paragraph{Exact balancing implies exact orthogonality on a working regression space.} Let ${\bm{\phi}}(X)=\left(\phi_1(X),\dots,\phi_p(X)\right)^\top$ be the basis used to estimate the representer, and define the working linear space \[ \Gamma_\phi \coloneqq \Big\{ \gamma:{\mathcal{X}}\to{\mathbb{R}} \mid \gamma(x)={\bm{\phi}}(x)^\top{\bm{\rho}}\text{ for some }{\bm{\rho}}\in{\mathbb{R}}^p \Big\}. \] Define the sample imbalance vector \[ \Delta\left(\widehat\alpha\right) \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat\alpha(X_i){\bm{\phi}}(X_i)-m\left(W_i,{\bm{\phi}}\right)\Big) \in{\mathbb{R}}^p, \] where $m\left(W,{\bm{\phi}}\right)\coloneqq\left(m\left(W,\phi_1\right),\dots,m\left(W,\phi_p\right)\right)t(W,\phi_p\right)}^\top$. Then for any $\gamma_\rho(x)={\bm{\phi}}(x)^\top{\bm{\rho}}\in\Gamma_\phi$, linearity gives the identity
Hence exact balance $\Delta\left(\widehat\alpha\right)=0$ implies \[ \frac{1}{n}\sum^n_{i=1}\widehat\alpha(X_i)\gamma(X_i) = \frac{1}{n}\sum^n_{i=1}m\left(W_i,\gamma\right) \qquad \forall \gamma\in\Gamma_\phi, \] which is precisely the empirical Riesz equation restricted to the sieve space $\Gamma_\phi$.
\paragraph{RW equals the orthogonal-score estimator on the working space.} If $\widehat\Delta_j\left(\widehat\alpha\right)=0$ for all $j$ (e.g., $\lambda=0$ in the convex and feasible regime), then by linearity it holds that
Consequently, for any $\gamma\in\Gamma_{\phi}$,
Identity ((ref)) is deterministic and requires no asymptotics: RW equals the orthogonal-score estimator for any $\gamma$ in the working space.
Equation ((ref)) shows that, on the working regression space spanned by ${\bm{\phi}}$, exact balancing enforces Neyman orthogonality. We refer to this mechanism as automatic Neyman orthogonalization.
\paragraph{Choice of loss--link pairs and automatic Neyman orthogonalization.} To attain automatic Neyman orthogonalization, two requirements must be met:
We emphasize two related points:
This suggests the following procedure:
\paragraph{What this does and does not mean.} Equation ((ref)) shows that the first-order plug-in bias coming from approximating $\gamma_0$ within $\Gamma_\phi$ is exactly removed once $\Delta\left(\widehat\alpha\right)=0$ and $\gamma_0\in\Gamma_\phi$.
Theorem (ref) is a working-model statement: exact balancing yields exact orthogonalization on the chosen span $\Gamma_\phi$. If $\gamma_0\notin\Gamma_\phi$, orthogonality does not hold exactly and approximation error must be controlled. Moreover, under cross fitting, balance constraints hold on training folds while the score is evaluated on held-out folds, so the finite-sample identity does not transfer directly.
In practice, balance is inexact due to regularization, cross fitting, finite-sample infeasibility, or optimization error. Let $\widehat\Delta_j\left(\widehat\alpha\right)$ be the imbalance defined in ((ref)). For any $\gamma(x)=\sum_{j=1}^p \rho_j \widetilde\phi_j(x)$ in the working span, linearity of $m$ yields the identity
Combining ((ref)) with Theorem (ref) gives explicit bounds. For example, with lasso ($a=1$), \[ \left| \frac{1}{n}\sum^n_{i=1}\widehat\alpha(X_i)\gamma(X_i) - \frac{1}{n}\sum^n_{i=1}m\left(W_i,\gamma\right) \right| \le \lambda \|{\bm{\rho}}\|_1 \qquad\text{for all }\gamma=\sum_j\rho_j\widetilde\phi_j. \] Thus, $\lambda$ directly controls the worst-case orthogonality defect over the working linear space.
\paragraph{A simple bound on imbalance over the working regression span.} Define the imbalance operator \[ \widehat\Delta\left(\alpha,f\right) \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\alpha(X_i)f(X_i)-m\left(W_i,f\right)\Big). \] When $f=\widetilde\phi_j$, this is exactly ((ref)). Under the KKT bounds, $\max_j\left|\widehat\Delta\left(\widehat\alpha,\widetilde\phi_j\right)\right|$ is controlled by $\lambda$ via Theorem (ref). If $\gamma=\sum_{j=1}^p\rho_j\widetilde\phi_j$, then
and analogous bounds hold for $a>1$ using ((ref)). Thus, inexact balancing yields explicit control of the score-equation error on the chosen span, with $\lambda$ acting as a tolerance parameter.
\paragraph{Inexact balancing, regularization, and cross-fitting.} Inexact balancing arises in at least two ubiquitous situations: $\lambda>0$ in the representer fit (stability or feasibility), and cross fitting, where exact KKT balance holds only on the training fold. In both cases, ((ref)) suggests a practical workflow: choose the dictionary and regularization so that $\gamma_0$ is well approximated by the working space and imbalance generalizes, meaning it remains small on held-out folds, and then use ARW or TMLE when needed for inference.
A particularly instructive regime is when both the representer and regression are modeled in the same linear span. This is the regime behind the augmented balancing equivalences emphasized by BrunsSmith2025augmentedbalancing.
Let $\Phi\in{\mathbb{R}}^{n\times p}$ be the design matrix with $\Phi_{ij}=\phi_j(X_i)$, and let $b\in{\mathbb{R}}^p$ be the target moment vector with \[ b_j=\frac{1}{n}\sum^n_{i=1}m\left(W_i,\phi_j\right). \] Exact balance in the linear representer model $\widehat\alpha(X)={\bm{\phi}}(X)^\top\widehat{\bm{\beta}}$ means $\frac{1}{n}\Phi^\top\Phi\,\widehat{\bm{\beta}}=b$ (when solvable). Let $\widehat{\bm{\rho}}^{\mathrm{OLS}}=\left(\Phi^\top\Phi\right)^\dagger\Phi^\top Y$ and $\widehat\gamma^{\mathrm{OLS}}(x)={\bm{\phi}}(x)^\top\widehat{\bm{\rho}}^{\mathrm{OLS}}$.
\paragraph{Why this matters.} Proposition (ref) shows that, in the linear--linear exact-balance regime, weighting and regression are two views of the same normal equations. This helps explain why some balancing-weight estimators collapse to regression estimators and why specific regularization choices can yield a single effective regression estimator. See BrunsSmith2025augmentedbalancing for general augmented-balancing identities and ridge and kernel-ridge special cases. This phenomenon is closely related to OLS is doubly robust arguments in Robins2007commentperformance. It is not a failure of orthogonalization: once orthogonality is enforced on the entire working regression space, the RW estimator is already the orthogonal estimator for that space.
\paragraph{Augmented balancing weights collapse to regression estimators.} BrunsSmith2025augmentedbalancing shows that when both models are linear, the augmented estimator is equivalent to a single linear regression whose coefficients are affine combinations of the original regression coefficients and unpenalized OLS coefficients. Under certain regularization choices, the augmented estimator collapses to OLS itself. This makes explicit that aggressive balancing can implicitly undo outcome regularization, which can reduce functional bias at the cost of higher variance.
\paragraph{Kernel ridge special case.} In RKHS settings, this algebraic collapse yields an interpretable undersmoothing rule. If both the outcome regression and the representer are fit by ridge in a common RKHS, the augmented estimator can be written as a single undersmoothed ridge regression with an effective regularization level, rather than as two separate nuisance estimators.
\paragraph{Special case: linear Riesz and linear regression under inexact balance.} When both $\alpha$ and $\gamma$ are fitted in the same linear span with $\ell_2$ penalties, BrunsSmith2025augmentedbalancing shows that the resulting ARW estimator can be written as a single ridge-type regression estimator with an effective regularization parameter. This makes explicit that augmentation can be interpreted as a data-dependent undersmoothing rule: regularization choices that are prediction-optimal for $\gamma$ can be too aggressive for inference on $\theta_0$, and the representer fit provides a principled correction.
If we cross-fit $\widehat\alpha$, the balancing identity typically holds on the training folds but not on the evaluation folds where the score is computed. In that regime, exact balance does not transfer as a finite-sample identity, and the drift term in ((ref)) must be controlled by convergence rates, as in standard DML logic. This is precisely where the augmented estimator $\widehat\theta^{\mathrm{ARW}}$ (or TMLE) is the safer default for inference.
\paragraph{Cross fitting and the Donsker trade-off.} Exact automatic regressor balancing can yield sharp finite-sample identities, but efficiency arguments without cross fitting typically require a Donsker-type condition. Cross fitting relaxes Donsker requirements, but it also turns balancing from an identity into a generalization property, so imbalance becomes a diagnostic rather than a deterministic guarantee. This trade-off must be addressed carefully in practice.
\paragraph{Basis functions depending only on $Z$.} In ATE estimation with standard propensity-score modeling, it is common to use basis functions depending only on $Z$, that is, ${\bm{\phi}}\colon {\mathcal{Z}} \to {\mathbb{R}}^p$. Under this choice, for ATE-type $m$ one typically has $m\left(W,\phi_j\right)=0$, so automatic covariate balancing controls only the weighted moments $\frac{1}{n}\sum^n_{i=1}\widehat\alpha(X_i)\phi_j(Z_i)$. Automatic Neyman orthogonalization of RW without regression adjustment is then strongest only for outcome-model components that lie in the same $Z$-only working span. If treatment effects are heterogeneous, the relevant regression components generally depend on $X=\left(D,Z\right)$, so one should use richer dictionaries, such as treatment-specific features $\left(D\phi(Z),(1-D)\phi(Z)\right)$ or other $X$-dependent bases, even though the Riesz representer corresponds to inverse propensity weights.
TMLE is another promising approach in debiased machine learning vanderLaan2006targetedmaximum. TMLE adds a perturbation to an initial estimate of $\gamma_0$ and targets an empirical score equation.
To make the contrast concrete, consider the sample Neyman error written as \[ \text{NeymanError} = \frac{1}{n}\sum^n_{i=1}\Big( \underbrace{\widehat{\alpha}(X_i)\big(Y_i-\widehat{\gamma}(X_i)\big)}_{\text{= $(\star)$}} + m\left(W_i,\widehat{\gamma}\right) - m\left(W_i,\gamma_0\right) \Big). \] TMLE updates $\widehat{\gamma}$ so that the empirical mean of $(\star)$ becomes zero, that is, it enforces \[ \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)\big(Y_i-\widehat{\gamma}^{(1)}(X_i)\big)=0. \]
In contrast, generalized Riesz regression targets the representer, and it controls the empirical drift term induced by imbalance. Equivalently, using $\widehat{\alpha}(X_i)Y_i=\widehat{\alpha}(X_i)\big(Y_i-\widehat{\gamma}(X_i)\big)+\widehat{\alpha}(X_i)\widehat{\gamma}(X_i)$, we can also rewrite NeymanError as \[ \text{NeymanError} = \frac{1}{n}\sum^n_{i=1}\Big( \widehat{\alpha}(X_i)Y_i + \underbrace{m\left(W_i,\widehat{\gamma}\right)-\widehat{\alpha}(X_i)\widehat{\gamma}(X_i)}_{\text{= $(\star\star)$}} - m\left(W_i,\gamma_0\right) \Big). \] Generalized Riesz regression works to make the sample average of $(\star\star)$ small by enforcing empirical Riesz equations, approximately, on a prescribed working class.
\paragraph{TMLE and generalized Riesz regression target different nuisances.}
These are complementary. In practice, one may estimate $\widehat{\alpha}$ by generalized Riesz regression and then apply a TMLE-type fluctuation to $\widehat{\gamma}$.
\paragraph{Automated and end-to-end variants.} Recent automatic debiasing work shows how to estimate $\alpha_0$ without requiring a closed-form representer formula, by directly minimizing an objective whose expectation is minimized at $\alpha_0$ and by using black-box machine learners. This includes Riesz regression and its extensions to generalized regressions Chernozhukov2022automaticdebiased, as well as multitask and targeted-regularization architectures Chernozhukov2022riesznet. These approaches move beyond finite-dimensional moment constraints: balancing becomes a first-order condition at the population level, and practical diagnostics such as imbalance, weight tails, and optimization residuals become central.
\paragraph{A useful comparison point.} Conceptually, TMLE automates orthogonality by updating $\gamma$ given $\alpha$, while generalized Riesz regression automates orthogonality by fitting $\alpha$ against a rich outcome space and then optionally augmenting. Hybrids exist, for example, adding TMLE-like targeted regularization terms while fitting $\alpha$ and $\gamma$ jointly, and they can be interpreted within our loss--link framework via ((ref)).
We close by emphasizing a modeling principle: the best feature map for balancing is the one that matches the approximation space for the regression components that matter for $\theta_0$.
\paragraph{Two complementary routes to a small remainder.} Orthogonal estimators have leading remainder controlled by \[ \|\widehat{\gamma}-\gamma_0\|_2\|\widehat{\alpha}-\alpha_0\|_2 \] under cross-fitting or Donsker-type conditions. This suggests two complementary strategies:
\paragraph{ARW as correction for estimation and approximation bias.} In modern applications, $\widehat{\gamma}$ is typically regularized (lasso, ridge, kernels, neural nets), so it can be biased for $\gamma_0$, and this bias can propagate to $\theta_0$ through $m\left(W,\widehat{\gamma}\right)$. The orthogonal ARW form corrects this bias to first order, leaving a second-order product remainder controlled by ((ref)). This clarifies the complementarity between outcome modeling and representer modeling.
\paragraph{A practical warning about misalignment.} If one balances only low-order moments of $Z$ but uses a rich learner for $\gamma_0(D,Z)$ with interactions and nonlinearities, then the balanced span and the regression approximation space are misaligned. Orthogonality defects can persist even if covariate balance diagnostics look excellent. This is a strong argument for regressor balancing on $X=\left(D,Z\right)$, or at least balancing a basis rich enough to approximate both treatment-specific regression components.
\paragraph{RKHS balancing and functional balance.} Kernel-based balancing methods can be interpreted as balancing an infinite-dimensional feature map, via RKHS mean embeddings, thereby avoiding ad hoc basis selection while controlling a rich class of covariate functions. This motivates implementing generalized Riesz regression in RKHS models, where balance becomes functional balance and the representer theorem yields finite-dimensional computations.
\paragraph{Representer-centric view and robustness under misspecification.} Complementing regression-centric undersmoothing, Singh2024kernelridge studies kernel ridge estimation of the Riesz representer and analyzes its generalization error in population $L_2$. A related message is robustness under misspecification in orthogonal constructions: for an orthogonal estimator built from a regression approximation and a representer approximation, a double-robust type statement persists in the sense that consistency can hold if either nuisance is sufficiently accurate, up to stochastic terms controlled by their convergence rates.
\paragraph{Implicit restriction in covariate balancing.} In ATE applications, it is common to set ${\bm{\phi}}={\bm{\phi}}(Z)$ so that balancing is interpreted as covariate balance. But if the regression function $\gamma_0(D,Z)$ contains heterogeneous components that are poorly approximated by a $Z$-only linear span, then ((ref)) does not directly control the orthogonality defect for those components. This motivates richer feature choices, such as separate dictionaries by treatment arm or $D\times f(Z)$ interactions, when the objective is debiasing rather than propensity interpretation.
This section provides advice for practitioners on how to use generalized Riesz regression. In generalized Riesz regression, the choice of the Bregman loss function $g$ is not merely a training loss choice. Appropriate loss--link pairs enforce regressor balancing automatically, which also reduces the error of the true Neyman orthogonal score. In addition, $g$ determines (i) the geometry of the fitted representer, (ii) the induced shape constraints through the link, and (iii) which pseudo-true limit is targeted under misspecification.
We have shown that generalized Riesz regression includes a broad class of objective functions for Riesz representer estimation. This section summarizes how we choose basis functions, link functions, loss functions, and the final estimator of the parameter of interest. These elements are closely related and should be chosen jointly from the following perspectives:
\paragraph{Exact balancing, inexact balancing, and when loss choice matters.} As shown in BrunsSmith2025augmentedbalancing, in some situations the choice of loss function in generalized Riesz regression does not affect the final estimator. If we do not use cross-fitting, $\lambda=0$, and exact balancing is feasible on the training sample, then the RW identity ((ref)) implies that the final estimator is loss-invariant on that sample. Moreover, in the linear--linear regime this estimator collapses to a regression-based estimator on the same working space, such as the sample average of the OLS estimator of $\gamma_0$ (Section (ref)). Under specific combinations of regularization for the representer and the regression nuisance, the final estimator can simplify further. For example, if we use an $\ell_2$-penalty for both estimators, the final estimator becomes the sample average of the ridge estimator of $\gamma_0$ Singh2024kernelridge. In contrast, under inexact balancing, different generators $g$ generally select different approximately balancing solutions in the same approximation space, so the loss choice can affect the final estimator (Figure (ref)).
\paragraph{Automatic regressor balancing determines the choice of loss and link.} From the viewpoint of constructing a Neyman-orthogonal final estimator, we aim to exploit automatic regressor balancing. As discussed in Section (ref), automatic balancing is a KKT implication when we model the dual coordinate $u=\partial g\circ\alpha$ linearly. Therefore, the loss--link pair should be chosen so that the link makes $u$ linear in parameters. Holding the link fixed while changing $g$ typically breaks the compatibility needed for ((ref)).
\paragraph{Sensitivity viewpoint for the loss--link pair.} The loss--link pair also controls sensitivity of the representer fit to the data, including outliers and tail observations. For ATE estimation, the following combinations are especially interpretable:
Related discussions in density ratio estimation appear in Menon2016linkinglosses and Zellinger2025binarylosses.
\paragraph{Regularization and the choice of final estimator.} If we do not use cross-fitting, $\lambda=0$, and exact balancing is feasible, then RW and ARW are equivalent on the training sample. If $\lambda>0$, RW and ARW generally differ. From the viewpoint of Neyman orthogonality under inexact balancing and cross-fitting, the ARW estimator or TMLE is the safer default.
\paragraph{ARW estimator and TMLE.} The ARW estimator shifts the difficulty of semiparametric inference toward Riesz representer estimation, while TMLE shifts it toward regression function estimation (Section (ref)).
\paragraph{Choice of basis functions.} Ideally, the regression function $\gamma_0$ lies in the linear span of ${\bm{\phi}}(X)$, as discussed in Section (ref). Under certain conditions, if we are interested only in minimax rates, overlap can be mitigated via outcome-modeling viewpoints, as discussed in Section (ref).
At the population level, for a convex model class $\mathcal H\subset L_2(P_X)$, generalized Riesz regression targets a Bregman projection \[ \alpha_g^\ast \in \arg\min_{\alpha\in\mathcal H}\ {\mathbb{E}}\!\left[\mathrm{BD}^\dagger_g\!\left(\alpha_0(X)\mid \alpha(X)\right)\right]. \] When $\alpha_0\notin\mathcal H$, different generators generally induce different pseudo-true limits $\alpha_g^\ast$. Thus, choosing $g$ is analogous to choosing a likelihood or criterion in a parametric model: it determines which aspects of $\alpha_0$ are prioritized by the fit.
A key message from Section (ref) is that once we model the dual coordinate $u=\partial g\circ\alpha$ linearly, the resulting KKT conditions enforce approximate moment equations that are the empirical analogue of sieve Riesz equations. In that regime, $g$ selects among approximately balancing solutions through its induced Bregman geometry (Appendix (ref)).
The loss choice becomes particularly interpretable when paired with a link that makes $u=\partial g\circ\alpha$ linear in a parameter vector, since KKT then yields balancing constraints. Below we summarize three representative pairings used throughout the paper.
We first introduce the combination of the squared-loss generator and a linear link. Let \[ g^{\mathrm{SQ}}(\alpha)=\tfrac12(\alpha-C)^2, \qquad \partial g^{\mathrm{SQ}}(\alpha)=\alpha-C, \] and consider the affine linear representer model \[ \alpha_{{\bm{\beta}}}(X) = C+{\bm{\phi}}(X)^\top{\bm{\beta}}, \] where ${\bm{\phi}} \colon {\mathcal{X}} \to {\mathbb{R}}^p$ is a basis function. This yields $u_{{\bm{\beta}}}(x)=(\partial g^{\mathrm{SQ}})(\alpha_{{\bm{\beta}}}(x))={\bm{\phi}}(x)^\top{\bm{\beta}}$ and hence automatic balancing by KKT.
Under the linear model, if we use an $\ell_1$ penalty in SQ-Riesz regression, the dual formulation implies that the representer fit is equivalent to solving the constrained quadratic program
This matches the “stable balancing weights” formulation in Zubizarreta2015stableweights. When $\lambda = 0$, the constraint enforces exact balancing: \[ \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}_i \phi_j(X_i) = \frac{1}{n}\sum^n_{i=1} m(W_i,\phi_j), \qquad j = 1,\dots,p, \] where $\widehat{\alpha}_i=\alpha_{\widehat{{\bm{\beta}}}}(X_i)$. In ATE estimation, this becomes \[ \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}_i \phi_j(D_i, Z_i) = \frac{1}{n}\sum^n_{i=1}\Big(\phi_j(1, Z_i) - \phi_j(0, Z_i)\Big), \qquad j = 1,\dots,p. \] An advantage of linear models is that we can express the entire ATE estimation problem using a single linear working space, as shown by BrunsSmith2025augmentedbalancing. This pairing is often attractive when weight stability is paramount, for example under weak overlap, because the induced geometry is quadratic.
We next describe the pairing of the UKL divergence generator and a log link, which connects generalized Riesz regression to entropy-type balancing and exponential tilting. Let $\xi:{\mathcal{X}}\to\{0,1\}$ be a known branch indicator and define $s(X)\coloneqq 2\xi(X)-1\in\{-1,1\}$. For a constant $C\ge 0$, consider the representer model
This is a log link model in the sense that, conditional on $s(X)$, the linear index is recovered by taking logarithms of $|\alpha_{{\bm{\beta}}}(X)|-C$.
If we use the UKL generator $g^{\mathrm{UKL}}$ in generalized Riesz regression with the model ((ref)) and an $\ell_1$ penalty, the dual formulation yields an entropy-type balancing program: it minimizes an entropic objective over nonnegative weights subject to linear moment constraints. A representative form is
Up to constants, the objective is the KL divergence from $\bm w$ to uniform weights and matches the entropy-balancing criterion of Hainmueller2012entropybalancing when $C=1$. When $s(X)$ varies across observations, branchwise balance can be imposed by including branch-specific features, such as $\xi(X)\phi_j(X)$ and $(1-\xi(X))\phi_j(X)$, so that the constraints act within each branch. These programs are closely related to stable weights Zubizarreta2015stableweights and overlap weights Li2018addressingextreme.
This loss--link pairing has two practical advantages. First, the model enforces the sign and positivity structure of the representer by construction because $|\alpha_{{\bm{\beta}}}(X)|\ge C$. Second, the induced objective is entropic and therefore shrinks weights toward uniformity, which often improves finite-sample stability, although the exponential form can still be sensitive to large linear indices.
\paragraph{Special case: density ratio estimation.} If $C=0$ and $\xi(X)=1$ for all $X$, then ((ref)) reduces to the positive log-linear model $\alpha_{{\bm{\beta}}}(X)=\exp\Big({\bm{\phi}}(X)^\top{\bm{\beta}}\Big)$. This is the exponential tilting form that arises in maximum-entropy estimation of density ratios Qin1998inferencesfor, Sugiyama2012densityratio.
\paragraph{Connection to logistic propensity models.} For ATE estimation, it is natural to take $\xi(X)=D$ and $C=1$. If the propensity score is modeled by a sigmoid link $e(Z)=\Lambda\Big(\eta_{{\bm{\beta}}}(Z)\Big)$, then the induced ATE representer has the log-link form ((ref)); see Appendix (ref) for details. This explains why UKL with a log link is estimand-consistent for ATE when one works with logistic propensity models.
We next introduce a specification that pairs the BP divergence generator with a link function that interpolates between the linear link used for SQ-Riesz regression and the log link used for UKL-Riesz regression. This specification is useful both as a robustness device and as a way to understand how automatic covariate balancing varies continuously with the choice of loss and link.
Let $\omega \in (0,\infty)$ and define $k \coloneqq 1 + 1/\omega$. Consider the following model for the Riesz representer:
where ${\bm{\phi}} \colon {\mathcal{X}} \to {\mathbb{R}}^p$ is a basis function and $\xi \colon {\mathcal{X}} \to \{0,1\}$ selects the branch. We call the link in ((ref)) a power link.
The choice of $(\xi,C)$ is application dependent. For example, in ATE estimation we typically set $\xi(X)=D$ and $C=1$, while in density ratio estimation we often set $\xi(X)=1$ and $C=0$ so that $\alpha_{\bm{\beta}}(X)$ is nonnegative by construction.
Under ((ref)), the dual characterization implies that BP-Riesz regression returns the minimum BP-loss solution among approximately balancing models. In particular, if we use an $\ell_1$ penalty in BP-Riesz regression, it is equivalent to solving a constrained problem of the form
with $\alpha_i$ restricted to the domain of $g^{\text{BP}}$, that is, $|\alpha_i| \ge C$. When $\lambda = 0$, the constraint ((ref)) enforces exact balancing: \[ \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}_i \phi_j(X_i) = \frac{1}{n}\sum^n_{i=1} m(W_i,\phi_j), \qquad j = 1,\dots,p, \] where $\widehat{\alpha}_i=\alpha_{\widehat{{\bm{\beta}}}}(X_i)$.
\paragraph{Relationship to the linear and log links.} The power link ((ref)) provides a continuous bridge between the linear and log specifications. As $\omega \to 0$, we have $k = 1 + 1/\omega \to \infty$ and \[ \Big(1+\frac{t}{k}\Big)^{1/\omega} = \Big(1+\omega t + o(\omega)\Big)^{1/\omega} \to \exp(t), \] so ((ref)) reduces to the log-link form used for UKL-Riesz regression. At $\omega=1$, the BP generator reduces to the squared-loss generator, and the link becomes an affine transformation of the linear index around the origin, which connects BP-Riesz regression to SQ-Riesz regression up to reparameterization.
This interpolation perspective is also consistent with the robustness interpretation of the BP divergence Basu1998robustandefficient,Sugiyama2012densityratio. Smaller $\omega$ makes the objective closer to a KL-type criterion, which can be efficient under correct specification, while larger $\omega$ yields behavior closer to squared loss and is typically more robust to misspecification and extreme weights.
Finally, the role of the loss choice becomes most visible in regimes where balance is inexact or where $\alpha_0$ is misspecified by the model class. At the population level, generalized Riesz regression targets a Bregman projection $\alpha_g^\ast$ of $\alpha_0$ onto the model class (Section (ref)). Different generators $g$ generally induce different pseudo-true limits and different stability properties of the resulting weights: squared-type losses tend to produce smoother and less extreme weights, while KL-type losses can be efficient under correct exponential-type specification but are often more sensitive to tail observations.
Moreover, as shown in Section (ref), in ATE estimation under a sigmoid propensity model the estimand-consistent choice is UKL-Riesz, whereas logistic MLE corresponds to a different loss--estimand pairing. This is why loss choice should be regarded as part of model specification rather than a mere computational detail.
Under exact balance, the RW identity ((ref)) is loss-invariant on the training sample. Under inexact balance, different generators $g$ select different Bregman-projection solutions within the same approximation space, producing different $\widehat{\alpha}$ even when the same dictionary is used. This affects stability of weights, how approximation error is distributed across the support, and the size of imbalance generalization. The density-ratio literature makes this sensitivity viewpoint explicit Menon2016linkinglosses,Zellinger2025binarylosses.
A pragmatic workflow is:
This section provides an estimation error analysis for generalized Riesz regression. We model the Riesz representer $\alpha_0$ by \[ \alpha_f(X) = \zeta^{-1}\Big(X, f(X)\Big), \] where $\zeta^{-1}$ is continuously differentiable and globally Lipschitz in its second argument, uniformly in $x \in {\mathcal{X}}$, and $f$ is a base model. Unlike Section (ref), we do not restrict $f$ to be a linear model. For example, in addition to linear models ${\bm{\phi}}(X)^\top {\bm{\beta}}$, we can use random forests, neural networks, and other models for $f$. In this section, we consider the case where we use RKHS methods and neural networks for $f$.
Throughout this section, we assume that the true Riesz representer is bounded.
This boundedness assumption holds in the standard ATE setting under common support and bounded outcomes. In many other applications, this assumption also holds. If we wish to allow unbounded support, we can develop an extension by imposing appropriate tail conditions. For example, density ratios between two Normal distributions may violate this assumption. In such cases, Zheng2022anerror presents a convergence rate analysis, and we can follow their approach. In practical data analysis, it is often reasonable to treat the Riesz representer as bounded.
First, we study the case with RKHS regression. Let ${\mathcal{F}}^{\text{RKHS}}$ be a class of RKHS functions and define \[ \widehat{f}^{\text{RKHS}} \coloneqq \operatorname*{arg\,min}_{f \in {\mathcal{F}}^{\text{RKHS}}}\left\{\widehat{\text{BD}}_{g}(\alpha_f) + \lambda \|f\|^2_{{\mathcal{F}}}\right\}, \] where $\|\cdot \|_{{\mathcal{F}}}$ is the RKHS norm and $\lambda > 0$ is a regularization parameter. We then define \[ \widehat{\alpha}^{\text{RKHS}}(x) \coloneqq \alpha_{\widehat{f}^{\text{RKHS}}}(x) \coloneqq \zeta^{-1}\left(x, \widehat{f}^{\text{RKHS}}(x)\right). \] We analyze the estimation error by adapting the approach of Kanamori2012statisticalanalysis, which studies RKHS-based LSIF for density ratio estimation.
For technical control of complexity, we use a localized class. Let $I(f)$ be a complexity measure on ${\mathcal{F}}^{\text{RKHS}}$, and define \[ {\mathcal{F}}^{\text{RKHS}}_M \coloneqq \big\{f \in {\mathcal{F}}^{\text{RKHS}}\colon I(f) \le M\big\}, \qquad {\mathcal{H}}^{\text{RKHS}} \coloneqq \left\{\zeta^{-1}\left(\cdot, f(\cdot)\right)\colon f\in {\mathcal{F}}^{\text{RKHS}}\right\}. \]
For bracketing entropy, see Definition 2.2 in VandeGeer2000empiricalprocesses and Appendix (ref).
The proof is provided in Appendix (ref), following the approach of Kanamori2012statisticalanalysis. The parameter $\tau$ is determined by the function class to which the true base model $f_0$ belongs.
Second, we provide an estimation error analysis when we use neural networks for ${\mathcal{H}}$. Our analysis follows Kato2021nonnegativebregman and Zheng2022anerror.
For ${\mathcal{F}}^{\text{FNN}}$, define \[ \widehat{f}^{\text{FNN}} \coloneqq \operatorname*{arg\,min}_{f \in {\mathcal{F}}^{\text{FNN}}}\left\{\widehat{\text{BD}}_{g}(\alpha_f)\right\}, \qquad \widehat{\alpha}^{\text{FNN}}(x) \coloneqq \zeta^{-1}\left(x, \widehat{f}^{\text{FNN}}(x)\right). \]
Let $\text{Pdim}({\mathcal{F}}^{\text{FNN}})$ be the pseudodimension of ${\mathcal{F}}^{\text{FNN}}$. For the definition, see Anthony1999neuralnetwork and Definition 3 in Zheng2022anerror.
The proof is provided in Appendix (ref), following Zheng2022anerror. This result implies minimax optimality of the proposed method when $f_0$ belongs to a H\"older class.
This subsection describes how we construct an efficient estimator for the parameter of interest $\theta_0$ using generalized Riesz regression. As discussed in Section (ref), we construct $\widehat{\theta}$ by solving \[ \frac{1}{n}\sum^n_{i=1}\psi\left(W_i;\widehat{\eta},\widehat{\theta}^{\text{ARW}}\right)=0, \] where the Neyman orthogonal score is \[ \psi(W;\eta,\theta) \coloneqq m(W,\gamma) + \alpha(X)\big(Y-\gamma(X)\big) - \theta, \qquad \eta \coloneqq (\alpha,\gamma). \] As introduced in Section (ref), we refer to this estimator as the ARW estimator.
For example, the Donsker condition holds when the bracketing entropy of ${\mathcal{H}}$ is finite. In contrast, it fails in high-dimensional or series regression settings where the model complexity diverges as $n \to \infty$. For neural networks, the condition can hold when both the number of layers and the width are fixed. If these quantities grow with the sample size, cross fitting is typically used instead.
Under these assumptions, asymptotic normality follows from standard debiased machine learning arguments.
Here, $V^*$ matches the semiparametric efficiency bound, the variance of the efficient influence function VanderVaart1998asymptoticstatistics,Hahn1998ontherole.
\paragraph{Automatic Neyman orthogonalization in the RW estimator.} A central theme in debiased machine learning is to construct estimators from Neyman orthogonal scores. In our setting, $\gamma_0(x)={\mathbb{E}}\left[Y\mid X=x\left(\right)\right]$, and $\alpha_0$ is the Riesz representer associated with the linear functional $\gamma \mapsto {\mathbb{E}}\left[m(W,\gamma)\right]$. While the ARW estimator requires an explicit regression estimator $\widehat{\gamma}$, the RW estimator \[ \widehat{\theta}^{\text{RW}} \coloneqq \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)Y_i \] uses only the estimated representer. The next theorem shows that if the representer fit achieves exact balancing on a linear working space that contains $\gamma_0$, then $\widehat{\theta}^{\text{RW}}$ admits an exact orthogonal-score representation with $\gamma_0$ plugged in.
This section provides applications of generalized Riesz regression: ATE estimation, AME estimation, and covariate shift adaptation (density ratio estimation). We introduce other applications such as difference-in-difference in Appendix (ref).
In ATE estimation, the linear functional is \[ m^{\text{ATE}}(W,\gamma)\coloneqq \gamma(1,Z)-\gamma(0,Z), \] and the Riesz representer is \[ \alpha^{\text{ATE}}_0(X)=\frac{D}{e_0(Z)}-\frac{1-D}{1-e_0(Z)}, \] where $e_0(Z) = P(D = 1\mid Z)$ is the propensity score. Let $r_0(1 , Z) \coloneqq \frac{1}{e_0(Z)}$ and $r_0(0 , Z) \coloneqq \frac{1}{1 - e_0(Z)}$ be the inverse propensity score, also called the density ratio. We estimate $\alpha^{\text{ATE}}_0$ by minimizing the empirical Bregman divergence objective $\widehat{\text{BD}}_g(\alpha)$ introduced in Section (ref), with $m=m^{\text{ATE}}$, an application-specific choice of $g$, and a model class for $\alpha$.
\paragraph{SQ-Riesz Regression.} We take the squared loss, \[ g^{\text{SQ}}(\alpha)=\alpha^2, \] and minimize the corresponding empirical Bregman objective. By substituting $g^{\text{SQ}}$ into ((ref)) and using $m^{\text{ATE}}(W,\gamma)=\gamma(1,Z)-\gamma(0,Z)$, we obtain, up to an additive constant that does not depend on $\alpha$, \[ \text{BD}_{g^{\text{SQ}}}(\alpha) = {\mathbb{E}}\left[ \alpha(D,Z)^2 -2\big(\alpha(1,Z)-\alpha(0,Z)\big) \right]. \] Thus, SQ-Riesz regression estimates $\alpha^{\text{ATE}}_0$ by \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \widehat{\text{BD}}_{g^{\text{SQ}}}(\alpha)+\lambda J(\alpha), \] where \[ \widehat{\text{BD}}_{g^{\text{SQ}}}(\alpha) \coloneqq \frac{1}{n}\sum_{i=1}^n \left( \alpha(D_i,Z_i)^2 -2\big(\alpha(1,Z_i)-\alpha(0,Z_i)\big) \right). \] This coincides with Riesz regression in Chernozhukov2021automaticdebiased and corresponds to LSIF in density ratio estimation Kanamori2009aleastsquares. With appropriate choices of ${\mathcal{H}}$, it also recovers nearest neighbor matching-based constructions, as discussed in Kato2025nearestneighbor.
\paragraph{UKL-Riesz Regression.} Consider Riesz representer models $\alpha$ such that $\alpha(1, x) > 1$ and $\alpha(0, x) < -1$ for all $x$. We next use the UKL divergence loss with $C=1$, \[ g^{\text{UKL}}(\alpha)=(|\alpha|-1)\log\left(|\alpha|-1\right)-|\alpha|. \] By substituting $g^{\text{UKL}}$ into ((ref)) and using $m^{\text{ATE}}(W,\gamma)=\gamma(1,Z)-\gamma(0,Z)$, we obtain
Thus, UKL-Riesz regression estimates $\alpha^{\text{ATE}}_0$ by \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \widehat{\text{BD}}_{g^{\text{UKL}}}(\alpha)+\lambda J(\alpha), \] where
This coincides with the tailored loss minimization with $\alpha = \beta = -1$ in Zhao2019covariatebalancing and corresponds to KLIEP in density ratio estimation Sugiyama2008directimportance.
\paragraph{BP-Riesz Regression.} BP-Riesz regression uses Basu's power divergence with $C=1$ and $\omega\in(0,\infty)$: \[ g^{\text{BP}}(\alpha) \coloneqq \frac{\big(|\alpha|-1\big)^{1+\omega}-\big(|\alpha|-1\big)}{\omega}-|\alpha|. \] Plugging $g^{\text{BP}}$ into ((ref)) and using $m^{\text{ATE}}$ yields the empirical objective
where $\upsilon \coloneqq 1+1/\omega$. We then estimate $\alpha^{\text{ATE}}_0$ by \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{BP}}}(\alpha)+\lambda J(\alpha). \]
\paragraph{BKL-Riesz Regression.} Consider Riesz representer models $\alpha$ such that $\alpha(1, x) > 1$ and $\alpha(0, x) < -1$ for all $x$. BKL-Riesz regression uses BKL divergence with $C=1$: \[g^{\text{BKL}}(\alpha) \coloneqq (|\alpha| - 1)\log \big(|\alpha| - 1\big) - (|\alpha| + 1)\log(|\alpha| + 1).\] By plugging $g^{\text{BKL}}$ into ((ref)) and using $m^{\text{ATE}}$, we have the following empirical objective function:
We then estimate $\alpha^{\text{ATE}}_0$ by \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}}\widehat{\text{BD}}_{g^{\text{BKL}}}(\alpha)+\lambda J(\alpha). \]
We consider the AME setup described in Section (ref). Let $X=(D,Z)$, where $D$ is a scalar continuous regressor and $Z$ is a vector of covariates. The target parameter is \[ \theta^{\text{AME}}_0 \coloneqq \mathbb{E}\Big[\partial_d \gamma_0(D,Z)\Big], \qquad m^{\text{AME}}(W,\gamma)\coloneqq \partial_d \gamma(D,Z). \] Assume that $X$ admits a density $f_0$ that is continuously differentiable and that an integration by parts argument is valid, for example, $\gamma(x)f_0(x)$ vanishes on the boundary of the support in the $d$ direction. Then \[ \mathbb{E}\Big[\partial_d \gamma(X)\Big] = \mathbb{E}\Big[\alpha^{\text{AME}}_0(X)\gamma(X)\Big], \qquad \alpha^{\text{AME}}_0(X)=-\partial_d \log f_0(X), \] so the AME Riesz representer is the negative score of the marginal density of $X$ with respect to $d$. Since $\partial_d \log f_0(D,Z)=\partial_d \log f_0(D\mid Z)$, we can equivalently view $\alpha^{\text{AME}}_0$ as the negative score of the conditional density of $D$ given $Z$.
To estimate $\alpha^{\text{AME}}_0$, we apply generalized Riesz regression with $m=m^{\text{AME}}$. The population objective in Section (ref) becomes \[ \text{BD}^{\text{AME}}_g(\alpha) \coloneqq \mathbb{E}\Big[ - g\big(\alpha(X)\big) + \partial g\big(\alpha(X)\big)\alpha(X) - \partial_d\Big(\partial g\big(\alpha(X)\big)\Big)\Big], \] and we minimize its empirical analogue over a differentiable model class ${\mathcal{H}}$ (so that $\partial_d \alpha(X)$ and $\partial_d\{\partial g(\alpha(X))\}$ are well defined), possibly with regularization.
\paragraph{SQ-Riesz Regression.} Let $g^{\text{SQ}}(\alpha)=(\alpha-C)^2$ for an arbitrary constant $C\in{\mathbb{R}}$, so that $\partial g^{\text{SQ}}(\alpha)=2(\alpha-C)$. Substituting into $\text{BD}^{\text{AME}}_g$ yields
where the constant does not depend on $\alpha$. Hence SQ-Riesz regression targets $\alpha^{\text{AME}}_0$ in $L_2$. This objective is also a score matching style criterion for estimating the score, written here in terms of the negative score $\alpha^{\text{AME}}_0$. The empirical estimator is \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \frac{1}{n}\sum_{i=1}^n\Big(\alpha(X_i)^2-2\partial_d \alpha(X_i)\Big) +\lambda J(\alpha). \] This method corresponds to Riesz regression for AME, as discussed in Chernozhukov2021automaticdebiased.
\paragraph{UKL-Riesz Regression.} To obtain a KL motivated loss that allows signed $\alpha$, we use the signed KL type convex function \[ g^{\text{UKL}}(\alpha)=|\alpha|\log|\alpha|-|\alpha|, \qquad \partial g^{\text{UKL}}(\alpha)=\operatorname{sign}(\alpha)\log|\alpha|, \] on a domain that excludes $\alpha=0$. Plugging into $\text{BD}^{\text{AME}}_g$ gives
Accordingly, we estimate $\alpha^{\text{AME}}_0$ by minimizing the empirical version over ${\mathcal{H}}$: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \frac{1}{n}\sum_{i=1}^n \Big( |\alpha(X_i)| -\partial_d\Big(\operatorname{sign}\big(\alpha(X_i)\big)\log|\alpha(X_i)|\Big)\Big)g)} +\lambda J(\alpha). \] In practice, one can use the shifted UKL loss in Section (ref) to avoid the singularity at zero and combine it with a branchwise link specification as in Section (ref).
\paragraph{BP-Riesz Regression.} BP-Riesz regression interpolates between squared distance and UKL divergence. For simplicity, we present the unshifted form with $C=0$: \[ g^{\text{BP}}(\alpha) \coloneqq \frac{|\alpha|^{1+\gamma}-|\alpha|}{\gamma}-|\alpha|, \qquad \partial g^{\text{BP}}(\alpha) = \left(1+\frac{1}{\gamma}\right)\operatorname{sign}(\alpha)\Big(|\alpha|^\gamma-1\Big), \qquad \gamma\in(0,\infty). \] Let $k\coloneqq 1+1/\gamma$. Then $\text{BD}^{\text{AME}}_g$ simplifies to \[ \text{BD}^{\text{AME}}_{g^{\text{BP}}}(\alpha) = \mathbb{E}\Big[ |\alpha(X)|^{1+\gamma} - \partial_d\Big(k\operatorname{sign}\big(\alpha(X)\big)\Big(|\alpha(X)|^\gamma-1\Big)\Big)g)\Big]} +\text{const}. \] We estimate $\alpha^{\text{AME}}_0$ by minimizing the empirical objective: \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \frac{1}{n}\sum_{i=1}^n \Big( |\alpha(X_i)|^{1+\gamma} - \partial_d\Big(k\operatorname{sign}\big(\alpha(X_i)\big)\Big(|\alpha(X_i)|^\gamma-1\Big)\Big)g)}\Big)g)g)}} +\lambda J(\alpha). \] As in Section (ref), $\gamma=1$ recovers the squared loss behavior (up to scaling), while $\gamma\to 0$ approaches a KL type criterion through the identity $\lim_{\gamma\to 0}(|\alpha|^\gamma-1)/\gamma=\log|\alpha|$.
\paragraph{BKL-Riesz Regression.} Finally, we can use the BKL loss from Section (ref) to obtain a logistic motivated criterion. Let $C>0$ and define \[ g^{\text{BKL}}(\alpha) \coloneqq (|\alpha|-C)\log\big(|\alpha|-C\big)-(|\alpha|+C)\log\big(|\alpha|+C\big), \qquad \partial g^{\text{BKL}}(\alpha) = \operatorname{sign}(\alpha)\log\left(\frac{|\alpha|-C}{|\alpha|+C}\right). \] Then the AME objective is \[ \text{BD}^{\text{AME}}_{g^{\text{BKL}}}(\alpha) = {\mathbb{E}}\left[ C\log\left(\frac{|\alpha(X)|-C}{|\alpha(X)|+C}\right) - \partial_d\left(\operatorname{sign}\big(\alpha(X)\big)\log\left(\frac{|\alpha(X)|-C}{|\alpha(X)|+C}\right)\right)\right]\right)}} +\text{const}, \] and the estimator minimizes its empirical counterpart over ${\mathcal{H}}$ with regularization. As in the ATE case, this loss is naturally paired with a logistic style link for the magnitude of $\alpha$, while sign changes can be handled via the branchwise constructions in Section (ref).
Once we obtain $\widehat{\alpha}^{\text{AME}}$ and an outcome regression estimator $\widehat{\gamma}$, we plug them into the Neyman orthogonal score in Section (ref) to form an estimator of $\theta_0^{\text{AME}}$.
We consider the covariate shift setting in Section (ref). Let $X$ be the source covariate distribution that generates labeled observations $\{(X_i,Y_i)\}^n_{i=1}$, and let $\widetilde{X}$ be the target covariate distribution that generates unlabeled observations $\{\widetilde{X}_j\}^m_{j=1}$, independent of the source sample. Let $p_0(x)$ and $p_1(x)$ be the pdfs of $X$ and $\widetilde{X}$, respectively. We assume that $p_0(x), p_1(x) > 0$ for all $x\in{\mathcal{X}}$. The Riesz representer for covariate shift adaptation is the density ratio \[ \alpha^{\text{CS}}_0(X)=r_0(X)\coloneqq \frac{p_1(X)}{p_0(X)}. \] We estimate $r_0$ directly by density ratio fitting under a Bregman divergence, avoiding separate density estimation for $p_0(X)$ and $p_1(X)$.
Let $g\colon {\mathbb{R}}_+\to{\mathbb{R}}$ be differentiable and strictly convex. The Bregman divergence between $r_0$ and a candidate ratio model $\alpha$ is \[ \text{BD}^\dagger_g\big(r_0\mid \alpha\big) \coloneqq {\mathbb{E}}_X\Big( g\big(r_0(X)\big)-g\big(\alpha(X)\big)-\partial g\big(\alpha(X)\big)\big(r_0(X)-\alpha(X)\big) \Big). \] Dropping the constant ${\mathbb{E}}_X\big[g(r_0(X))\big]$ and using the identity ${\mathbb{E}}_X\big[r_0(X)h(X)\big]={\mathbb{E}}_{\widetilde{X}}\big[h(X)\big]$, we obtain the equivalent population objective
Given samples $\{X_i\}_{i\in{\mathcal{I}}_S}$ and $\{\widetilde{X}_j\}_{j\in{\mathcal{I}}_T}$, the empirical objective is
We estimate the density ratio by \[ \widehat{\alpha} \coloneqq \operatorname*{arg\,min}_{\alpha\in{\mathcal{H}}} \widehat{\text{BD}}^{\text{CS}}_g(\alpha)+\lambda J(\alpha), \] where ${\mathcal{H}}$ is a model class and $J$ is a regularizer. A convenient way to enforce $\alpha(x)\ge 0$ is to use a link specification such as $\alpha(x)=\exp\big(f(x)\big)$ with a flexible regression model $f$.
\paragraph{SQ-Riesz Regression.} For the squared loss, take \[ g^{\text{SQ}}(\alpha)=(\alpha-1)^2, \qquad \partial g^{\text{SQ}}(\alpha)=2(\alpha-1). \] Substituting into ((ref)) and dropping constants that do not depend on $\alpha$, we obtain \[ \widehat{\text{BD}}^{\text{CS}}_{g^{\text{SQ}}}(\alpha) = \frac{1}{n}\sum^n_{i=1}\alpha(X_i)^2 - \frac{2}{m}\sum^m_{j=1}\alpha(\widetilde{X}_j). \] This is the classical least-squares importance fitting (LSIF) objective in density ratio estimation Kanamori2009aleastsquares. In the debiased machine learning literature, the same squared loss criterion is also used as Riesz regression for covariate shift adaptation Chernozhukov2025automaticdebiased. Related extensions include doubly robust covariate shift adaptation schemes that combine density ratio estimation with regression adjustment Kato2024doubledebiasedcovariateshift.
\paragraph{UKL-Riesz Regression.} For a KL motivated objective on ${\mathbb{R}}_+$, take \[ g^{\text{UKL}}(\alpha)=\alpha\log\alpha-\alpha, \qquad \partial g^{\text{UKL}}(\alpha)=\log\alpha. \] Then ((ref)) becomes, up to an additive constant that does not depend on $\alpha$, \[ \widehat{\text{BD}}^{\text{CS}}_{g^{\text{UKL}}}(\alpha) = \frac{1}{n}\sum^n_{i=1}\alpha(X_i) - \frac{1}{m}\sum^m_{j=1}\log \alpha(\widetilde{X}_j). \] A standard implementation imposes the normalization constraint $\frac{1}{n}\sum^n_{i=1}\alpha(X_i)=1$, in which case minimizing $\widehat{\text{BD}}^{\text{CS}}_{g^{\text{UKL}}}$ is equivalent to maximizing the target log-likelihood $\frac{1}{m}\sum^m_{j=1}\log \alpha(\widetilde{X}_j)$ subject to normalization and nonnegativity, which yields KLIEP style procedures Sugiyama2008directimportance. This constrained view is also useful for understanding the dual characterization and the associated moment matching property.
\paragraph{BP-Riesz Regression.} BP-Riesz regression interpolates between squared loss and KL type objectives. For $\gamma\in(0,\infty)$, consider the BP choice on ${\mathbb{R}}_+$ with $C=0$, \[ g^{\text{BP}}(\alpha) \coloneqq \frac{\alpha^{1+\gamma}-\alpha}{\gamma}-\alpha, \qquad \partial g^{\text{BP}}(\alpha) = \left(1+\frac{1}{\gamma}\right)\Big(\alpha^\gamma-1\Big). \] A useful simplification is that $\partial g^{\text{BP}}(\alpha)\alpha-g^{\text{BP}}(\alpha)=\alpha^{1+\gamma}$, so ((ref)) reduces, up to constants, to \[ \widehat{\text{BD}}^{\text{CS}}_{g^{\text{BP}}}(\alpha) = \frac{1}{n}\sum^n_{i=1}\alpha(X_i)^{1+\gamma} - \left(1+\frac{1}{\gamma}\right)\frac{1}{m}\sum^m_{j=1}\alpha(\widetilde{X}_j)^{\gamma}. \] When $\gamma=1$, this objective coincides with the SQ-Riesz objective, up to scaling and constants. As $\gamma\to 0$, it approaches a KL flavored criterion via the expansion $\alpha^\gamma=1+\gamma\log \alpha+o(\gamma)$, providing a continuous bridge between LSIF and KLIEP, and offering a robustness device against extreme ratios Basu1998robustandefficient,Sugiyama2012densityratio.
\paragraph{BKL-Riesz Regression.} BKL-Riesz regression corresponds to probabilistic classification based density ratio estimation, which estimates the log density ratio by fitting a classifier that discriminates target covariates from source covariates Qin1998inferencesfor,Cheng2004semiparametricdensity. Let $S\in\{0,1\}$ denote a domain indicator, where $S=1$ for target and $S=0$ for source, and let $\pi\coloneqq P(S=1)$ denote the mixture class prior. Under Bayes' rule, the density ratio satisfies \[ r_0(x)=\frac{p_1(x)}{p_0(x)} = \frac{1-\pi}{\pi}\frac{P(S=1\mid X=x)}{P(S=0\mid X=x)}. \] We model $P(S=1\mid X=x)$ by a logistic specification \[ p_{\bm{\beta}}(S=1\mid X=x)\coloneqq \frac{1}{1+\exp\big(-{\bm{\phi}}(x)^\top{\bm{\beta}}\big)}, \] and estimate ${\bm{\beta}}$ by regularized Bernoulli likelihood on the pooled sample: \[ \widehat{{\bm{\beta}}} \coloneqq \operatorname*{arg\,min}_{{\bm{\beta}}} -\frac{1}{n+m} \left( \sum^n_{i=1}\log\big(1-p_{\bm{\beta}}(S=1\mid X_i)\big) + \sum^m_{j=1}\log p_{\bm{\beta}}(S=1\mid \widetilde{X}_j) \right) +\lambda\|{\bm{\beta}}\|_2^2. \] With $\widehat{\pi}\coloneqq \frac{m}{n+m}$, we then set \[ \widehat{\alpha}(x) \coloneqq \frac{1-\widehat{\pi}}{\widehat{\pi}} \frac{p_{\widehat{{\bm{\beta}}}}(S=1\mid X=x)}{1-p_{\widehat{{\bm{\beta}}}}(S=1\mid X=x)} = \frac{1-\widehat{\pi}}{\widehat{\pi}}\exp\big({\bm{\phi}}(x)^\top\widehat{{\bm{\beta}}}\big). \] This construction enforces nonnegativity by design and connects density ratio estimation to standard classification tools.
We evaluate generalized Riesz regression as a building block for debiased machine learning, focusing on average treatment effect (ATE) estimation. Across all experiments, we compare three ways of estimating the ATE Riesz representer (bias-correction term) introduced in Section (ref): SQ-Riesz (squared-loss objective), UKL-Riesz (unnormalized-KL objective), and BKL-Riesz (binary-KL objective). In the ATE setting, BKL-Riesz coincides with estimating the propensity score by Bernoulli likelihood (logistic MLE) and then plugging it into the closed-form ATE Riesz representer; we therefore refer to it as “BKL-Riesz = MLE.”
Given an estimate of the outcome regression $\widehat{\gamma}$ and an estimate of the Riesz representer $\widehat{\alpha}$, we report three ATE estimators:
We quantify accuracy by the mean squared error (MSE) of the ATE estimate and quantify uncertainty by the empirical coverage ratio (CR) of nominal $95$
\paragraph{Design} The covariates are three-dimensional, $K = 3$, and we fix the sample size at $n = 3000$. In each Monte Carlo replication, we generate covariates $Z_i \in {\mathbb{R}}^3$ from a multivariate normal distribution $\mathcal{N}(0, I_3)$ and construct a nonlinear propensity score model with polynomial and interaction terms as \[e_0(Z_i) = \frac{1}{1 + \exp\bigl(-h(Z_i)\bigr)},\] where \[h(Z_i) = \sum_{j=1}^3 a_j Z_{i,j} + \sum_{j=1}^3 b_j Z_{i,j}^2 + c_{1} Z_{i,1} Z_{i,2} + c_{2} Z_{i,2} Z_{i,3} + c_{3} Z_{i,1} Z_{i,3}.\] The coefficients \(a_j\), \(b_j\), and \(c_j\) are independently drawn from \(\mathcal{N}(0,0.5)\). Given these propensity scores, the treatment assignment \(D_i\) is sampled accordingly. We then generate the outcome as \[Y_i = 1.0 + \left(\sum^3_{j=1}Z_{i,j} \widetilde{a}_j\right)^2 + 1/\left(1 + \exp\left(-\left(\sum^3_{j=1}Z^2_{i,j} \widetilde{b}_j\right) \right)\right) }\right) \right)\right) }} + \tau_0 D_i + \varepsilon_i,\] where \(\varepsilon_i \sim \mathcal{N}(0,1)\) and \(\tau_0 = 5.0\).
\paragraph{Estimators and implementation} We estimate the Riesz representer using the following variants, matched to Table (ref).
For the Riesz representer and regression models, we separately use a neural network with one hidden layer consisting of $100$ nodes. To avoid relying on the Donsker condition, we estimate all nuisance functions using two-fold cross fitting. In each replication, we split the sample into two folds, estimate the nuisance functions on one fold, evaluate the corresponding scores on the other fold, and then swap the roles of the folds. The final estimators aggregate the two cross-fitted scores.
This experiment does not guarantee automatic Neyman orthogonalization, since we use cross fitting and do not use the same basis functions for outcome modeling. However, this implementation is standard in debiased machine learning; therefore, we adopt it.
We repeat the simulation 100 times. The “True” columns in Table (ref) report infeasible oracle performance using the true nuisance functions.
\paragraph{Results} Table (ref) highlights three robust patterns. First, oracle baselines separate estimation error from intrinsic variance. The oracle ARW estimator is close to the efficiency benchmark (MSE $=0.01$) and achieves near-nominal coverage (CR $=0.98$). In contrast, even with the true propensity score, oracle RW remains noisy (MSE $=1.44$) and undercovers (CR $=0.84$), reflecting the well-known finite-sample instability of pure weighting in challenging overlap regimes.
Second, the plug-in RA estimator is not reliable for inference in this design. Across feasible implementations, RA has moderate MSE (about $0.38$–$0.40$) but extremely poor coverage (CR $=0.06$–$0.12$). This indicates that the outcome regression learner, while not catastrophically inaccurate in MSE, does not deliver a reliable uncertainty estimate when used without orthogonalization, and the resulting Wald intervals are severely miscalibrated.
Third, how we fit the Riesz representer matters substantially for IPW, and ARW mitigates (but does not eliminate) this sensitivity. For IPW, SQ-Riesz (Linear) is the best-performing option in Table (ref) (MSE $=0.49$) and yields near-nominal coverage (CR $=0.98$). In contrast, RW based on UKL-Riesz has larger MSE (about $1.50$) and noticeably lower coverage (CR $=0.68$–$0.73$), while BKL-Riesz (= MLE) performs worst (MSE $=3.79$, CR $=0.32$), consistent with propensity-score MLE producing more extreme effective weights in this design.
The ARW estimator is uniformly more stable than RW and RA in terms of MSE, but calibration still depends on the Riesz-representer fit. SQ-Riesz (Logit) attains the best ARW MSE (0.08) with CR 0.89. UKL-Riesz achieves similarly small ARW MSE (0.10) but exhibits undercoverage (CR $=0.77$–$0.81$). BKL-Riesz (= MLE) improves substantially over its RW counterpart (ARW MSE $=0.23$), yet its coverage remains poor (CR $=0.60$). Overall, directly fitting the Riesz representer via generalized Riesz regression can materially improve finite-sample performance relative to the MLE plug-in baseline, and the combination of objective and link specification plays a first-order role, especially for RW and for the calibration of ARW intervals.
We next evaluate the same family of estimators on the semi-synthetic IHDP benchmark, following Chernozhukov2022riesznet. We use the standard setting “A” in the npci package and generate 1000 replications. Each replication contains $n=747$ observations with a binary treatment, an outcome, and $p=25$ covariates. The estimand is the ATE.
We report DM, IPW, and ARW for each Riesz-representer estimator (SQ-Riesz, UKL-Riesz, and BKL-Riesz (= MLE)). We consider two nuisance-learner families:
For each configuration, we compute the MSE of the ATE estimate and the empirical coverage ratio (CR) of nominal $95$% Wald-type confidence intervals across the 1000 replications; CR close to $0.95$ indicates well-calibrated uncertainty quantification. Results appear in Table (ref).
Two findings stand out. With neural networks, ARW is consistently accurate (MSE around $0.31$–$0.43$) and well calibrated for UKL-Riesz and BKL-Riesz (CR $=0.94$ and $0.90$), while SQ-Riesz yields overly conservative ARW intervals (CR $=1.00$). In contrast, RA has very low coverage (CR near zero) and RW exhibits large error, especially for SQ-Riesz (MSE $=6.82$), reinforcing that orthogonalization is essential in this benchmark.
With RKHS, performance becomes much more sensitive to the particular objective: SQ-Riesz deteriorates sharply (MSE around $20$ with CR $=0$ for both RA and ARW), whereas UKL-Riesz and BKL-Riesz remain substantially more stable. In particular, UKL-Riesz attains strong RW calibration under RKHS (CR $=0.93$) with comparatively low MSE (1.78), while BKL-Riesz provides a competitive alternative (IPW MSE $=1.22$ with CR $=0.81$). These results underscore that, in finite samples, the interaction between the Riesz-representer objective and the nuisance-function learner can be decisive, and that UKL-type objectives can offer noticeably more robust behavior than squared-loss fitting in this semi-synthetic setting.
Density ratio estimation often suffers from a characteristic form of overfitting. Kato2021nonnegativebregman refers to this issue as train-loss hacking and shows that the empirical objective can be artificially reduced by inflating $r(X^{(\text{nu})})$ at the training points. Rhodes2020telescopingdensityratio highlights a related mechanism: when $p_{\text{nu}}$ and $p_{\text{de}}$ are far apart, for example when $\mathrm{KL}(p_{\text{nu}}\|p_{\text{de}})$ is on the order of tens of nats, the estimation problem enters a large-gap regime that exacerbates overfitting. They refer to this phenomenon as the density chasm. Although the two papers emphasize different viewpoints, both point to the same underlying difficulty: finite samples provide weak control of the ratio in regions where the two distributions have little overlap.
\paragraph{Non-negative Bregman divergence} Kato2021nonnegativebregman proposes a modification of the Bregman divergence objective that isolates the problematic component and applies a non-negative correction under a mild boundedness condition on $r_0$. Specifically, choose $0<C<1/R$ with $R\coloneqq \sup r_0$. The population objective decomposes, up to an additive constant, into a non-negative term plus a bounded residual. At the sample level, the method replaces the non-negative component with its positive part $[\cdot]_+$. This yields an objective that curbs train-loss hacking while remaining within the Bregman divergence framework Kiryo2017positiveunlabeledlearning,Kato2021nonnegativebregman.
\paragraph{Telescoping density ratio estimation} Rhodes2020telescopingdensityratio proposes telescoping density ratio estimation, which targets overfitting in large-gap regimes by introducing intermediate waymark distributions $p_0 = p_{\text{nu}},p_1,\dots,p_m = p_{\text{de}}$. The method estimates local ratios $p_k/p_{k+1}$ and combines them through the identity \[ \frac{p_0(x)}{p_m(x)} = \prod_{k=0}^{m-1}\frac{p_k(x)}{p_{k+1}(x)}. \] Each local ratio corresponds to a smaller distributional gap, which makes perfect classification harder and typically stabilizes ratio estimation at finite sample sizes. As a result, telescoping can improve robustness and generalization in practice.
Telescoping density ratio estimation is also closely connected to score matching. As the number of intermediate ratios tends to infinity, the log density ratio can be expressed as an integral of time scores along a continuum of bridge distributions, and can be approximated by aggregating these score functions Choi2022densityratio. Building on this idea, Choi2022densityratio proposes density ratio estimation via infinitesimal classification. See Appendix (ref) for details.
We provide a Python package called genriesz.\footnote{Code: \url{https://github.com/MasaKat0/genriesz}. Document: \url{https://genriesz.readthedocs.io/en/latest/}.} The package implements generalized Riesz regression (GRR) under Bregman divergences, providing an end-to-end interface for fitting the Riesz representer and constructing debiased estimators. At a high level, users specify (i) a linear functional $m(X,\gamma)$, (ii) a basis map $\phi(X)$, and (iii) a Bregman generator $g(X,\alpha)$. The package then (a) constructs the automatic covariate balancing link function induced by $g$, (b) estimates the Riesz representer $\widehat{\alpha}(X)$ by solving a penalized GRR problem, and (c) optionally fits an outcome model $\widehat{\gamma}(X)$. The main entry point grr_functional returns direct-method, inverse-probability-weighting, and augmented estimators (when requested), together with standard errors, confidence intervals, and $p$-values, and supports cross-fitting. For lower-level control, genriesz also exposes solvers (GRR, GRRGLM) with ridge, lasso, and general $\ell_p$ regularization options.
The core installation is lightweight (Python 3.10 or later with NumPy and SciPy) and is available on PyPI via pip install genriesz. Optional extras provide integrations with scikit-learn and PyTorch, enabling additional feature maps such as random forest leaf encodings and neural network feature maps. The library includes a collection of ready-to-use components, including Bregman generators (e.g., squared and Kullback--Leibler-type generators) and basis functions such as polynomial bases, treatment-interaction bases, RKHS-style random Fourier features and Nystr\"om bases, and a $k$NN catchment-area basis for matching-style estimators. Custom generators and bases can also be supplied as Python callables; when analytic derivatives are not provided, the package can fall back on numerical differentiation and scalar root-finding to evaluate the induced link.
For reproducibility and ease of use, the repository includes runnable scripts under examples/ and an end-to-end Jupyter notebook (notebooks/GRR_end_to_end_examples.ipynb). These examples illustrate how to apply genriesz to common causal estimands such as the average treatment effect (grr_ate), average marginal effects (grr_ame), and average policy effects (grr_policy_effect).
There are other implementations of debiased machine learning. See also DoubleML Bach2022doublemlpython,Bach2024doublemlr and EconML Battocchi2019econmla.
This paper develops a unified perspective on estimating the Riesz representer, namely, the bias correction term that appears in Neyman orthogonal scores for a broad class of causal and structural parameters. We formulate Riesz representer estimation as fitting a model to the unknown representer under a Bregman divergence, which yields an empirical risk minimization objective that depends only on observed data. This generalized Riesz regression recovers Riesz regression and least squares importance fitting under squared loss, it recovers KL based tailored loss minimization and its dual entropy balancing weights under a KL type loss, and it connects logistic likelihood based propensity modeling with classification based density ratio estimation through a binary KL criterion. By pairing the loss with an appropriate link function, we make explicit a dual characterization that delivers automatic covariate balancing or moment matching, which clarifies when popular balancing schemes arise as primal or dual solutions. We provide convergence rate results for kernel methods and neural networks, including minimax optimality under standard smoothness classes, and we show how the framework instantiates in ATE, AME, APE, and covariate shift adaptation. Our experiments suggest that directly estimating the bias correction term can be competitive with common propensity score based baselines and can be stable across divergence choices when combined with cross fitting. Overall, the proposed framework bridges density ratio estimation and causal inference, and it offers a single set of tools for designing, analyzing, and implementing Riesz representer estimators, while motivating extensions such as nearest neighbor and score matching based constructions.